1 /** 2 * @file websock.c Implementation of The WebSocket Protocol 3 * 4 * Copyright (C) 2010 Creytiv.com 5 */ 6 7 #include <re_types.h> 8 #include <re_fmt.h> 9 #include <re_mem.h> 10 #include <re_mbuf.h> 11 #include <re_sa.h> 12 #include <re_list.h> 13 #include <re_tmr.h> 14 #include <re_srtp.h> 15 #include <re_tcp.h> 16 #include <re_tls.h> 17 #include <re_msg.h> 18 #include <re_http.h> 19 #include <re_base64.h> 20 #include <re_sha.h> 21 #include <re_sys.h> 22 #include <re_websock.h> 23 24 25 enum { 26 TIMEOUT_CLOSE = 10000, 27 BUFSIZE_MAX = 131072, 28 }; 29 30 enum websock_state { 31 ACCEPTING = 0, 32 CONNECTING, 33 OPEN, 34 CLOSING, 35 CLOSED, 36 }; 37 38 struct websock { 39 websock_shutdown_h *shuth; 40 void *arg; 41 bool shutdown; 42 }; 43 44 struct websock_conn { 45 struct tmr tmr; 46 struct sa peer; 47 char nonce[24]; 48 struct websock *sock; 49 struct tcp_conn *tc; 50 struct tls_conn *sc; 51 struct mbuf *mb; 52 struct http_req *req; 53 websock_estab_h *estabh; 54 websock_recv_h *recvh; 55 websock_close_h *closeh; 56 void *arg; 57 enum websock_state state; 58 unsigned kaint; 59 bool active; 60 }; 61 62 63 static const char magic[] = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"; 64 65 66 static void timeout_handler(void *arg); 67 68 69 static void dummy_recv_handler(const struct websock_hdr *hdr, struct mbuf *mb, 70 void *arg) 71 { 72 (void)hdr; 73 (void)mb; 74 (void)arg; 75 } 76 77 78 static void internal_close_handler(int err, void *arg) 79 { 80 struct websock_conn *conn = arg; 81 (void)err; 82 83 mem_deref(conn); 84 } 85 86 87 static void sock_destructor(void *arg) 88 { 89 struct websock *sock = arg; 90 91 if (sock->shutdown) { 92 sock->shutdown = false; 93 mem_ref(sock); 94 if (sock->shuth) 95 sock->shuth(sock->arg); 96 return; 97 } 98 } 99 100 101 static void conn_destructor(void *arg) 102 { 103 struct websock_conn *conn = arg; 104 105 if (conn->state == OPEN) 106 (void)websock_close(conn, WEBSOCK_GOING_AWAY, "Going Away"); 107 108 if (conn->state == CLOSING) { 109 110 conn->recvh = dummy_recv_handler; 111 conn->closeh = internal_close_handler; 112 conn->arg = conn; 113 114 tmr_start(&conn->tmr, TIMEOUT_CLOSE, timeout_handler, conn); 115 116 /* important: the hack below depends on this */ 117 mem_ref(conn); 118 return; 119 } 120 121 tmr_cancel(&conn->tmr); 122 mem_deref(conn->sc); 123 mem_deref(conn->tc); 124 mem_deref(conn->mb); 125 mem_deref(conn->req); 126 mem_deref(conn->sock); 127 } 128 129 130 static void conn_close(struct websock_conn *conn, int err) 131 { 132 tmr_cancel(&conn->tmr); 133 conn->sc = mem_deref(conn->sc); 134 conn->tc = mem_deref(conn->tc); 135 conn->state = CLOSED; 136 137 conn->closeh(err, conn->arg); 138 } 139 140 141 static void timeout_handler(void *arg) 142 { 143 struct websock_conn *conn = arg; 144 145 conn_close(conn, ETIMEDOUT); 146 } 147 148 149 static void keepalive_handler(void *arg) 150 { 151 struct websock_conn *conn = arg; 152 153 tmr_start(&conn->tmr, conn->kaint, keepalive_handler, conn); 154 155 (void)websock_send(conn, WEBSOCK_PING, NULL); 156 } 157 158 159 static enum websock_scode websock_err2scode(int err) 160 { 161 switch (err) { 162 163 case EOVERFLOW: return WEBSOCK_MESSAGE_TOO_BIG; 164 case EPROTO: return WEBSOCK_PROTOCOL_ERROR; 165 case EBADMSG: return WEBSOCK_PROTOCOL_ERROR; 166 default: return WEBSOCK_INTERNAL_ERROR; 167 } 168 } 169 170 171 static int websock_decode(struct websock_hdr *hdr, struct mbuf *mb) 172 { 173 uint8_t v, *p; 174 size_t i; 175 176 if (mbuf_get_left(mb) < 2) 177 return ENODATA; 178 179 v = mbuf_read_u8(mb); 180 hdr->fin = v>>7 & 0x1; 181 hdr->rsv1 = v>>6 & 0x1; 182 hdr->rsv2 = v>>5 & 0x1; 183 hdr->rsv3 = v>>4 & 0x1; 184 hdr->opcode = v & 0x0f; 185 186 v = mbuf_read_u8(mb); 187 hdr->mask = v>>7 & 0x1; 188 hdr->len = v & 0x7f; 189 190 if (hdr->len == 126) { 191 192 if (mbuf_get_left(mb) < 2) 193 return ENODATA; 194 195 hdr->len = ntohs(mbuf_read_u16(mb)); 196 } 197 else if (hdr->len == 127) { 198 199 if (mbuf_get_left(mb) < 8) 200 return ENODATA; 201 202 hdr->len = sys_ntohll(mbuf_read_u64(mb)); 203 } 204 205 if (hdr->mask) { 206 207 if (mbuf_get_left(mb) < (4 + hdr->len)) 208 return ENODATA; 209 210 hdr->mkey[0] = mbuf_read_u8(mb); 211 hdr->mkey[1] = mbuf_read_u8(mb); 212 hdr->mkey[2] = mbuf_read_u8(mb); 213 hdr->mkey[3] = mbuf_read_u8(mb); 214 215 for (i=0, p=mbuf_buf(mb); i<hdr->len; i++) 216 p[i] = p[i] ^ hdr->mkey[i%4]; 217 } 218 else { 219 if (mbuf_get_left(mb) < hdr->len) 220 return ENODATA; 221 } 222 223 return 0; 224 } 225 226 227 static void recv_handler(struct mbuf *mb, void *arg) 228 { 229 struct websock_conn *conn = arg; 230 int err = 0; 231 232 if (conn->mb) { 233 234 const size_t len = mbuf_get_left(mb), pos = conn->mb->pos; 235 236 if ((mbuf_get_left(conn->mb) + len) > BUFSIZE_MAX) { 237 err = EOVERFLOW; 238 goto out; 239 } 240 241 conn->mb->pos = conn->mb->end; 242 243 err = mbuf_write_mem(conn->mb, mbuf_buf(mb), len); 244 if (err) 245 goto out; 246 247 conn->mb->pos = pos; 248 } 249 else { 250 conn->mb = mem_ref(mb); 251 } 252 253 while (conn->mb) { 254 255 struct websock_hdr hdr; 256 size_t pos, end; 257 258 pos = conn->mb->pos; 259 260 err = websock_decode(&hdr, conn->mb); 261 if (err) { 262 if (err == ENODATA) { 263 conn->mb->pos = pos; 264 err = 0; 265 break; 266 } 267 268 goto out; 269 } 270 271 if (conn->active == hdr.mask) { 272 err = EPROTO; 273 goto out; 274 } 275 276 if (hdr.rsv1 || hdr.rsv2 || hdr.rsv3) { 277 err = EPROTO; 278 goto out; 279 } 280 281 mb = conn->mb; 282 283 end = mb->end; 284 mb->end = mb->pos + (size_t)hdr.len; 285 286 if (end > mb->end) { 287 struct mbuf *mbn = mbuf_alloc(end - mb->end); 288 if (!mbn) { 289 err = ENOMEM; 290 goto out; 291 } 292 293 (void)mbuf_write_mem(mbn, mb->buf + mb->end, 294 end - mb->end); 295 mbn->pos = 0; 296 297 conn->mb = mbn; 298 } 299 else { 300 conn->mb = NULL; 301 } 302 303 switch (hdr.opcode) { 304 305 case WEBSOCK_CONT: 306 case WEBSOCK_TEXT: 307 case WEBSOCK_BIN: 308 mem_ref(conn); 309 conn->recvh(&hdr, mb, conn->arg); 310 311 if (mem_nrefs(conn) == 1) { 312 313 if (conn->state == OPEN) 314 (void)websock_close(conn, 315 WEBSOCK_GOING_AWAY, 316 "Going Away"); 317 318 /* 319 * This is a hack. We enforce CLOSING 320 * state so we know the connection will 321 * continue to live. 322 */ 323 conn->state = CLOSING; 324 } 325 mem_deref(conn); 326 break; 327 328 case WEBSOCK_CLOSE: 329 if (conn->state == OPEN) 330 (void)websock_send(conn, WEBSOCK_CLOSE, "%b", 331 mbuf_buf(mb), mbuf_get_left(mb)); 332 conn_close(conn, 0); 333 mem_deref(mb); 334 return; 335 336 case WEBSOCK_PING: 337 (void)websock_send(conn, WEBSOCK_PONG, "%b", 338 mbuf_buf(mb), mbuf_get_left(mb)); 339 break; 340 341 case WEBSOCK_PONG: 342 break; 343 344 default: 345 mem_deref(mb); 346 err = EPROTO; 347 goto out; 348 } 349 350 mem_deref(mb); 351 } 352 353 out: 354 if (err) { 355 (void)websock_close(conn, websock_err2scode(err), NULL); 356 conn_close(conn, err); 357 } 358 } 359 360 361 static void close_handler(int err, void *arg) 362 { 363 struct websock_conn *conn = arg; 364 365 conn_close(conn, err); 366 } 367 368 369 static int accept_print(struct re_printf *pf, const struct pl *key) 370 { 371 uint8_t digest[SHA_DIGEST_LENGTH]; 372 SHA_CTX ctx; 373 374 SHA1_Init(&ctx); 375 SHA1_Update(&ctx, key->p, key->l); 376 SHA1_Update(&ctx, magic, sizeof(magic)-1); 377 SHA1_Final(digest, &ctx); 378 379 return base64_print(pf, digest, sizeof(digest)); 380 } 381 382 383 static void http_resp_handler(int err, const struct http_msg *msg, void *arg) 384 { 385 struct websock_conn *conn = arg; 386 const struct http_hdr *hdr; 387 struct pl key; 388 char buf[32]; 389 390 if (err || msg->scode != 101) 391 goto fail; 392 393 if (!http_msg_hdr_has_value(msg, HTTP_HDR_UPGRADE, "websocket")) 394 goto fail; 395 396 if (!http_msg_hdr_has_value(msg, HTTP_HDR_CONNECTION, "Upgrade")) 397 goto fail; 398 399 hdr = http_msg_hdr(msg, HTTP_HDR_SEC_WEBSOCKET_ACCEPT); 400 if (!hdr) 401 goto fail; 402 403 key.p = conn->nonce; 404 key.l = sizeof(conn->nonce); 405 406 if (re_snprintf(buf, sizeof(buf), "%H", accept_print, &key) < 0) 407 goto fail; 408 409 if (pl_strcmp(&hdr->val, buf)) 410 goto fail; 411 412 /* here we are ok */ 413 414 conn->state = OPEN; 415 (void)tcp_conn_peer_get(conn->tc, &conn->peer); 416 417 if (conn->kaint) 418 tmr_start(&conn->tmr, conn->kaint, keepalive_handler, conn); 419 420 conn->estabh(conn->arg); 421 return; 422 423 fail: 424 conn_close(conn, err ? err : EPROTO); 425 } 426 427 428 static void http_conn_handler(struct tcp_conn *tc, struct tls_conn *sc, 429 void *arg) 430 { 431 struct websock_conn *conn = arg; 432 433 conn->tc = mem_ref(tc); 434 conn->sc = mem_ref(sc); 435 436 tcp_set_handlers(conn->tc, NULL, recv_handler, close_handler, conn); 437 } 438 439 440 int websock_connect(struct websock_conn **connp, struct websock *sock, 441 struct http_cli *cli, const char *uri, unsigned kaint, 442 websock_estab_h *estabh, websock_recv_h *recvh, 443 websock_close_h *closeh, void *arg, 444 const char *fmt, ...) 445 { 446 struct websock_conn *conn; 447 uint8_t nonce[16]; 448 va_list ap; 449 size_t len; 450 int err; 451 452 if (!connp || !sock || !cli || !uri || !estabh || !recvh || !closeh) 453 return EINVAL; 454 455 conn = mem_zalloc(sizeof(*conn), conn_destructor); 456 if (!conn) 457 return ENOMEM; 458 459 /* The nonce MUST be selected randomly for each connection */ 460 rand_bytes(nonce, sizeof(nonce)); 461 462 len = sizeof(conn->nonce); 463 464 err = base64_encode(nonce, sizeof(nonce), conn->nonce, &len); 465 if (err) 466 goto out; 467 468 conn->sock = mem_ref(sock); 469 conn->kaint = kaint; 470 conn->estabh = estabh; 471 conn->recvh = recvh; 472 conn->closeh = closeh; 473 conn->arg = arg; 474 conn->state = CONNECTING; 475 conn->active = true; 476 477 /* Protocol Handshake */ 478 va_start(ap, fmt); 479 err = http_request(&conn->req, cli, "GET", uri, 480 http_resp_handler, NULL, conn, 481 "Upgrade: websocket\r\n" 482 "Connection: upgrade\r\n" 483 "Sec-WebSocket-Key: %b\r\n" 484 "Sec-WebSocket-Version: 13\r\n" 485 "%v" 486 "\r\n", 487 conn->nonce, sizeof(conn->nonce), 488 fmt, &ap); 489 va_end(ap); 490 if (err) 491 goto out; 492 493 http_req_set_conn_handler(conn->req, http_conn_handler); 494 495 out: 496 if (err) 497 mem_deref(conn); 498 else 499 *connp = conn; 500 501 return err; 502 } 503 504 505 int websock_accept(struct websock_conn **connp, struct websock *sock, 506 struct http_conn *htconn, const struct http_msg *msg, 507 unsigned kaint, websock_recv_h *recvh, 508 websock_close_h *closeh, void *arg) 509 { 510 const struct http_hdr *key; 511 struct websock_conn *conn; 512 int err; 513 514 if (!connp || !sock || !htconn || !msg || !recvh || !closeh) 515 return EINVAL; 516 517 if (!http_msg_hdr_has_value(msg, HTTP_HDR_UPGRADE, "websocket")) 518 return EBADMSG; 519 520 if (!http_msg_hdr_has_value(msg, HTTP_HDR_CONNECTION, "Upgrade")) 521 return EBADMSG; 522 523 if (!http_msg_hdr_has_value(msg, HTTP_HDR_SEC_WEBSOCKET_VERSION, "13")) 524 return EBADMSG; 525 526 key = http_msg_hdr(msg, HTTP_HDR_SEC_WEBSOCKET_KEY); 527 if (!key) 528 return EBADMSG; 529 530 conn = mem_zalloc(sizeof(*conn), conn_destructor); 531 if (!conn) 532 return ENOMEM; 533 534 err = http_reply(htconn, 101, "Switching Protocols", 535 "Upgrade: websocket\r\n" 536 "Connection: Upgrade\r\n" 537 "Sec-WebSocket-Accept: %H\r\n" 538 "\r\n", 539 accept_print, &key->val); 540 if (err) 541 goto out; 542 543 sa_cpy(&conn->peer, http_conn_peer(htconn)); 544 conn->sock = mem_ref(sock); 545 conn->tc = mem_ref(http_conn_tcp(htconn)); 546 conn->sc = mem_ref(http_conn_tls(htconn)); 547 conn->kaint = kaint; 548 conn->recvh = recvh; 549 conn->closeh = closeh; 550 conn->arg = arg; 551 conn->state = OPEN; 552 conn->active = false; 553 554 tcp_set_handlers(conn->tc, NULL, recv_handler, close_handler, conn); 555 http_conn_close(htconn); 556 557 if (conn->kaint) 558 tmr_start(&conn->tmr, conn->kaint, keepalive_handler, conn); 559 560 out: 561 if (err) 562 mem_deref(conn); 563 else 564 *connp = conn; 565 566 return err; 567 } 568 569 570 static int websock_encode(struct mbuf *mb, bool fin, 571 enum websock_opcode opcode, bool mask, size_t len) 572 { 573 int err; 574 575 err = mbuf_write_u8(mb, (fin<<7) | (opcode & 0x0f)); 576 577 if (len > 0xffff) { 578 err |= mbuf_write_u8(mb, (mask<<7) | 127); 579 err |= mbuf_write_u64(mb, sys_htonll(len)); 580 } 581 else if (len > 125) { 582 err |= mbuf_write_u8(mb, (mask<<7) | 126); 583 err |= mbuf_write_u16(mb, htons(len)); 584 } 585 else { 586 err |= mbuf_write_u8(mb, (mask<<7) | len); 587 } 588 589 if (mask) { 590 uint8_t mkey[4]; 591 uint8_t *p; 592 size_t i; 593 594 rand_bytes(mkey, sizeof(mkey)); 595 596 err |= mbuf_write_mem(mb, mkey, sizeof(mkey)); 597 598 for (i=0, p=mbuf_buf(mb); i<len; i++) 599 p[i] = p[i] ^ mkey[i%4]; 600 } 601 602 return err; 603 } 604 605 606 static int websock_vsend(struct websock_conn *conn, enum websock_opcode opcode, 607 enum websock_scode scode, const char *fmt, va_list ap) 608 { 609 const size_t hsz = conn->active ? 14 : 10; 610 size_t len, start; 611 struct mbuf *mb; 612 int err = 0; 613 614 if (conn->state != OPEN) 615 return ENOTCONN; 616 617 mb = mbuf_alloc(2048); 618 if (!mb) 619 return ENOMEM; 620 621 mb->pos = hsz; 622 623 if (scode) 624 err |= mbuf_write_u16(mb, htons(scode)); 625 if (fmt) 626 err |= mbuf_vprintf(mb, fmt, ap); 627 if (err) 628 goto out; 629 630 len = mb->pos - hsz; 631 632 if (len > 0xffff) 633 start = mb->pos = 0; 634 else if (len > 125) 635 start = mb->pos = 6; 636 else 637 start = mb->pos = 8; 638 639 err = websock_encode(mb, true, opcode, conn->active, len); 640 if (err) 641 goto out; 642 643 mb->pos = start; 644 645 err = tcp_send(conn->tc, mb); 646 if (err) 647 goto out; 648 649 out: 650 mem_deref(mb); 651 652 return err; 653 } 654 655 656 int websock_send(struct websock_conn *conn, enum websock_opcode opcode, 657 const char *fmt, ...) 658 { 659 va_list ap; 660 int err; 661 662 if (!conn) 663 return EINVAL; 664 665 va_start(ap, fmt); 666 err = websock_vsend(conn, opcode, 0, fmt, ap); 667 va_end(ap); 668 669 return err; 670 } 671 672 673 int websock_close(struct websock_conn *conn, enum websock_scode scode, 674 const char *fmt, ...) 675 { 676 va_list ap; 677 int err; 678 679 if (!conn) 680 return EINVAL; 681 682 if (!scode) 683 fmt = NULL; 684 685 va_start(ap, fmt); 686 err = websock_vsend(conn, WEBSOCK_CLOSE, scode, fmt, ap); 687 va_end(ap); 688 689 if (!err) 690 conn->state = CLOSING; 691 692 return err; 693 } 694 695 696 const struct sa *websock_peer(const struct websock_conn *conn) 697 { 698 return conn ? &conn->peer : NULL; 699 } 700 701 702 int websock_alloc(struct websock **sockp, websock_shutdown_h *shuth, void *arg) 703 { 704 struct websock *sock; 705 706 if (!sockp) 707 return EINVAL; 708 709 sock = mem_zalloc(sizeof(*sock), sock_destructor); 710 if (!sock) 711 return ENOMEM; 712 713 sock->shuth = shuth; 714 sock->arg = arg; 715 716 *sockp = sock; 717 718 return 0; 719 } 720 721 722 void websock_shutdown(struct websock *sock) 723 { 724 if (!sock || sock->shutdown) 725 return; 726 727 sock->shutdown = true; 728 mem_deref(sock); 729 } 730