基于 POCO 框架的 TCP 连接分流程序

介绍

下面的程序实现了对 TCP 连接的分流，即将一个 TCP 连接的流量分布到多个 TCP 连接上进行传输。

本程序的主要作用是在特定网络环境下提升通过 TCP 连接的 OpenVPN 服务的速率，使之充分利用带宽。

程序的主要复杂之处在于单生产者—多消费者、多生产者—单消费者两种同步方式的实现。

代码 

[Cpp]  view plain copy print ?

1. // multisocks, copyright (c) 2013 coolypf  
2.   
3. #include <stdio.h>  
4. #include <stdlib.h>  
5. #include <string.h>  
6. #include <string>  
7. #include <map>  
8. #include <vector>  
9. #include <Poco/Thread.h>  
10. #include <Poco/Event.h>  
11. #include <Poco/ErrorHandler.h>  
12. #include <Poco/FIFOBuffer.h>  
13. #include <Poco/Net/StreamSocket.h>  
14. #include <Poco/Net/ServerSocket.h>  
15.   
16. using namespace std;  
17. using namespace Poco;  
18. using namespace Poco::Net;  
19.   
20. class error_handler : public ErrorHandler  
21. {  
22. public:  
23.     virtual void exception(const Exception& exc)  
24.     {  
25.         printf("[!] Thread %u: Caught exception: %s\n", (unsigned int)Thread::currentTid(), exc.displayText().c_str());  
26.         fflush(stdout);  
27.         exit(1);  
28.     }  
29.   
30.     virtual void exception(const std::exception& exc)  
31.     {  
32.         printf("[!] Thread %u: Caught exception: %s\n", (unsigned int)Thread::currentTid(), exc.what());  
33.         fflush(stdout);  
34.         exit(2);  
35.     }  
36.   
37.     virtual void exception()  
38.     {  
39.         printf("[!] Thread %u: Caught unknown exception\n", (unsigned int)Thread::currentTid());  
40.         fflush(stdout);  
41.         exit(3);  
42.     }  
43. };  
44.   
45. class  
46. {  
47.     map<string, string> config;  
48.     string empty;  
49.   
50.     int read_file(const char *filename, vector<char> &out)  
51.     {  
52.         FILE *fp = fopen(filename, "rb");  
53.         if (!fp)  
54.         {  
55.             printf("[!] Fail to load: %s\n", filename);  
56.             return 1;  
57.         }  
58.         char buf[4096];  
59.         size_t sz;  
60.         while ((sz = fread(buf, 1, 4096, fp)))  
61.             out.insert(out.end(), buf, buf + sz);  
62.         fclose(fp);  
63.         for (size_t i = 0; i < out.size(); ++i)  
64.             if (!out[i])  
65.                 out[i] = ' ';  
66.         return 0;  
67.     }  
68.   
69.     void split_lines(const vector<char> &in, vector<string> &out)  
70.     {  
71.         size_t pos = 0;  
72.         for (size_t i = 0; i < in.size(); ++i)  
73.         {  
74.             if (in[i] != '\r' && in[i] != '\n')  
75.                 continue;  
76.             out.push_back(string(in.begin() + pos, in.begin() + i));  
77.             if (in[i] == '\r' && i + 1 < in.size() && in[i + 1] == '\n')  
78.                 i++;  
79.             pos = i + 1;  
80.         }  
81.         if (pos < in.size())  
82.             out.push_back(string(in.begin() + pos, in.end()));  
83.     }  
84.   
85.     void trim_string(string &str)  
86.     {  
87.         const char *spaces = " \r\n\t\xb\xc";  
88.         int first = 0, last = (int)str.size() - 1;  
89.         while (first <= last && strchr(spaces, str[first]))  
90.             ++first;  
91.         while (last >= first && strchr(spaces, str[last]))  
92.             --last;  
93.         str.resize(last + 1);  
94.         str.erase(0, first);  
95.     }  
96.   
97.     void split_string(const string &in, vector<string> &out, const char *delims)  
98.     {  
99.         size_t pos = 0;  
100.         for (size_t i = 0; i < in.size(); ++i)  
101.         {  
102.             if (!strchr(delims, in[i]))  
103.                 continue;  
104.             string part(in.begin() + pos, in.begin() + i);  
105.             trim_string(part);  
106.             out.push_back(part);  
107.             pos = i + 1;  
108.         }  
109.         {  
110.             string part(in.begin() + pos, in.end());  
111.             trim_string(part);  
112.             out.push_back(part);  
113.         }  
114.     }  
115.   
116. public:  
117.     int load(const char *filename)  
118.     {  
119.         vector<char> content;  
120.         if (read_file(filename, content))  
121.             return 1;  
122.         vector<string> lines;  
123.         split_lines(content, lines);  
124.         for (size_t i = 0; i < lines.size(); ++i)  
125.         {  
126.             if (lines[i].empty() || lines[i][0] == '#')  
127.                 continue;  
128.             vector<string> parts;  
129.             split_string(lines[i], parts, "=");  
130.             if (parts.size() != 2 || parts[0].empty())  
131.             {  
132.                 printf("[!] Invalid config: %s (%d)\n", filename, (int)i + 1);  
133.                 continue;  
134.             }  
135.             if (config.find(parts[0]) != config.end())  
136.                 printf("[!] Override config: %s (%d)\n", filename, (int)i + 1);  
137.             config[parts[0]] = parts[1];  
138.         }  
139.         return 0;  
140.     }  
141.   
142.     int i(const char *key) const  
143.     {  
144.         map<string, string>::const_iterator iter = config.find(string(key));  
145.         if (iter == config.end())  
146.         {  
147.             printf("[!] No config: %s\n", key);  
148.             return 0;  
149.         }  
150.         int ret = 0;  
151.         if (sscanf(iter->second.c_str(), "%d", &ret) != 1)  
152.             printf("[!] Invalid int: %s\n", key);  
153.         return ret;  
154.     }  
155.   
156.     const string & s(const char *key) const  
157.     {  
158.         map<string, string>::const_iterator iter = config.find(string(key));  
159.         if (iter == config.end())  
160.         {  
161.             printf("[!] No config: %s\n", key);  
162.             return empty;  
163.         }  
164.         return iter->second;  
165.     }  
166. } config;  
167.   
168. // Read from out buffer and send to remote  
169. class reader : public Runnable  
170. {  
171.     int index;  
172.     StreamSocket socket;  
173.     Event &readable, &writable;  
174.     FIFOBuffer &buffer;  
175.     bool &self_stopped, &peer_stopped;  
176. public:  
177.     reader(int i, StreamSocket s, Event &r, Event &w, FIFOBuffer &b, bool &self, bool &peer)  
178.         : index(i), socket(s), readable(r), writable(w), buffer(b), self_stopped(self), peer_stopped(peer)  
179.     {  
180.     }  
181.   
182.     virtual void run()  
183.     {  
184.         printf("[.] Reader %d, tid = %u\n", index, (unsigned int)Thread::currentTid());  
185.         fflush(stdout);  
186.         char buf[8192];  
187.         bool closed = false;  
188.         while (true)  
189.         {  
190.             if (buffer.isReadable())  
191.             {  
192.                 int sz = (int)buffer.read(buf, 8192);  
193.                 writable.set();  
194.                 int sent = 0;  
195.                 while (sent < sz)  
196.                 {  
197.                     int sz2;  
198.                     try  
199.                     {  
200.                         sz2 = socket.sendBytes(buf + sent, sz - sent);  
201.                     }  
202.                     catch (Exception &exc)  
203.                     {  
204.                         printf("[!] Reader %d: %s\n", index, exc.displayText().c_str());  
205.                         fflush(stdout);  
206.                         sz2 = 0;  
207.                     }  
208.                     if (sz2 <= 0)  
209.                     {  
210.                         closed = true;  
211.                         break;  
212.                     }  
213.                     sent += sz2;  
214.                 }  
215.                 if (closed)  
216.                     break;  
217.             }  
218.             else  
219.             {  
220.                 if (peer_stopped)  
221.                     break;  
222.                 readable.wait();  
223.             }  
224.         }  
225.         try { socket.shutdownSend(); } catch (...) {}  
226.         self_stopped = true;  
227.         writable.set();  
228.     }  
229. };  
230.   
231. // Receive from remote and write to in buffer  
232. class writer : public Runnable  
233. {  
234.     int index;  
235.     StreamSocket socket;  
236.     Event &readable, &writable;  
237.     FIFOBuffer &buffer;  
238.     bool &self_stopped, &peer_stopped;  
239. public:  
240.     writer(int i, StreamSocket s, Event &r, Event &w, FIFOBuffer &b, bool &self, bool &peer)  
241.         : index(i), socket(s), readable(r), writable(w), buffer(b), self_stopped(self), peer_stopped(peer)  
242.     {  
243.     }  
244.   
245.     virtual void run()  
246.     {  
247.         printf("[.] Writer %d, tid = %u\n", index, (unsigned int)Thread::currentTid());  
248.         fflush(stdout);  
249.         char buf[8192];  
250.         while (true)  
251.         {  
252.             int sz;  
253.             try  
254.             {  
255.                 sz = socket.receiveBytes(buf, 8192);  
256.             }  
257.             catch (Exception &exc)  
258.             {  
259.                 printf("[!] Writer %d: %s\n", index, exc.displayText().c_str());  
260.                 fflush(stdout);  
261.                 sz = 0;  
262.             }  
263.             if (sz <= 0)  
264.                 break;  
265.             int written = 0;  
266.             while (written < sz)  
267.             {  
268.                 if (buffer.isWritable())  
269.                 {  
270.                     written += buffer.write(buf + written, sz - written);  
271.                     readable.set();  
272.                 }  
273.                 else  
274.                 {  
275.                     if (peer_stopped)  
276.                         break;  
277.                     writable.wait();  
278.                 }  
279.             }  
280.             if (peer_stopped)  
281.                 break;  
282.         }  
283.         try { socket.shutdownReceive(); } catch (...) {}  
284.         self_stopped = true;  
285.         readable.set();  
286.     }  
287. };  
288.   
289. // Receive from local and write to out buffers  
290. class divider : public Runnable  
291. {  
292.     int nr_conn;  
293.     StreamSocket socket;  
294.     vector<Event *> &readables;  
295.     Event &writable;  
296.     vector<FIFOBuffer *> buffers;  
297.     bool &self_stopped, *peers_stopped;  
298. public:  
299.     divider(int n, StreamSocket s, vector<Event *> &vr, Event &w, vector<FIFOBuffer *> &b, bool &self, bool *peers)  
300.         : nr_conn(n), socket(s), readables(vr), writable(w), buffers(b), self_stopped(self), peers_stopped(peers)  
301.     {  
302.     }  
303.   
304.     virtual void run()  
305.     {  
306.         printf("[.] Divider, tid = %u\n", (unsigned int)Thread::currentTid());  
307.         fflush(stdout);  
308.         int sz0 = 8192 * nr_conn;  
309.         char *buf0 = new char[sz0];  
310.         char **bufv = new char *[nr_conn];  
311.         for (int i = 0; i < nr_conn; ++i)  
312.             bufv[i] = new char[8192];  
313.         int *szv = new int[nr_conn];  
314.         int *writtenv = new int[nr_conn];  
315.         long long total = 0;  
316.         bool peer_stopped = false;  
317.         while (true)  
318.         {  
319.             int sz;  
320.             try  
321.             {  
322.                 sz = socket.receiveBytes(buf0, sz0);  
323.             }  
324.             catch (Exception &exc)  
325.             {  
326.                 printf("[!] Divider: %s\n", exc.displayText().c_str());  
327.                 fflush(stdout);  
328.                 sz = 0;  
329.             }  
330.             if (sz <= 0)  
331.                 break;  
332.             for (int i = 0; i < nr_conn; ++i)  
333.             {  
334.                 szv[i] = 0;  
335.                 writtenv[i] = 0;  
336.             }  
337.             for (int i = 0; i < sz; ++i)  
338.                 bufv[(total + i) % nr_conn][szv[(total + i) % nr_conn]++] = buf0[i];  
339.             total += sz;  
340.             while (true)  
341.             {  
342.                 bool done = true;  
343.                 for (int i = 0; i < nr_conn; ++i)  
344.                     if (writtenv[i] < szv[i])  
345.                         done = false;  
346.                 if (done)  
347.                     break;  
348.                 bool written = false;  
349.                 for (int i = 0; i < nr_conn; ++i)  
350.                 {  
351.                     if (buffers[i]->isWritable())  
352.                     {  
353.                         writtenv[i] += buffers[i]->write(bufv[i] + writtenv[i], szv[i] - writtenv[i]);  
354.                         readables[i]->set();  
355.                         written = true;  
356.                     }  
357.                 }  
358.                 if (!written)  
359.                 {  
360.                     for (int i = 0; i < nr_conn; ++i)  
361.                         if (peers_stopped[i])  
362.                             peer_stopped = true;  
363.                     if (peer_stopped)  
364.                         break;  
365.                     writable.wait();  
366.                     writable.reset();  
367.                 }  
368.                 if (peer_stopped)  
369.                     break;  
370.             }  
371.         }  
372.         try { socket.shutdownReceive(); } catch (...) {}  
373.         self_stopped = true;  
374.         for (int i = 0; i < nr_conn; ++i)  
375.             readables[i]->set();  
376.     }  
377. };  
378.   
379. // Read from in buffers and send to local  
380. class combiner : public Runnable  
381. {  
382.     int nr_conn;  
383.     StreamSocket socket;  
384.     Event &readable;  
385.     vector<Event *> &writables;  
386.     vector<FIFOBuffer *> buffers;  
387.     bool &self_stopped, *peers_stopped;  
388. public:  
389.     combiner(int n, StreamSocket s, Event &r, vector<Event *> &vw, vector<FIFOBuffer *> &b, bool &self, bool *peers)  
390.         : nr_conn(n), socket(s), readable(r), writables(vw), buffers(b), self_stopped(self), peers_stopped(peers)  
391.     {  
392.     }  
393.   
394.     virtual void run()  
395.     {  
396.         printf("[.] Combiner, tid = %u\n", (unsigned int)Thread::currentTid());  
397.         fflush(stdout);  
398.         char *buf0 = new char[8192 * nr_conn];  
399.         char **bufv = new char *[nr_conn];  
400.         for (int i = 0; i < nr_conn; ++i)  
401.             bufv[i] = new char[8192];  
402.         int *szv = new int[nr_conn];  
403.         int *readv = new int[nr_conn];  
404.         long long total = 0;  
405.         bool peer_stopped = false, closed = false;  
406.         while (true)  
407.         {  
408.             for (int i = 0; i < nr_conn; ++i)  
409.             {  
410.                 szv[i] = buffers[i]->peek(bufv[i], 8192);  
411.                 readv[i] = 0;  
412.             }  
413.             int sz = 0;  
414.             while (true)  
415.             {  
416.                 int i = (total + sz) % nr_conn;  
417.                 if (readv[i] >= szv[i])  
418.                     break;  
419.                 buf0[sz++] = bufv[i][readv[i]++];  
420.             }  
421.             total += sz;  
422.             if (sz > 0)  
423.             {  
424.                 for (int i = 0; i < nr_conn; ++i)  
425.                 {  
426.                     if (readv[i] > 0)  
427.                     {  
428.                         buffers[i]->read(bufv[i], readv[i]);  
429.                         writables[i]->set();  
430.                     }  
431.                 }  
432.                 int sent = 0;  
433.                 while (sent < sz)  
434.                 {  
435.                     int sz1;  
436.                     try  
437.                     {  
438.                         sz1 = socket.sendBytes(buf0 + sent, sz - sent);  
439.                     }  
440.                     catch (Exception &exc)  
441.                     {  
442.                         printf("[!] Combiner: %s\n", exc.displayText().c_str());  
443.                         fflush(stdout);  
444.                         sz1 = 0;  
445.                     }  
446.                     if (sz1 <= 0)  
447.                     {  
448.                         closed = true;  
449.                         break;  
450.                     }  
451.                     sent += sz1;  
452.                 }  
453.                 if (closed)  
454.                     break;  
455.             }  
456.             else  
457.             {  
458.                 for (int i = 0; i < nr_conn; ++i)  
459.                     if (peers_stopped[i])  
460.                         peer_stopped = true;  
461.                 if (peer_stopped)  
462.                     break;  
463.                 readable.wait();  
464.                 readable.reset();  
465.             }  
466.         }  
467.         try { socket.shutdownSend(); } catch (...) {}  
468.         self_stopped = true;  
469.         for (int i = 0; i < nr_conn; ++i)  
470.             writables[i]->set();  
471.     }  
472. };  
473.   
474. int main(int argc, char **argv)  
475. {  
476.     config.load("multisocks.txt");  
477.     for (int i = 1; i < argc; ++i)  
478.         if (config.load(argv[i]))  
479.             return 1;  
480.     if (!config.s("log").empty())  
481.         freopen(config.s("log").c_str(), "w", stdout);  
482.     ErrorHandler::set(new error_handler);  
483.     fflush(stdout);  
484.     try  
485.     {  
486.         int nr_conn = config.i("nr_conn");  
487.         bool divider_stopped = false, combiner_stopped = false;  
488.         bool *readers_stopped = new bool[nr_conn](), *writers_stopped = new bool[nr_conn]();  
489.         Event in_readable(false), out_writable(false);  
490.         in_readable.reset();  
491.         out_writable.set();  
492.         vector<Event *> in_writables, out_readables;  
493.         for (int i = 0; i < nr_conn; ++i)  
494.         {  
495.             in_writables.push_back(new Event);  
496.             in_writables.back()->set();  
497.             out_readables.push_back(new Event);  
498.             out_readables.back()->reset();  
499.         }  
500.         vector<StreamSocket> remotes;  
501.         if (config.s("remote") == "listen")  
502.         {  
503.             ServerSocket server;  
504.             if (config.s("remote.protocol") == "ipv6")  
505.                 server.bind6(config.i("remote.port"), true, true);  
506.             else  
507.                 server.bind(config.i("remote.port"), true);  
508.             server.listen();  
509.             for (int i = 0; i < nr_conn; ++i)  
510.             {  
511.                 SocketAddress client_addr;  
512.                 remotes.push_back(server.acceptConnection(client_addr));  
513.                 remotes.back().setNoDelay(true);  
514.                 printf("[.] Remote incoming: %s\n", client_addr.toString().c_str());  
515.                 fflush(stdout);  
516.                 if (client_addr.host() != remotes.front().peerAddress().host())  
517.                     throw Exception(string("client addresses mismatch"));  
518.             }  
519.             server.close();  
520.         }  
521.         else  
522.         {  
523.             SocketAddress server_addr(config.s("remote.host"), config.i("remote.port"));  
524.             for (int i = 0; i < nr_conn; ++i)  
525.             {  
526.                 StreamSocket remote;  
527.                 if (!config.s("remote.bind").empty())  
528.                     remote.impl()->bind(SocketAddress(config.s("remote.bind"), 0), true);  
529.                 remote.connect(server_addr);  
530.                 remote.setNoDelay(true);  
531.                 remotes.push_back(remote);  
532.                 printf("[.] Connected to remote from: %s\n", remote.address().toString().c_str());  
533.                 fflush(stdout);  
534.             }  
535.         }  
536.         vector<FIFOBuffer *> in_buffers, out_buffers;  
537.         vector<reader *> readers;  
538.         vector<writer *> writers;  
539.         Thread *reader_threads = new Thread[nr_conn], *writer_threads = new Thread[nr_conn];  
540.         for (int i = 0; i < nr_conn; ++i)  
541.         {  
542.             in_buffers.push_back(new FIFOBuffer(config.i("buffer_size")));  
543.             out_buffers.push_back(new FIFOBuffer(config.i("buffer_size")));  
544.             readers.push_back(new reader(i, remotes[i], *out_readables[i], out_writable, *out_buffers.back(), readers_stopped[i], divider_stopped));  
545.             writers.push_back(new writer(i, remotes[i], in_readable, *in_writables[i], *in_buffers.back(), writers_stopped[i], combiner_stopped));  
546.             reader_threads[i].setStackSize(65536);  
547.             writer_threads[i].setStackSize(65536);  
548.             reader_threads[i].start(*readers.back());  
549.             writer_threads[i].start(*writers.back());  
550.         }  
551.         StreamSocket local;  
552.         if (config.s("local") == "listen")  
553.         {  
554.             ServerSocket server;  
555.             if (config.s("local.protocol") == "ipv6")  
556.                 server.bind6(config.i("local.port"), true, true);  
557.             else  
558.                 server.bind(config.i("local.port"), true);  
559.             server.listen();  
560.             local = server.acceptConnection();  
561.             printf("[.] Local incoming: %s\n", local.peerAddress().toString().c_str());  
562.             fflush(stdout);  
563.             local.setNoDelay(true);  
564.             server.close();  
565.         }  
566.         else  
567.         {  
568.             SocketAddress server_addr(config.s("local.host"), config.i("local.port"));  
569.             local.connect(server_addr);  
570.             local.setNoDelay(true);  
571.             printf("[.] Connected to local from: %s\n", local.address().toString().c_str());  
572.             fflush(stdout);  
573.         }  
574.         divider *pdivider = new divider(nr_conn, local, out_readables, out_writable, out_buffers, divider_stopped, readers_stopped);  
575.         combiner *pcombiner = new combiner(nr_conn, local, in_readable, in_writables, in_buffers, combiner_stopped, writers_stopped);  
576.         Thread divider_thread, combiner_thread;  
577.         divider_thread.setStackSize(65536);  
578.         combiner_thread.setStackSize(65536);  
579.         divider_thread.start(*pdivider);  
580.         combiner_thread.start(*pcombiner);  
581.         printf("[.] Connection established\n");  
582.         fflush(stdout);  
583.         divider_thread.join();  
584.         combiner_thread.join();  
585.         local.close();  
586.         for (int i = 0; i < nr_conn; ++i)  
587.         {  
588.             reader_threads[i].join();  
589.             writer_threads[i].join();  
590.             remotes[i].close();  
591.         }  
592.         printf("[.] Connection closed\n");  
593.     }  
594.     catch (Exception &exc)  
595.     {  
596.         ErrorHandler::handle(exc);  
597.     }  
598.     catch (exception &exc)  
599.     {  
600.         ErrorHandler::handle(exc);  
601.     }  
602.     catch (...)  
603.     {  
604.         ErrorHandler::handle();  
605.     }  
606.     return 0;  
607. }  

测试

客户端网络环境为 CERNET ，服务器是位于达拉斯的 BurstNET VPS ，连接方式是 TCP/IPv6 。

直接连接， OpenVPN 速率为 341 KB/s 。使用 5 个连接进行分流， OpenVPN 速率可达 1.4 MB/s 。