#define _GNU_SOURCE
#include <event2/buffer.h>
#include <event2/bufferevent.h>
#include <event2/event.h>
#include <event2/listener.h>
#include <event2/util.h>
#include <inttypes.h>
#include <openssl/evp.h>
#include <openssl/md5.h>
#include <openssl/rand.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#ifdef _WIN32
#ifndef _WIN32_WINNT
#define _WIN32_WINNT 0x0A00
#endif
#define _WINSOCK_DEPRECATED_NO_WARNINGS
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "Ws2_32.lib")
typedef int socklen_t;
#ifndef ssize_t
#define ssize_t SSIZE_T
#endif
#define CLOSESOCK(s) closesocket((s))
#define SOCK_ERR() WSAGetLastError()
#define WOULD_BLOCK(e) ((e)==WSAEWOULDBLOCK)
#else
#include <arpa/inet.h>
#include <errno.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#define CLOSESOCK(s) close((s))
#define SOCK_ERR() errno
#define WOULD_BLOCK(e) ((e)==EAGAIN || (e)==EWOULDBLOCK)
#endif

#define RECV_BUF 8192
#define DNS_CACHE_SIZE 256
#define UDP_MAP_CAP 128
#define DOMAIN_MAX 255

static const char *k_domains[] = {"google.com", "youtube.com", "github.com"};
static const size_t k_domains_cnt = 3;

static bool ends_with(const char *s, const char *suffix) {
  size_t ls = strlen(s), lt = strlen(suffix);
  return lt <= ls && strcasecmp(s + (ls - lt), suffix) == 0;
}

static bool host_in_domain_list(const char *host) {
  for (size_t i = 0; i < k_domains_cnt; ++i) {
    if (ends_with(host, k_domains[i])) return true;
  }
  return false;
}

static bool resolve_host_port(const char *host, uint16_t port,
                              struct sockaddr_storage *ss, socklen_t *sslen) {
  struct addrinfo hints, *res = NULL;
  char portstr[8];
  snprintf(portstr, sizeof(portstr), "%u", (unsigned)port);
  memset(&hints, 0, sizeof(hints));
  hints.ai_socktype = SOCK_STREAM;
  hints.ai_family = AF_UNSPEC;
  if (getaddrinfo(host, portstr, &hints, &res) != 0 || !res) {
    if (res) freeaddrinfo(res);
    return false;
  }
  memcpy(ss, res->ai_addr, res->ai_addrlen);
  *sslen = (socklen_t)res->ai_addrlen;
  freeaddrinfo(res);
  return true;
}

static int host_literal_ip(const char *host, uint8_t *packed16, size_t *packed_len) {
  struct in_addr a4; struct in6_addr a6;
  if (inet_pton(AF_INET, host, &a4) == 1) {
    memcpy(packed16, &a4, 4); *packed_len = 4; return AF_INET;
  }
  if (inet_pton(AF_INET6, host, &a6) == 1) {
    memcpy(packed16, &a6, 16); *packed_len = 16; return AF_INET6;
  }
  return 0;
}

static void kdf_evp_bytes_to_key(const uint8_t *pw, size_t pwlen, uint8_t out32[32]) {
  uint8_t h1[16]; uint8_t h2[16];
  MD5_CTX ctx;
  MD5_Init(&ctx);
  MD5_Update(&ctx, pw, pwlen);
  MD5_Final(h1, &ctx);
  MD5_Init(&ctx);
  MD5_Update(&ctx, h1, 16);
  MD5_Update(&ctx, pw, pwlen);
  MD5_Final(h2, &ctx);
  memcpy(out32, h1, 16);
  memcpy(out32 + 16, h2, 16);
}

typedef struct {
  EVP_CIPHER_CTX *ctx;
} aes_cfb_stream_t;

static void aes_cfb_init(aes_cfb_stream_t *s, const uint8_t key[32],
                         const uint8_t iv[16], int enc) {
  s->ctx = EVP_CIPHER_CTX_new();
  EVP_CipherInit_ex(s->ctx, EVP_aes_256_cfb128(), NULL, key, iv, enc);
}

static void aes_cfb_free(aes_cfb_stream_t *s) {
  if (s->ctx) EVP_CIPHER_CTX_free(s->ctx);
}

static void aes_cfb_update(aes_cfb_stream_t *s, const uint8_t *in,
                           size_t inlen, uint8_t *out, int *outlen) {
  int n = 0;
  EVP_CipherUpdate(s->ctx, out, &n, in, (int)inlen);
  *outlen = n;
}

static int aes_cfb_one_shot(const uint8_t key[32], const uint8_t iv[16],
                            int enc, const uint8_t *in, size_t inlen,
                            uint8_t *out) {
  EVP_CIPHER_CTX *c = EVP_CIPHER_CTX_new();
  if (!c) return -1;
  int n1 = 0, n2 = 0;
  EVP_CipherInit_ex(c, EVP_aes_256_cfb128(), NULL, key, iv, enc);
  if (!EVP_CipherUpdate(c, out, &n1, in, (int)inlen)) { EVP_CIPHER_CTX_free(c); return -1; }
  if (!EVP_CipherFinal_ex(c, out + n1, &n2)) { EVP_CIPHER_CTX_free(c); return -1; }
  EVP_CIPHER_CTX_free(c);
  return n1 + n2;
}

