xref: /unit/src/test/nxt_unit_websocket_chat.c (revision 1710:e598cd15bd91)
1 
2 /*
3  * Copyright (C) NGINX, Inc.
4  */
5 
6 #include <fcntl.h>
7 #include <stdio.h>
8 #include <string.h>
9 #include <errno.h>
10 
11 #include <sys/mman.h>
12 #include <sys/stat.h>
13 
14 #include <nxt_unit.h>
15 #include <nxt_unit_request.h>
16 #include <nxt_clang.h>
17 #include <nxt_websocket.h>
18 #include <nxt_unit_websocket.h>
19 #include <nxt_main.h>
20 
21 
22 #define CONTENT_TYPE    "Content-Type"
23 #define CONTENT_LENGTH  "Content-Length"
24 #define TEXT_HTML       "text/html"
25 
26 typedef struct {
27     nxt_queue_link_t  link;
28     int               id;
29 } ws_chat_request_data_t;
30 
31 
32 static int ws_chat_root(nxt_unit_request_info_t *req);
33 static void ws_chat_broadcast(const char *buf, size_t size);
34 
35 
36 static const char     ws_chat_index_html[];
37 static const int      ws_chat_index_html_size;
38 
39 static char           ws_chat_index_content_length[34];
40 static int            ws_chat_index_content_length_size;
41 
42 static nxt_queue_t    ws_chat_sessions;
43 static int            ws_chat_next_id = 0;
44 
45 
46 static void
ws_chat_request_handler(nxt_unit_request_info_t * req)47 ws_chat_request_handler(nxt_unit_request_info_t *req)
48 {
49     static char             buf[1024];
50     int                     buf_size;
51     int                     rc = NXT_UNIT_OK;
52     nxt_unit_request_t      *r;
53     ws_chat_request_data_t  *data;
54 
55     r = req->request;
56 
57     const char* target = nxt_unit_sptr_get(&r->target);
58 
59     if (strcmp(target, "/") == 0) {
60         rc = ws_chat_root(req);
61         goto fail;
62     }
63 
64     if (strcmp(target, "/chat") == 0) {
65         if (!nxt_unit_request_is_websocket_handshake(req)) {
66             goto notfound;
67         }
68 
69         rc = nxt_unit_response_init(req, 101, 0, 0);
70         if (nxt_slow_path(rc != NXT_UNIT_OK)) {
71             goto fail;
72         }
73 
74         data = req->data;
75         nxt_queue_insert_tail(&ws_chat_sessions, &data->link);
76 
77         data->id = ws_chat_next_id++;
78 
79         nxt_unit_response_upgrade(req);
80         nxt_unit_response_send(req);
81 
82 
83         buf_size = snprintf(buf, sizeof(buf), "Guest #%d has joined.", data->id);
84 
85         ws_chat_broadcast(buf, buf_size);
86 
87         return;
88     }
89 
90 notfound:
91 
92     rc = nxt_unit_response_init(req, 404, 0, 0);
93 
94 fail:
95 
96     nxt_unit_request_done(req, rc);
97 }
98 
99 
100 static int
ws_chat_root(nxt_unit_request_info_t * req)101 ws_chat_root(nxt_unit_request_info_t *req)
102 {
103     int rc;
104 
105     rc = nxt_unit_response_init(req, 200 /* Status code. */,
106                                 2 /* Number of response headers. */,
107                                 nxt_length(CONTENT_TYPE)
108                                 + nxt_length(TEXT_HTML)
109                                 + nxt_length(CONTENT_LENGTH)
110                                 + ws_chat_index_content_length_size
111                                 + ws_chat_index_html_size);
112     if (nxt_slow_path(rc != NXT_UNIT_OK)) {
113         return rc;
114     }
115 
116     rc = nxt_unit_response_add_field(req,
117                                      CONTENT_TYPE, nxt_length(CONTENT_TYPE),
118                                      TEXT_HTML, nxt_length(TEXT_HTML));
119     if (nxt_slow_path(rc != NXT_UNIT_OK)) {
120         return rc;
121     }
122 
123     rc = nxt_unit_response_add_field(req,
124                                      CONTENT_LENGTH, nxt_length(CONTENT_LENGTH),
125                                      ws_chat_index_content_length,
126                                      ws_chat_index_content_length_size);
127     if (nxt_slow_path(rc != NXT_UNIT_OK)) {
128         return rc;
129     }
130 
131     rc = nxt_unit_response_add_content(req, ws_chat_index_html,
132                                        ws_chat_index_html_size);
133     if (nxt_slow_path(rc != NXT_UNIT_OK)) {
134         return rc;
135     }
136 
137     return nxt_unit_response_send(req);
138 }
139 
140 
141 static void
ws_chat_broadcast(const char * buf,size_t size)142 ws_chat_broadcast(const char *buf, size_t size)
143 {
144     ws_chat_request_data_t   *data;
145     nxt_unit_request_info_t  *req;
146 
147     nxt_unit_debug(NULL, "broadcast: %*.s", (int) size, buf);
148 
149     nxt_queue_each(data, &ws_chat_sessions, ws_chat_request_data_t, link) {
150 
151         req = nxt_unit_get_request_info_from_data(data);
152 
153         nxt_unit_req_debug(req, "send: %*.s", (int) size, buf);
154 
155         nxt_unit_websocket_send(req, NXT_WEBSOCKET_OP_TEXT, 1, buf, size);
156     } nxt_queue_loop;
157 }
158 
159 
160 static void
ws_chat_websocket_handler(nxt_unit_websocket_frame_t * ws)161 ws_chat_websocket_handler(nxt_unit_websocket_frame_t *ws)
162 {
163     int                     buf_size;
164     static char             buf[1024];
165     ws_chat_request_data_t  *data;
166 
167     if (ws->header->opcode != NXT_WEBSOCKET_OP_TEXT) {
168         return;
169     }
170 
171     data = ws->req->data;
172 
173     buf_size = snprintf(buf, sizeof(buf), "Guest #%d: ", data->id);
174 
175     buf_size += nxt_unit_websocket_read(ws, buf + buf_size,
176                                         nxt_min(sizeof(buf),
177                                                 ws->content_length));
178 
179     ws_chat_broadcast(buf, buf_size);
180 
181     nxt_unit_websocket_done(ws);
182 }
183 
184 
185 static void
ws_chat_close_handler(nxt_unit_request_info_t * req)186 ws_chat_close_handler(nxt_unit_request_info_t *req)
187 {
188     int                     buf_size;
189     static char             buf[1024];
190     ws_chat_request_data_t  *data;
191 
192     data = req->data;
193     buf_size = snprintf(buf, sizeof(buf), "Guest #%d has disconnected.",
194                         data->id);
195 
196     nxt_queue_remove(&data->link);
197     nxt_unit_request_done(req, NXT_UNIT_OK);
198 
199     ws_chat_broadcast(buf, buf_size);
200 }
201 
202 
203 int
main()204 main()
205 {
206     nxt_unit_ctx_t   *ctx;
207     nxt_unit_init_t  init;
208 
209     ws_chat_index_content_length_size =
210         snprintf(ws_chat_index_content_length,
211                  sizeof(ws_chat_index_content_length), "%d",
212                  ws_chat_index_html_size);
213 
214     nxt_queue_init(&ws_chat_sessions);
215 
216     memset(&init, 0, sizeof(nxt_unit_init_t));
217 
218     init.callbacks.request_handler = ws_chat_request_handler;
219     init.callbacks.websocket_handler = ws_chat_websocket_handler;
220     init.callbacks.close_handler = ws_chat_close_handler;
221 
222     init.request_data_size = sizeof(ws_chat_request_data_t);
223 
224     ctx = nxt_unit_init(&init);
225     if (ctx == NULL) {
226         return 1;
227     }
228 
229     nxt_unit_run(ctx);
230 
231     nxt_unit_done(ctx);
232 
233     return 0;
234 }
235 
236 
237 static const char ws_chat_index_html[] =
238 "<html>\n"
239 "<head>\n"
240 "    <title>WebSocket Chat Examples</title>\n"
241 "    <style type=\"text/css\">\n"
242 "        input#chat {\n"
243 "            width: 410px\n"
244 "        }\n"
245 "\n"
246 "        #container {\n"
247 "            width: 400px;\n"
248 "        }\n"
249 "\n"
250 "        #console {\n"
251 "            border: 1px solid #CCCCCC;\n"
252 "            border-right-color: #999999;\n"
253 "            border-bottom-color: #999999;\n"
254 "            height: 170px;\n"
255 "            overflow-y: scroll;\n"
256 "            padding: 5px;\n"
257 "            width: 100%;\n"
258 "        }\n"
259 "\n"
260 "        #console p {\n"
261 "            padding: 0;\n"
262 "            margin: 0;\n"
263 "        }\n"
264 "    </style>\n"
265 "    <script>\n"
266 "        \"use strict\";\n"
267 "\n"
268 "        var Chat = {};\n"
269 "\n"
270 "        Chat.socket = null;\n"
271 "\n"
272 "        Chat.connect = (function(host) {\n"
273 "            if ('WebSocket' in window) {\n"
274 "                Chat.socket = new WebSocket(host);\n"
275 "            } else if ('MozWebSocket' in window) {\n"
276 "                Chat.socket = new MozWebSocket(host);\n"
277 "            } else {\n"
278 "                Console.log('Error: WebSocket is not supported by this browser.');\n"
279 "                return;\n"
280 "            }\n"
281 "\n"
282 "            Chat.socket.onopen = function () {\n"
283 "                Console.log('Info: WebSocket connection opened.');\n"
284 "                document.getElementById('chat').onkeydown = function(event) {\n"
285 "                    if (event.keyCode == 13) {\n"
286 "                        Chat.sendMessage();\n"
287 "                    }\n"
288 "                };\n"
289 "            };\n"
290 "\n"
291 "            Chat.socket.onclose = function () {\n"
292 "                document.getElementById('chat').onkeydown = null;\n"
293 "                Console.log('Info: WebSocket closed.');\n"
294 "            };\n"
295 "\n"
296 "            Chat.socket.onmessage = function (message) {\n"
297 "                Console.log(message.data);\n"
298 "            };\n"
299 "        });\n"
300 "\n"
301 "        Chat.initialize = function() {\n"
302 "            var proto = 'ws://';\n"
303 "            if (window.location.protocol == 'https:') {\n"
304 "                proto = 'wss://'\n"
305 "            }\n"
306 "            Chat.connect(proto + window.location.host + '/chat');\n"
307 "        };\n"
308 "\n"
309 "        Chat.sendMessage = (function() {\n"
310 "            var message = document.getElementById('chat').value;\n"
311 "            if (message != '') {\n"
312 "                Chat.socket.send(message);\n"
313 "                document.getElementById('chat').value = '';\n"
314 "            }\n"
315 "        });\n"
316 "\n"
317 "        var Console = {};\n"
318 "\n"
319 "        Console.log = (function(message) {\n"
320 "            var console = document.getElementById('console');\n"
321 "            var p = document.createElement('p');\n"
322 "            p.style.wordWrap = 'break-word';\n"
323 "            p.innerHTML = message;\n"
324 "            console.appendChild(p);\n"
325 "            while (console.childNodes.length > 25) {\n"
326 "                console.removeChild(console.firstChild);\n"
327 "            }\n"
328 "            console.scrollTop = console.scrollHeight;\n"
329 "        });\n"
330 "\n"
331 "        Chat.initialize();\n"
332 "\n"
333 "      </script>\n"
334 "</head>\n"
335 "<body>\n"
336 "<noscript><h2 style=\"color: #ff0000\">Seems your browser doesn't support Javascript! Websockets rely on Javascript being enabled. Please enable\n"
337 "    Javascript and reload this page!</h2></noscript>\n"
338 "<div>\n"
339 "    <p><input type=\"text\" placeholder=\"type and press enter to chat\" id=\"chat\" /></p>\n"
340 "    <div id=\"container\">\n"
341 "        <div id=\"console\"/>\n"
342 "    </div>\n"
343 "</div>\n"
344 "</body>\n"
345 "</html>\n"
346 ;
347 
348 static const int  ws_chat_index_html_size = nxt_length(ws_chat_index_html);
349