utils.hpp

#include <sys/socket.h>
#include <sys/epoll.h>

int create_and_bind_udp(uint16_t port);
int create_bind_and_listen_tcp(uint16_t port);
void print_addr(const sockaddr* peer_addr, socklen_t peer_addr_len);

// Base class for handling notifications from epoll.
class EPollHandlerInterface {
protected:
  const int sock_;

  EPollHandlerInterface(const int sock) : sock_(sock) {}

  int epoll_add(const int epollfd);

public:
  virtual ~EPollHandlerInterface();

  // If the result of `process_read` is negative, it means that the socket
  // should be closed and the epoll handler deleted.
  virtual int process_read(const epoll_event* event) = 0;

  // If the result of `process_write` is negative, it means that the socket
  // should be closed and the epoll handler deleted.
  virtual int process_write(const epoll_event* event) = 0;
};

// Rather than passing the UDP socket's file descriptor to
// RecvHandlerUDP, we pass this abstract interface so that we can
// restrict what the RecvHandlerUDP is able to do with the socket.
class SocketHandlerUDP {
public:
  virtual ~SocketHandlerUDP() {}

  // Send a reply. (Wrapper around sendto().)
  virtual ssize_t replyto(
    const char* buf, size_t buflen,
    const sockaddr *dest_addr, socklen_t addrlen
  ) = 0;
};

class RecvHandlerUDP {
public:
  virtual ~RecvHandlerUDP() {};

  // This method is called by `EpollRecvHandlerUDP::process_read`
  // when a message is received. Return 0 to keep receiving messages,
  // or a negative number to close the socket and delete the epoll
  // handler.
  virtual int receive(
    const uint8_t* buf, ssize_t len,
    SocketHandlerUDP& sock,
    const sockaddr* peer_addr, socklen_t peer_addr_len
  ) = 0;
};

// Rather than passing the TCP socket's file descriptor to
// RecvHandlerTCP, we pass this abstract interface so that we can
// restrict what the RecvHandlerTCP is able to do with the socket.
class SocketHandlerTCP {
public:
  virtual ~SocketHandlerTCP() {}

  // Send a reply. (Wrapper around send().)
  virtual ssize_t reply(const char* buf, size_t buflen) = 0;
};

class RecvHandlerTCP {
public:
  virtual ~RecvHandlerTCP() {};

  // This is method is called immediately after the connection is
  // accepted. The return value is the number of bytes expected on the next
  // message. But if the return value is negative then it means that the
  // socket should be closed and the epoll handler deleted.
  virtual ssize_t accept(SocketHandlerTCP& sock) = 0;

  // `EpollRecvHandlerTCP` calls this method when it has received the
  // number of bytes that we asked for. Unlike the corresponding method in
  // `RecvHandlerUDP`, this method does not need a `len` parameter, because
  // we will get the exact number of bytes that we asked for. (We always
  // have to tell `EpollRecvHandlerTCP` how big we are expecting the next
  // message to be.) The return value is the number of bytes expected on
  // the next message. But if the return value is negative then it means
  // that the socket should be closed and the epoll handler deleted.
  virtual ssize_t receive(SocketHandlerTCP& sock, const uint8_t* buf) = 0;

  // The socket disconnected before we were finished sending/receiving.
  virtual void disconnect() = 0;
};

// Factory class for constructing instances of RecvHandlerTCP.
class BuildRecvHandlerTCP {
public:
  virtual ~BuildRecvHandlerTCP() {};

  virtual RecvHandlerTCP* build(
    sockaddr* peer_addr, socklen_t peer_addr_len
  ) = 0;
};

// Implementation of EPollHandlerInterface which reads data from
// a file descriptor. The data that was read is passed to the `handler_`
// field for further processing.
class EpollRecvHandlerUDP :
  public EPollHandlerInterface,
  virtual private SocketHandlerUDP
{
  // Handler pointer is owned by this class.
  RecvHandlerUDP* handler_;

  // Takes ownership of `handler`.
  EpollRecvHandlerUDP(const int sock, RecvHandlerUDP* handler)
    : EPollHandlerInterface(sock)
    , handler_(handler)
  {}

  // Implements SocketHandlerUDP interface.
  ssize_t replyto(
    const char* buf, size_t buflen,
    const sockaddr *dest_addr, socklen_t addrlen
  ) override;

public:
  // Factory method. Constructs the class and registers the file descriptor
  // with epoll. Returns zero on success, otherwise negative.
  static int build(
    const int epollfd, const int sock, RecvHandlerUDP* handler
  );

  virtual ~EpollRecvHandlerUDP();

  int process_read(const epoll_event *) override;

  int process_write(const epoll_event *) override {
    // TODO: implement buffering for UDP sends, similar
    // to how it is done in EpollRecvHandlerTCP.
    return 0;
  }
};