static bool encode_addr(uint8_t atyp, const char *host, uint16_t port,
                        struct evbuffer *out, bool tcp) {
  uint8_t packed[16]; size_t plen = 0;
  switch (atyp) {
    case 0x01: {
      if (host_literal_ip(host, packed, &plen) != AF_INET) return false;
      evbuffer_add_printf(out, "%c", tcp ? (char)0x61 : (char)0x01);
      evbuffer_add(out, packed, 4);
      break;
    }
    case 0x03: {
      size_t len = strlen(host);
      if (len > 255) return false;
      evbuffer_add_printf(out, "%c", tcp ? (char)0x63 : (char)0x03);
      if (!tcp) evbuffer_add_printf(out, "%c", (char)len);
      else evbuffer_add_printf(out, "%c", (char)len);
      evbuffer_add(out, host, len);
      break;
    }
    case 0x04: {
      if (host_literal_ip(host, packed, &plen) != AF_INET6) return false;
      evbuffer_add_printf(out, "%c", tcp ? (char)0x64 : (char)0x04);
      evbuffer_add(out, packed, 16);
      break;
    }
    default: return false;
  }
  uint16_t nport = htons(port);
  evbuffer_add(out, &nport, 2);
  return true;
}

static int dns_build_query(const char *host, uint8_t *out, size_t cap, size_t *outlen, uint16_t *out_id) {
  if (!host || !out || cap < 12) return -1;
  uint16_t id; RAND_bytes((unsigned char*)&id, sizeof(id));
  *out_id = id;

  size_t off = 0;
  uint16_t flags = htons(0x0100);
  uint16_t qd = htons(1), an = 0, ns = 0, ar = 0;
  if (cap < 12) return -1;
  *(uint16_t*)(out+off) = htons(id); off+=2;
  *(uint16_t*)(out+off) = flags; off+=2;
  *(uint16_t*)(out+off) = qd; off+=2;
  *(uint16_t*)(out+off) = an; off+=2;
  *(uint16_t*)(out+off) = ns; off+=2;
  *(uint16_t*)(out+off) = ar; off+=2;

  const char *p = host, *dot = NULL;
  while (*p) {
    dot = strchr(p, '.');
    size_t lab_len = dot ? (size_t)(dot - p) : strlen(p);
    if (lab_len == 0 || lab_len > 63) return -1;
    if (off + 1 + lab_len >= cap) return -1;
    out[off++] = (uint8_t)lab_len;
    memcpy(out + off, p, lab_len); off += lab_len;
    if (!dot) break;
    p = dot + 1;
  }
  if (off + 1 + 4 > cap) return -1;
  out[off++] = 0;
  *(uint16_t*)(out+off) = htons(1); off+=2;
  *(uint16_t*)(out+off) = htons(1); off+=2;

  *outlen = off;
  return 0;
}

static int dns_skip_name(const uint8_t *buf, size_t len, size_t *off) {
  size_t i = *off;
  while (i < len) {
    uint8_t c = buf[i++];
    if ((c & 0xC0) == 0xC0) {
      if (i >= len) return -1;
      i++; break;
    } else if (c == 0) {
      break;
    } else {
      if (i + c > len) return -1;
      i += c;
    }
  }
  *off = i;
  return 0;
}

static int dns_parse_a(const uint8_t *buf, size_t len, char *ip_str, size_t ip_cap) {
  if (len < 12) return -1;
  uint16_t qd = ntohs(*(const uint16_t *)(buf + 4));
  uint16_t an = ntohs(*(const uint16_t *)(buf + 6));
  size_t off = 12;
  for (uint16_t i = 0; i < qd; ++i) {
    if (dns_skip_name(buf, len, &off) != 0) return -1;
    if (off + 4 > len) return -1;
    off += 4;
  }
  for (uint16_t i = 0; i < an; ++i) {
    if (dns_skip_name(buf, len, &off) != 0) return -1;
    if (off + 10 > len) return -1;
    uint16_t type  = ntohs(*(const uint16_t *)(buf + off)); off += 2;
    uint16_t class_= ntohs(*(const uint16_t *)(buf + off)); off += 2;
    off += 4;
    uint16_t rdlen = ntohs(*(const uint16_t *)(buf + off)); off += 2;
    if (off + rdlen > len) return -1;
    if (type == 1 && class_ == 1 && rdlen == 4) {
      struct in_addr a4; memcpy(&a4, buf + off, 4);
      const char *ret = inet_ntop(AF_INET, &a4, ip_str, (socklen_t)ip_cap);
      return ret ? 0 : -1;
    }
    off += rdlen;
  }
  return -1;
}

typedef struct {
  bool used;
  time_t expire;
  char host[64];
  char ip[16];
} dns_cache_entry_t;

typedef struct {
  dns_cache_entry_t tab[DNS_CACHE_SIZE];
} dns_cache_t;

static void dns_cache_init(dns_cache_t *c) { memset(c, 0, sizeof(*c)); }

static const char* dns_cache_get(dns_cache_t *c, const char *host) {
  time_t now = time(NULL);
  for (size_t i = 0; i < DNS_CACHE_SIZE; ++i) {
    if (c->tab[i].used && strcmp(c->tab[i].host, host) == 0) {
      if (c->tab[i].expire > now) return c->tab[i].ip;
      c->tab[i].used = false;
    }
  }
  return NULL;
}

static void dns_cache_put(dns_cache_t *c, const char *host, const char *ip, int ttl_sec) {
  time_t now = time(NULL);
  size_t slot = DNS_CACHE_SIZE; time_t oldest = now;
  for (size_t i = 0; i < DNS_CACHE_SIZE; ++i) {
    if (!c->tab[i].used) { slot = i; break; }
    if (c->tab[i].expire < oldest) { oldest = c->tab[i].expire; slot = i; }
  }
  if (slot >= DNS_CACHE_SIZE) slot = 0;
  dns_cache_entry_t *e = &c->tab[slot];
  e->used = true;
  e->expire = now + ttl_sec;
  strncpy(e->host, host, sizeof(e->host)-1); e->host[sizeof(e->host)-1] = 0;
  strncpy(e->ip, ip, sizeof(e->ip)-1); e->ip[sizeof(e->ip)-1] = 0;
}

