1#include "utils.h"
2#include <pico_ipv4.h>
3#include <pico_ipv6.h>
4#include <pico_socket.h>
5/*** START TCP BENCH ***/
6#define TCP_BENCH_TX  1
7#define TCP_BENCH_RX  2
8#define TCP_BENCH_TX_FOREVER 3
9static char *buffer1;
10static char *buffer0;
11
12int tcpbench_mode = 0;
13struct pico_socket *tcpbench_sock = NULL;
14static pico_time tcpbench_time_start, tcpbench_time_end;
15
16void cb_tcpbench(uint16_t ev, struct pico_socket *s)
17{
18    static int closed = 0;
19    static unsigned long count = 0;
20    uint8_t recvbuf[1500];
21    uint16_t port;
22    char peer[200];
23    /* struct pico_socket *sock_a; */
24
25    static int tcpbench_wr_size = 0;
26    static int tcpbench_rd_size = 0;
27    int tcpbench_w = 0;
28    int tcpbench_r = 0;
29    double tcpbench_time = 0;
30
31    count++;
32
33    if (ev & PICO_SOCK_EV_RD) {
34        do {
35            /* read data, but discard */
36            tcpbench_r = pico_socket_read(s, recvbuf, 1500);
37            if (tcpbench_r > 0) {
38                tcpbench_rd_size += tcpbench_r;
39            }
40        } while (tcpbench_r > 0);
41        if (tcpbench_time_start == 0)
42            tcpbench_time_start = PICO_TIME_MS();
43
44        printf("tcpbench_rd_size = %d      \r", tcpbench_rd_size);
45    }
46
47    if (ev & PICO_SOCK_EV_CONN) {
48        if (!IPV6_MODE) {
49            struct pico_ip4 orig;
50            if (tcpbench_mode == TCP_BENCH_TX || tcpbench_mode == TCP_BENCH_TX_FOREVER) {
51                printf("tcpbench> Connection established with server.\n");
52            } else if (tcpbench_mode == TCP_BENCH_RX) {
53                /* sock_a = pico_socket_accept(s, &orig, &port); */
54                pico_socket_accept(s, &orig, &port);
55                pico_ipv4_to_string(peer, orig.addr);
56                printf("tcpbench> Connection established with %s:%d.\n", peer, short_be(port));
57            }
58        } else {
59            struct pico_ip6 orig;
60            if (tcpbench_mode == TCP_BENCH_TX || tcpbench_mode == TCP_BENCH_TX_FOREVER) {
61                printf("tcpbench> Connection established with server.\n");
62            } else if (tcpbench_mode == TCP_BENCH_RX) {
63                /* sock_a = pico_socket_accept(s, &orig, &port); */
64                pico_socket_accept(s, &orig, &port);
65#ifdef PICO_SUPPORT_IPV6
66                pico_ipv6_to_string(peer, orig.addr);
67                printf("tcpbench> Connection established with [%s]:%d.\n", peer, short_be(port));
68#endif
69            }
70        }
71    }
72
73    if (ev & PICO_SOCK_EV_FIN) {
74        printf("tcpbench> Socket closed. Exit normally. \n");
75        if (tcpbench_mode == TCP_BENCH_RX) {
76            tcpbench_time_end = PICO_TIME_MS();
77            tcpbench_time = (tcpbench_time_end - tcpbench_time_start) / 1000.0; /* get number of seconds */
78            printf("tcpbench> received %d bytes in %lf seconds\n", tcpbench_rd_size, tcpbench_time);
79            printf("tcpbench> average read throughput %lf kbit/sec\n", ((tcpbench_rd_size * 8.0) / tcpbench_time) / 1000);
80            pico_socket_shutdown(s, PICO_SHUT_WR);
81            printf("tcpbench> Called shutdown write, ev = %d\n", ev);
82        }
83
84        if (!pico_timer_add(5000, deferred_exit, NULL)) {
85            printf("tcpbench> Failed to start exit timer, exiting now\n");
86            exit(1);
87        }
88    }
89
90    if (ev & PICO_SOCK_EV_ERR) {
91        printf("tcpbench> ---- Socket Error received: %s. Bailing out.\n", strerror(pico_err));
92        if (!pico_err == PICO_ERR_ECONNRESET) {
93            if (pico_timer_add(5000, deferred_exit, NULL)) {
94                printf("tcpbench> Failed to start exit timer, exiting now\n");
95                exit(1);
96            }
97        }
98        else {
99            printf("tcpbench> ---- Socket Error: '%s'. Was unexpected! Something went wrong.\n", strerror(pico_err));
100            exit(2);
101        }
102    }
103
104    if (ev & PICO_SOCK_EV_CLOSE) {
105        printf("tcpbench> event close\n");
106        if (tcpbench_mode == TCP_BENCH_RX) {
107            pico_socket_close(s);
108            printf("tcpbench> Called shutdown write, ev = %d\n", ev);
109        } else if (tcpbench_mode == TCP_BENCH_TX || tcpbench_mode == TCP_BENCH_TX_FOREVER) {
110            pico_socket_close(s);
111            return;
112        }
113    }
114
115    if (ev & PICO_SOCK_EV_WR) {
116        if (((tcpbench_wr_size < TCPSIZ) && (tcpbench_mode == TCP_BENCH_TX)) || tcpbench_mode == TCP_BENCH_TX_FOREVER) {
117            do {
118                tcpbench_w = pico_socket_write(tcpbench_sock, buffer0 + (tcpbench_wr_size % TCPSIZ), TCPSIZ - (tcpbench_wr_size % TCPSIZ));
119                if (tcpbench_w > 0) {
120                    tcpbench_wr_size += tcpbench_w;
121                    /* printf("tcpbench> SOCKET WRITTEN - %d\n",tcpbench_w); */
122                } else {
123                    /* printf("pico_socket_write returned %d\n", tcpbench_w); */
124                }
125
126                if (tcpbench_time_start == 0)
127                    tcpbench_time_start = PICO_TIME_MS();
128            } while(tcpbench_w > 0);
129            printf("tcpbench_wr_size = %d      \r", tcpbench_wr_size);
130        } else {
131            if (!closed && tcpbench_mode == TCP_BENCH_TX) {
132                tcpbench_time_end = PICO_TIME_MS();
133                pico_socket_shutdown(s, PICO_SHUT_WR);
134                printf("tcpbench> TCPSIZ written\n");
135                printf("tcpbench> Called shutdown()\n");
136                tcpbench_time = (tcpbench_time_end - tcpbench_time_start) / 1000.0; /* get number of seconds */
137                printf("tcpbench> Transmitted %u bytes in %lf seconds\n", TCPSIZ, tcpbench_time);
138                printf("tcpbench> average write throughput %lf kbit/sec\n", ((TCPSIZ * 8.0) / tcpbench_time) / 1000);
139                closed = 1;
140            }
141        }
142    }
143}
144
145void app_tcpbench(char *arg)
146{
147    struct pico_socket *s;
148    char *dport = NULL;
149    char *dest = NULL;
150    char *mode = NULL;
151    char *nagle = NULL;
152    int port = 0, i;
153    uint16_t port_be = 0;
154    char *nxt;
155    char *sport = NULL;
156    int nagle_off = 1;
157    union {
158        struct pico_ip4 ip4;
159        struct pico_ip6 ip6;
160    } inaddr_any = {
161        .ip4 = {0}, .ip6 = {{0}}
162    };
163
164    nxt = cpy_arg(&mode, arg);
165
166    if ((*mode == 't') || (*mode == 'f')) { /* TEST BENCH SEND MODE */
167        if (*mode == 't')
168            tcpbench_mode = TCP_BENCH_TX;
169        else
170            tcpbench_mode = TCP_BENCH_TX_FOREVER;
171
172        printf("tcpbench> TX\n");
173
174        nxt = cpy_arg(&dest, nxt);
175        if (!dest) {
176            fprintf(stderr, "tcpbench send needs the following format: tcpbench:tx:dst_addr[:dport][:n] -- 'n' is for nagle\n");
177            exit(255);
178        }
179
180        printf ("+++ Dest is %s\n", dest);
181        if (nxt) {
182            printf("Next arg: %s\n", nxt);
183            nxt = cpy_arg(&dport, nxt);
184            printf("Dport: %s\n", dport);
185        }
186
187        if (nxt) {
188            printf("Next arg: %s\n", nxt);
189            nxt = cpy_arg(&nagle, nxt);
190            printf("nagle: %s\n", nagle);
191            if (strlen(nagle) == 1 && nagle[0] == 'n') {
192                nagle_off = 0;
193                printf("Nagle algorithm enabled\n");
194            }
195        }
196
197        if (dport) {
198            port = atoi(dport);
199            port_be = short_be((uint16_t)port);
200        }
201
202        if (port == 0) {
203            port_be = short_be(5555);
204        }
205
206        buffer0 = malloc(TCPSIZ);
207        buffer1 = malloc(TCPSIZ);
208        printf("Buffer1 (%p)\n", buffer1);
209        for (i = 0; i < TCPSIZ; i++) {
210            char c = (i % 26) + 'a';
211            buffer0[i] = c;
212        }
213        memset(buffer1, 'a', TCPSIZ);
214        printf("tcpbench> Connecting to: %s:%d\n", dest, short_be(port_be));
215
216        if (!IPV6_MODE) {
217            struct pico_ip4 server_addr;
218            s = pico_socket_open(PICO_PROTO_IPV4, PICO_PROTO_TCP, &cb_tcpbench);
219            if (!s)
220                exit(1);
221
222            pico_socket_setoption(s, PICO_TCP_NODELAY, &nagle_off);
223
224            /* NOTE: used to set a fixed local port and address
225               local_port = short_be(6666);
226               pico_string_to_ipv4("10.40.0.11", &local_addr.addr);
227               pico_socket_bind(s, &local_addr, &local_port);*/
228
229            pico_string_to_ipv4(dest, &server_addr.addr);
230            pico_socket_connect(s, &server_addr, port_be);
231        } else {
232            struct pico_ip6 server_addr;
233            s = pico_socket_open(PICO_PROTO_IPV6, PICO_PROTO_TCP, &cb_tcpbench);
234            if (!s)
235                exit(1);
236
237            pico_socket_setoption(s, PICO_TCP_NODELAY, &nagle_off);
238
239            /* NOTE: used to set a fixed local port and address
240               local_port = short_be(6666);
241               pico_string_to_ipv4("10.40.0.11", &local_addr.addr);
242               pico_socket_bind(s, &local_addr, &local_port);*/
243#ifdef PICO_SUPPORT_IPV6
244            pico_string_to_ipv6(dest, server_addr.addr);
245            pico_socket_connect(s, &server_addr, port_be);
246#endif
247
248        }
249
250    } else if (*mode == 'r') { /* TEST BENCH RECEIVE MODE */
251        int ret;
252        tcpbench_mode = TCP_BENCH_RX;
253        printf("tcpbench> RX\n");
254
255        cpy_arg(&sport, nxt);
256        if (!sport) {
257            fprintf(stderr, "tcpbench receive needs the following format: tcpbench:rx[:dport]\n");
258            exit(255);
259        }
260
261        if (sport) {
262            printf("s-port is %s\n", sport);
263            port = atoi(sport);
264            port_be = short_be((uint16_t)port);
265            printf("tcpbench> Got port %d\n", port);
266            free(sport);
267        }
268
269        if (port == 0) {
270            port_be = short_be(5555);
271        }
272
273        printf("tcpbench> OPEN\n");
274        if (!IPV6_MODE)
275            s = pico_socket_open(PICO_PROTO_IPV4, PICO_PROTO_TCP, &cb_tcpbench);
276        else
277            s = pico_socket_open(PICO_PROTO_IPV6, PICO_PROTO_TCP, &cb_tcpbench);
278
279        if (!s)
280            exit(1);
281
282        printf("tcpbench> BIND\n");
283        if (!IPV6_MODE)
284            ret = pico_socket_bind(s, &inaddr_any.ip4, &port_be);
285        else
286            ret = pico_socket_bind(s, &inaddr_any.ip6, &port_be);
287
288        if (ret < 0) {
289            printf("tcpbench> BIND failed because %s\n", strerror(pico_err));
290            exit(1);
291        }
292
293        printf("tcpbench> LISTEN\n");
294        if (pico_socket_listen(s, 40) != 0)
295            exit(1);
296
297        printf("tcpbench> listening port %u ...\n", short_be(port_be));
298    } else {
299        printf("tcpbench> wrong mode argument\n");
300        exit(1);
301    }
302
303    tcpbench_sock = s;
304
305    /* free strdups */
306    if (dport)
307        free(dport);
308
309    if (dest)
310        free (dest);
311
312    if (mode)
313        free (mode);
314
315    if (nagle)
316        free (nagle);
317
318    return;
319}
320/*** END TCP BENCH ***/
321