Linux C/C++网络编程实战-陈硕-笔记4-TTCP代码概述

TCP block - c语言

• 代码

  #include "examples/ace/ttcp/common.h"
  #include "muduo/base/Timestamp.h"
  #include "muduo/base/Types.h"
  
  #undef NDEBUG
  
  #include <assert.h>
  #include <errno.h>
  #include <stdio.h>
  #include <unistd.h>
  
  #include <netinet/in.h>
  #include <arpa/inet.h>
  
  static int acceptOrDie(uint16_t port)
  {
    int listenfd = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    assert(listenfd >= 0);
  
    int yes = 1;
    if (setsockopt(listenfd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes)))
    {
      perror("setsockopt");
      exit(1);
    }
  
    struct sockaddr_in addr;
    muduo::memZero(&addr, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = INADDR_ANY;
    if (bind(listenfd, reinterpret_cast<struct sockaddr*>(&addr), sizeof(addr)))
    {
      perror("bind");
      exit(1);
    }
  
    if (listen(listenfd, 5))
    {
      perror("listen");
      exit(1);
    }
  
    struct sockaddr_in peer_addr;
    muduo::memZero(&peer_addr, sizeof(peer_addr));
    socklen_t addrlen = 0;
    int sockfd = ::accept(listenfd, reinterpret_cast<struct sockaddr*>(&peer_addr), &addrlen);
    if (sockfd < 0)
    {
      perror("accept");
      exit(1);
    }
    ::close(listenfd);
    return sockfd;
  }
  
  static int write_n(int sockfd, const void* buf, int length)
  {
    int written = 0;
    while (written < length)
    {
      ssize_t nw = ::write(sockfd, static_cast<const char*>(buf) + written, length - written);
      if (nw > 0)
      {
        written += static_cast<int>(nw);
      }
      else if (nw == 0)
      {
        break;  // EOF
      }
      else if (errno != EINTR)
      {
        perror("write");
        break;
      }
    }
    return written;
  }
  
  static int read_n(int sockfd, void* buf, int length)
  {
    int nread = 0;
    while (nread < length)
    {
      ssize_t nr = ::read(sockfd, static_cast<char*>(buf) + nread, length - nread);
      if (nr > 0)
      {
        nread += static_cast<int>(nr);
      }
      else if (nr == 0)
      {
        break;  // EOF
      }
      else if (errno != EINTR)
      {
        perror("read");
        break;
      }
    }
    return nread;
  }
  
  void transmit(const Options& opt)
  {
    struct sockaddr_in addr = resolveOrDie(opt.host.c_str(), opt.port);
    printf("connecting to %s:%d\n", inet_ntoa(addr.sin_addr), opt.port);
  
    int sockfd = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    assert(sockfd >= 0);
    int ret = ::connect(sockfd, reinterpret_cast<struct sockaddr*>(&addr), sizeof(addr));
    if (ret)
    {
      perror("connect");
      printf("Unable to connect %s\n", opt.host.c_str());
      ::close(sockfd);
      return;
    }
  
    printf("connected\n");
    muduo::Timestamp start(muduo::Timestamp::now());
    struct SessionMessage sessionMessage = { 0, 0 };
    sessionMessage.number = htonl(opt.number);
    sessionMessage.length = htonl(opt.length);
    if (write_n(sockfd, &sessionMessage, sizeof(sessionMessage)) != sizeof(sessionMessage))
    {
      perror("write SessionMessage");
      exit(1);
    }
  
    const int total_len = static_cast<int>(sizeof(int32_t) + opt.length);
    PayloadMessage* payload = static_cast<PayloadMessage*>(::malloc(total_len));
    assert(payload);
    payload->length = htonl(opt.length);
    for (int i = 0; i < opt.length; ++i)
    {
      payload->data[i] = "0123456789ABCDEF"[i % 16];
    }
  
    double total_mb = 1.0 * opt.length * opt.number / 1024 / 1024;
    printf("%.3f MiB in total\n", total_mb);
  
    for (int i = 0; i < opt.number; ++i)
    {
      int nw = write_n(sockfd, payload, total_len);
      assert(nw == total_len);
  
      int ack = 0;
      int nr = read_n(sockfd, &ack, sizeof(ack));
      assert(nr == sizeof(ack));
      ack = ntohl(ack);
      assert(ack == opt.length);
    }
  
    ::free(payload);
    ::close(sockfd);
    double elapsed = timeDifference(muduo::Timestamp::now(), start);
    printf("%.3f seconds\n%.3f MiB/s\n", elapsed, total_mb / elapsed);
  }
  
  
  /* 服务端：接收 */
  void receive(const Options& opt)
  {
    int sockfd = acceptOrDie(opt.port);                                                    // 接收连接
  
    struct SessionMessage sessionMessage = { 0, 0 };                                       // 接收消息，comment.h 中定义 8 字节
    if (read_n(sockfd, &sessionMessage, sizeof(sessionMessage)) != sizeof(sessionMessage)) // 保证完整的读取 8 字节
    {
      perror("read SessionMessage");
      exit(1);
    }
  
    /* sessionMessage.number 条消息，每条消息长 sessionMessage.length */ 
    sessionMessage.number = ntohl(sessionMessage.number);                                 // 消息由网络字节序 转至 本机字节序
    sessionMessage.length = ntohl(sessionMessage.length);
    printf("receive number = %d\nreceive length = %d\n",                                  // 服务端打印预计接收的数据大小
           sessionMessage.number, sessionMessage.length);
    const int total_len = static_cast<int>(sizeof(int32_t) + sessionMessage.length);      // 准备缓冲区接收
                                                                                          // 安全漏洞，对方发送的length值很大，造成缓冲区。 这里可以进行一个判断，限制最大length长度
    PayloadMessage* payload = static_cast<PayloadMessage*>(::malloc(total_len));
    assert(payload);
  
    for (int i = 0; i < sessionMessage.number; ++i)                                       //  number 次消息循环读取
    {
      payload->length = 0;
      if (read_n(sockfd, &payload->length, sizeof(payload->length)) != sizeof(payload->length)) // 先读4字节的， 这里存放每条消息的长度
      {
        perror("read length");
        exit(1);
      }
      payload->length = ntohl(payload->length);
      assert(payload->length == sessionMessage.length);                                   // 每次消息的 len 应该与 sessionMessage.len 一致
      
      /* 读消息 */
      if (read_n(sockfd, payload->data, payload->length) != payload->length)              // 读剩余的消息
      {
        perror("read payload data");
        exit(1);
      }
      
      /* 响应字段，返回本次收到的字节数 */
      int32_t ack = htonl(payload->length);
      if (write_n(sockfd, &ack, sizeof(ack)) != sizeof(ack))
      {
        perror("write ack");
        exit(1);
      }
    }
    
    ::free(payload);
    ::close(sockfd);
  }