#include "socks5.h"

void socks5_free_ctx(void* elem) {
  struct socks5_ctx* ctx = elem;
  free(ctx);
}

void socks5_create_dns_client(struct evt_core_ctx* ctx, char* proxy_host, char* proxy_port, char* addr, uint16_t port) {
  struct evt_core_fdinfo fdinfo;
  struct evt_core_cat cat;
  struct socks5_ctx* s5ctx;
  struct evt_core_fdinfo* reg_fdinfo;
  char url[1024];

  // 0. Compute domain length and enforce an upper bound on its size
  size_t domainLength = strlen(addr);
  if (domainLength > 255) {
    fprintf(stderr, "domain is too long\n");
    exit(EXIT_FAILURE);
  }

  // 1. Open connection
  int sock = create_tcp_client (proxy_host, proxy_port);
  if (sock < 0) {
    fprintf(stderr, "Unable to connect to proxy %s:%s\n", proxy_host, proxy_port);
    exit(EXIT_FAILURE);
  }

  // 2. Create fdinfo
  fdinfo.cat = &cat;
  fdinfo.cat->name = "socks5-send-handshake";
  fdinfo.fd = sock;
  fdinfo.other = malloc(sizeof(struct socks5_ctx));
  if (fdinfo.other == NULL) {
    perror("malloc failed");
    exit(EXIT_FAILURE);
  }
  memset(fdinfo.other, 0, sizeof(struct socks5_ctx));
  fdinfo.free_other = socks5_free_ctx;
  sprintf(url, "socks5:send-hs:%s:%d", addr, port);
  fdinfo.url = url;

  // 3. Fill socks5_ctx structures
  s5ctx = fdinfo.other;
  s5ctx->port = port;
  s5ctx->addr = strdup(addr);

  // 3.1 Client handshake to send
  s5ctx->ch.ver = VER_SOCKS5;
  s5ctx->ch.nmethods = 0x01;
  s5ctx->ch.methods[0] = METHOD_NOAUTH;
  s5ctx->ch_size = sizeof(uint8_t) * (2 + s5ctx->ch.nmethods);

  // 3.2 Client request to send
  s5ctx->cr.ver = VER_SOCKS5;
  s5ctx->cr.cmd = CMD_CONNECT;
  s5ctx->cr.rsv = 0x00;
  s5ctx->cr.atyp = ATYP_DOMAINNAME;
  s5ctx->cr.dst_addr_len = domainLength;
  s5ctx->cr.dst_addr = addr;
  s5ctx->cr.port = htons(port);

  // 3.3 Generate client request buffer
  s5ctx->cr_size = 0;
  fill_buffer(&s5ctx->cr_size, s5ctx->cr_buffer, &s5ctx->cr, 5*sizeof(uint8_t));
  fill_buffer(&s5ctx->cr_size, s5ctx->cr_buffer, s5ctx->cr.dst_addr, s5ctx->cr.dst_addr_len*sizeof(char));
  fill_buffer(&s5ctx->cr_size, s5ctx->cr_buffer, &s5ctx->cr.port, sizeof(uint16_t));

  reg_fdinfo = evt_core_add_fd (ctx, &fdinfo);
}

int on_socks5_send_handshake(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
  struct socks5_ctx* s5ctx = fdinfo->other;

  size_t written = write(fdinfo->fd, (char*)&s5ctx->ch + s5ctx->ch_cursor, s5ctx->ch_size - s5ctx->ch_cursor);
  if (written == -1 && errno == EAGAIN) return 1;
  if (written < 0) {
    perror("write failed on tcp socket in socks5");
    evt_core_mv_fd2(ctx, fdinfo, "socks5-failed");
    return 1;
  }
  s5ctx->ch_cursor += written;
  if (s5ctx->ch_cursor < s5ctx->ch_size) return 0;

  evt_core_mv_fd2(ctx, fdinfo, "socks5-recv-handshake");
  return 1;
}

int on_socks5_recv_handshake(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
  struct socks5_ctx* s5ctx = fdinfo->other;
  int readn = 0;

  readn = read(fdinfo->fd, (char*)&s5ctx->sh + s5ctx->sh_cursor, sizeof(s5ctx->sh) - s5ctx->sh_cursor);
  if (readn == -1 && errno == EAGAIN) return 1;
  if (readn < 0) {
    perror("sock5 handshake failed read");
    evt_core_mv_fd2(ctx, fdinfo, "socks5-failed");
    return 1;
  }

  s5ctx->sh_cursor += readn;
  if (s5ctx->sh_cursor < sizeof(s5ctx->sh)) return 0;

  if (s5ctx->ch.ver != s5ctx->sh.ver || s5ctx->sh.method != s5ctx->ch.methods[0]) {
    fprintf(stderr, "Protocol error: client asks for ver=%d, method=%d and server answers with ver=%d, method=%d\n",
              s5ctx->ch.ver, s5ctx->ch.methods[0], s5ctx->sh.ver, s5ctx->sh.method);
    evt_core_mv_fd2(ctx, fdinfo, "socks5-failed");
    return 1;
  }
  printf("[socks5_server_handshake] fd=%d, ver=%d, method=%d\n", fdinfo->fd, s5ctx->sh.ver, s5ctx->sh.method);
  evt_core_mv_fd2(ctx, fdinfo, "socks5-send-client-req");
  return 1;
}