typedef struct {
  char remote_host[64];
  int remote_port;
  char password[64];
  char mx_head[128];
  char listen_host[16];
  int listen_port;
  int recv_buf;
  int connect_timeout;
  int udp_timeout;
} config_t;

static void config_default(config_t *cfg) {
  memset(cfg, 0, sizeof(*cfg));
  strcpy(cfg->remote_host, "121.14.152.149");
  cfg->remote_port = 10004;
  strcpy(cfg->password, "dwz1GtF7");
  strcpy(cfg->mx_head, "com.win64.oppc.game.common:22021709,102024080020541279");
  strcpy(cfg->listen_host, "0.0.0.0");
  cfg->listen_port = 10807;
  cfg->recv_buf = RECV_BUF;
  cfg->connect_timeout = 10;
  cfg->udp_timeout = 180;
}

struct bridge_s;
typedef struct bridge_s bridge_t;

typedef struct {
  bool used;
  time_t ts;
  uint8_t key[128]; size_t key_len;
  struct sockaddr_storage client_ss; socklen_t client_len;
} udp_map_entry_t;

typedef struct {
  bridge_t *owner;
  int fd;
  struct event *ev_read;
  struct event *ev_timer;
  struct sockaddr_storage remote_ss; socklen_t remote_len;
  uint8_t key32[32];
  udp_map_entry_t map[UDP_MAP_CAP];
  struct sockaddr_storage last_client; socklen_t last_client_len; bool has_last;
} udp_assoc_t;

static int udp_map_find(udp_assoc_t *ua, const uint8_t *k, size_t klen) {
  for (int i = 0; i < UDP_MAP_CAP; ++i) {
    if (ua->map[i].used && ua->map[i].key_len == klen &&
        memcmp(ua->map[i].key, k, klen) == 0) return i;
  }
  return -1;
}

static int udp_map_alloc_slot(udp_assoc_t *ua) {
  for (int i = 0; i < UDP_MAP_CAP; ++i) if (!ua->map[i].used) return i;
  time_t oldest = time(NULL); int idx = 0;
  for (int i = 0; i < UDP_MAP_CAP; ++i) {
    if (ua->map[i].ts < oldest) { oldest = ua->map[i].ts; idx = i; }
  }
  return idx;
}

static void udp_assoc_close(udp_assoc_t *ua) {
  if (!ua) return;
  if (ua->ev_read) event_free(ua->ev_read);
  if (ua->ev_timer) event_free(ua->ev_timer);
  if (ua->fd >= 0) CLOSESOCK(ua->fd);
  free(ua);
}

typedef enum {
  ST_GREETING = 0, ST_REQUEST = 1, ST_CONNECTING_REMOTE = 2,
  ST_STREAM = 3, ST_UDP_ASSOC = 4, ST_CLOSED = 5
} bridge_state_t;

typedef struct {
  uint8_t cmd, atyp;
  char host[128];
  uint16_t port;
} socks5_req_t;

struct bridge_s {
  struct event_base *base;
  config_t *cfg;
  struct bufferevent *cli_bev;
  evutil_socket_t cli_fd;
  struct bufferevent *srv_bev;
  evutil_socket_t srv_fd;
  bridge_state_t state;
  uint8_t key32[32];
  aes_cfb_stream_t up_enc; bool up_ready;
  aes_cfb_stream_t down_dec; bool down_ready;
  uint8_t iv_server[16]; size_t iv_have;
  udp_assoc_t *udp_assoc;
  dns_cache_t dns_cache;
  socks5_req_t req;
};

static void socks5_send_reply(struct bufferevent *bev, uint8_t rep,
                              const char *bind_host, uint16_t bind_port) {
  struct evbuffer *out = evbuffer_new();
  uint8_t ver = 5, rsv = 0;
  evbuffer_add(out, &ver, 1);
  evbuffer_add(out, &rep, 1);
  evbuffer_add(out, &rsv, 1);

  uint8_t packed[16]; size_t plen = 0;
  int af = host_literal_ip(bind_host, packed, &plen);
  if (af == AF_INET) {
    uint8_t t = 0x01; evbuffer_add(out, &t, 1);
    evbuffer_add(out, packed, 4);
  } else if (af == AF_INET6) {
    uint8_t t = 0x04; evbuffer_add(out, &t, 1);
    evbuffer_add(out, packed, 16);
  } else {
    size_t len = strlen(bind_host);
    if (len > 255) len = 255;
    uint8_t t = 0x03; uint8_t l = (uint8_t)len;
    evbuffer_add(out, &t, 1);
    evbuffer_add(out, &l, 1);
    evbuffer_add(out, bind_host, len);
  }
  uint16_t p = htons(bind_port);
  evbuffer_add(out, &p, 2);
  bufferevent_write_buffer(bev, out);
  evbuffer_free(out);
}