// Implementation of EPollHandlerInterface which reads data from
// a file descriptor. The data that was read is passed to the `handler_`
// field for further processing.
class EpollRecvHandlerTCP :
  public EPollHandlerInterface,
  virtual private SocketHandlerTCP
{
  // Handler pointer is owned by this class.
  RecvHandlerTCP* handler_;

  uint8_t* recvbuf_;  // Buffer for incoming data.
  size_t received_;   // Number of bytes received so far. (Stored in `recvbuf_`.)
  size_t remaining_;  // Number of bytes we are still waiting for.

  uint8_t* sendbuf_;    // Buffer for outgoing data.
  size_t sendbufsize_;  // Total size of sendbuf_
  size_t sendbufpos_;   // Current position in sendbuf_

  // Takes ownership of `handler`.
  EpollRecvHandlerTCP(const int sock, RecvHandlerTCP* handler);

  // This method is called after the class has been constructed and
  // registered with epoll. Returns zero on success, otherwise negative.
  int accept();

  // Implements SocketHandlerTCP interface.
  ssize_t reply(const char* buf, size_t buflen) override;

public:
  virtual ~EpollRecvHandlerTCP();

  // Factory method. Constructs the class and registers the file descriptor
  // with epoll. Returns zero on success, otherwise negative.
  static int build(
    const int epollfd, const int sock, RecvHandlerTCP* handler
  );

  int process_read(const epoll_event *) override;
  int process_write(const epoll_event *) override;
};

// Implementation of EPollHandlerInterface which handles TCP connection
// request.
class EpollTcpConnectHandler : public EPollHandlerInterface {
  const int epollfd_;

  // Factory pointer is owned by this class.
  BuildRecvHandlerTCP* factory_;

  // Takes ownership of `factory`.
  EpollTcpConnectHandler(
    const int epollfd, const int sock, BuildRecvHandlerTCP* factory
  );

public:
  // Factory method. Constructs the class and registers the file descriptor
  // with epoll. Returns zero on success, otherwise negative.
  static int build(
    const int epollfd, const int sock, BuildRecvHandlerTCP* factory
  );

  virtual ~EpollTcpConnectHandler();

  int process_read(const epoll_event *) override;
  int process_write(const epoll_event *) override;
};

// If a TCP conversation is always exactly the same every time, then we can
// make it run faster by caching it. We use this class to cache the data
// from the first run, so that we can just play it back on subsequent
// runs. The proxy class TCP_Cache_Record is used to initialize the cache
// and TCP_Cache_Playback is used for playback.
class TCP_Cache {
  char* sendbuf_;
  size_t total_send_;
  size_t total_receive_;

public:
  TCP_Cache();
  ~TCP_Cache();

  bool isInitialized() const;

  // Called by EpsonHandlerCacheProxyTCP after it has collected
  // the relevant data.
  void init(char* sendbuf, size_t total_send, size_t total_receive);

  ssize_t accept(SocketHandlerTCP& sock);
};

// Implementation of RecvHandlerTCP which just runs a previously cached
// conversation: it sends the same sequence of bytes that were sent before
// and doesn't bother to read the incoming messages. It just waits until is
// has received the expected number of bytes and then closes the
// connection.
class TCP_Cache_Playback : public RecvHandlerTCP {
  TCP_Cache& cache_;

public:
  explicit TCP_Cache_Playback(TCP_Cache& cache);

  ssize_t accept(SocketHandlerTCP& sock) override;

  ssize_t receive(SocketHandlerTCP&, const uint8_t*) override {
    return 0;
  }

  void disconnect() override {}
};

// This proxy class is used to initialize a TCP_Cache.  It intercepts all
// the data sent and received by `handler_`, which is the RecvHandlerTCP
// which we are proxying.
class TCP_Cache_Record : public RecvHandlerTCP {
  // The object that we are proxying. We own this pointer.
  RecvHandlerTCP* handler_;

  // If caching is successful, then we call the `init()`
  // method on this object.
  TCP_Cache& cache_;

  char* sendbuf_;
  size_t total_send_;
  size_t total_receive_;

  class SocketHandlerProxy;

public:
  TCP_Cache_Record(RecvHandlerTCP* handler, TCP_Cache& cache);
  ~TCP_Cache_Record();

  ssize_t accept(SocketHandlerTCP& sock) override;

  ssize_t receive(SocketHandlerTCP& sock, const uint8_t* buf) override;

  void disconnect() override {}
};

// This function runs an infinite loop, listening for epoll events. You
// need to use `epoll_create1` to initialize epoll before calling this
// function. You also need to open the relevant sockets and register them
// with epoll before calling this function, for example by calling
// `EpollTcpConnectHandler::build()`.
[[ noreturn ]] void epoll_main_loop(const int epollfd);