1/* Copyright (c) 2017 - 2022 LiteSpeed Technologies Inc.  See LICENSE. */
2/*
3 * echo_server.c -- QUIC server that echoes back input line by line
4 */
5
6#include <assert.h>
7#include <signal.h>
8#include <stdio.h>
9#include <stdlib.h>
10#include <string.h>
11#include <sys/queue.h>
12#include <time.h>
13#ifndef WIN32
14#include <unistd.h>
15#include <netinet/in.h>
16#else
17#include "vc_compat.h"
18#include "getopt.h"
19#endif
20
21#include "lsquic.h"
22#include "test_common.h"
23#include "../src/liblsquic/lsquic_hash.h"
24#include "test_cert.h"
25#include "prog.h"
26
27#include "../src/liblsquic/lsquic_logger.h"
28
29
30struct lsquic_conn_ctx;
31
32struct echo_server_ctx {
33    TAILQ_HEAD(, lsquic_conn_ctx)   conn_ctxs;
34    unsigned max_reqs;
35    int n_conn;
36    struct sport_head sports;
37    struct prog *prog;
38};
39
40struct lsquic_conn_ctx {
41    TAILQ_ENTRY(lsquic_conn_ctx)    next_connh;
42    lsquic_conn_t       *conn;
43    struct echo_server_ctx   *server_ctx;
44};
45
46
47static lsquic_conn_ctx_t *
48echo_server_on_new_conn (void *stream_if_ctx, lsquic_conn_t *conn)
49{
50    struct echo_server_ctx *server_ctx = stream_if_ctx;
51    lsquic_conn_ctx_t *conn_h = calloc(1, sizeof(*conn_h));
52    conn_h->conn = conn;
53    conn_h->server_ctx = server_ctx;
54    TAILQ_INSERT_TAIL(&server_ctx->conn_ctxs, conn_h, next_connh);
55    LSQ_NOTICE("New connection!");
56    print_conn_info(conn);
57    return conn_h;
58}
59
60
61static void
62echo_server_on_conn_closed (lsquic_conn_t *conn)
63{
64    lsquic_conn_ctx_t *conn_h = lsquic_conn_get_ctx(conn);
65    if (conn_h->server_ctx->n_conn)
66    {
67        --conn_h->server_ctx->n_conn;
68        LSQ_NOTICE("Connection closed, remaining: %d", conn_h->server_ctx->n_conn);
69        if (0 == conn_h->server_ctx->n_conn)
70            prog_stop(conn_h->server_ctx->prog);
71    }
72    else
73        LSQ_NOTICE("Connection closed");
74    TAILQ_REMOVE(&conn_h->server_ctx->conn_ctxs, conn_h, next_connh);
75    free(conn_h);
76}
77
78
79struct lsquic_stream_ctx {
80    lsquic_stream_t     *stream;
81    struct echo_server_ctx   *server_ctx;
82    char                 buf[0x100];
83    size_t               buf_off;
84};
85
86
87static lsquic_stream_ctx_t *
88echo_server_on_new_stream (void *stream_if_ctx, lsquic_stream_t *stream)
89{
90    lsquic_stream_ctx_t *st_h = malloc(sizeof(*st_h));
91    st_h->stream = stream;
92    st_h->server_ctx = stream_if_ctx;
93    st_h->buf_off = 0;
94    lsquic_stream_wantread(stream, 1);
95    return st_h;
96}
97
98
99static struct lsquic_conn_ctx *
100find_conn_h (const struct echo_server_ctx *server_ctx, lsquic_stream_t *stream)
101{
102    struct lsquic_conn_ctx *conn_h;
103    lsquic_conn_t *conn;
104
105    conn = lsquic_stream_conn(stream);
106    TAILQ_FOREACH(conn_h, &server_ctx->conn_ctxs, next_connh)
107        if (conn_h->conn == conn)
108            return conn_h;
109    return NULL;
110}
111
112
113static void
114echo_server_on_read (lsquic_stream_t *stream, lsquic_stream_ctx_t *st_h)
115{
116    struct lsquic_conn_ctx *conn_h;
117    size_t nr;
118
119    nr = lsquic_stream_read(stream, st_h->buf + st_h->buf_off++, 1);
120    if (0 == nr)
121    {
122        LSQ_NOTICE("EOF: closing connection");
123        lsquic_stream_shutdown(stream, 2);
124        conn_h = find_conn_h(st_h->server_ctx, stream);
125        lsquic_conn_close(conn_h->conn);
126    }
127    else if ('\n' == st_h->buf[ st_h->buf_off - 1 ])
128    {
129        /* Found end of line: echo it back */
130        lsquic_stream_wantwrite(stream, 1);
131        lsquic_stream_wantread(stream, 0);
132    }
133    else if (st_h->buf_off == sizeof(st_h->buf))
134    {
135        /* Out of buffer space: line too long */
136        LSQ_NOTICE("run out of buffer space");
137        lsquic_stream_shutdown(stream, 2);
138    }
139    else
140    {
141        /* Keep reading */;
142    }
143}
144
145
146static void
147echo_server_on_write (lsquic_stream_t *stream, lsquic_stream_ctx_t *st_h)
148{
149    lsquic_stream_write(stream, st_h->buf, st_h->buf_off);
150    st_h->buf_off = 0;
151    lsquic_stream_flush(stream);
152    lsquic_stream_wantwrite(stream, 0);
153    lsquic_stream_wantread(stream, 1);
154}
155
156
157static void
158echo_server_on_stream_close (lsquic_stream_t *stream, lsquic_stream_ctx_t *st_h)
159{
160    struct lsquic_conn_ctx *conn_h;
161    LSQ_NOTICE("%s called", __func__);
162    conn_h = find_conn_h(st_h->server_ctx, stream);
163    LSQ_WARN("%s: TODO: free connection handler %p", __func__, conn_h);
164    free(st_h);
165}
166
167
168const struct lsquic_stream_if server_echo_stream_if = {
169    .on_new_conn            = echo_server_on_new_conn,
170    .on_conn_closed         = echo_server_on_conn_closed,
171    .on_new_stream          = echo_server_on_new_stream,
172    .on_read                = echo_server_on_read,
173    .on_write               = echo_server_on_write,
174    .on_close               = echo_server_on_stream_close,
175};
176
177
178static void
179usage (const char *prog)
180{
181    const char *const slash = strrchr(prog, '/');
182    if (slash)
183        prog = slash + 1;
184    printf(
185"Usage: %s [opts]\n"
186"\n"
187"Options:\n"
188                , prog);
189}
190
191
192int
193main (int argc, char **argv)
194{
195    int opt, s;
196    struct prog prog;
197    struct echo_server_ctx server_ctx;
198
199    memset(&server_ctx, 0, sizeof(server_ctx));
200    server_ctx.prog = &prog;
201    TAILQ_INIT(&server_ctx.sports);
202    TAILQ_INIT(&server_ctx.conn_ctxs);
203
204    prog_init(&prog, LSENG_SERVER, &server_ctx.sports,
205                                        &server_echo_stream_if, &server_ctx);
206
207    while (-1 != (opt = getopt(argc, argv, PROG_OPTS "hn:")))
208    {
209        switch (opt) {
210        case 'n':
211            server_ctx.n_conn = atoi(optarg);
212            break;
213        case 'h':
214            usage(argv[0]);
215            prog_print_common_options(&prog, stdout);
216            exit(0);
217        default:
218            if (0 != prog_set_opt(&prog, opt, optarg))
219                exit(1);
220        }
221    }
222
223    add_alpn("echo");
224    if (0 != prog_prep(&prog))
225    {
226        LSQ_ERROR("could not prep");
227        exit(EXIT_FAILURE);
228    }
229
230    LSQ_DEBUG("entering event loop");
231
232    s = prog_run(&prog);
233    prog_cleanup(&prog);
234
235    exit(0 == s ? EXIT_SUCCESS : EXIT_FAILURE);
236}
237