static int parse_socks5_request(struct evbuffer *in, socks5_req_t *req, size_t *consumed) {
  size_t len = evbuffer_get_length(in);
  if (len < 4) return 0;
  unsigned char *p = evbuffer_pullup(in, len);
  if (!p) return 0;

  if (p[0] != 5) return -2;
  req->cmd = p[1];
  req->atyp = p[3];
  size_t off = 4;
  if (req->atyp == 0x01) {
    if (len < off + 4 + 2) return 0;
    char ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, p + off, ip, sizeof(ip));
    strncpy(req->host, ip, sizeof(req->host)-1); req->host[127]=0;
    off += 4;
    req->port = ntohs(*(uint16_t*)(p + off)); off += 2;
  } else if (req->atyp == 0x03) {
    if (len < off + 1) return 0;
    uint8_t l = p[off++]; if (len < off + l + 2) return 0;
    size_t cplen = (l > 127) ? 127 : l;
    memcpy(req->host, p + off, cplen); req->host[cplen]=0;
    off += l;
    req->port = ntohs(*(uint16_t*)(p + off)); off += 2;
  } else if (req->atyp == 0x04) {
    if (len < off + 16 + 2) return 0;
    char ip6[INET6_ADDRSTRLEN];
    inet_ntop(AF_INET6, p + off, ip6, sizeof(ip6));
    strncpy(req->host, ip6, sizeof(req->host)-1); req->host[127]=0;
    off += 16;
    req->port = ntohs(*(uint16_t*)(p + off)); off += 2;
  } else {
    return -3;
  }
  *consumed = off;
  return 1;
}

static void udp_assoc_on_timer(evutil_socket_t fd, short what, void *arg);
static void udp_assoc_on_read(evutil_socket_t fd, short what, void *arg);

