#include "socks5.h"

struct socks5_ctx {
  struct client_handshake ch;
  struct server_handshake sh;
  struct client_request cr;
  struct server_reply sr;
  uint64_t ch_cursor;
  uint64_t sh_cursor;
  uint64_t cr_cursor;
  uint64_t sr_cursor;
  char cr_buffer[262];
  size_t ch_size;
  size_t cr_size;
};

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

void create_socks5_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;
  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 = strdup(url);

  // 3. Fill socks5_ctx structures
  s5ctx = fdinfo.other;

  // 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));

  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, &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 0;
  }
  s5ctx->ch_cursor += written;
  if (s5ctx->ch_cursor < s5ctx->ch_size) return 0;

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

int on_socks5_recv_handshake(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {

    if (err == -1) {
        perror("sock5 handshake failed read");
        return -1;
    }
    if (ch.ver != sh.ver || sh.method != ch.methods[0]) {
      fprintf(stderr, "Protocol error: client asks for ver=%d, method=%d and server answers with ver=%d, method=%d\n",
              ch.ver, ch.methods[0], sh.ver, sh.method);
      return -1;
    }
    printf("[server_handshake] ver=%d, method=%d\n", sh.ver, sh.method);
    return 0;
}

int socks5_reply(int sock) {
    int res;
    struct server_reply sr = {0};
    res = read_entity(sock, &sr, sizeof(uint8_t) * 4);
    if (res == -1) goto read_error;

    switch(sr.atyp) {
        case ATYP_IPV4:
            if (read_entity(sock, sr.bind_addr.ipv4, sizeof(uint8_t) * 4) == -1)
                goto read_error;
            break;
        case ATYP_DOMAINNAME:
            if (read_entity(sock, &sr.bind_addr.dns.len, sizeof(uint8_t) * 4) == -1)
                goto read_error;
            if (read_entity(sock, sr.bind_addr.dns.str, sizeof(char) * sr.bind_addr.dns.len) == -1)
                goto read_error;
            break;
        case ATYP_IPV6:
            if (read_entity(sock, sr.bind_addr.ipv6, sizeof(uint8_t) * 16) == -1)
                goto read_error;
            break;
        default:
            fprintf(stderr, "Unsupported ATYP in server reply\n");
            return -128;
    }
    res = read_entity(sock, &sr.port, sizeof(uint16_t));
    if (res == -1) {
      perror("read_entity");
      exit(EXIT_FAILURE);
    }

    if (sr.rep < 0 || sr.rep > 0x08) {
      fprintf(stderr, "Invalid reply field\n");
      return -128;
    }

    printf("[server_reply] ver=%d, rep=%s, atyp=%d, port=%d\n", sr.ver, rep_msg[sr.rep], sr.atyp, sr.port);
    return -sr.rep;

read_error:
    fprintf(stderr, "Unable to read ATYP\n");
    return -128;
}

int socks5_connect_dns(int sock, char* addr, uint16_t port) {
    char buffer[262];
    size_t domainLength = strlen(addr);
    if (domainLength > 255) {
        fprintf(stderr, "domain is too long\n");
        exit(EXIT_FAILURE);
    }
    struct client_request cr = { 0x05, CMD_CONNECT, 0x00, ATYP_DOMAINNAME, (uint8_t)domainLength, addr, htons(port) };
    size_t written = 0;
    fill_buffer(&written, buffer, &cr, 5*sizeof(uint8_t));
    fill_buffer(&written, buffer, cr.dst_addr, cr.dst_addr_len*sizeof(char));
    fill_buffer(&written, buffer, &cr.port, sizeof(uint16_t));

    if (written != write(sock, buffer, written)) {
        fprintf(stderr, "partial/failed write\n");
        return -1;
    }
    return 0;
}

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