int on_socks5_send_client_req(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
  struct socks5_ctx* s5ctx = fdinfo->other;
  int written = 0;

  written = write(fdinfo->fd, (char*)s5ctx->cr_buffer + s5ctx->cr_cursor, s5ctx->cr_size - s5ctx->cr_cursor);
  if (written == -1 && errno == EAGAIN) return 1;
  if (written < 0) {
    fprintf(stderr, "socks5 send client request failed\n");
    evt_core_mv_fd2 (ctx, fdinfo, "socks5-failed");
    return 1;
  }
  s5ctx->cr_cursor += written;
  if (s5ctx->cr_cursor < s5ctx->cr_size) return 0;

  evt_core_mv_fd2 (ctx, fdinfo, "socks5-recv-server-reply");
  return 1;
}

int socks5_server_reply_atyp_ipv4(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
  struct socks5_ctx* s5ctx = fdinfo->other;
  size_t fixed_headers_size = (char*)&s5ctx->sr.bind_addr.ipv4 - (char*)&s5ctx->sr;
  size_t host_size = sizeof(s5ctx->sr.bind_addr.ipv4);
  uint64_t relative_cursor = (s5ctx->sr_cursor - fixed_headers_size);
  int nread = 0;

  nread = read(fdinfo->fd,
       (char*)s5ctx->sr.bind_addr.ipv4 + relative_cursor,
       host_size - relative_cursor);

  if (nread == -1 && errno == EAGAIN) return 1;
  if (nread < 0) {
    perror("write failed on tcp socket in socks5");
    evt_core_mv_fd2(ctx, fdinfo, "socks5-failed");
    return 1;
  }

  s5ctx->sr_cursor += nread;
  if (s5ctx->sr_cursor < fixed_headers_size + host_size) return 0;

  s5ctx->sr_host_read = 1;
  return 0;
}

int socks5_server_reply_atyp_ipv6(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
  struct socks5_ctx* s5ctx = fdinfo->other;
  size_t fixed_headers_size = (char*)&s5ctx->sr.bind_addr.ipv6 - (char*)&s5ctx->sr;
  size_t host_size = sizeof(s5ctx->sr.bind_addr.ipv6);
  uint64_t relative_cursor = (s5ctx->sr_cursor - fixed_headers_size);
  int nread = 0;

  nread = read(fdinfo->fd,
       (char*)s5ctx->sr.bind_addr.ipv6 + relative_cursor,
       host_size - relative_cursor);

  if (nread == -1 && errno == EAGAIN) return 1;
  if (nread < 0) {
    perror("write failed on tcp socket in socks5");
    evt_core_mv_fd2(ctx, fdinfo, "socks5-failed");
    return 1;
  }
  s5ctx->sr_cursor += nread;
  if (s5ctx->sr_cursor < fixed_headers_size + host_size) return 0;

  s5ctx->sr_host_read = 1;
  return 0;
}

int socks5_server_reply_atyp_dn(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
  struct socks5_ctx* s5ctx = fdinfo->other;
  size_t fixed_headers_size = (char*)&s5ctx->sr.bind_addr.ipv6 - (char*)&s5ctx->sr;
  size_t dn_size_size = sizeof(s5ctx->sr.bind_addr.dns.len);
  uint64_t relative_cursor = 0;
  int nread = 0;

  if (s5ctx->sr_cursor < fixed_headers_size + dn_size_size) {
    relative_cursor = (s5ctx->sr_cursor - fixed_headers_size);
    nread = read(fdinfo->fd, (char*)&s5ctx->sr.bind_addr.dns.len + relative_cursor, dn_size_size - relative_cursor);
    if (nread == -1 && errno == EAGAIN) return 1;
    if (nread < 0) {
      perror("write failed on tcp socket in socks5");
      evt_core_mv_fd2(ctx, fdinfo, "socks5-failed");
      return 1;
    }
    s5ctx->sr_cursor += nread;
    return 0;
  }

  relative_cursor = s5ctx->sr_cursor - fixed_headers_size - sizeof(s5ctx->sr.bind_addr.dns.len);
  nread = read(fdinfo->fd, (char*)&s5ctx->sr.bind_addr.dns.str + relative_cursor, s5ctx->sr.bind_addr.dns.len - relative_cursor);
  if (nread == -1 && errno == EAGAIN) return 1;
  if (nread < 0) {
    perror("write failed on tcp socket in socks5");
    evt_core_mv_fd2(ctx, fdinfo, "socks5-failed");
    return 1;
  }
  s5ctx->sr_cursor += nread;
  if (s5ctx->sr_cursor < fixed_headers_size + dn_size_size + s5ctx->sr.bind_addr.dns.len) return 0;

  s5ctx->sr_host_read = 1;
  return 0;
}