static udp_assoc_t* udp_assoc_create(bridge_t *b) {
  udp_assoc_t *ua = (udp_assoc_t*)calloc(1, sizeof(*ua));
  ua->owner = b;
  ua->fd = socket(AF_INET, SOCK_DGRAM, 0);
  if (ua->fd < 0) { free(ua); return NULL; }
  int on = 1;
  setsockopt(ua->fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
#ifdef SO_REUSEPORT
  setsockopt(ua->fd, SOL_SOCKET, SO_REUSEPORT, (const char*)&on, sizeof(on));
#endif
  int sz = 1<<18;
  setsockopt(ua->fd, SOL_SOCKET, SO_SNDBUF, (const char*)&sz, sizeof(sz));
  setsockopt(ua->fd, SOL_SOCKET, SO_RCVBUF, (const char*)&sz, sizeof(sz));
  struct sockaddr_in sin; memset(&sin,0,sizeof(sin));
  sin.sin_family = AF_INET; sin.sin_port = htons(0);
  inet_pton(AF_INET, b->cfg->listen_host, &sin.sin_addr);
  if (bind(ua->fd, (struct sockaddr*)&sin, sizeof(sin)) != 0) {
    CLOSESOCK(ua->fd); free(ua); return NULL;
  }
  evutil_make_socket_nonblocking(ua->fd);

  if (!resolve_host_port(b->cfg->remote_host, (uint16_t)b->cfg->remote_port,
                         &ua->remote_ss, &ua->remote_len)) {
    CLOSESOCK(ua->fd); free(ua); return NULL;
  }
  memcpy(ua->key32, b->key32, 32);

  ua->ev_read = event_new(b->base, ua->fd, EV_READ|EV_PERSIST, udp_assoc_on_read, ua);
  event_add(ua->ev_read, NULL);

  struct timeval tv = { .tv_sec = 5, .tv_usec = 0 };
  ua->ev_timer = event_new(b->base, -1, EV_PERSIST, udp_assoc_on_timer, ua);
  event_add(ua->ev_timer, &tv);
  return ua;
}

static void udp_assoc_on_timer(evutil_socket_t fd, short what, void *arg) {
  (void)fd; (void)what;
  udp_assoc_t *ua = (udp_assoc_t*)arg;
  time_t now = time(NULL);
  for (int i = 0; i < UDP_MAP_CAP; ++i) {
    if (ua->map[i].used && now - ua->map[i].ts > ua->owner->cfg->udp_timeout) {
      ua->map[i].used = false;
    }
  }
}

static bool udp_parse_client_packet(const uint8_t *buf, size_t len,
                                    size_t *addr_off, size_t *payload_off) {
  if (len < 3) return false;
  if (buf[0] != 0 || buf[1] != 0 || buf[2] != 0) return false;
  size_t off = 3;
  if (len < off + 1) return false;
  uint8_t atyp = buf[off++];
  if (atyp == 0x01) {
    if (len < off + 4 + 2) return false;
    off += 4 + 2;
  } else if (atyp == 0x03) {
    if (len < off + 1) return false;
    uint8_t l = buf[off++];
    if (len < off + l + 2) return false;
    off += l + 2;
  } else if (atyp == 0x04) {
    if (len < off + 16 + 2) return false;
    off += 16 + 2;
  } else return false;

  *addr_off = 3;
  *payload_off = off;
  return true;
}

static bool udp_clone_addr_block(const uint8_t *buf, size_t addr_off,
                                 size_t payload_off, uint8_t *out, size_t *outlen) {
  size_t n = payload_off - addr_off;
  if (n > 128) return false;
  memcpy(out, buf + addr_off, n);
  *outlen = n;
  return true;
}

static void udp_assoc_sendto(udp_assoc_t *ua, const void *data, size_t len) {
  sendto(ua->fd, (const char*)data, (int)len, 0,
         (struct sockaddr*)&ua->remote_ss, ua->remote_len);
}

static void udp_assoc_on_read(evutil_socket_t fd, short what, void *arg) {
  (void)what;
  udp_assoc_t *ua = (udp_assoc_t*)arg;
  uint8_t buf[8192];
  struct sockaddr_storage from; socklen_t fromlen = sizeof(from);

  for (;;) {
    int n = (int)recvfrom(fd, (char*)buf, (int)sizeof(buf), 0,
                          (struct sockaddr*)&from, &fromlen);
    if (n < 0) {
      int e = SOCK_ERR();
      if (WOULD_BLOCK(e)) return;
      return;
    }

    bool from_remote = false;
    if (from.ss_family == ua->remote_ss.ss_family) {
      if (from.ss_family == AF_INET) {
        struct sockaddr_in *a=(struct sockaddr_in*)&from, *b=(struct sockaddr_in*)&ua->remote_ss;
        from_remote = (a->sin_port==b->sin_port && a->sin_addr.s_addr==b->sin_addr.s_addr);
      } else if (from.ss_family == AF_INET6) {
        struct sockaddr_in6 *a=(struct sockaddr_in6*)&from, *b=(struct sockaddr_in6*)&ua->remote_ss;
        from_remote = (a->sin6_port==b->sin6_port && memcmp(&a->sin6_addr,&b->sin6_addr,sizeof(a->sin6_addr))==0);
      }
    }

    if (from_remote) {
      if (n < 16) continue;
      const uint8_t *iv = buf;
      const uint8_t *ct = buf + 16; size_t ctlen = (size_t)n - 16;
      uint8_t plain[8192];
      int m = aes_cfb_one_shot(ua->key32, iv, 0, ct, ctlen, plain);
      if (m <= 0) continue;

      size_t off = 0; if (m < 1) continue;
      uint8_t atyp = plain[off++];
      if (atyp == 0x01) off += 4 + 2;
      else if (atyp == 0x03) { if (off>= (size_t)m) continue; uint8_t l=plain[off++]; off += l + 2; }
      else if (atyp == 0x04) off += 16 + 2;
      else continue;
      if ((size_t)m < off) continue;
      size_t addrlen = off; size_t payload_len = (size_t)m - off;

      int idx = udp_map_find(ua, plain, addrlen);
      struct sockaddr_storage dst; socklen_t dstlen = 0;
      if (idx >= 0) {
        dst = ua->map[idx].client_ss; dstlen = ua->map[idx].client_len;
      } else if (ua->has_last) {
        dst = ua->last_client; dstlen = ua->last_client_len;
      } else {
        continue;
      }

      uint8_t out[8192];
      size_t outlen = 3 + addrlen + payload_len;
      if (outlen > sizeof(out)) continue;
      out[0] = 0; out[1] = 0; out[2] = 0;
      memcpy(out + 3, plain, addrlen);
      memcpy(out + 3 + addrlen, plain + addrlen, payload_len);
      sendto(ua->fd, (const char*)out, (int)outlen, 0,
             (struct sockaddr*)&dst, dstlen);

    } else {
      size_t addr_off=0, payload_off=0;
      if (!udp_parse_client_packet(buf, (size_t)n, &addr_off, &payload_off)) continue;

      uint8_t keybuf[128]; size_t klen=0;
      if (!udp_clone_addr_block(buf, addr_off, payload_off, keybuf, &klen)) continue;
      int idx = udp_map_find(ua, keybuf, klen);
      if (idx < 0) idx = udp_map_alloc_slot(ua);
      ua->map[idx].used = true; ua->map[idx].ts = time(NULL);
      memcpy(ua->map[idx].key, keybuf, klen); ua->map[idx].key_len = klen;
      ua->map[idx].client_ss = from; ua->map[idx].client_len = fromlen;
      ua->last_client = from; ua->last_client_len = fromlen; ua->has_last = true;

      uint8_t iv[16]; RAND_bytes(iv, sizeof(iv));
      size_t plain_len = (size_t)n - 3;
      const uint8_t *plain = buf + 3;
      uint8_t ct[8192];
      int m = aes_cfb_one_shot(ua->key32, iv, 1, plain, plain_len, ct);
      if (m <= 0) continue;

      uint8_t out[8192];
      if ((size_t)m + 16 > sizeof(out)) continue;
      memcpy(out, iv, 16);
      memcpy(out + 16, ct, m);
      udp_assoc_sendto(ua, out, (size_t)m + 16);
    }
  }
}

static void bridge_free(bridge_t *b) {
  if (!b) return;
  if (b->cli_bev) bufferevent_free(b->cli_bev);
  if (b->srv_bev) bufferevent_free(b->srv_bev);
  if (b->udp_assoc) udp_assoc_close(b->udp_assoc);
  if (b->up_ready) aes_cfb_free(&b->up_enc);
  if (b->down_ready) aes_cfb_free(&b->down_dec);
  free(b);
}

static int handle_socks5_greeting(struct bufferevent *bev) {
  struct evbuffer *in = bufferevent_get_input(bev);
  size_t len = evbuffer_get_length(in);
  if (len < 2) return 0;
  unsigned char hdr[2];
  evbuffer_copyout(in, hdr, 2);
  if (hdr[0] != 5) return -2;
  size_t need = 2 + hdr[1];
  if (len < need) return 0;
  evbuffer_drain(in, need);
  uint8_t resp[2] = {0x05, 0x00};
  bufferevent_write(bev, resp, 2);
  return 1;
}

static void on_client_read(struct bufferevent *bev, void *ctx);
static void on_server_read(struct bufferevent *bev, void *ctx);
static void on_client_event(struct bufferevent *bev, short events, void *ctx);
static void on_server_event(struct bufferevent *bev, short events, void *ctx);

static void bridge_start_stream(bridge_t *b) {
  b->state = ST_STREAM;
  bufferevent_setcb(b->cli_bev, on_client_read, NULL, on_client_event, b);
  bufferevent_setcb(b->srv_bev, on_server_read, NULL, on_server_event, b);
  bufferevent_enable(b->cli_bev, EV_READ|EV_WRITE);
  bufferevent_enable(b->srv_bev, EV_READ|EV_WRITE);
}

static void after_remote_connected(bridge_t *b) {
  struct evbuffer *addr = evbuffer_new();
  if (!encode_addr(b->req.atyp, b->req.host, b->req.port, addr, true)) {
    socks5_send_reply(b->cli_bev, 0x04, "0.0.0.0", 0);
    evbuffer_free(addr); bridge_free(b); return;
  }
  size_t mx_len = strlen(b->cfg->mx_head);
  if (mx_len > 255) { evbuffer_free(addr); bridge_free(b); return; }
  uint8_t ml = (uint8_t)mx_len;
  evbuffer_add(addr, &ml, 1);
  evbuffer_add(addr, b->cfg->mx_head, mx_len);

  size_t plain_len = evbuffer_get_length(addr);
  uint8_t *plain = (uint8_t*)malloc(plain_len);
  evbuffer_copyout(addr, plain, plain_len);

  uint8_t iv[16]; RAND_bytes(iv, sizeof(iv));
  aes_cfb_init(&b->up_enc, b->key32, iv, 1); b->up_ready = true;

  uint8_t *ct = (uint8_t*)malloc(plain_len);
  int m = 0; aes_cfb_update(&b->up_enc, plain, plain_len, ct, &m);

  bufferevent_write(b->srv_bev, iv, 16);
  bufferevent_write(b->srv_bev, ct, (size_t)m);

  socks5_send_reply(b->cli_bev, 0x00, "0.0.0.0", 0);

  free(plain); free(ct); evbuffer_free(addr);
  bridge_start_stream(b);
}

static void connect_remote(bridge_t *b) {
  b->srv_bev = bufferevent_socket_new(b->base, -1, BEV_OPT_CLOSE_ON_FREE|BEV_OPT_DEFER_CALLBACKS);
  bufferevent_setcb(b->srv_bev, NULL, NULL, on_server_event, b);
  bufferevent_enable(b->srv_bev, EV_READ|EV_WRITE);

  struct sockaddr_storage ss; socklen_t sslen;
  if (!resolve_host_port(b->cfg->remote_host, (uint16_t)b->cfg->remote_port, &ss, &sslen)) {
    socks5_send_reply(b->cli_bev, 0x04, "0.0.0.0", 0);
    bridge_free(b); return;
  }

  if (bufferevent_socket_connect(b->srv_bev, (struct sockaddr*)&ss, sslen) != 0) {
    socks5_send_reply(b->cli_bev, 0x05, "0.0.0.0", 0);
    bridge_free(b); return;
  }

  struct timeval tv = { .tv_sec = b->cfg->connect_timeout, .tv_usec = 0 };
  bufferevent_set_timeouts(b->srv_bev, &tv, &tv);
  b->state = ST_CONNECTING_REMOTE;
}

typedef struct {
  bridge_t *b;
  int fd;
  struct event *ev_read;
  struct event *ev_timer;
  uint16_t id;
  char host[64];
} dns_tunnel_t;

static void dns_tunnel_free(dns_tunnel_t *dt) {
  if (!dt) return;
  if (dt->ev_read) event_free(dt->ev_read);
  if (dt->ev_timer) event_free(dt->ev_timer);
  if (dt->fd >= 0) CLOSESOCK(dt->fd);
  free(dt);
}

static void dns_tunnel_on_timeout(evutil_socket_t fd, short what, void *arg) {
  (void)fd; (void)what;
  dns_tunnel_t *dt = (dns_tunnel_t*)arg;
  socks5_send_reply(dt->b->cli_bev, 0x04, "0.0.0.0", 0);
  dns_tunnel_free(dt);
  bridge_free(dt->b);
}

static void dns_tunnel_on_read(evutil_socket_t fd, short what, void *arg) {
  (void)what;
  dns_tunnel_t *dt = (dns_tunnel_t*)arg;
  uint8_t buf[2048];
  struct sockaddr_storage from; socklen_t fromlen = sizeof(from);
  int n = (int)recvfrom(fd, (char*)buf, (int)sizeof(buf), 0,
                        (struct sockaddr*)&from, &fromlen);
  if (n < 0) return;
  if (n < 16) { dns_tunnel_free(dt); return; }

  uint8_t plain[2048];
  int m = aes_cfb_one_shot(dt->b->key32, buf, 0, buf + 16, (size_t)n - 16, plain);
  if (m <= 0) { dns_tunnel_free(dt); return; }

  size_t off = 0;
  if (plain[0] == 0x01) off = 1 + 4 + 2;
  else if (plain[0] == 0x04) off = 1 + 16 + 2;
  else { dns_tunnel_free(dt); return; }
  if ((size_t)m <= off) { dns_tunnel_free(dt); return; }

  char ip[INET_ADDRSTRLEN];
  if (dns_parse_a(plain + off, (size_t)m - off, ip, sizeof(ip)) == 0) {
    dns_cache_put(&dt->b->dns_cache, dt->host, ip, 300);
    strncpy(dt->b->req.host, ip, sizeof(dt->b->req.host)-1);
    dt->b->req.atyp = 0x01;
    dns_tunnel_free(dt);
    connect_remote(dt->b);
    return;
  }
  socks5_send_reply(dt->b->cli_bev, 0x04, "0.0.0.0", 0);
  dns_tunnel_free(dt);
  bridge_free(dt->b);
}

static void start_dns_over_tunnel(bridge_t *b, const char *host) {
  const char *cached = dns_cache_get(&b->dns_cache, host);
  if (cached) {
    strncpy(b->req.host, cached, sizeof(b->req.host)-1);
    b->req.atyp = 0x01;
    connect_remote(b);
    return;
  }

  dns_tunnel_t *dt = (dns_tunnel_t*)calloc(1, sizeof(*dt));
  dt->b = b; dt->fd = socket(AF_INET, SOCK_DGRAM, 0);
  if (dt->fd < 0) { free(dt); socks5_send_reply(b->cli_bev, 0x04, "0.0.0.0", 0); bridge_free(b); return; }
  evutil_make_socket_nonblocking(dt->fd);
  strncpy(dt->host, host, sizeof(dt->host)-1);

  struct sockaddr_storage r; socklen_t rlen;
  if (!resolve_host_port(b->cfg->remote_host, (uint16_t)b->cfg->remote_port, &r, &rlen)) {
    CLOSESOCK(dt->fd); free(dt); socks5_send_reply(b->cli_bev, 0x04, "0.0.0.0", 0); bridge_free(b); return;
  }

  uint8_t q[512]; size_t qlen=0; uint16_t qid=0;
  if (dns_build_query(host, q, sizeof(q), &qlen, &qid) != 0) {
    CLOSESOCK(dt->fd); free(dt); socks5_send_reply(b->cli_bev, 0x04, "0.0.0.0", 0); bridge_free(b); return;
  }
  dt->id = qid;

  struct evbuffer *ab = evbuffer_new();
  encode_addr(0x01, "8.8.8.8", 53, ab, false);
  size_t ablen = evbuffer_get_length(ab);
  uint8_t *plain = (uint8_t*)malloc(ablen + qlen);
  evbuffer_copyout(ab, plain, ablen);
  memcpy(plain + ablen, q, qlen);

  uint8_t iv[16]; RAND_bytes(iv, sizeof(iv));
  uint8_t ct[1024]; int m = aes_cfb_one_shot(b->key32, iv, 1, plain, ablen + qlen, ct);
  uint8_t out[1024];
  memcpy(out, iv, 16);
  memcpy(out + 16, ct, m);

  sendto(dt->fd, out, 16 + m, 0, (struct sockaddr*)&r, rlen);

  dt->ev_read = event_new(b->base, dt->fd, EV_READ|EV_PERSIST, dns_tunnel_on_read, dt);
  event_add(dt->ev_read, NULL);
  struct timeval tv = { .tv_sec = b->cfg->connect_timeout, .tv_usec = 0 };
  dt->ev_timer = event_new(b->base, -1, 0, dns_tunnel_on_timeout, dt);
  evtimer_add(dt->ev_timer, &tv);

  evbuffer_free(ab); free(plain);
}

static void handle_udp_associate(bridge_t *b) {
  b->udp_assoc = udp_assoc_create(b);
  if (!b->udp_assoc) {
    socks5_send_reply(b->cli_bev, 0x01, "0.0.0.0", 0);
    bridge_free(b); return;
  }
  struct sockaddr_in sin; socklen_t slen = sizeof(sin);
  getsockname(b->udp_assoc->fd, (struct sockaddr*)&sin, &slen);
  char bind_ip[INET_ADDRSTRLEN]; inet_ntop(AF_INET, &sin.sin_addr, bind_ip, sizeof(bind_ip));
  uint16_t bind_port = ntohs(sin.sin_port);
  socks5_send_reply(b->cli_bev, 0x00, bind_ip, bind_port);
  b->state = ST_UDP_ASSOC;
}

static void on_client_read(struct bufferevent *bev, void *ctx) {
  bridge_t *b = (bridge_t*)ctx;

  if (b->state == ST_GREETING) {
    int r = handle_socks5_greeting(bev);
    if (r < 0) { bridge_free(b); return; }
    if (r == 1) b->state = ST_REQUEST;
    if (r == 0) return;
  }

  if (b->state == ST_REQUEST) {
    size_t consumed = 0;
    int r = parse_socks5_request(bufferevent_get_input(bev), &b->req, &consumed);
    if (r < 0) { bridge_free(b); return; }
    if (r == 0) return;
    evbuffer_drain(bufferevent_get_input(bev), consumed);

    if (b->req.cmd == 0x01) {
      if (b->req.atyp == 0x03 && host_in_domain_list(b->req.host)) {
        start_dns_over_tunnel(b, b->req.host);
        return;
      }
      connect_remote(b);
    } else if (b->req.cmd == 0x03) {
      handle_udp_associate(b);
    } else {
      socks5_send_reply(b->cli_bev, 0x07, "0.0.0.0", 0);
      bridge_free(b);
    }
    return;
  }

  if (b->state == ST_STREAM) {
    struct evbuffer *in = bufferevent_get_input(b->cli_bev);
    size_t len = evbuffer_get_length(in);
    if (!b->up_ready || len == 0) return;

    while (len) {
      size_t chunk = len > (size_t)b->cfg->recv_buf ? (size_t)b->cfg->recv_buf : len;
      unsigned char *p = evbuffer_pullup(in, chunk);
      if (!p) break;
      uint8_t *ct = (uint8_t*)malloc(chunk);
      int m = 0; aes_cfb_update(&b->up_enc, p, chunk, ct, &m);
      bufferevent_write(b->srv_bev, ct, (size_t)m);
      evbuffer_drain(in, chunk);
      free(ct);
      len -= chunk;
    }
  }

  if (b->state == ST_UDP_ASSOC) {
    struct evbuffer *in = bufferevent_get_input(b->cli_bev);
    size_t len = evbuffer_get_length(in);
    if (len) evbuffer_drain(in, len);
  }
}

static void on_server_read(struct bufferevent *bev, void *ctx) {
  bridge_t *b = (bridge_t*)ctx;
  struct evbuffer *in = bufferevent_get_input(b->srv_bev);

  if (!b->down_ready) {
    size_t len = evbuffer_get_length(in);
    if (len < 16) return;
    evbuffer_remove(in, b->iv_server, 16);
    aes_cfb_init(&b->down_dec, b->key32, b->iv_server, 0);
    b->down_ready = true;
  }

  size_t len = evbuffer_get_length(in);
  while (len) {
    size_t chunk = len > (size_t)b->cfg->recv_buf ? (size_t)b->cfg->recv_buf : len;
    unsigned char *p = evbuffer_pullup(in, chunk);
    if (!p) break;
    uint8_t *pt = (uint8_t*)malloc(chunk);
    int m = 0; aes_cfb_update(&b->down_dec, p, chunk, pt, &m);
    bufferevent_write(b->cli_bev, pt, (size_t)m);
    evbuffer_drain(in, chunk);
    free(pt);
    len -= chunk;
  }
}

static void on_client_event(struct bufferevent *bev, short events, void *ctx) {
  bridge_t *b = (bridge_t*)ctx;
  if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
    bridge_free(b);
  }
}

