gproxy_protocol/adapt/responses_ws/
connect.rs

1use super::*;
2use crate::{
3    HttpBody, WireRequest, WireResponse,
4    capability::{CapabilityLimits, Upstream, UpstreamConnection},
5    codec,
6    wire::openai::responses::websocket::ClientEvent,
7};
8
9#[derive(Debug)]
10#[allow(clippy::large_enum_variant)]
11pub enum ResponsesWsConnect {
12    Connected {
13        handshake: WireResponse<()>,
14        session: ResponsesWsSession,
15    },
16    Rejected(WireResponse<HttpBody>),
17}
18
19/// Failed preflight/transport, or contradictory connected-handshake metadata.
20/// A normal HTTP rejection is the untouched `ResponsesWsConnect::Rejected`.
21#[derive(Debug)]
22pub struct ResponsesWsConnectError {
23    pub error: TransformError,
24    pub handshake: Option<Box<WireResponse<()>>>,
25}
26
27impl ResponsesWsConnectError {
28    pub fn kind(&self) -> TransformErrorKind {
29        self.error.kind()
30    }
31}
32
33impl From<TransformError> for ResponsesWsConnectError {
34    fn from(error: TransformError) -> Self {
35        Self {
36            error,
37            handshake: None,
38        }
39    }
40}
41
42impl std::fmt::Display for ResponsesWsConnectError {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        self.error.fmt(f)
45    }
46}
47
48impl std::error::Error for ResponsesWsConnectError {
49    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
50        Some(&self.error)
51    }
52}
53
54/// The host supplies the actual upgrade, including authentication and deadlines.
55/// HTTP/1.1 101 and a successful HTTP/2 CONNECT status remain unmodified.
56pub async fn connect<U: Upstream>(
57    upstream: &U,
58    target: &U::Target,
59    request: WireRequest<()>,
60    limits: ResponsesWsLimits,
61) -> Result<ResponsesWsConnect, ResponsesWsConnectError> {
62    let bounds = bound(limits, upstream.limits())?;
63    // Even an empty typed response.create must fit before opening a connection.
64    codec::encode_json(
65        &ClientEvent::ResponseCreate(GenerateContentRequestBody::builder().build()),
66        codec_limits(bounds.send_event.min(bounds.send_bytes)),
67    )
68    .map_err(|e| codec_error(e, true))?;
69    match upstream
70        .connect(target, request)
71        .await
72        .map_err(TransformError::from)?
73    {
74        UpstreamConnection::Rejected(response) => Ok(ResponsesWsConnect::Rejected(response)),
75        UpstreamConnection::Connected { handshake, socket } => {
76            if handshake.status != http::StatusCode::SWITCHING_PROTOCOLS
77                && !handshake.status.is_success()
78            {
79                drop(socket);
80                return Err(ResponsesWsConnectError {
81                    error: invalid(
82                        "host returned a connected socket with unsuccessful handshake status",
83                    ),
84                    handshake: Some(Box::new(handshake)),
85                });
86            }
87            Ok(ResponsesWsConnect::Connected {
88                handshake,
89                session: ResponsesWsSession::connected(socket, bounds),
90            })
91        }
92    }
93}
94
95fn host_usize(value: u64) -> usize {
96    usize::try_from(value).unwrap_or(usize::MAX)
97}
98
99fn bound(limits: ResponsesWsLimits, host: CapabilityLimits) -> Result<Bounds, TransformError> {
100    let send_frame = limits
101        .max_frame_bytes
102        .min(host_usize(host.ws_frame_bytes))
103        .min(host_usize(host.write_bytes));
104    let receive_frame = limits
105        .max_frame_bytes
106        .min(host_usize(host.ws_frame_bytes))
107        .min(host_usize(host.read_bytes));
108    let send_event = limits.max_event_bytes.min(send_frame);
109    let receive_event = limits.max_event_bytes.min(receive_frame);
110    let send_bytes = limits.max_send_bytes.min(host_usize(host.write_bytes));
111    let receive_bytes = limits.max_receive_bytes.min(host_usize(host.read_bytes));
112    let mut collector = limits.collector;
113
114    collector.max_bytes = collector.max_bytes.min(receive_bytes);
115    collector.max_text_bytes = collector.max_text_bytes.min(receive_bytes);
116    collector.max_json_bytes = collector.max_json_bytes.min(receive_bytes);
117    if [
118        send_frame,
119        receive_frame,
120        send_event,
121        receive_event,
122        send_bytes,
123        receive_bytes,
124        collector.max_bytes,
125    ]
126    .contains(&0)
127    {
128        return Err(limit(
129            "responses.websocket.limits",
130            "positive usable byte limits required",
131        ));
132    }
133    Ok(Bounds {
134        send_frame,
135        receive_frame,
136        send_event,
137        receive_event,
138        send_bytes,
139        receive_bytes,
140
141        allow_binary: limits.allow_binary,
142        collector,
143    })
144}
gproxy_protocol/adapt/responses_ws/connect.rs

gproxy_protocol/adapt/responses_ws/
connect.rs
