1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

//! # Client handshake request
//!
//! A client sends a handshake request to the server. It includes the following information:
//!
//! ```yml
//! GET /chat HTTP/1.1
//! Host: example.com:8000
//! Upgrade: websocket
//! Connection: Upgrade
//! Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
//! Sec-WebSocket-Version: 13
//! ```
//!
//! The server must be careful to understand everything the client asks for, otherwise security issues can occur.
//! If any header is not understood or has an incorrect value, the server should send a 400 ("Bad Request")} response and immediately close the socket.
//!
//! ### Tips
//!
//! All browsers send an Origin header.
//! You can use this header for security (checking for same origin, automatically allowing or denying, etc.) and send a 403 Forbidden if you don't like what you see.
//! However, be warned that non-browser agents can send a faked Origin. Most applications reject requests without this header.
//!
//! Any http headers is allowed. (Do whatever you want with them)
//!
//! ### Note
//!
//! -  HTTP version must be `1.1` or greater, and method must be `GET`
//! - `Host` header field containing the server's authority.
//! - `Upgrade` header field containing the value `"websocket"`
//! - `Connection` header field that includes the token `"Upgrade"`
//! - `Sec-WebSocket-Version` header field containing the value `13`
//! - `Sec-WebSocket-Key` header field with a base64-encoded value that, when decoded, is 16 bytes in length.
//! -  Request may include any other header fields, for example, cookies and/or authentication-related header fields.
//! -  Optionally, `Origin` header field.  This header field is sent by all browser clients.

use crate::http::Header;
use sha1::{Digest, Sha1};

/// WebSocket magic string used during the WebSocket handshake
pub const MAGIC_STRING: &[u8; 36] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

/// ### Example
///
/// ```rust
/// use web_socket::handshake::accept_key_from;
/// assert_eq!(accept_key_from("dGhlIHNhbXBsZSBub25jZQ=="), "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=");
/// ```
#[inline]
pub fn accept_key_from(sec_ws_key: impl AsRef<[u8]>) -> String {
    let mut sha1 = Sha1::new();
    sha1.update(sec_ws_key.as_ref());
    sha1.update(MAGIC_STRING);
    base64_encode(sha1.finalize())
}

/// ## Server handshake response
///
/// When the server receives the handshake request,
/// It should send back a special response that indicates that the protocol will be changing from HTTP to WebSocket.
///
/// The `Sec-WebSocket-Accept` header is important in that the server must derive it from the `Sec-WebSocket-Key` that the client sent to it.
///
/// ### Example
///
/// ```rust
/// let res = [
///     "HTTP/1.1 101 Switching Protocols",
///     "Upgrade: websocket",
///     "Connection: Upgrade",
///     "Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=",
///     "",
///     ""
/// ];
/// let field: Option<(&str, &str)> = None;
/// assert_eq!(web_socket::handshake::response("dGhlIHNhbXBsZSBub25jZQ==", field), res.join("\r\n"));
/// ```
///
/// To get it, concatenate the client's `Sec-WebSocket-Key` and the string _"258EAFA5-E914-47DA-95CA-C5AB0DC85B11"_ together (it's a [Magic string](https://en.wikipedia.org/wiki/Magic_string)), take the SHA-1 hash of the result, and return the base64 encoding of that hash.
///
///
/// 1. If the connection is happening on an HTTPS (HTTP-over-TLS) port,
///    perform a TLS handshake over the connection.  If this fails
///    (e.g., the client indicated a host name in the extended client
///    hello "server_name" extension that the server does not host),
///    then close the connection.
///
/// 2. The server can perform additional client authentication, Or The server MAY redirect the client.
///
///
/// ### Note
///
/// - Regular HTTP status codes can be used only before the handshake. After the handshake succeeds, you have to use a different set of codes (defined in section 7.4 of the spec)
pub fn response(
    sec_ws_key: impl AsRef<str>,
    headers: impl IntoIterator<Item = impl Header>,
) -> String {
    let key = accept_key_from(sec_ws_key.as_ref());
    let headers: String = headers.into_iter().map(|f| Header::fmt(&f)).collect();
    format!("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: {key}\r\n{headers}\r\n")
}

/// ### Example
///
/// ```no_run
/// use web_socket::handshake::request;
/// let _ = request("example.com", "/path", [("key", "value")]);
/// ```
///
/// ### Output
///
/// ```yaml
/// GET /path HTTP/1.1
/// Host: example.com
/// Upgrade: websocket
/// Connection: Upgrade
/// Sec-WebSocket-Version: 13
/// Sec-WebSocket-Key: D3E1sFZlZfeZgNXtVHfhKg== # randomly generated
/// key: value
/// ...
/// ```
pub fn request(
    host: impl AsRef<str>,
    path: impl AsRef<str>,
    headers: impl IntoIterator<Item = impl Header>,
) -> (String, String) {
    let host = host.as_ref();
    let path = path.as_ref().trim_start_matches('/');
    let sec_key = base64_encode(crate::utils::rand_u128().to_ne_bytes());
    let headers: String = headers.into_iter().map(|f| Header::fmt(&f)).collect();
    (format!("GET /{path} HTTP/1.1\r\nHost: {host}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Version: 13\r\nSec-WebSocket-Key: {sec_key}\r\n{headers}\r\n"),sec_key)
}

#[inline]
fn base64_encode(string: impl AsRef<[u8]>) -> String {
    base64::Engine::encode(&base64::prelude::BASE64_STANDARD, string)
}