// iocp_s5.c  --  single-file, WinSock IOCP + CNG(AES-256-CFB) + DNS-over-tunnel
// Build: cl /nologo /O2 /DNDEBUG /MT /utf-8 iocp_s5.c /link /MACHINE:X64 /OPT:REF /OPT:ICF ws2_32.lib mswsock.lib bcrypt.lib

#define _CRT_SECURE_NO_WARNINGS
#include <winsock2.h>
#include <mswsock.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <bcrypt.h>
#include <stdio.h>
#include <stdint.h>
#include <stdbool.h>
#include <time.h>

#pragma comment(lib, "ws2_32.lib")
#pragma comment(lib, "mswsock.lib")
#pragma comment(lib, "bcrypt.lib")

//===================== Config =====================

typedef struct {
    char  remote_host[64];
    int   remote_port;
    char  password[64];
    char  mx_head[128];
    char  listen_host[32];
    int   listen_port;
    int   recv_buf;
    int   connect_timeout_sec;
    int   udp_timeout_sec;
    int   verbose;
    int   dns_on; // 1=enable DNS over tunnel for domain CONNECT
} config_t;

static void config_default(config_t* c){
    memset(c,0,sizeof(*c));
    strcpy(c->remote_host, "121.14.152.149");
    c->remote_port = 10004;
    strcpy(c->password, "dwz1GtF7");
    strcpy(c->mx_head,  "com.win64.oppc.game.common:22021709,102024080020541279");
    strcpy(c->listen_host, "0.0.0.0");
    c->listen_port = 10807;
    c->recv_buf = 8192;
    c->connect_timeout_sec = 10;
    c->udp_timeout_sec = 180;
    c->verbose = 1;
    c->dns_on = 1;
}

static void parse_args(int argc, char** argv, config_t* c){
    config_default(c);
    for (int i=1;i<argc;i++){
        if (!strcmp(argv[i],"--remote-host") && i+1<argc) { strncpy(c->remote_host, argv[++i], sizeof(c->remote_host)-1); }
        else if (!strcmp(argv[i],"--remote-port") && i+1<argc) { c->remote_port = atoi(argv[++i]); }
        else if (!strcmp(argv[i],"--password") && i+1<argc) { strncpy(c->password, argv[++i], sizeof(c->password)-1); }
        else if (!strcmp(argv[i],"--mx") && i+1<argc) { strncpy(c->mx_head, argv[++i], sizeof(c->mx_head)-1); }
        else if (!strcmp(argv[i],"--listen") && i+1<argc) { strncpy(c->listen_host, argv[++i], sizeof(c->listen_host)-1); }
        else if (!strcmp(argv[i],"--port") && i+1<argc) { c->listen_port = atoi(argv[++i]); }
        else if (!strcmp(argv[i],"--recv-buf") && i+1<argc) { c->recv_buf = atoi(argv[++i]); }
        else if (!strcmp(argv[i],"--connect-timeout") && i+1<argc) { c->connect_timeout_sec = atoi(argv[++i]); }
        else if (!strcmp(argv[i],"--udp-timeout") && i+1<argc) { c->udp_timeout_sec = atoi(argv[++i]); }
        else if (!strcmp(argv[i],"--verbose") && i+1<argc) { c->verbose = atoi(argv[++i]); }
        else if (!strcmp(argv[i],"--dns-on") && i+1<argc) { c->dns_on = atoi(argv[++i]); }
    }
}

//===================== Log =====================

static DWORD g_main_tid = 0;
static int g_verbose = 1;

static uint64_t now_ms(){
    FILETIME ft; GetSystemTimeAsFileTime(&ft);
    ULARGE_INTEGER u; u.LowPart=ft.dwLowDateTime; u.HighPart=ft.dwHighDateTime;
    return (u.QuadPart/10000ULL);
}
static void logv(char lv, const char* fmt, ...){
    if (lv=='D' && !g_verbose) return;
    va_list ap; va_start(ap, fmt);
    SYSTEMTIME st; GetLocalTime(&st);
    DWORD tid = GetCurrentThreadId();
    char buf[1024];
    int n = vsnprintf(buf,sizeof(buf),fmt,ap);
    va_end(ap);
    if (n<0) n=0; if (n> (int)sizeof(buf)-1) n=(int)sizeof(buf)-1;
    printf("[%02d:%02d:%02d.%03d][%c][T%u] %s\n",
        st.wHour,st.wMinute,st.wSecond,st.wMilliseconds, lv, (unsigned)tid, buf);
    fflush(stdout);
}

//===================== Helpers =====================

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

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

//===================== KDF: EVP_BytesToKey(md5) 32B =====================

