xref: /unit/src/test/nxt_unit_websocket_echo.c (revision 2229:3a230013e58a)
1 
2 /*
3  * Copyright (C) NGINX, Inc.
4  */
5 
6 #include <string.h>
7 #include <stdlib.h>
8 
9 #include <nxt_unit.h>
10 #include <nxt_unit_request.h>
11 #include <nxt_clang.h>
12 #include <nxt_websocket.h>
13 #include <nxt_unit_websocket.h>
14 
15 
16 static void
ws_echo_request_handler(nxt_unit_request_info_t * req)17 ws_echo_request_handler(nxt_unit_request_info_t *req)
18 {
19     int         rc;
20     const char  *target;
21 
22     rc = NXT_UNIT_OK;
23     target = nxt_unit_sptr_get(&req->request->target);
24 
25     if (strcmp(target, "/") == 0) {
26         if (!nxt_unit_request_is_websocket_handshake(req)) {
27             goto notfound;
28         }
29 
30         rc = nxt_unit_response_init(req, 101, 0, 0);
31         if (nxt_slow_path(rc != NXT_UNIT_OK)) {
32             goto fail;
33         }
34 
35         nxt_unit_response_upgrade(req);
36         nxt_unit_response_send(req);
37 
38         return;
39     }
40 
41 notfound:
42 
43     rc = nxt_unit_response_init(req, 404, 0, 0);
44 
45 fail:
46 
47     nxt_unit_request_done(req, rc);
48 }
49 
50 
51 static void
ws_echo_websocket_handler(nxt_unit_websocket_frame_t * ws)52 ws_echo_websocket_handler(nxt_unit_websocket_frame_t *ws)
53 {
54     uint8_t                  opcode;
55     ssize_t                  size;
56     nxt_unit_request_info_t  *req;
57 
58     static size_t            buf_size = 0;
59     static uint8_t           *buf = NULL;
60 
61     if (buf_size < ws->content_length) {
62         buf = realloc(buf, ws->content_length);
63         buf_size = ws->content_length;
64     }
65 
66     req = ws->req;
67     opcode = ws->header->opcode;
68 
69     if (opcode == NXT_WEBSOCKET_OP_PONG) {
70         nxt_unit_websocket_done(ws);
71         return;
72     }
73 
74     size = nxt_unit_websocket_read(ws, buf, ws->content_length);
75 
76     nxt_unit_websocket_send(req, opcode, ws->header->fin, buf, size);
77     nxt_unit_websocket_done(ws);
78 
79     if (opcode == NXT_WEBSOCKET_OP_CLOSE) {
80         nxt_unit_request_done(req, NXT_UNIT_OK);
81     }
82 }
83 
84 
85 int
main(void)86 main(void)
87 {
88     nxt_unit_ctx_t   *ctx;
89     nxt_unit_init_t  init;
90 
91     memset(&init, 0, sizeof(nxt_unit_init_t));
92 
93     init.callbacks.request_handler = ws_echo_request_handler;
94     init.callbacks.websocket_handler = ws_echo_websocket_handler;
95 
96     ctx = nxt_unit_init(&init);
97     if (ctx == NULL) {
98         return 1;
99     }
100 
101     nxt_unit_run(ctx);
102     nxt_unit_done(ctx);
103 
104     return 0;
105 }
106