static void on_server_event(struct bufferevent *bev, short events, void *ctx) {
  bridge_t *b = (bridge_t*)ctx;
  if (events & BEV_EVENT_CONNECTED) {
    evutil_socket_t fd = bufferevent_getfd(bev);
    int on = 1, sz = 1<<18;
    setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, (const char*)&on, sizeof(on));
    setsockopt(fd, SOL_SOCKET, SO_SNDBUF, (const char*)&sz, sizeof(sz));
    setsockopt(fd, SOL_SOCKET, SO_RCVBUF, (const char*)&sz, sizeof(sz));
#ifdef _WIN32
    DWORD bytes = 0, enabled = 1;
    WSAIoctl(fd, 0x98000010, &enabled, sizeof(enabled), NULL, 0, &bytes, NULL, NULL);
#endif
    after_remote_connected(b);
    return;
  }
  if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR|BEV_EVENT_TIMEOUT)) {
    socks5_send_reply(b->cli_bev, 0x05, "0.0.0.0", 0);
    bridge_free(b);
  }
}

static void on_accept(struct evconnlistener *lev, evutil_socket_t fd,
                      struct sockaddr *addr, int socklen, void *ctx) {
  (void)addr; (void)socklen;
  config_t *cfg = (config_t*)ctx;

  bridge_t *b = (bridge_t*)calloc(1, sizeof(*b));
  b->base = evconnlistener_get_base(lev);
  b->cfg = cfg;
  b->cli_fd = fd;
  b->cli_bev = bufferevent_socket_new(b->base, fd, BEV_OPT_CLOSE_ON_FREE|BEV_OPT_DEFER_CALLBACKS);
#ifdef _WIN32
  {
    DWORD bytes=0, enabled=1;
    WSAIoctl(fd, 0x98000010, &enabled, sizeof(enabled), NULL, 0, &bytes, NULL, NULL);
  }
#endif
  bufferevent_setcb(b->cli_bev, on_client_read, NULL, on_client_event, b);
  bufferevent_enable(b->cli_bev, EV_READ|EV_WRITE);

  kdf_evp_bytes_to_key((const uint8_t*)cfg->password, strlen(cfg->password), b->key32);
  dns_cache_init(&b->dns_cache);
  b->state = ST_GREETING;
}