static bool md5_once(const uint8_t* in, DWORD inlen, uint8_t out16[16]){
    BCRYPT_ALG_HANDLE hAlg=NULL; BCRYPT_HASH_HANDLE hHash=NULL;
    NTSTATUS s=BCryptOpenAlgorithmProvider(&hAlg, BCRYPT_MD5_ALGORITHM, NULL, 0); if (s) return false;
    s=BCryptCreateHash(hAlg, &hHash, NULL, 0, NULL, 0, 0); if (s){ BCryptCloseAlgorithmProvider(hAlg,0); return false; }
    s=BCryptHashData(hHash, (PUCHAR)in, inlen, 0); if (s){ BCryptDestroyHash(hHash); BCryptCloseAlgorithmProvider(hAlg,0); return false; }
    s=BCryptFinishHash(hHash, out16, 16, 0);
    BCryptDestroyHash(hHash); BCryptCloseAlgorithmProvider(hAlg,0);
    return s==0;
}
static bool kdf_key32_from_password(const char* pw, uint8_t out32[32]){
    uint8_t h1[16], h2[16];
    if (!md5_once((const uint8_t*)pw, (DWORD)strlen(pw), h1)) return false;
    uint8_t tmp[16+64]; memcpy(tmp,h1,16); memcpy(tmp+16,pw,strlen(pw));
    if (!md5_once(tmp, (DWORD)(16+strlen(pw)), h2)) return false;
    memcpy(out32,h1,16); memcpy(out32+16,h2,16);
    return true;
}

//===================== AES-256-CFB via CNG =====================

typedef struct {
    BCRYPT_ALG_HANDLE hAlg;
    BCRYPT_KEY_HANDLE hKey;
    uint8_t iv[16];       // updated in place per call
    uint8_t key[32];
    BOOL    is_enc;       // 1=encrypt, 0=decrypt
    BOOL    init_ok;
} cfb_ctx_t;

static bool cfb_init(cfb_ctx_t* c, const uint8_t key[32], const uint8_t iv16[16], BOOL enc){
    memset(c,0,sizeof(*c));
    memcpy(c->key, key, 32);
    memcpy(c->iv,  iv16, 16);
    c->is_enc = enc;

    NTSTATUS s=BCryptOpenAlgorithmProvider(&c->hAlg, BCRYPT_AES_ALGORITHM, NULL, 0);
    if (s){ logv('E',"BCryptOpenAlgorithmProvider AES failed: 0x%08X", s); return false; }

    s=BCryptSetProperty(c->hAlg, BCRYPT_CHAINING_MODE, (PUCHAR)BCRYPT_CHAIN_MODE_CFB, (ULONG)sizeof(BCRYPT_CHAIN_MODE_CFB), 0);
    if (s){ logv('E',"Set ChainingModeCFB failed: 0x%08X", s); BCryptCloseAlgorithmProvider(c->hAlg,0); return false; }

    s=BCryptGenerateSymmetricKey(c->hAlg, &c->hKey, NULL, 0, (PUCHAR)c->key, 32, 0);
    if (s){ logv('E',"GenerateSymmetricKey failed: 0x%08X", s); BCryptCloseAlgorithmProvider(c->hAlg,0); return false; }

    c->init_ok = TRUE;
    return true;
}
static void cfb_free(cfb_ctx_t* c){
    if (c->hKey){ BCryptDestroyKey(c->hKey); c->hKey=NULL; }
    if (c->hAlg){ BCryptCloseAlgorithmProvider(c->hAlg,0); c->hAlg=NULL; }
    c->init_ok=FALSE;
}

// in==out allowed? 为安全起见用不同 buffer（上层已分配）
static bool cfb_update(cfb_ctx_t* c, const uint8_t* in, DWORD inlen, uint8_t* out, DWORD* outlen){
    if (!c->init_ok) return false;
    ULONG cb=0; NTSTATUS s;
    if (c->is_enc)
        s=BCryptEncrypt(c->hKey,(PUCHAR)in,inlen,NULL,(PUCHAR)c->iv,16,(PUCHAR)out,inlen,&cb,0);
    else
        s=BCryptDecrypt(c->hKey,(PUCHAR)in,inlen,NULL,(PUCHAR)c->iv,16,(PUCHAR)out,inlen,&cb,0);
    if (s){ logv('E',"CFB update failed 0x%08X", s); return false; }
    *outlen = cb;
    return true;
}

//===================== DNS minimal =====================

