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    pub(super) bridge: Option<B>,
80    pub(super) flow: IdentityFlow,
81    pub(super) source: Option<<B::NativeEvent as NativeEvent>::Collector>,
82    pub(super) client: Option<<B::ClientEvent as NativeEvent>::Collector>,
83    pub(super) reader: Option<NativeReader>,
84    pub(super) encoder: output::Encoder,
85    pub(super) queued: VecDeque<(B::ClientEvent, u64)>,
86    pub(super) held: Vec<(B::ClientEvent, u64)>,
87    pub(super) holding: bool,
88    pub(super) queued_bytes: u64,
89    pub(super) ready: Option<ReadyChunk<B::ClientEvent>>,
90    pub(super) last_native_event: Option<B::NativeEvent>,
91    pub(super) native_final: Option<Collected<NativeFull<B>>>,
92    pub(super) client_final: Option<ClientFull<B>>,
93    pub(super) final_saved: bool,
94    pub(super) report: Report,
95    pub(super) emit_usage: bool,
96    pub(super) metadata: Option<WireResponse<()>>,
97    pub(super) rejected: Option<WireResponse<Bytes>>,
98    pub(super) sent: bool,
99    pub(super) eof: bool,
100    pub(super) websocket_terminal: bool,
101    pub(super) finished: bool,
102    pub(super) failed: bool,
103}
104
105impl<B: StreamBridge> StreamInvocation<B> {
106    pub(super) async fn new<S: StateStore>(
107        original: B::ClientRequest,
108        target: B::NativeRequest,
109        selected: StreamTarget,
110        settings: StreamSettings,
111        bridge: B,
112        report: Report,
113        state: &GenerationStateAccess<'_, S>,
114    ) -> Result<Self, TransformError> {
115        let original = original.into_declared();
116        let target = target.into_declared();
117        let flow = bridge.identities().clone();
118        if flow.namespace() != selected.identities.response.namespace() {
119            return Err(super::invalid("bridge response namespace mismatch"));
120        }
121        let source = B::NativeEvent::collector(
122            IdentityFlow::new(flow.namespace()),
123            TargetIdPolicy::new(B::NativeEvent::DIALECT),
124            settings.events,
125        );
126        let client = B::ClientEvent::collector(
127            IdentityFlow::new(flow.namespace()),
128            selected.identities.response_policy.clone(),
129            settings.events,
130        );
131        Ok(Self {
132            history: None,
133            pending_native: None,
134            pending_converted: None,
135            finishing_source: false,
136            image_resources_required: false,
137            resource_revision: 0,
138            original,
139            target,
140            selected,
141            settings,
142            binding: StateBinding::new(state),
143            websocket_lane: None,
144            bridge: Some(bridge),
145            flow,
146            source: Some(source),
147            client: Some(client),
148            reader: None,
149            encoder: output::Encoder::new(settings.client_framing, settings.codec),
150            queued: VecDeque::new(),
151            held: Vec::new(),
152            holding: false,
153            queued_bytes: 0,
154            ready: None,
155            last_native_event: None,
156            native_final: None,
157            client_final: None,
158            final_saved: false,
159            report,
160            emit_usage: true,
161            metadata: None,
162            rejected: None,
163            sent: false,
164            eof: false,
165            websocket_terminal: false,
166            finished: false,
167            failed: false,
168        })
169    }
170    pub fn original_request(&self) -> &B::ClientRequest {
171        &self.original
172    }
173    pub fn target_request(&self) -> &B::NativeRequest {
174        &self.target
175    }
176    pub fn send_started(&self) -> bool {
177        self.sent
178    }
179    pub fn report(&self) -> &Report {
180        &self.report
181    }
182    pub fn native_metadata(&self) -> Option<&WireResponse<()>> {
183        self.metadata.as_ref()
184    }
185    pub fn last_native_event(&self) -> Option<&B::NativeEvent> {
186        self.last_native_event.as_ref()
187    }
188    pub fn native_result(&self) -> Option<&NativeFull<B>> {
189        self.native_final.as_ref().map(|v| &v.value)
190    }
191    pub fn client_result(&self) -> Option<&ClientFull<B>> {
192        if self.final_saved {
193            self.client_final.as_ref()
194        } else {
195            None
196        }
197    }
198    pub fn rejected_response(&self) -> Option<&WireResponse<Bytes>> {
199        self.rejected.as_ref()
200    }
201    pub async fn start<U: Upstream, S: StateStore>(
202        &mut self,
203        upstream: &U,
204        target: &U::Target,
205        state: &GenerationStateAccess<'_, S>,
206    ) -> Result<StreamStart, TransformError> {
207        self.binding.check(state)?;
208        if self.sent || self.failed || self.websocket_lane.is_some() {
209            return Err(super::conflict(
210                "stream send already started or failed; cannot replay POST",
211            ));
212        }
213        let mut request_limits = self.settings.codec;
214        request_limits.max_body_bytes = request_limits
215            .max_body_bytes
216            .min(upstream.limits().write_bytes);
217        let body = codec::encode_json(&self.target, request_limits).map_err(super::codec_error)?;
218        let mut headers = self.selected.endpoint.headers.clone();
219        for name in [
220            http::header::CONTENT_LENGTH,
221            http::header::CONTENT_ENCODING,
222            http::header::TRANSFER_ENCODING,
223        ] {
224            headers.remove(name);
225        }
226        headers.insert(
227            http::header::CONTENT_TYPE,
228            http::HeaderValue::from_static("application/json"),
229        );
230        self.sent = true;
231        let response = upstream
232            .send(
233                target,
234                WireRequest {
235                    method: http::Method::POST,
236                    path: self.selected.endpoint.path.clone(),
237                    query: self.selected.endpoint.query.clone(),
238                    headers,
239                    body: HttpBody::Bytes(body),
240                },
241            )
242            .await?;
243        let mut read_limits = self.settings.codec;
244        read_limits.max_body_bytes = read_limits.max_body_bytes.min(upstream.limits().read_bytes);
245        self.metadata = Some(WireResponse {
246            status: response.status,
247            headers: response.headers.clone(),
248            body: (),
249        });
250        if !response.status.is_success() {
251            let raw = WireResponse {
252                status: response.status,
253                headers: response.headers,
254                body: codec::read_http_body(response.body, read_limits)
255                    .await
256                    .map_err(super::codec_error)?,
257            };
258            self.rejected = Some(WireResponse {
259                status: raw.status,
260                headers: raw.headers.clone(),
261                body: raw.body.clone(),
262            });
263            return Ok(StreamStart::Rejected(raw));
264        }
265        let accepted = WireResponse {
266            status: response.status,
267            headers: output::headers(response.headers.clone(), self.settings.client_framing),
268            body: (),
269        };
270        self.reader = Some(NativeReader::new(
271            response.body,
272            self.settings.source_framing,
273            read_limits,
274        ));
275        if let Some(value) = response.headers.get(http::header::CONTENT_TYPE) {
276            let valid = value
277                .to_str()
278                .ok()
279                .and_then(|v| v.split(';').next())
280                .is_some_and(|v| match self.settings.source_framing {
281                    SourceFraming::Sse => v.trim().eq_ignore_ascii_case("text/event-stream"),
282                    SourceFraming::JsonArray => v.trim().eq_ignore_ascii_case("application/json"),
283                    SourceFraming::Ndjson => {
284                        matches!(v.trim(), "application/x-ndjson" | "application/ndjson")
285                    }
286                });
287            if !valid {
288                self.failed = true;
289                self.reader = None;
290                return Err(super::invalid(
291                    "native content type disagrees with selected stream framing",
292                ));
293            }
294        }
295        Ok(StreamStart::Streaming(accepted))
296    }
297}