Skip to content

Commit 72886f2

Browse files
matyhtfmatyhtf
committed
Optimize Client::socks5_handshake(), fix core tests
1 parent e5b1c90 commit 72886f2

File tree

2 files changed

+109
-69
lines changed

2 files changed

+109
-69
lines changed

‎core-tests/src/network/client.cpp

+98-63
Original file line numberDiff line numberDiff line change
@@ -217,9 +217,19 @@ TEST(client, shutdown_all) {
217217
}
218218

219219
#ifdef SW_USE_OPENSSL
220-
TEST(client, ssl_1) {
221-
int ret;
222220

221+
static const char *request_baidu = "GET / HTTP/1.1\r\n"
222+
"Host: www.baidu.com\r\n"
223+
"Connection: close\r\n"
224+
"User-Agent: Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) "
225+
"Chrome/51.0.2704.106 Safari/537.36"
226+
"\r\n\r\n";
227+
228+
static const char *domain_baidu = "www.baidu.com";
229+
230+
#define SOCKS5_WITH_AUTH 1
231+
232+
TEST(client, ssl_1) {
223233
bool connected = false;
224234
bool closed = false;
225235
swoole::String buf(65536);
@@ -230,21 +240,14 @@ TEST(client, ssl_1) {
230240
client.enable_ssl_encrypt();
231241
client.onConnect = [&connected](Client *cli) {
232242
connected = true;
233-
cli->send(cli,
234-
SW_STRL("GET / HTTP/1.1\r\n"
235-
"Host: www.baidu.com\r\n"
236-
"Connection: close\r\n"
237-
"User-Agent: Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) "
238-
"Chrome/51.0.2704.106 Safari/537.36"
239-
"\r\n\r\n"),
240-
0);
243+
cli->send(cli, request_baidu, strlen(request_baidu), 0);
241244
};
242245

243246
client.onError = [](Client *cli) {};
244247
client.onClose = [&closed](Client *cli) { closed = true; };
245248
client.onReceive = [&buf](Client *cli, const char *data, size_t length) { buf.append(data, length); };
246-
ret = client.connect(&client, "www.baidu.com", 443, -1, 0);
247-
ASSERT_EQ(ret, 0);
249+
250+
ASSERT_EQ(client.connect(&client, domain_baidu, 443, -1, 0), 0);
248251

249252
swoole_event_wait();
250253

@@ -253,89 +256,121 @@ TEST(client, ssl_1) {
253256
ASSERT_TRUE(buf.contains("Baidu"));
254257
}
255258

256-
257-
TEST(client, http_proxy) {
259+
static void proxy_async_test(Client &client, bool https) {
258260
int ret;
259261

262+
swoole_event_init(SW_EVENTLOOP_WAIT_EXIT);
263+
260264
bool connected = false;
261265
bool closed = false;
262266
swoole::String buf(65536);
263267

264-
swoole_event_init(SW_EVENTLOOP_WAIT_EXIT);
265-
266-
Client client(SW_SOCK_TCP, true);
267-
client.enable_ssl_encrypt();
268-
client.http_proxy = new HttpProxy();
269-
client.http_proxy->proxy_host = std::string(TEST_HTTP_PROXY_HOST);
270-
client.http_proxy->proxy_port = TEST_HTTP_PROXY_PORT;
268+
if (https) {
269+
client.enable_ssl_encrypt();
270+
}
271271

272272
client.onConnect = [&connected](Client *cli) {
273273
connected = true;
274-
cli->send(cli,
275-
SW_STRL("GET / HTTP/1.1\r\n"
276-
"Host: www.baidu.com\r\n"
277-
"Connection: close\r\n"
278-
"User-Agent: Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) "
279-
"Chrome/51.0.2704.106 Safari/537.36"
280-
"\r\n\r\n"),
281-
0);
274+
cli->send(cli, request_baidu, strlen(request_baidu), 0);
282275
};
283276

284277
client.onError = [](Client *cli) {};
285278
client.onClose = [&closed](Client *cli) { closed = true; };
286279
client.onReceive = [&buf](Client *cli, const char *data, size_t length) { buf.append(data, length); };
287-
ret = client.connect(&client, "www.baidu.com", 443, -1, 0);
288-
ASSERT_EQ(ret, 0);
280+
281+
ASSERT_EQ(client.connect(&client, domain_baidu, https ? 443 : 80, -1, 0), 0);
289282

290283
swoole_event_wait();
291284

292285
ASSERT_TRUE(connected);
293286
ASSERT_TRUE(closed);
294-
ASSERT_TRUE(buf.contains("Baidu"));
287+
ASSERT_TRUE(buf.contains("www.baidu.com"));
295288
}
296289

297-
TEST(client, socks5_proxy) {
298-
int ret;
299-
300-
bool connected = false;
301-
bool closed = false;
290+
static void proxy_sync_test(Client &client, bool https) {
302291
swoole::String buf(65536);
292+
if (https) {
293+
client.enable_ssl_encrypt();
294+
}
295+
296+
ASSERT_EQ(client.connect(&client, domain_baidu, https ? 443 : 80, -1, 0), 0);
297+
ASSERT_GT(client.send(&client, request_baidu, strlen(request_baidu), 0), 0);
298+
299+
while(true) {
300+
char rbuf[4096];
301+
auto nr = client.recv(&client, rbuf, sizeof(rbuf), 0);
302+
if (nr <= 0) {
303+
break;
304+
}
305+
buf.append(rbuf, nr);
306+
}
307+
308+
ASSERT_TRUE(buf.contains("www.baidu.com"));
309+
}
303310

304-
swoole_event_init(SW_EVENTLOOP_WAIT_EXIT);
305-
306-
Client client(SW_SOCK_TCP, true);
307-
client.enable_ssl_encrypt();
308-
311+
static void proxy_set_socks5_proxy(Client &client) {
309312
client.socks5_proxy = new Socks5Proxy();
310313
client.socks5_proxy->host = std::string("127.0.0.1");
311314
client.socks5_proxy->port = 1080;
312315
client.socks5_proxy->dns_tunnel = 1;
313-
client.socks5_proxy->method = 0x02;
316+
#if SOCKS5_WITH_AUTH
317+
client.socks5_proxy->method = SW_SOCKS5_METHOD_AUTH;
314318
client.socks5_proxy->username = std::string("user");
315319
client.socks5_proxy->password = std::string("password");
320+
#endif
321+
}
316322

317-
client.onConnect = [&connected](Client *cli) {
318-
connected = true;
319-
cli->send(cli,
320-
SW_STRL("GET / HTTP/1.1\r\n"
321-
"Host: www.baidu.com\r\n"
322-
"Connection: close\r\n"
323-
"User-Agent: Mozilla/5.0 (Windows NT 10.0; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) "
324-
"Chrome/51.0.2704.106 Safari/537.36"
325-
"\r\n\r\n"),
326-
0);
327-
};
323+
static void proxy_set_http_proxy(Client &client) {
324+
client.http_proxy = new HttpProxy();
325+
client.http_proxy->proxy_host = std::string(TEST_HTTP_PROXY_HOST);
326+
client.http_proxy->proxy_port = TEST_HTTP_PROXY_PORT;
327+
}
328328

329-
client.onError = [](Client *cli) {};
330-
client.onClose = [&closed](Client *cli) { closed = true; };
331-
client.onReceive = [&buf](Client *cli, const char *data, size_t length) { buf.append(data, length); };
332-
ret = client.connect(&client, "www.baidu.com", 443, -1, 0);
333-
ASSERT_EQ(ret, 0);
329+
TEST(client, https_get_async_with_http_proxy) {
330+
Client client(SW_SOCK_TCP, true);
331+
proxy_set_http_proxy(client);
332+
proxy_async_test(client, true);
333+
}
334334

335-
swoole_event_wait();
335+
TEST(client, https_get_async_with_socks5_proxy) {
336+
Client client(SW_SOCK_TCP, true);
337+
proxy_set_socks5_proxy(client);
338+
proxy_async_test(client, true);
339+
}
336340

337-
ASSERT_TRUE(connected);
338-
ASSERT_TRUE(closed);
339-
ASSERT_TRUE(buf.contains("Baidu"));
341+
TEST(client, https_get_sync_with_http_proxy) {
342+
Client client(SW_SOCK_TCP, false);
343+
proxy_set_http_proxy(client);
344+
proxy_sync_test(client, true);
345+
}
346+
347+
TEST(client, https_get_sync_with_socks5_proxy) {
348+
Client client(SW_SOCK_TCP, false);
349+
proxy_set_socks5_proxy(client);
350+
proxy_sync_test(client, true);
351+
}
352+
353+
TEST(client, http_get_async_with_http_proxy) {
354+
Client client(SW_SOCK_TCP, true);
355+
proxy_set_http_proxy(client);
356+
proxy_async_test(client, false);
357+
}
358+
359+
TEST(client, http_get_async_with_socks5_proxy) {
360+
Client client(SW_SOCK_TCP, true);
361+
proxy_set_socks5_proxy(client);
362+
proxy_async_test(client, false);
363+
}
364+
365+
TEST(client, http_get_sync_with_http_proxy) {
366+
Client client(SW_SOCK_TCP, false);
367+
proxy_set_http_proxy(client);
368+
proxy_sync_test(client, false);
369+
}
370+
371+
TEST(client, http_get_sync_with_socks5_proxy) {
372+
Client client(SW_SOCK_TCP, false);
373+
proxy_set_socks5_proxy(client);
374+
proxy_sync_test(client, false);
340375
}
341376
#endif

‎src/network/client.cc

+11-6
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ int Client::socks5_handshake(const char *recv_data, size_t length) {
193193

194194
ctx->state = SW_SOCKS5_STATE_AUTH;
195195

196-
return send(this, ctx->buf, ctx->username.length() + ctx->password.length() + 3, 0);
196+
return send(this, ctx->buf, ctx->username.length() + ctx->password.length() + 3, 0) > 0 ? SW_OK : SW_ERR;
197197
}
198198
// send connect request
199199
else {
@@ -211,14 +211,14 @@ int Client::socks5_handshake(const char *recv_data, size_t length) {
211211
memcpy(buf, ctx->target_host.c_str(), ctx->target_host.length());
212212
buf += ctx->target_host.length();
213213
*(uint16_t *) buf = htons(ctx->target_port);
214-
return send(this, ctx->buf, ctx->target_host.length() + 7, 0);
214+
return send(this, ctx->buf, ctx->target_host.length() + 7, 0) > 0 ? SW_OK : SW_ERR;
215215
} else {
216216
buf[3] = 0x01;
217217
buf += 4;
218218
*(uint32_t *) buf = htons(ctx->target_host.length());
219219
buf += 4;
220220
*(uint16_t *) buf = htons(ctx->target_port);
221-
return send(this, ctx->buf, ctx->target_host.length() + 7, 0);
221+
return send(this, ctx->buf, ctx->target_host.length() + 7, 0) > 0 ? SW_OK : SW_ERR;
222222
}
223223
}
224224
} else if (ctx->state == SW_SOCKS5_STATE_AUTH) {
@@ -249,13 +249,14 @@ int Client::socks5_handshake(const char *recv_data, size_t length) {
249249
#endif
250250
if (result == 0) {
251251
ctx->state = SW_SOCKS5_STATE_READY;
252+
return SW_OK;
252253
} else {
253254
swoole_error_log(SW_LOG_NOTICE,
254255
SW_ERROR_SOCKS5_SERVER_ERROR,
255256
"Socks5 server error, reason :%s",
256257
Socks5Proxy::strerror(result));
258+
return SW_ERR;
257259
}
258-
return result;
259260
}
260261
return SW_OK;
261262
}
@@ -940,9 +941,13 @@ static int Client_onStreamRead(Reactor *reactor, Event *event) {
940941
goto _connect_fail;
941942
}
942943
cli->buffer->length += n;
943-
if (cli->socks5_handshake(buf, buf_size) < 0 || cli->socks5_proxy->state != SW_SOCKS5_STATE_READY) {
944+
if (cli->socks5_handshake(buf, buf_size) < 0) {
944945
goto _connect_fail;
945946
}
947+
if (cli->socks5_proxy->state != SW_SOCKS5_STATE_READY) {
948+
return SW_OK;
949+
}
950+
cli->buffer->clear();
946951
if (!do_ssl_handshake) {
947952
execute_onConnect(cli);
948953
return SW_OK;
@@ -1119,7 +1124,7 @@ static int Client_onWrite(Reactor *reactor, Event *event) {
11191124
// socks5 proxy
11201125
if (cli->socks5_proxy && cli->socks5_proxy->state == SW_SOCKS5_STATE_WAIT) {
11211126
char buf[3];
1122-
Socks5Proxy::pack(buf, cli->socks5_proxy->username.empty() ? 0x00 : 0x02);
1127+
Socks5Proxy::pack(buf, cli->socks5_proxy->username.empty() ? 0 : SW_SOCKS5_METHOD_AUTH);
11231128
cli->socks5_proxy->state = SW_SOCKS5_STATE_HANDSHAKE;
11241129
return cli->send(cli, buf, sizeof(buf), 0);
11251130
}

0 commit comments

Comments
 (0)