static int dns_build_query_a(const char* host, uint8_t* out, size_t cap, size_t* outlen){
    if (!host || !out || cap<12) return -1;
    uint16_t id = (uint16_t)rand();
    size_t off=0;
    if (cap<12) return -1;
    *(uint16_t*)(out+off) = htons(id); off+=2;
    *(uint16_t*)(out+off) = htons(0x0100); off+=2; // rd
    *(uint16_t*)(out+off) = htons(1); off+=2;      // qdcount
    *(uint16_t*)(out+off) = 0; off+=2;             // an
    *(uint16_t*)(out+off) = 0; off+=2;             // ns
    *(uint16_t*)(out+off) = 0; off+=2;             // ar
    // qname
    const char* p=host;
    while (*p){
        const char* dot = strchr(p,'.');
        size_t lab = dot ? (size_t)(dot-p) : strlen(p);
        if (lab==0 || lab>63) return -1;
        if (off+1+lab>=cap) return -1;
        out[off++] = (uint8_t)lab;
        memcpy(out+off, p, lab); off+=lab;
        if (!dot) break;
        p = dot+1;
    }
    if (off+1+4>cap) return -1;
    out[off++] = 0;
    *(uint16_t*)(out+off) = htons(1); off+=2; // QTYPE A
    *(uint16_t*)(out+off) = htons(1); off+=2; // QCLASS IN
    *outlen = off;
    return 0;
}
static int dns_skip_name(const uint8_t* b,size_t len,size_t* off){
    size_t i=*off;
    while (i<len){
        uint8_t c=b[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_ip(const uint8_t* b,size_t len,char ip[16]){
    if (len<12) return -1;
    uint16_t qd = ntohs(*(uint16_t*)(b+4));
    uint16_t an = ntohs(*(uint16_t*)(b+6));
    size_t off=12;
    for (uint16_t i=0;i<qd;i++){ if (dns_skip_name(b,len,&off)) return -1; if (off+4>len) return -1; off+=4; }
    for (uint16_t i=0;i<an;i++){
        if (dns_skip_name(b,len,&off)) return -1;
        if (off+10>len) return -1;
        uint16_t type = ntohs(*(uint16_t*)(b+off)); off+=2;
        uint16_t cls  = ntohs(*(uint16_t*)(b+off)); off+=2;
        off+=4; // ttl
        uint16_t rdlen= ntohs(*(uint16_t*)(b+off)); off+=2;
        if (off+rdlen>len) return -1;
        if (type==1 && cls==1 && rdlen==4){
            const struct in_addr* a4=(const struct in_addr*)(b+off);
            inet_ntop(AF_INET, a4, ip, 16);
            return 0;
        }
        off+=rdlen;
    }
    return -1;
}

//===================== protocol: encode_addr =====================
// tcp=true : first byte 'a'(0x61)/'c'(0x63)/'d'(0x64)
// tcp=false: first byte 0x01/0x03/0x04
static size_t encode_addr_block(bool tcp, uint8_t atyp, const char* host, uint16_t port, uint8_t* out, size_t cap){
    uint8_t* p=out; size_t left=cap;
    uint8_t packed[16]; size_t plen=0;
    int afip = host_literal_ip(host,packed,&plen);

    if (tcp){
        if (afip==AF_INET && atyp==0x01){ if (left<1+4+2) return 0; *p++=0x61; memcpy(p,packed,4); p+=4; }
        else if (afip==AF_INET6 && atyp==0x04){ if (left<1+16+2) return 0; *p++=0x64; memcpy(p,packed,16); p+=16; }
        else if (atyp==0x03){
            size_t L=strlen(host); if (L>255) return 0;
            if (left<1+1+L+2) return 0;
            *p++=0x63; *p++=(uint8_t)L; memcpy(p,host,L); p+=L;
        } else return 0;
    }else{
        if (afip==AF_INET && atyp==0x01){ if (left<1+4+2) return 0; *p++=0x01; memcpy(p,packed,4); p+=4; }
        else if (afip==AF_INET6 && atyp==0x04){ if (left<1+16+2) return 0; *p++=0x04; memcpy(p,packed,16); p+=16; }
        else if (atyp==0x03){
            size_t L=strlen(host); if (L>255) return 0;
            if (left<1+1+L+2) return 0;
            *p++=0x03; *p++=(uint8_t)L; memcpy(p,host,L); p+=L;
        } else return 0;
    }
    uint16_t np = htons(port);
    memcpy(p,&np,2); p+=2;
    return (size_t)(p - out);
}

//===================== IOCP Types =====================

typedef enum { OP_NONE=0, OP_CLI_RECV, OP_CLI_SEND, OP_SRV_RECV, OP_SRV_SEND, OP_CONNECT } opkind_t;

typedef struct iocp_op_s {
    OVERLAPPED ov;
    WSABUF     buf;
    opkind_t   kind;
    char       storage[1]; // flexible for debugging
} iocp_op_t;

static iocp_op_t* op_alloc(opkind_t k, DWORD cap){
    iocp_op_t* op = (iocp_op_t*)calloc(1, sizeof(iocp_op_t) + (cap?cap-1:0));
    op->kind = k; op->buf.buf = cap? op->storage : NULL; op->buf.len = cap;
    return op;
}
static void op_free(iocp_op_t* op){ free(op); }

//===================== Bridge =====================

typedef struct {
    SOCKET s_cli, s_srv;
    HANDLE iocp;
    config_t* cfg;

    // ConnectEx
    LPFN_CONNECTEX pConnectEx;

    // state
    enum { ST_GREETING=0, ST_REQUEST=1, ST_CONNECTING=2, ST_STREAM=3, ST_CLOSED=4 } st;

    // SOCKS5 request
    uint8_t  req_cmd, req_atyp;
    char     req_host[128];
    uint16_t req_port;

    // crypto
    uint8_t  key32[32];
    uint8_t  iv_up[16];   // client->server: sent once at beginning
    cfb_ctx_t c_up;       // encryptor  (iv_up evolves in place)
    BOOL     up_inited;
    uint8_t  iv_dn[16];   // server->client: received once from remote
    size_t   ivdn_have;   // bytes collected for dn iv
    cfb_ctx_t c_dn;       // decryptor
    BOOL     dn_inited;

    LONG     pending_ops; // refcnt-like
    BOOL     closing;

    // buffers
    DWORD    recv_cap;

} bridge_t;

static void bridge_start_close(bridge_t* b);

//============= small utils =============

static void rand_bytes(uint8_t* p, size_t n){
    for (size_t i=0;i<n;i++) p[i] = (uint8_t)(rand() & 0xFF);
}

static bool s5_greeting(uint8_t* in, DWORD inlen, DWORD* consumed){
    if (inlen<2) return false;
    if (in[0]!=5) return false;
    DWORD need = 2 + in[1];
    if (inlen<need) return false;
    *consumed = need;
    return true;
}
static int s5_parse_request(uint8_t* in, DWORD inlen, uint8_t* pcmd, uint8_t* patyp, char* host, uint16_t* pport, DWORD* consumed){
    if (inlen<4 || in[0]!=5) return 0;
    DWORD off=4;
    *pcmd = in[1]; *patyp=in[3];
    if (*patyp==0x01){
        if (inlen<off+4+2) return 0;
        char ip[16]; InetNtopA(AF_INET, in+off, ip, sizeof(ip));
        strncpy(host, ip, 127); host[127]=0;
        off+=4; *pport = ntohs(*(uint16_t*)(in+off)); off+=2;
    }else if (*patyp==0x03){
        if (inlen<off+1) return 0;
        uint8_t L = in[off++]; if (inlen<off+L+2) return 0;
        size_t cp = L>127?127:L; memcpy(host,in+off,cp); host[cp]=0;
        off+=L; *pport = ntohs(*(uint16_t*)(in+off)); off+=2;
    }else if (*patyp==0x04){
        if (inlen<off+16+2) return 0;
        char ip6[64]; InetNtopA(AF_INET6, in+off, ip6, sizeof(ip6));
        strncpy(host, ip6, 127); host[127]=0;
        off+=16; *pport = ntohs(*(uint16_t*)(in+off)); off+=2;
    }else return -1;
    *consumed=off; return 1;
}
static void s5_reply(SOCKET s, uint8_t rep, const char* bind_host, uint16_t bind_port){
    uint8_t out[4+1+255+2]; uint8_t* p=out;
    *p++=5; *p++=rep; *p++=0; // RSV
    uint8_t ipbuf[16]; size_t plen=0; int af=host_literal_ip(bind_host, ipbuf, &plen);
    if (af==AF_INET){ *p++=0x01; memcpy(p,ipbuf,4); p+=4; }
    else if (af==AF_INET6){ *p++=0x04; memcpy(p,ipbuf,16); p+=16; }
    else{ size_t L=strlen(bind_host); if (L>255) L=255; *p++=0x03; *p++=(uint8_t)L; memcpy(p,bind_host,L); p+=L; }
    uint16_t np = htons(bind_port); memcpy(p,&np,2); p+=2;
    DWORD sent=0; WSASend(s, (WSABUF*)&(WSABUF){ .buf=(CHAR*)out, .len=(ULONG)(p-out) }, 1, &sent, 0, NULL, NULL);
}

//===================== GQCS loop =====================

typedef struct {
    HANDLE iocp;
    config_t* cfg;
} worker_arg_t;

static void post_recv(bridge_t* b, SOCKET s, opkind_t kind){
    iocp_op_t* op=op_alloc(kind, b->recv_cap);
    InterlockedIncrement(&b->pending_ops);
    DWORD fl=0; DWORD recvd=0;
    int rc = WSARecv(s, &op->buf, 1, &recvd, &fl, &op->ov, NULL);
    if (rc==SOCKET_ERROR){
        int e=WSAGetLastError();
        if (e!=WSA_IO_PENDING){
            logv('E',"WSARecv immediate fail e=%d", e);
            InterlockedDecrement(&b->pending_ops);
            op_free(op);
            bridge_start_close(b);
        }else{
            if (g_verbose) logv('D',"post RECV kind=%d cap=%u (pending_ops=%ld)", kind, op->buf.len, b->pending_ops);
        }
    }
}
static void post_send(bridge_t* b, SOCKET s, const void* data, DWORD len, opkind_t kind){
    iocp_op_t* op=op_alloc(kind, len);
    memcpy(op->buf.buf, data, len);
    op->buf.len=len;
    InterlockedIncrement(&b->pending_ops);
    DWORD sent=0;
    int rc=WSASend(s, &op->buf, 1, &sent, 0, &op->ov, NULL);
    if (rc==SOCKET_ERROR){
        int e=WSAGetLastError();
        if (e!=WSA_IO_PENDING){
            logv('E',"WSASend immediate fail e=%d", e);
            InterlockedDecrement(&b->pending_ops);
            op_free(op);
            bridge_start_close(b);
        }else{
            if (g_verbose) logv('D',"post SEND kind=%d len=%u (pending_ops=%ld)", kind, len, b->pending_ops);
        }
    }
}

static void after_remote_connected(bridge_t* b){
    // 1) 生成上行IV并初始化加密流
    rand_bytes(b->iv_up, 16);
    if (!cfb_init(&b->c_up, b->key32, b->iv_up, TRUE)){ bridge_start_close(b); return; }
    b->up_inited=TRUE;

    // 2) 组首包：encode_addr(tcp=true) + MX
    uint8_t addr[256]; size_t alen = encode_addr_block(true, b->req_atyp, b->req_host, b->req_port, addr, sizeof(addr));
    if (!alen){ logv('E',"encode_addr(tcp) failed"); bridge_start_close(b); return; }

    size_t mxL = strlen(b->cfg->mx_head); if (mxL>255) mxL=255;
    uint8_t plain[256+1+255]; memcpy(plain, addr, alen); plain[alen]=(uint8_t)mxL; memcpy(plain+alen+1, b->cfg->mx_head, mxL);
    DWORD plain_len = (DWORD)(alen + 1 + mxL);

    // 3) 加密首包（不含 IV）
    uint8_t ct[512]; DWORD ctlen=0;
    if (!cfb_update(&b->c_up, plain, plain_len, ct, &ctlen)){ bridge_start_close(b); return; }

    // 4) 发送：IV + CT
    uint8_t out[16+512]; memcpy(out, b->iv_up, 16); memcpy(out+16, ct, ctlen);
    post_send(b, b->s_srv, out, 16+ctlen, OP_SRV_SEND);

    // 5) 给客户端回复 OK，并进入流转发
    s5_reply(b->s_cli, 0x00, "0.0.0.0", 0);
    b->st = ST_STREAM;

    // 6) 分别贴上读取
    post_recv(b, b->s_cli, OP_CLI_RECV);
    post_recv(b, b->s_srv, OP_SRV_RECV);

    logv('I',"enter STREAM (+posted cli/srv recv)");
}

//===================== DNS over tunnel (UDP) =====================

static bool dns_over_tunnel_query(const config_t* cfg, const uint8_t key32[32], const char* host, char out_ip[16]){
    // 明文: encode_addr(tcp=false) to 8.8.8.8:53 + DNS query
    uint8_t addr[64]; size_t alen = encode_addr_block(false, 0x01, "8.8.8.8", 53, addr, sizeof(addr));
    if (!alen) return false;

    uint8_t q[512]; size_t qlen=0;
    if (dns_build_query_a(host, q, sizeof(q), &qlen)!=0) return false;

    uint8_t plain[800]; if (alen+qlen>sizeof(plain)) return false;
    memcpy(plain, addr, alen); memcpy(plain+alen, q, qlen);
    uint8_t iv[16]; rand_bytes(iv,16);
    cfb_ctx_t enc; if (!cfb_init(&enc, key32, iv, TRUE)) return false;
    uint8_t ct[800]; DWORD ctlen=0;
    bool ok = cfb_update(&enc, plain, (DWORD)(alen+qlen), ct, &ctlen);
    cfb_free(&enc); if (!ok) return false;

    uint8_t out[16+800]; memcpy(out, iv,16); memcpy(out+16, ct, ctlen);
    // UDP sendto
    SOCKET su = WSASocket(AF_INET, SOCK_DGRAM, IPPROTO_UDP, NULL,0,0);
    if (su==INVALID_SOCKET) return false;

    struct sockaddr_storage rs; int rslen=0;
    if (!resolve_host_port(cfg->remote_host, (uint16_t)cfg->remote_port, &rs, &rslen)){ closesocket(su); return false; }

    DWORD to=2000; setsockopt(su, SOL_SOCKET, SO_RCVTIMEO, (const char*)&to, sizeof(to));

    int sn = sendto(su, (const char*)out, (int)(16+ctlen), 0, (struct sockaddr*)&rs, rslen);
    if (sn<=0){ closesocket(su); return false; }

    uint8_t buf[1200];
    struct sockaddr_storage from; int fromlen=sizeof(from);
    int rn = recvfrom(su, (char*)buf, sizeof(buf), 0, (struct sockaddr*)&from, &fromlen);
    closesocket(su);
    if (rn<16) return false;

    cfb_ctx_t dec; if (!cfb_init(&dec, key32, buf, FALSE)) return false;
    uint8_t plain2[1200]; DWORD p2=0;
    ok = cfb_update(&dec, buf+16, rn-16, plain2, &p2);
    cfb_free(&dec);
    if (!ok || p2==0) return false;

    // 跳过 addr 块
    size_t off=0;
    uint8_t first = plain2[0];
    if (first==0x01){ off = 1+4+2; }
    else if (first==0x03){ if (p2<2) return false; uint8_t L=plain2[1]; off = 1+1+L+2; if (off>p2) return false; }
    else if (first==0x04){ off = 1+16+2; }
    else return false;
    if (off>=p2) return false;

    return dns_parse_a_ip(plain2+off, p2-off, out_ip)==0;
}

//===================== ConnectEx helper =====================

static bool load_connectex(SOCKET s, LPFN_CONNECTEX* pOut){
    GUID g = WSAID_CONNECTEX; DWORD bytes=0;
    int rc = WSAIoctl(s, SIO_GET_EXTENSION_FUNCTION_POINTER, &g, sizeof(g), pOut, sizeof(*pOut), &bytes, NULL, NULL);
    return rc==0 && *pOut!=NULL;
}
static bool bridge_connect_remote(bridge_t* b){
    // 建 srv socket & 绑定本地
    b->s_srv = WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL, 0, WSA_FLAG_OVERLAPPED);
    if (b->s_srv==INVALID_SOCKET){ logv('E',"WSASocket srv fail %d", WSAGetLastError()); return false; }
    CreateIoCompletionPort((HANDLE)b->s_srv, b->iocp, (ULONG_PTR)b, 0);

    SOCKADDR_STORAGE local; int llen=0;
    // 绑定任意本地
    struct sockaddr_in sin={0}; sin.sin_family=AF_INET; sin.sin_addr.s_addr=0; sin.sin_port=0;
    if (bind(b->s_srv, (struct sockaddr*)&sin, sizeof(sin))!=0){
        logv('E',"bind srv fail %d", WSAGetLastError()); return false;
    }

    if (!load_connectex(b->s_srv, &b->pConnectEx)){ logv('E',"load ConnectEx fail"); return false; }

    SOCKADDR_STORAGE rs; int rslen=0;
    if (!resolve_host_port(b->cfg->remote_host, (uint16_t)b->cfg->remote_port, &rs, &rslen)){
        logv('E',"resolve remote fail");
        return false;
    }

    iocp_op_t* op = op_alloc(OP_CONNECT, 0);
    InterlockedIncrement(&b->pending_ops);

    BOOL ok = b->pConnectEx(b->s_srv, (SOCKADDR*)&rs, rslen, NULL, 0, NULL, &op->ov);
    if (!ok){
        int e=WSAGetLastError();
        if (e!=ERROR_IO_PENDING){
            logv('E',"ConnectEx immediate fail %d", e);
            InterlockedDecrement(&b->pending_ops);
            op_free(op);
            return false;
        }
    }
    logv('I',"ConnectEx to remote server %s:%d", b->cfg->remote_host, b->cfg->remote_port);
    return true;
}

//===================== Close =====================

static void bridge_free(bridge_t* b){
    if (!b) return;
    if (b->s_cli!=INVALID_SOCKET){ closesocket(b->s_cli); b->s_cli=INVALID_SOCKET; }
    if (b->s_srv!=INVALID_SOCKET){ closesocket(b->s_srv); b->s_srv=INVALID_SOCKET; }
    if (b->up_inited){ cfb_free(&b->c_up); b->up_inited=FALSE; }
    if (b->dn_inited){ cfb_free(&b->c_dn); b->dn_inited=FALSE; }
    free(b);
}
static void bridge_start_close(bridge_t* b){
    if (b->closing) return;
    b->closing=TRUE;
    logv('I',"bridge %p start_close (pending_ops=%ld)", b, b->pending_ops);
    shutdown(b->s_srv, SD_BOTH);
    shutdown(b->s_cli, SD_BOTH);
}

//===================== Worker (GQCS) =====================

static void handle_cli_data(bridge_t* b, const uint8_t* data, DWORD n){
    if (b->st!=ST_STREAM) return;
    if (!b->up_inited){ logv('E',"up cipher not ready"); bridge_start_close(b); return; }

    // 加密并发给远端（不再发送IV）
    uint8_t* out = (uint8_t*)malloc(n);
    DWORD m=0; if (!cfb_update(&b->c_up, data, n, out, &m)){ free(out); bridge_start_close(b); return; }
    if (g_verbose) logv('D',"STREAM up: in=%u enc=%u", n, m);
    post_send(b, b->s_srv, out, m, OP_SRV_SEND);
    free(out);
}

static void handle_srv_data(bridge_t* b, uint8_t* data, DWORD n){
    DWORD off=0;
    // 下行首次收到时，先收满 16B IV
    if (!b->dn_inited){
        DWORD need = (DWORD)(16 - b->ivdn_have);
        if (n < need){
            memcpy(b->iv_dn + b->ivdn_have, data, n);
            b->ivdn_have += n;
            return;
        }else{
            // 完成 IV 收集
            memcpy(b->iv_dn + b->ivdn_have, data, need);
            b->ivdn_have = 16;
            off = need;
            if (!cfb_init(&b->c_dn, b->key32, b->iv_dn, FALSE)){ bridge_start_close(b); return; }
            b->dn_inited=TRUE;
            logv('I',"downstream IV received");
        }
    }
    if (off < n){
        uint8_t* out=(uint8_t*)malloc(n - off);
        DWORD m=0;
        if (!cfb_update(&b->c_dn, data+off, n-off, out, &m)){ free(out); bridge_start_close(b); return; }
        if (g_verbose) logv('D',"STREAM down: in=%u dec=%u", n-off, m);
        post_send(b, b->s_cli, out, m, OP_CLI_SEND);
        free(out);
    }
}

static DWORD WINAPI worker_thread(LPVOID arg_){
    worker_arg_t* arg=(worker_arg_t*)arg_;
    HANDLE iocp=arg->iocp;
    for (;;){
        DWORD bytes=0; ULONG_PTR key=0; OVERLAPPED* pov=NULL;
        BOOL ok = GetQueuedCompletionStatus(iocp, &bytes, &key, &pov, INFINITE);
        bridge_t* b=(bridge_t*)key;
        iocp_op_t* op = (iocp_op_t*)pov;
        if (!pov){ // shutdown signal?
            break;
        }
        if (!ok){
            DWORD gle=GetLastError();
            // 某些取消/超时也会走这里
            // 统一当作该 op 完成，交给状态机关闭
        }

        if (op->kind==OP_CLI_RECV){
            if (g_verbose) logv('D',"GQCS: CLI_RECV %u bytes", bytes);
            if (bytes==0){ // 客户端关闭
                InterlockedDecrement(&b->pending_ops);
                op_free(op);
                bridge_start_close(b);
                continue;
            }

            // 状态机：GREETING 或 REQUEST 或 STREAM
            if (b->st==ST_GREETING){
                DWORD need=0;
                if (!s5_greeting((uint8_t*)op->buf.buf, bytes, &need)){
                    // 不足，理论上应缓存；这里简单化，认为一次收齐（curl等会在一个包里）
                    if (bytes<2){ // 再贴收
                        post_recv(b, b->s_cli, OP_CLI_RECV);
                        InterlockedDecrement(&b->pending_ops); op_free(op);
                        continue;
                    }
                    logv('E',"S5 greeting malformed");
                    InterlockedDecrement(&b->pending_ops); op_free(op);
                    bridge_start_close(b);
                    continue;
                }
                // 丢掉 greeting，并回 05 00
                uint8_t ok2[2]={0x05,0x00};
                post_send(b, b->s_cli, ok2, 2, OP_CLI_SEND);
                if (g_verbose) logv('I',"S5 greeting ok");
                b->st = ST_REQUEST;
                // 再收 request
                post_recv(b, b->s_cli, OP_CLI_RECV);
                InterlockedDecrement(&b->pending_ops); op_free(op);
                continue;

            }else if (b->st==ST_REQUEST){
                uint8_t cmd,atyp; char host[128]; uint16_t port=0; DWORD used=0;
                int r = s5_parse_request((uint8_t*)op->buf.buf, bytes, &cmd,&atyp,host,&port,&used);
                if (r<=0){
                    // 继续收
                    post_recv(b, b->s_cli, OP_CLI_RECV);
                    InterlockedDecrement(&b->pending_ops); op_free(op);
                    continue;
                }
                b->req_cmd=cmd; b->req_atyp=atyp; strncpy(b->req_host,host,127); b->req_port=port;
                if (g_verbose) logv('I',"S5 request cmd=0x%02X atyp=0x%02X host=%s port=%u", cmd,atyp,host,port);

                if (cmd==0x01){ // CONNECT
                    // 先派生密钥
                    if (!kdf_key32_from_password(b->cfg->password, b->key32)){
                        logv('E',"KDF failed"); bridge_start_close(b);
                        InterlockedDecrement(&b->pending_ops); op_free(op); continue;
                    }

                    // 如果是域名且启用 DNS over tunnel，先尝试查询 A
                    if (b->cfg->dns_on && atyp==0x03){
                        char ip[16]={0};
                        if (dns_over_tunnel_query(b->cfg, b->key32, host, ip)){
                            strncpy(b->req_host, ip, 127); b->req_atyp = 0x01;
                            logv('I',"DOT result %s -> %s", host, ip);
                        }else{
                            logv('W',"DOT fail, fallback to remote resolve (keep domain)");
                        }
                    }

                    // 连接远端中继
                    if (!bridge_connect_remote(b)){
                        s5_reply(b->s_cli, 0x05, "0.0.0.0", 0);
                        bridge_start_close(b);
                        InterlockedDecrement(&b->pending_ops); op_free(op); continue;
                    }
                    b->st = ST_CONNECTING;

                }else if (cmd==0x03){
                    // UDP ASSOC（此版本先不实现 SOCKS UDP 转发，只回 OK，防吞包）
                    s5_reply(b->s_cli, 0x00, "0.0.0.0", 0);
                }else{
                    s5_reply(b->s_cli, 0x07, "0.0.0.0", 0);
                    bridge_start_close(b);
                }
                InterlockedDecrement(&b->pending_ops); op_free(op);
                continue;

            }else if (b->st==ST_STREAM){
                handle_cli_data(b, (uint8_t*)op->buf.buf, bytes);
                // 继续收客户端
                post_recv(b, b->s_cli, OP_CLI_RECV);
                InterlockedDecrement(&b->pending_ops); op_free(op);
                continue;
            }

            // 其它状态
            InterlockedDecrement(&b->pending_ops); op_free(op);

        }else if (op->kind==OP_SRV_RECV){
            if (g_verbose) logv('D',"GQCS: SRV_RECV %u bytes", bytes);
            if (bytes==0){
                InterlockedDecrement(&b->pending_ops);
                op_free(op);
                bridge_start_close(b);
                continue;
            }
            handle_srv_data(b, (uint8_t*)op->buf.buf, bytes);
            // 继续收远端
            post_recv(b, b->s_srv, OP_SRV_RECV);
            InterlockedDecrement(&b->pending_ops); op_free(op);

        }else if (op->kind==OP_CLI_SEND || op->kind==OP_SRV_SEND){
            if (g_verbose) logv('D',"GQCS: %s %u bytes", op->kind==OP_CLI_SEND?"CLI_SEND":"SRV_SEND", bytes);
            InterlockedDecrement(&b->pending_ops); op_free(op);

        }else if (op->kind==OP_CONNECT){
            DWORD gle = 0;
            BOOL ok2 = WSAGetOverlappedResult(b->s_srv, &op->ov, &bytes, FALSE, &gle);
            if (!ok2){
                logv('E',"ConnectEx complete err gle=%u", gle);
                s5_reply(b->s_cli, 0x05, "0.0.0.0", 0);
                InterlockedDecrement(&b->pending_ops); op_free(op);
                bridge_start_close(b); continue;
            }
            // 必须调用 setsockopt/更新为连接态
            setsockopt(b->s_srv, IPPROTO_TCP, TCP_NODELAY, (const char*)&(int){1}, sizeof(int));
            // 将 socket 转换为连接态（微软文档：ConnectEx 完成后需调用 setsockopt 或 WSAIoctl SIO_KEEPALIVE_VALS/或者调用 getsockname 等）
            struct sockaddr_in name; int namelen=sizeof(name);
            getsockname(b->s_srv,(struct sockaddr*)&name,&namelen);

            if (g_verbose) logv('I',"ConnectEx completed ok=1 gle=0");
            InterlockedDecrement(&b->pending_ops); op_free(op);

            after_remote_connected(b);
        }

        // 统一关闭判定
        if (b->closing && InterlockedCompareExchange(&b->pending_ops,0,0)==0){
            logv('I',"bridge %p free", b);
            bridge_free(b);
        }
    }
    return 0;
}

//===================== Accept loop =====================

static DWORD WINAPI accept_loop(LPVOID arg_){
    worker_arg_t* wa=(worker_arg_t*)arg_;
    config_t* cfg = wa->cfg;

    SOCKET ls = WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL,0, WSA_FLAG_OVERLAPPED);
    if (ls==INVALID_SOCKET){ logv('E',"listen socket fail %d", WSAGetLastError()); return 1; }

    struct sockaddr_in sin={0}; sin.sin_family=AF_INET; sin.sin_port=htons((u_short)cfg->listen_port);
    InetPtonA(AF_INET, cfg->listen_host, &sin.sin_addr);
    if (bind(ls,(struct sockaddr*)&sin,sizeof(sin))!=0){ logv('E',"bind fail %d", WSAGetLastError()); return 1; }
    if (listen(ls, SOMAXCONN)!=0){ logv('E',"listen fail %d", WSAGetLastError()); return 1; }

    logv('I',"ready.");

    for (;;){
        struct sockaddr_storage cs; int clen=sizeof(cs);
        SOCKET s = accept(ls, (struct sockaddr*)&cs, &clen);
        if (s==INVALID_SOCKET){ int e=WSAGetLastError(); if (e==WSAEINTR) break; continue; }

        char ipbuf[64]; uint16_t cport=0;
        if (cs.ss_family==AF_INET){
            struct sockaddr_in* a=(struct sockaddr_in*)&cs;
            inet_ntop(AF_INET, &a->sin_addr, ipbuf, sizeof(ipbuf)); cport=ntohs(a->sin_port);
        }else{
            strcpy(ipbuf,"::1"); cport=0;
        }
        logv('I',"accepted socket=%llu from %s:%u", (unsigned long long)s, ipbuf, cport);

        // 建 bridge
        bridge_t* b=(bridge_t*)calloc(1,sizeof(bridge_t));
        b->s_cli = s; b->s_srv=INVALID_SOCKET; b->cfg = cfg; b->recv_cap = (DWORD)cfg->recv_buf; b->st=ST_GREETING;
        b->iocp = wa->iocp;
        CreateIoCompletionPort((HANDLE)b->s_cli, wa->iocp, (ULONG_PTR)b, 0);

        // 贴收客户端
        post_recv(b, b->s_cli, OP_CLI_RECV);
    }
    return 0;
}

//===================== main =====================

int main(int argc, char** argv){
    config_t cfg; parse_args(argc,argv,&cfg);
    g_verbose = cfg.verbose;
    srand((unsigned)time(NULL));
    g_main_tid = GetCurrentThreadId();

    WSADATA w; if (WSAStartup(MAKEWORD(2,2), &w)!=0){ return 1; }

    HANDLE iocp = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
    logv('I',"=== mx iocp socks5 ===");
    logv('I',"listen %s:%d  remote %s:%d  verbose=%d", cfg.listen_host, cfg.listen_port, cfg.remote_host, cfg.remote_port, cfg.verbose);
    logv('I',"udp_timeout=%ds connect_timeout=%ds", cfg.udp_timeout_sec, cfg.connect_timeout_sec);

    worker_arg_t wa = { .iocp=iocp, .cfg=&cfg };

    // workers
    const int N=3; HANDLE th[N];
    for (int i=0;i<N;i++) th[i]=CreateThread(NULL,0,worker_thread,&wa,0,NULL);

    // accept
    HANDLE thacc = CreateThread(NULL,0,accept_loop,&wa,0,NULL);

    WaitForSingleObject(thacc, INFINITE);
    // signal workers stop
    for (int i=0;i<N;i++) PostQueuedCompletionStatus(iocp, 0, 0, NULL);
    WaitForMultipleObjects(N, th, TRUE, INFINITE);

    CloseHandle(iocp);
    WSACleanup();
    return 0;
}