Oregami
Repositories/oxedyne/fe2o3

oxedyne/fe2o3/fe2o3_steel/src/srv/wsproxy.rs

7.8 KiB, 29 runs

created by r1870400018:20228, which is this file's identity for as long as the history lasts, whatever it is later renamed to

download · who wrote it · its history

1//! Relaying a WebSocket upgrade to an upstream server.
2//!
3//! Two kinds of route need the same thing done. A [`ProxyRoute`](crate::srv::cfg::ProxyRoute)
4//! carries a whole application, WebSocket endpoints included; a
5//! [`WsRoute`](crate::srv::cfg::WsRoute) carries one path to a server that speaks its own
6//! protocol. Both forward the handshake and then get out of the way.
7//!
8//! Getting out of the way is the point. After the `101` the connection is no longer HTTP, and the
9//! frames on it are between the client and the upstream: this module never parses one. It copies
10//! bytes in both directions until an end closes, which is also why nothing here needs to know
11//! which sub-protocol, extensions or message sizes the two ends agreed on.
12//!
13//! [Written with AI entirely](https://need2know.ai/entirely-ai/code)\
14//! Anthropic Claude
15
16use oxedyne_fe2o3_core::prelude::*;
17use oxedyne_fe2o3_net::http::{
18 fwd::{
19 self,
20 ForwardedPolicy,
21 },
22 header::HttpHeadline,
23 msg::HttpMessage,
24};
25
26use std::net::SocketAddr;
27
28use tokio::{
29 io::{
30 AsyncRead,
31 AsyncReadExt,
32 AsyncWrite,
33 AsyncWriteExt,
34 },
35 net::TcpStream,
36};
37
38
39const MAX_RESPONSE_HEADER_BYTES: usize = 64 * 1024; // upstream block awaiting the `101`
40
41/// The upstream path a request is forwarded to, query included.
42///
43/// The query must ride through verbatim: an upstream that dispatches on a query parameter never
44/// sees it otherwise, and silently gets the default.
45pub fn upstream_target(request: &HttpMessage, base_path: &str) -> Outcome<String> {
46 let query = match &request.header.headline {
47 HttpHeadline::Request { loc, .. } => loc.query.clone(),
48 _ => return Err(err!(
49 "A websocket upgrade must be an HTTP request, not a response.";
50 Invalid, Bug)),
51 };
52 Ok(match query.is_empty() {
53 true => base_path.to_string(),
54 false => fmt!("{}?{}", base_path, query),
55 })
56}
57
58/// Forward `request` to the upstream WebSocket server at `host:port` as a `GET` of
59/// `upstream_path`, relay the response to `client`, and then copy bytes both ways until either end
60/// closes.
61///
62/// Every header the client sent is forwarded except the ones this hop owns -- see
63/// [`fwd::MANAGED_HEADERS`] and [`fwd::FORWARDED_HEADERS`] -- so `Sec-WebSocket-Key`,
64/// `Sec-WebSocket-Version`, `Sec-WebSocket-Protocol` and any cookies reach the upstream untouched.
65/// The upstream therefore computes the `Sec-WebSocket-Accept` the client will check, and this hop
66/// never has to.
67///
68/// `policy` decides whether the caller was entitled to speak the forwarding headers. With the
69/// default empty policy the caller's copies are dropped, so the values this hop appends are the
70/// only ones the upstream sees.
71///
72/// Returns once a direction closes. A client that goes away closes the upstream's write half, and
73/// an upstream that goes away closes the client's, so neither end is left holding a socket the
74/// other has abandoned.
75pub async fn tunnel_upgrade<S>(
76 client: &mut S,
77 request: &HttpMessage,
78 host: &str,
79 port: u16,
80 upstream_path: &str,
81 src_addr: SocketAddr,
82 policy: &ForwardedPolicy,
83 id: &str,
84)
85 -> Outcome<()>
86 where S: AsyncRead + AsyncWrite + Unpin,
87{
88 // Connect to the upstream.
89 let mut upstream = match TcpStream::connect((host, port)).await {
90 Ok(s) => s,
91 Err(e) => return Err(err!(e,
92 "{}: ws relay: failed to connect to {}:{}.", id, host, port;
93 IO, Network, Init)),
94 };
95
96 // Reconstruct the upgrade request for the upstream. The header work is shared with the HTTP
97 // proxy path in `fe2o3_net::http::fwd`: two copies of it is how one gets fixed and the other
98 // does not.
99 let req = fwd::build_upgrade_request_head(
100 upstream_path,
101 host,
102 request,
103 &src_addr,
104 policy,
105 );
106
107 match upstream.write_all(req.as_bytes()).await {
108 Ok(()) => (),
109 Err(e) => return Err(err!(e,
110 "{}: ws relay: failed to send the upgrade request upstream.", id;
111 IO, Network, Wire, Write)),
112 }
113 match upstream.flush().await {
114 Ok(()) => (),
115 Err(e) => return Err(err!(e,
116 "{}: ws relay: failed to flush the upstream connection.", id;
117 IO, Network, Wire, Write)),
118 }
119
120 // Read the upstream's response -- a 101 if it accepted -- and forward it verbatim, whatever it
121 // says. A refusal is the upstream's answer to give, and the client is entitled to read it.
122 let mut buf = vec![0u8; 8192];
123 let mut accum: Vec<u8> = Vec::new();
124 loop {
125 let n = match upstream.read(&mut buf).await {
126 Ok(0) => {
127 return Err(err!(
128 "{}: ws relay: the upstream closed before answering the upgrade.", id;
129 IO, Network, Wire, Read, Missing));
130 }
131 Ok(n) => n,
132 Err(e) => return Err(err!(e,
133 "{}: ws relay: error reading the upstream response.", id;
134 IO, Network, Wire, Read)),
135 };
136 accum.extend_from_slice(&buf[..n]);
137 if let Some(pos) = accum.windows(4).position(|w| w == b"\r\n\r\n") {
138 let header_end = pos + 4;
139 let response_bytes = &accum[..header_end];
140 let extra_bytes = &accum[header_end..];
141
142 match client.write_all(response_bytes).await {
143 Ok(()) => (),
144 Err(e) => return Err(err!(e,
145 "{}: ws relay: failed to forward the upgrade response.", id;
146 IO, Network, Wire, Write)),
147 }
148 // Anything the upstream sent after its headers is already a frame, and belongs to the
149 // client as much as the headers did.
150 if !extra_bytes.is_empty() {
151 match client.write_all(extra_bytes).await {
152 Ok(()) => (),
153 Err(e) => return Err(err!(e,
154 "{}: ws relay: failed to forward the upstream's first frames.", id;
155 IO, Network, Wire, Write)),
156 }
157 }
158 match client.flush().await {
159 Ok(()) => (),
160 Err(e) => return Err(err!(e,
161 "{}: ws relay: failed to flush the client connection.", id;
162 IO, Network, Wire, Write)),
163 }
164 break;
165 }
166 if accum.len() > MAX_RESPONSE_HEADER_BYTES {
167 return Err(err!(
168 "{}: ws relay: the upstream response headers exceed {} bytes.",
169 id, MAX_RESPONSE_HEADER_BYTES;
170 IO, Network, Input, TooBig));
171 }
172 }
173
174 // Copy bytes both ways until a direction ends.
175 let (mut client_r, mut client_w) = tokio::io::split(client);
176 let (mut upstream_r, mut upstream_w) = upstream.into_split();
177
178 log!(log_get_level!(), "{}: ws relay: tunnel to {}:{} established.", id, host, port);
179
180 tokio::select! {
181 // Client -> upstream.
182 res = tokio::io::copy(&mut client_r, &mut upstream_w) => {
183 match res {
184 Ok(_) => log!(log_get_level!(),
185 "{}: ws relay: client -> upstream closed.", id),
186 Err(e) => log!(log_get_level!(),
187 "{}: ws relay: client -> upstream error: {}", id, e),
188 }
189 let _ = upstream_w.shutdown().await;
190 }
191 // Upstream -> client.
192 res = tokio::io::copy(&mut upstream_r, &mut client_w) => {
193 match res {
194 Ok(_) => log!(log_get_level!(),
195 "{}: ws relay: upstream -> client closed.", id),
196 Err(e) => log!(log_get_level!(),
197 "{}: ws relay: upstream -> client error: {}", id, e),
198 }
199 let _ = client_w.shutdown().await;
200 }
201 }
202
203 log!(log_get_level!(), "{}: ws relay: tunnel closed.", id);
204 Ok(())
205}