to inspire confidence in somebody.

0%

也许是一个最简单的poll实现Socket代理

这几天翻了翻项目的代码,看到了一个非常简单的代理程序,使用poll实现,可以在代理过程中输出数据流,基本上算是教科书级别的poll使用例子了,分享一下:

1
#include <cstdio>
2
#include <cstdlib>
3
#include <cstring>
4
#include <iterator>
5
#include <iostream>
6
#include <unistd.h>
7
#include <errno.h>
8
#include <sys/time.h>
9
#include <sys/types.h>
10
#include <sys/socket.h>
11
#include <arpa/inet.h>
12
#include <fcntl.h>
13
#include <poll.h>
14
15
#define PROXY_PORT  8888
16
#define BIN_TXT     0
17
#define RAW_TXT     1
18
19
static const char *backendHost = 0;
20
static int backendPort;
21
static int oflag = RAW_TXT;
22
23
bool proxy(int cfd);
24
25
// g++ -o proxy proxy.cpp -I. -Wall
26
int main(int argc, char *argv[])
27
{
28
    if (argc == 3 || argc == 4) {
29
        backendHost = argv[1];  
30
        backendPort = atoi(argv[2]);    
31
        if (argc == 4) oflag = atoi(argv[3]);
32
    } else {
33
        fprintf(stderr, "usage: %s BackendHost BackendPort oflag\n", argv[0]);  
34
        exit(-1);
35
    }
36
37
    int fd = socket(AF_INET, SOCK_STREAM, 0);
38
39
    struct sockaddr_in addr;
40
    memset(&addr, 0x00, sizeof(addr));
41
    addr.sin_family = AF_INET;
42
    addr.sin_port = htons(PROXY_PORT);
43
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
44
45
    int flags = 1;
46
    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (void *)&flags, sizeof(flags));
47
48
    if (bind(fd, (struct sockaddr*)&addr, sizeof(addr)) != 0) {
49
        fprintf(stderr, "[error] bind error, %s\n", strerror(errno));
50
        return EXIT_FAILURE;
51
    }
52
53
    listen(fd, 10);
54
55
    int cfd;
56
    while ((cfd = accept(fd, NULL, NULL)) > 0) {
57
#ifdef USE_FORK
58
        pid_t pid = fork();
59
        if (pid == 0) {
60
#endif
61
            proxy(cfd);
62
            close(cfd);
63
#ifdef USE_FORK
64
            exit(0);
65
        } else if (pid == -1) {
66
            fprintf(stderr, "[error] fork error, %s\n", strerror(errno));   
67
        } else {
68
            close(cfd); 
69
        }
70
#endif
71
    }
72
    close(fd);
73
74
    return EXIT_SUCCESS;
75
}
76
77
static int connectBackend()
78
{
79
    int sfd = socket(AF_INET, SOCK_STREAM, 0);
80
    if (sfd == -1) {
81
        fprintf(stderr, "socket error, %s\n", strerror(errno)); 
82
        return -1;
83
    }
84
85
    struct sockaddr_in addr;
86
    memset(&addr, 0x00, sizeof(addr));
87
    addr.sin_family = AF_INET;
88
    addr.sin_port = htons(backendPort);
89
    inet_pton(AF_INET, backendHost, &addr.sin_addr);
90
91
    if (connect(sfd, (struct sockaddr*)&addr, sizeof(addr)) != 0) {
92
        fprintf(stderr, "connect backed error, %s\n", strerror(errno));
93
        close(sfd);
94
        return -1;
95
    }
96
97
    fprintf(stdout, "[debug %d] connect backend ok\n", getpid());
98
    return sfd;
99
}
100
101
static void txtPrint(const char *buffer, size_t size)
102
{
103
    /* stderr unbuffered */
104
    if (oflag == BIN_TXT) {
105
        for (size_t i = 0; i < size; ++i) {
106
            fprintf(stderr, "%02x", (unsigned char) buffer[i]);
107
        }
108
    } else {
109
        for (size_t i = 0; i < size; ++i) {
110
            if (isprint((unsigned char) buffer[i])) {
111
                fprintf(stderr, "%c", buffer[i]);   
112
            } else {
113
                fprintf(stderr, "%%%02x", (unsigned char) buffer[i]);
114
            }
115
        }
116
    }
117
}
118
119
bool proxy(int cfd)
120
{
121
    const size_t len = 128;
122
    char buf[len];
123
    ssize_t n;
124
125
    fcntl(cfd, F_SETFL, fcntl(cfd, F_GETFL) | O_NONBLOCK);
126
127
    int sfd = connectBackend();
128
    if (sfd == -1) return false;
129
130
    fcntl(sfd, F_SETFL, fcntl(sfd, F_GETFL) | O_NONBLOCK);
131
132
    struct pollfd pfd[2] = {
133
        {cfd, POLLIN, 0},   
134
        {sfd, POLLIN, 0}
135
    };
136
137
    int nfds = 2;
138
139
    bool stop = false;
140
    while (!stop) {
141
        int nready = poll(pfd, nfds, -1);
142
        if (nready == -1) {
143
            fprintf(stderr, "[error] poll, %s\n", strerror(errno)); 
144
            close(sfd);
145
            return false;
146
        }
147
148
        for (int i = 0; i < nfds; ++i) {
149
            if (pfd[i].revents & POLLIN) {
150
                while ((n = recv(pfd[i].fd, buf, len, 0)) > 0) {
151
                    int fd = (pfd[i].fd == cfd) ? sfd : cfd;
152
                    ssize_t nn;
153
                    size_t offset = 0;
154
                    while ((nn = send(fd, buf + offset, n - offset, 0)) > 0) {
155
                        offset += nn;
156
                        if (offset == (size_t) n) break;
157
                    }
158
                    struct timeval tv;
159
                    gettimeofday(&tv, 0);
160
                    if (pfd[i].fd == cfd) {
161
                        fprintf(stdout, "[debug] %lu:%lu client read\n", (unsigned long) tv.tv_sec, (unsigned long) tv.tv_usec);
162
                    } else {
163
                        fprintf(stdout, "[debug] %lu:%lu server read\n", (unsigned long) tv.tv_sec, (unsigned long) tv.tv_usec);
164
                    }
165
                    txtPrint(buf, n);
166
                }
167
                if (n == 0) {
168
                    if (pfd[i].fd == cfd) {
169
                        fprintf(stdout, "[debug] client closed\n");
170
                    } else {
171
                        fprintf(stdout, "[debug] server closed\n");
172
                    }
173
                    stop = true;
174
                } else if (n == -1 && errno != EAGAIN) {
175
                    fprintf(stderr, "[error] read %s error, %s\n", pfd[i].fd == cfd ? "client" : "backend",
176
                                                                 strerror(errno));  
177
                    close(sfd);
178
                    return false;
179
                }
180
            }
181
        }
182
    }
183
184
    fprintf(stdout, "[debug %d] disconnect\n", getpid());
185
186
    close(sfd);
187
    return true;
188
}

使用g++ -o proxy proxy.cpp -I. -Wall 命令编译,运行时直接指定需要代理的后端IP和Port即可,也支持fork以支持多条链接。