From a2bf7ddec148ec0a26d7b8462ca5d48c074baadb Mon Sep 17 00:00:00 2001
From: Quentin Dufour <quentin.dufour@inria.fr>
Date: Fri, 22 Mar 2019 17:14:35 +0100
Subject: [PATCH] WIP rewrite socks5 file (broken)

---
 src/evt_core.c  |   7 +-
 src/evt_core.h  |   2 +-
 src/net_tools.c |  10 ++
 src/net_tools.h |   1 +
 src/socks5.c    | 237 ++++++++++++++++++++++++++++++++++++------------
 src/socks5.h    |   2 +-
 6 files changed, 198 insertions(+), 61 deletions(-)

diff --git a/src/evt_core.c b/src/evt_core.c
index edd2b7a..ff6f942 100644
--- a/src/evt_core.c
+++ b/src/evt_core.c
@@ -77,6 +77,9 @@ void evt_core_mv_fd(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo, st
 
   // 3. Update new category
   g_array_append_val (fdinfo->cat->socklist, fdinfo);
+
+  // 4. Update epoll flags
+  update_fd_epoll (ctx->epollfd, fdinfo->fd, fdinfo->cat->flags);
 }
 
 void evt_core_mv_fd2(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo, char* to_cat) {
@@ -88,7 +91,7 @@ void evt_core_mv_fd2(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo, c
   evt_core_mv_fd (ctx, fdinfo, cat);
 }
 
-void evt_core_add_fd(struct evt_core_ctx* ctx, struct evt_core_fdinfo* user_data) {
+struct evt_core_fdinfo* evt_core_add_fd(struct evt_core_ctx* ctx, struct evt_core_fdinfo* user_data) {
   // 1. Fetch fd category
   struct evt_core_cat* cat = g_hash_table_lookup(ctx->catlist, user_data->cat->name);
   if (cat == NULL) {
@@ -123,6 +126,8 @@ void evt_core_add_fd(struct evt_core_ctx* ctx, struct evt_core_fdinfo* user_data
   // 5. Add file descriptor to epoll
   add_fd_to_epoll(ctx->epollfd, user_data->fd, cat->flags);
   printf("Added fd=%d with url=%s in cat=%s\n", fdinfo->fd, fdinfo->url, fdinfo->cat->name);
+
+  return fdinfo;
 }
 
 struct evt_core_cat* evt_core_rm_fd(struct evt_core_ctx* ctx, int fd) {
diff --git a/src/evt_core.h b/src/evt_core.h
index 6aaf4cc..0b5b2ac 100644
--- a/src/evt_core.h
+++ b/src/evt_core.h
@@ -47,7 +47,7 @@ void evt_core_init(struct evt_core_ctx* ctx);
 void evt_core_add_cat(struct evt_core_ctx* ctx, struct evt_core_cat* cat);
 void evt_core_mv_fd(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo, struct evt_core_cat* to_cat);
 void evt_core_mv_fd2(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo, char* to_cat);
-void evt_core_add_fd(struct evt_core_ctx* ctx, struct evt_core_fdinfo* user_data);
+struct evt_core_fdinfo* evt_core_add_fd(struct evt_core_ctx* ctx, struct evt_core_fdinfo* user_data);
 struct evt_core_cat* evt_core_rm_fd(struct evt_core_ctx* ctx, int fd);
 void evt_core_free(struct evt_core_ctx* ctx);
 void evt_core_loop(struct evt_core_ctx* ctx);
diff --git a/src/net_tools.c b/src/net_tools.c
index dbee61d..cf1c419 100644
--- a/src/net_tools.c
+++ b/src/net_tools.c
@@ -165,3 +165,13 @@ void add_fd_to_epoll(int epollfd, int fd, uint32_t flags) {
     exit(EXIT_FAILURE);
   }
 }
+
+void update_fd_epoll(int epollfd, int fd, uint32_t flags) {
+  struct epoll_event current_event = {0};
+  current_event.events = flags;
+  current_event.data.fd = fd;
+  if (epoll_ctl (epollfd, EPOLL_CTL_MOD, fd, &current_event) == -1) {
+    perror("Failed to update a file descriptor to epoll with epoll_ctl");
+    exit(EXIT_FAILURE);
+  }
+}
diff --git a/src/net_tools.h b/src/net_tools.h
index 9756c12..a847b9a 100644
--- a/src/net_tools.h
+++ b/src/net_tools.h
@@ -14,5 +14,6 @@ int create_tcp_server(char* host, char* service);
 int create_udp_server(char* host, char* service);
 int make_socket_non_blocking(int fd);
 void add_fd_to_epoll(int epollfd, int fd, uint32_t flags);
+void update_fd_epoll(int epollfd, int fd, uint32_t flags);
 int read_entity(int fd, void* entity, int size);
 void fill_buffer(size_t* written, char* dest, void *src, size_t n);
diff --git a/src/socks5.c b/src/socks5.c
index 6d36599..61c76c1 100644
--- a/src/socks5.c
+++ b/src/socks5.c
@@ -9,11 +9,14 @@ struct socks5_ctx {
   uint64_t sh_cursor;
   uint64_t cr_cursor;
   uint64_t sr_cursor;
+  uint8_t sr_atyp_read;
   char cr_buffer[262];
   size_t ch_size;
   size_t cr_size;
 };
 
+int on_socks5_send_handshake(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo);
+
 void socks5_free_ctx(void* elem) {
   struct socks5_ctx* ctx = elem;
   free(ctx);
@@ -23,6 +26,7 @@ void create_socks5_dns_client(struct evt_core_ctx* ctx, char* proxy_host, char*
   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
@@ -77,7 +81,8 @@ void create_socks5_dns_client(struct evt_core_ctx* ctx, char* proxy_host, char*
   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);
+  reg_fdinfo = evt_core_add_fd (ctx, &fdinfo);
+  reg_fdinfo->cat->cb(ctx, reg_fdinfo);
 }
 
 int on_socks5_send_handshake(struct evt_core_ctx* ctx, struct evt_core_fdinfo* fdinfo) {
@@ -88,54 +93,190 @@ int on_socks5_send_handshake(struct evt_core_ctx* ctx, struct evt_core_fdinfo* f
   if (written < 0) {
     perror("write failed on tcp socket in socks5");
     evt_core_mv_fd2(ctx, fdinfo, "socks5-failed");
-    return 0;
+    fdinfo->cat->cb(ctx, fdinfo);
+    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 0;
+  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;
 
-    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;
+  readn = read(fdinfo->fd, &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");
+    fdinfo->cat->cb(ctx, fdinfo);
+    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");
+    fdinfo->cat->cb(ctx, fdinfo);
+    return 1;
+  }
+  printf("[server_handshake] ver=%d, method=%d\n", s5ctx->sh.ver, s5ctx->sh.method);
+  evt_core_mv_fd2(ctx, fdinfo, "socks5-send-client-req");
+  fdinfo->cat->cb(ctx, fdinfo);
+  return 1;
 }
 
-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;
+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;
 
-    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;
+  written = write(fdinfo->fd, 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");
+    fdinfo->cat->cb(ctx, fdinfo);
+    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 = (void*)&s5ctx->sr.bind_addr.ipv4 - (void*)&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,
+       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");
+    fdinfo->cat->cb(ctx, fdinfo);
+    return 1;
+  }
+
+  s5ctx->sr_cursor += nread;
+  if (s5ctx->sr_cursor < fixed_headers_size + host_size) return 0;
+
+  s5ctx->sr_atyp_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 = (void*)&s5ctx->sr.bind_addr.ipv6 - (void*)&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,
+       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");
+    fdinfo->cat->cb(ctx, fdinfo);
+    return 1;
+  }
+  s5ctx->sr_cursor += nread;
+  if (s5ctx->sr_cursor < fixed_headers_size + host_size) return 0;
+
+  s5ctx->sr_atyp_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 = (void*)&s5ctx->sr.bind_addr.ipv6 - (void*)&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, &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");
+      fdinfo->cat->cb(ctx, fdinfo);
+      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, &s5ctx->sr.bind_addr.dns.str + relative_cursor, s5ctx->sr.bind_addr.dns.len - relative_cursor);
+
+  return 0;
+}
+
+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 = (void*)&s5ctx->sr.bind_addr - (void*)&s5ctx->sr;
+  size_t fixed_tail_size = (void*)(&s5ctx->sr + 1) - (void*)&s5ctx->sr.port ;
+
+  if (s5ctx->sr_cursor < fixed_headers_size) {
+    readn = read(fdinfo->fd, &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
+  }
+
+  if (!s5ctx->sr_atyp_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;
+  }
+
+move_to_failed:
+  evt_core_mv_fd2 (ctx, fdinfo, "socks5-failed");
+  fdinfo->cat->cb(ctx, fdinfo);
+  return 1;
+  /*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) {
@@ -153,27 +294,7 @@ int socks5_reply(int sock) {
 
 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;
+    return -128;*/
 }
 
 char* socks5_rep (enum socks5_rep rep) {
diff --git a/src/socks5.h b/src/socks5.h
index a5f8b9a..6966b09 100644
--- a/src/socks5.h
+++ b/src/socks5.h
@@ -42,7 +42,7 @@ enum methods {
 union socks5_addr {
     struct {
         uint8_t len;
-        char* str;
+        char str[256];
     } dns;
     uint8_t ipv4[4];
     uint8_t ipv6[16];