SOCKS代理网络操作中未知数据大小的缓冲区处理及poll用法咨询

问题描述 #

1. 用简单回显服务器测试SOCKS4代理时数据转发正常，但切换到其他协议（不同数据包大小）时，不确定如何判断传输完成，以及确保所有数据正确转发到远程对等端和客户端。
2. 尝试用poll()替代select()，确认poll()无需FD_SET、FD_ISSET等操作且代码更简洁是否属于正常情况。

原始线程处理代码（同步读写） #

// Contains the client socket and remote peer socket after successful connect:
typedef struct socks_connection {
    int client_sock;
    int target_sock;
} socks_connection;


void *socks_connection_thread(void *sockets) {
    printf("Thread started\n");
    socks_connection conn = *(socks_connection*) sockets;

    uint8_t buf[512];
    int rbytes;
    int wbytes;

    for(;;) {

        rbytes = recv(conn.client_sock, buf, sizeof(buf), 0);
        printf("read: %d\n", rbytes);
        if(rbytes < 0) {
            perror("read");
        }

        wbytes = send(conn.target_sock, buf, rbytes, 0);
        printf("send: %d\n", wbytes);
        if(rbytes < 0) {
            perror("send");
        }

        rbytes = recv(conn.target_sock, buf, sizeof(buf), 0);
        if(rbytes < 0) {
            perror("read");
        }
        printf("read: %d\n", rbytes);

        wbytes = send(conn.client_sock, buf, rbytes, 0);
        printf("send: %d\n", wbytes);
        if(rbytes < 0) {
            perror("send");
        }
    }

}

改进后的poll()版本代码 #

#define MAX_SOCKETS 2
#define DEFAULT_TIMEOUT (3 * 60 * 1000)
#define CLIENT_POLL 0
#define REMOTE_POLL 1

void *socks_connection_thread(void *pipefd) {

    pthread_detach(pthread_self());
    socks_connection *conn = pipefd;
    int rc = 0;
    int timeout = DEFAULT_TIMEOUT;

    struct pollfd pfds[MAX_SOCKETS];
    nfds_t nfds = MAX_SOCKETS;

    uint8_t client_buf[1024];
    size_t client_buf_size = 0;

    uint8_t target_buf[1024];
    size_t target_buf_size = 0;
    
    ssize_t num_bytes;

    memset(&pfds, 0, sizeof(pfds));

    int opt = 1;
    ioctl(conn->client_sock, FIONBIO, &opt);
    ioctl(conn->target_sock, FIONBIO, &opt);

    pfds[CLIENT_POLL].fd = conn->client_sock;
    pfds[CLIENT_POLL].events = POLLIN;

    pfds[REMOTE_POLL].fd = conn->target_sock;
    pfds[REMOTE_POLL].events = POLLIN;

    for(;;) {

        if(socksshutdown) break;

        /* waiting for some events */
        rc = poll(pfds, MAX_SOCKETS, timeout);
        if(rc < 0) {
            fprintf(stderr, "poll() failed: %s\n", strerror(errno));
            break;
        }

        if(rc == 0) {
            fprintf(stderr, "poll() timed out. End Connection\n");
            break;
        }

        /* there is something to read form the client side */
        if(pfds[CLIENT_POLL].revents & POLLIN)
        {

            num_bytes = readFromSocket(conn->client_sock, client_buf, sizeof(client_buf));
            if(num_bytes < 0) break; // client connection lost
            if(num_bytes > 0) {
                printf("read from client: %s (%ld)\n", client_buf, num_bytes);
                client_buf_size += num_bytes;
            }
            num_bytes = sendToSocket(conn->target_sock, client_buf, num_bytes);
            if(num_bytes < 0) break;
            if(num_bytes > 0) {
                printf("forward to remote peer: %s (%ld)\n", client_buf, num_bytes);
            }
            pfds[CLIENT_POLL].revents = 0;

        }
        /* there is something to read from the remote side */
        else if(pfds[REMOTE_POLL].revents & POLLIN)
        {
            //printf("Got data from remote.\n");
            num_bytes = readFromSocket(conn->target_sock, target_buf, sizeof(target_buf));
            if (num_bytes < 0) break; // remote connection lost
            if (num_bytes > 0) {
                printf("read from client: %s (%ld)\n", target_buf, num_bytes);
                target_buf_size += num_bytes;
            }
            num_bytes = sendToSocket(conn->client_sock, target_buf, target_buf_size);
                if (num_bytes < 0) break;
                if (num_bytes > 0) {
                    printf("forward to client: %s (%ld)\n", target_buf, num_bytes);
                    // remove the sent bytes...
                } 
            pfds[REMOTE_POLL].revents = 0;
        } else {
            // unexpected event result appeared so close the connection
            break;
        }

    }

    // all done
    close(conn->client_sock);
    close(conn->target_sock);
    printf("Thread terminating\n");

}

解决方案与说明 #

一、不同协议下的数据转发处理 #

1. 传输完成的判断

   • 当recv()返回0时，表示对端已主动关闭连接（发送FIN包），此时应关闭对应socket并终止转发循环。
   • 当recv()返回-1时，需区分错误类型：如果是EAGAIN或EWOULDBLOCK（非阻塞IO下无数据可读），属于正常情况，无需断开；其他错误（如ECONNRESET）则表示连接异常，需终止转发。
   • 原同步代码的核心问题是强制交替读写，会导致线程阻塞在某一个recv()调用上，无法处理另一端的新数据，必须用poll()/select()这类多路复用IO监听两个socket的可读事件，谁有数据就处理谁。

2. 确保数据完整转发

   • send()可能无法一次性发送所有字节（比如内核缓冲区满），因此需要循环调用send()，直到所有待发送字节都被写入，或者遇到错误。示例封装函数：

     ssize_t full_send(int sock, const void *buf, size_t len) {
         size_t sent = 0;
         while (sent < len) {
             ssize_t ret = send(sock, (const uint8_t*)buf + sent, len - sent, 0);
             if (ret < 0) {
                 if (errno == EAGAIN || errno == EWOULDBLOCK) {
                     // 非阻塞下需等待可写，可结合poll监听POLLOUT事件
                     continue;
                 }
                 return ret;
             }
             sent += ret;
         }
         return sent;
     }
     

   • 读取数据时，非阻塞IO下recv()返回EAGAIN意味着当前无数据，无需处理，等待下一次poll()触发POLLIN事件即可。

二、关于poll()的使用确认 #

poll()确实比select()更简洁，无需FD_SET、FD_ISSET等宏操作，这是正常的：

• poll()通过struct pollfd数组管理多个socket，直接设置需要监听的事件（如POLLIN），返回后检查revents字段即可判断事件类型。
• 相比select()，poll()没有FD_SETSIZE的限制，能管理更多socket，代码逻辑更清晰，是多路复用IO的更优选择。

三、poll()版本代码的优化点 #

1. 缓冲区管理问题：代码中client_buf_size和target_buf_size的累积逻辑错误，会导致重复发送旧数据。代理无需累积数据，读到多少就转发多少即可，直接使用num_bytes作为发送长度，去掉缓冲区大小累积的逻辑。
2. 打印信息错误：处理远程socket读取时，打印的是read from client，应改为read from remote。
3. readFromSocket/sendToSocket实现：需确保这两个函数正确处理非阻塞IO的EAGAIN错误，比如readFromSocket遇到EAGAIN时应返回0（表示当前无数据）而非-1，避免误判连接断开。
4. 事件重置：无需手动设置pfds[...].revents = 0，poll()每次调用会覆盖revents字段。

内容的提问来源于stack exchange，提问作者mortytheshorty