Skip to main content

gproxy_protocol/adapt/generate/stream/
invoke.rs

1use super::super::{Endpoint, GenerationIdentity, GenerationStateAccess};
2use super::{
3    binding::StateBinding,
4    bridge::StreamBridge,
5    event::{Collected, EventLimits, NativeEvent},
6    output,
7    reader::{NativeReader, SourceFraming},
8};
9use crate::{
10    HttpBody, WireRequest, WireResponse,
11    capability::{StateStore, Upstream},
12    codec::{self, CodecLimits},
13    transform::{
14        Report, TransformError,
15        identity::{IdentityFlow, TargetIdPolicy},
16    },
17    wire::DeclaredFields,
18};
19use bytes::Bytes;
20use std::collections::VecDeque;
21
22pub(super) type NativeFull<B> = <<B as StreamBridge>::NativeEvent as NativeEvent>::Full;
23
24pub(super) type ClientFull<B> = <<B as StreamBridge>::ClientEvent as NativeEvent>::Full;
25
26/// Selected operation and per-invocation identity namespaces, supplied by the host.
27pub struct StreamTarget {
28    pub endpoint: Endpoint,
29    pub identities: GenerationIdentity,
30}
31
32#[derive(Debug, Clone, Copy)]
33pub struct StreamSettings {
34    pub codec: CodecLimits,
35    pub events: EventLimits,
36    pub source_framing: SourceFraming,
37    pub client_framing: SourceFraming,
38}
39
40impl StreamSettings {}
41
42#[derive(Debug)]
43pub enum StreamStart {
44    Streaming(WireResponse<()>),
45    Rejected(WireResponse<Bytes>),
46}
47
48/// A client event and its already bounded native transport encoding. The final
49/// framing chunk (`[DONE]` or `]`) has `event = None` and `finished = true`.
50#[derive(Debug)]
51pub struct StreamChunk<E> {
52    pub event: Option<E>,
53    pub bytes: Bytes,
54    pub finished: bool,
55}
56
57pub(super) struct ReadyChunk<E> {
58    pub chunk: StreamChunk<E>,
59}
60
61/// Caller-owned invocation and progress. Cancellation retains pending state
62/// writes/events; a started upstream send is never implicitly repeated.
63pub struct StreamInvocation<B: StreamBridge> {
64    pub(super) original: B::ClientRequest,
65    pub(super) history: Option<super::history::History>,
66    pub(super) pending_native: Option<B::NativeEvent>,
67    pub(super) pending_converted: Option<Vec<B::ClientEvent>>,
68    pub(super) finishing_source: bool,
69    pub(super) image_resources_required: bool,
70    pub(super) resource_revision: u64,
71    pub(super) target: B::NativeRequest,
72    pub(super) selected: StreamTarget,
73    pub(super) settings: StreamSettings,
74    /// The state this invocation was prepared against; every step must use it.
75    pub(super) binding: StateBinding,
76    /// The output lane a WebSocket start chose, so an HTTP start is refused
77    /// afterwards and a retried WebSocket start cannot switch lanes.
78    pub(super) websocket_lane: Option<Option<String>>,
79    // Keep protocol collectors and conversion state off the caller's stack.
80    // The invocation crosses several async preparation layers before boxing
81    // by a host, and debug builds otherwise duplicate these large values.
82    pub(super) bridge: Option<Box<B>>,
83    pub(super) flow: IdentityFlow,
84    pub(super) source: Option<Box<<B::NativeEvent as NativeEvent>::Collector>>,
85    pub(super) client: Option<Box<<B::ClientEvent as NativeEvent>::Collector>>,
86    pub(super) reader: Option<NativeReader>,
87    pub(super) encoder: output::Encoder,
88    pub(super) queued: VecDeque<(B::ClientEvent, u64)>,
89    pub(super) held: Vec<(B::ClientEvent, u64)>,
90    pub(super) holding: bool,
91    pub(super) queued_bytes: u64,
92    pub(super) ready: Option<ReadyChunk<B::ClientEvent>>,
93    pub(super) last_native_event: Option<B::NativeEvent>,
94    pub(super) native_final: Option<Collected<NativeFull<B>>>,
95    pub(super) client_final: Option<ClientFull<B>>,
96    pub(super) final_saved: bool,
97    pub(super) report: Report,
98    pub(super) emit_usage: bool,
99    pub(super) metadata: Option<WireResponse<()>>,
100    pub(super) rejected: Option<WireResponse<Bytes>>,
101    pub(super) sent: bool,
102    pub(super) eof: bool,
103    pub(super) websocket_terminal: bool,
104    pub(super) finished: bool,
105    pub(super) failed: bool,
106}
107
108impl<B: StreamBridge> StreamInvocation<B> {
109    pub(super) async fn new<S: StateStore>(
110        original: B::ClientRequest,
111        target: B::NativeRequest,
112        selected: StreamTarget,
113        settings: StreamSettings,
114        bridge: B,
115        report: Report,
116        state: &GenerationStateAccess<'_, S>,
117    ) -> Result<Self, TransformError> {
118        let original = original.into_declared();
119        let target = target.into_declared();
120        let flow = bridge.identities().clone();
121        if flow.namespace() != selected.identities.response.namespace() {
122            return Err(super::invalid("bridge response namespace mismatch"));
123        }
124        let source = B::NativeEvent::collector(
125            IdentityFlow::new(flow.namespace()),
126            TargetIdPolicy::new(B::NativeEvent::DIALECT),
127            settings.events,
128        );
129        let client = B::ClientEvent::collector(
130            IdentityFlow::new(flow.namespace()),
131            selected.identities.response_policy.clone(),
132            settings.events,
133        );
134        Ok(Self {
135            history: None,
136            pending_native: None,
137            pending_converted: None,
138            finishing_source: false,
139            image_resources_required: false,
140            resource_revision: 0,
141            original,
142            target,
143            selected,
144            settings,
145            binding: StateBinding::new(state),
146            websocket_lane: None,
147            bridge: Some(Box::new(bridge)),
148            flow,
149            source: Some(Box::new(source)),
150            client: Some(Box::new(client)),
151            reader: None,
152            encoder: output::Encoder::new(settings.client_framing, settings.codec),
153            queued: VecDeque::new(),
154            held: Vec::new(),
155            holding: false,
156            queued_bytes: 0,
157            ready: None,
158            last_native_event: None,
159            native_final: None,
160            client_final: None,
161            final_saved: false,
162            report,
163            emit_usage: true,
164            metadata: None,
165            rejected: None,
166            sent: false,
167            eof: false,
168            websocket_terminal: false,
169            finished: false,
170            failed: false,
171        })
172    }
173    pub fn original_request(&self) -> &B::ClientRequest {
174        &self.original
175    }
176    pub fn target_request(&self) -> &B::NativeRequest {
177        &self.target
178    }
179    pub fn send_started(&self) -> bool {
180        self.sent
181    }
182    pub fn report(&self) -> &Report {
183        &self.report
184    }
185    pub fn native_metadata(&self) -> Option<&WireResponse<()>> {
186        self.metadata.as_ref()
187    }
188    pub fn last_native_event(&self) -> Option<&B::NativeEvent> {
189        self.last_native_event.as_ref()
190    }
191    pub fn native_result(&self) -> Option<&NativeFull<B>> {
192        self.native_final.as_ref().map(|v| &v.value)
193    }
194    pub fn client_result(&self) -> Option<&ClientFull<B>> {
195        if self.final_saved {
196            self.client_final.as_ref()
197        } else {
198            None
199        }
200    }
201    pub fn rejected_response(&self) -> Option<&WireResponse<Bytes>> {
202        self.rejected.as_ref()
203    }
204    pub async fn start<U: Upstream, S: StateStore>(
205        &mut self,
206        upstream: &U,
207        target: &U::Target,
208        state: &GenerationStateAccess<'_, S>,
209    ) -> Result<StreamStart, TransformError> {
210        self.binding.check(state)?;
211        if self.sent || self.failed || self.websocket_lane.is_some() {
212            return Err(super::conflict(
213                "stream send already started or failed; cannot replay POST",
214            ));
215        }
216        let mut request_limits = self.settings.codec;
217        request_limits.max_body_bytes = request_limits
218            .max_body_bytes
219            .min(upstream.limits().write_bytes);
220        let body = codec::encode_json(&self.target, request_limits).map_err(super::codec_error)?;
221        let mut headers = self.selected.endpoint.headers.clone();
222        for name in [
223            http::header::CONTENT_LENGTH,
224            http::header::CONTENT_ENCODING,
225            http::header::TRANSFER_ENCODING,
226        ] {
227            headers.remove(name);
228        }
229        headers.insert(
230            http::header::CONTENT_TYPE,
231            http::HeaderValue::from_static("application/json"),
232        );
233        self.sent = true;
234        let response = upstream
235            .send(
236                target,
237                WireRequest {
238                    method: http::Method::POST,
239                    path: self.selected.endpoint.path.clone(),
240                    query: self.selected.endpoint.query.clone(),
241                    headers,
242                    body: HttpBody::Bytes(body),
243                },
244            )
245            .await?;
246        let mut read_limits = self.settings.codec;
247        read_limits.max_body_bytes = read_limits.max_body_bytes.min(upstream.limits().read_bytes);
248        self.metadata = Some(WireResponse {
249            status: response.status,
250            headers: response.headers.clone(),
251            body: (),
252        });
253        if !response.status.is_success() {
254            let raw = WireResponse {
255                status: response.status,
256                headers: response.headers,
257                body: codec::read_http_body(response.body, read_limits)
258                    .await
259                    .map_err(super::codec_error)?,
260            };
261            self.rejected = Some(WireResponse {
262                status: raw.status,
263                headers: raw.headers.clone(),
264                body: raw.body.clone(),
265            });
266            return Ok(StreamStart::Rejected(raw));
267        }
268        let accepted = WireResponse {
269            status: response.status,
270            headers: output::headers(response.headers.clone(), self.settings.client_framing),
271            body: (),
272        };
273        self.reader = Some(NativeReader::new(
274            response.body,
275            self.settings.source_framing,
276            read_limits,
277        ));
278        if let Some(value) = response.headers.get(http::header::CONTENT_TYPE) {
279            let valid = value
280                .to_str()
281                .ok()
282                .and_then(|v| v.split(';').next())
283                .is_some_and(|v| match self.settings.source_framing {
284                    SourceFraming::Sse => v.trim().eq_ignore_ascii_case("text/event-stream"),
285                    SourceFraming::JsonArray => v.trim().eq_ignore_ascii_case("application/json"),
286                    SourceFraming::Ndjson => {
287                        matches!(v.trim(), "application/x-ndjson" | "application/ndjson")
288                    }
289                });
290            if !valid {
291                self.failed = true;
292                self.reader = None;
293                return Err(super::invalid(
294                    "native content type disagrees with selected stream framing",
295                ));
296            }
297        }
298        Ok(StreamStart::Streaming(accepted))
299    }
300}