static void on_accept_error(struct evconnlistener *lev, void *ctx) {
  (void)ctx; (void)lev;
}

static void parse_args(int argc, char **argv, config_t *cfg) {
  config_default(cfg);
  for (int i = 1; i < argc; ++i) {
    if (!strcmp(argv[i], "--remote-host") && i+1 < argc) {
      strncpy(cfg->remote_host, argv[++i], sizeof(cfg->remote_host)-1);
    } else if (!strcmp(argv[i], "--remote-port") && i+1 < argc) {
      cfg->remote_port = atoi(argv[++i]);
    } else if (!strcmp(argv[i], "--password") && i+1 < argc) {
      strncpy(cfg->password, argv[++i], sizeof(cfg->password)-1);
    } else if (!strcmp(argv[i], "--mx") && i+1 < argc) {
      strncpy(cfg->mx_head, argv[++i], sizeof(cfg->mx_head)-1);
    } else if (!strcmp(argv[i], "--listen") && i+1 < argc) {
      strncpy(cfg->listen_host, argv[++i], sizeof(cfg->listen_host)-1);
    } else if (!strcmp(argv[i], "--port") && i+1 < argc) {
      cfg->listen_port = atoi(argv[++i]);
    }
  }
}

int main(int argc, char **argv) {
  config_t cfg; parse_args(argc, argv, &cfg);

  struct event_base *base = event_base_new();
  if (!base) return 1;

  struct sockaddr_in sin; memset(&sin,0,sizeof(sin));
  sin.sin_family = AF_INET;
  sin.sin_port = htons((uint16_t)cfg.listen_port);
  inet_pton(AF_INET, cfg.listen_host, &sin.sin_addr);

  struct evconnlistener *lev =
      evconnlistener_new_bind(base, on_accept, &cfg,
          LEV_OPT_CLOSE_ON_FREE|LEV_OPT_REUSEABLE|LEV_OPT_REUSEABLE_PORT, -1,
          (struct sockaddr*)&sin, sizeof(sin));
  if (!lev) { event_base_free(base); return 1; }

  evconnlistener_set_error_cb(lev, on_accept_error);

  int rc = event_base_dispatch(base);

  evconnlistener_free(lev);
  event_base_free(base);
  return rc ? 1 : 0;
}