size_t socks5_server_reply_size(struct server_reply* sr) {
  size_t fixed_headers_size = (char*)&sr->bind_addr - (char*)sr;
  size_t fixed_tail_size = (char*)(sr + 1) - (char*)&sr->port;
  size_t host_size = 0;

  if (sr->atyp == ATYP_IPV4) {
    host_size = sizeof(sr->bind_addr.ipv4);
  } else if (sr->atyp == ATYP_IPV6) {
    host_size = sizeof(sr->bind_addr.ipv6);
  } else if (sr->atyp == ATYP_DOMAINNAME) {
    host_size = sizeof(sr->bind_addr.dns.len) + sr->bind_addr.dns.len;
  } else {
    fprintf(stderr, "Unsupported ATYP for SOCK5\n");
    exit(EXIT_FAILURE);
  }

  return fixed_headers_size + host_size + fixed_tail_size;
}

int on_socks5_recv_server_reply(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
  struct socks5_ctx* s5ctx = fdinfo->other;
  int readn = 0;
  size_t fixed_headers_size = (char*)&s5ctx->sr.bind_addr - (char*)&s5ctx->sr;

  // Read headers
  if (s5ctx->sr_cursor < fixed_headers_size) {
    readn = read(fdinfo->fd, (char*)&s5ctx->sr + s5ctx->sr_cursor, fixed_headers_size - s5ctx->sr_cursor);
    if (readn == -1 && errno == EAGAIN) return 1;
    if (readn < 0) goto move_to_failed;
    s5ctx->sr_cursor += readn;
    return 0; // Needed as we might have not read enough bytes and free us from writing a loop
  }

  // Read host
  if (!s5ctx->sr_host_read) {
    if (s5ctx->sr.atyp == ATYP_IPV4) return socks5_server_reply_atyp_ipv4(ctx, fdinfo);
    else if (s5ctx->sr.atyp == ATYP_IPV6) return socks5_server_reply_atyp_ipv6(ctx, fdinfo);
    else if (s5ctx->sr.atyp == ATYP_DOMAINNAME) return socks5_server_reply_atyp_dn(ctx, fdinfo);
    else goto move_to_failed;
  }

  // Read port
  size_t final_size = socks5_server_reply_size(&s5ctx->sr);
  if (s5ctx->sr_cursor < final_size) {
    size_t relative_cursor = s5ctx->sr_cursor - (final_size - sizeof(s5ctx->sr.port));
    readn = read(fdinfo->fd, (char*)&s5ctx->sr.port + relative_cursor, sizeof(s5ctx->sr.port) - relative_cursor);
    if (readn == -1 && errno == EAGAIN) return 1;
    if (readn < 0) goto move_to_failed;
    s5ctx->sr_cursor += readn;
    return 0; // Needed as we might have not read enough bytes and free us from writing a loop
  }

  // Do some checks
  if (s5ctx->sr.rep > 0x08) goto move_to_failed;
  printf("[socks5_server_reply] fd=%d, ver=%d, rep=%s, atyp=%d, port=%d\n", fdinfo->fd, s5ctx->sr.ver, rep_msg[s5ctx->sr.rep], s5ctx->sr.atyp, s5ctx->sr.port);

  if (s5ctx->sr.rep != SOCKS5_REP_SUCCESS) goto move_to_failed;

  evt_core_mv_fd2 (ctx, fdinfo, "socks5-success");
  return 1;
move_to_failed:
  evt_core_mv_fd2 (ctx, fdinfo, "socks5-failed");
  return 1;
}

int on_socks5_err(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
  evt_core_mv_fd2 (ctx, fdinfo, "socks5-failed");
  return 1;
}

void socks5_init(struct evt_core_ctx* ctx) {
  struct evt_core_cat template = {0};

  template.cb = on_socks5_send_handshake;
  template.err_cb = on_socks5_err;
  template.name = "socks5-send-handshake";
  template.flags = EPOLLOUT | EPOLLET;
  evt_core_add_cat (ctx, &template);

  template.cb = on_socks5_recv_handshake;
  template.err_cb = on_socks5_err;
  template.name = "socks5-recv-handshake";
  template.flags = EPOLLIN | EPOLLET;
  evt_core_add_cat (ctx, &template);

  template.cb = on_socks5_send_client_req;
  template.err_cb = on_socks5_err;
  template.name = "socks5-send-client-req";
  template.flags = EPOLLOUT | EPOLLET;
  evt_core_add_cat(ctx, &template);

  template.cb = on_socks5_recv_server_reply;
  template.err_cb = on_socks5_err;
  template.name =  "socks5-recv-server-reply";
  template.flags = EPOLLIN | EPOLLET;
  evt_core_add_cat(ctx, &template);
}

char* socks5_rep (enum socks5_rep rep) {
  return rep_msg[rep];
}