TcpStream.cpp

#include "TcpStream.hpp"
#include <cstring>
#include <iostream>
#include <stdexcept>
#include <string>

#ifdef _WIN32
typedef int RECV_SEND_T;
#else
typedef ssize_t RECV_SEND_T;
#endif

TcpStream::TcpStream(const std::string &host, int port) {
#ifdef _WIN32
  WSADATA wsa_data;
  if (WSAStartup(MAKEWORD(1, 1), &wsa_data) != 0) {
    throw std::runtime_error("Failed to initialize sockets");
  }
#endif
  sock = socket(AF_INET, SOCK_STREAM, 0);
  if (sock == -1) {
    throw std::runtime_error("Failed to create socket");
  }
  int yes = 1;
  if (setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, (char *)&yes, sizeof(int)) <
      0) {
    throw std::runtime_error("Failed to set TCP_NODELAY");
  }
  addrinfo hints, *servinfo;
  std::memset(&hints, 0, sizeof(hints));
  hints.ai_family = AF_INET;
  hints.ai_socktype = SOCK_STREAM;
  if (getaddrinfo(host.c_str(), std::to_string(port).c_str(), &hints,
                  &servinfo) != 0) {
    throw std::runtime_error("Failed to get addr info");
  }
  if (connect(sock, servinfo->ai_addr, servinfo->ai_addrlen) == -1) {
    throw std::runtime_error("Failed to connect");
  }
  freeaddrinfo(servinfo);
}

class TcpInputStream : public InputStream {
public:
  TcpInputStream(std::shared_ptr<TcpStream> tcpStream)
      : tcpStream(tcpStream), bufferPos(0), bufferSize(0) {}
  void readBytes(char *buffer, size_t byteCount) {
    while (byteCount > 0) {
      if (bufferSize > 0) {
        if (bufferSize >= byteCount) {
          memcpy(buffer, this->buffer + bufferPos, byteCount);
          bufferPos += byteCount;
          bufferSize -= byteCount;
          return;
        }
        memcpy(buffer, this->buffer + bufferPos, bufferSize);
        buffer += bufferSize;
        byteCount -= bufferSize;
        bufferPos += bufferSize;
        bufferSize = 0;
      }
      if (bufferPos == BUFFER_CAPACITY) {
        bufferPos = 0;
      }
      RECV_SEND_T received =
          recv(tcpStream->sock, this->buffer + bufferPos + bufferSize,
               BUFFER_CAPACITY - bufferPos - bufferSize, 0);
      if (received < 0) {
        throw std::runtime_error("Failed to read from socket");
      }
      bufferSize += received;
    }
  }

private:
  static const size_t BUFFER_CAPACITY = 8 * 1024;
  char buffer[BUFFER_CAPACITY];
  size_t bufferPos;
  size_t bufferSize;
  std::shared_ptr<TcpStream> tcpStream;
};

class TcpOutputStream : public OutputStream {
public:
  TcpOutputStream(std::shared_ptr<TcpStream> tcpStream)
      : tcpStream(tcpStream), bufferPos(0), bufferSize(0) {}
  void writeBytes(const char *buffer, size_t byteCount) {
    while (byteCount > 0) {
      size_t capacity = BUFFER_CAPACITY - bufferPos - bufferSize;
      if (capacity >= byteCount) {
        memcpy(this->buffer + bufferPos + bufferSize, buffer, byteCount);
        bufferSize += byteCount;
        return;
      }
      memcpy(this->buffer + bufferPos + bufferSize, buffer, capacity);
      bufferSize += capacity;
      byteCount -= capacity;
      buffer += capacity;
      flush();
    }
  }
  void flush() {
    while (bufferSize > 0) {
      RECV_SEND_T sent =
          send(tcpStream->sock, buffer + bufferPos, bufferSize, 0);
      if (sent < 0) {
        throw std::runtime_error("Failed to write to socket");
      }
      bufferPos += sent;
      bufferSize -= sent;
    }
    bufferPos = 0;
  }

private:
  static const size_t BUFFER_CAPACITY = 8 * 1024;
  char buffer[BUFFER_CAPACITY];
  size_t bufferPos;
  size_t bufferSize;
  std::shared_ptr<TcpStream> tcpStream;
};

std::shared_ptr<InputStream>
getInputStream(std::shared_ptr<TcpStream> tcpStream) {
  return std::shared_ptr<TcpInputStream>(new TcpInputStream(tcpStream));
}

std::shared_ptr<OutputStream>
getOutputStream(std::shared_ptr<TcpStream> tcpStream) {
  return std::shared_ptr<TcpOutputStream>(new TcpOutputStream(tcpStream));
}