krabka_app_sdk/
connect_client.rs

1//! Minimal Connect client pieces used by the Rust app SDK.
2
3use bytes::{Buf as _, Bytes, BytesMut};
4use futures_util::{StreamExt as _, stream};
5use http_body_util::{BodyExt as _, Full, StreamBody};
6use hyper::{
7    Method, Request, Uri, body::Frame as BodyFrame, client::conn::http2, header::CONTENT_TYPE,
8};
9use hyper_util::rt::{TokioExecutor, TokioIo};
10use serde::Deserialize;
11use tokio::sync::mpsc;
12
13const STREAM_CHANNEL_CAPACITY: usize = 32;
14
15/// Connect envelope codec.
16pub mod envelope {
17    use super::{Bytes, BytesMut, ConnectClientError, EndStream, TrailerBody};
18
19    /// Message frame flag byte.
20    pub const MESSAGE_FLAGS: u8 = 0x00;
21    /// `EndStream` trailer frame flag byte.
22    pub const END_STREAM_FLAGS: u8 = 0x02;
23
24    /// Decoded Connect frame.
25    #[derive(Debug, Clone, PartialEq, Eq)]
26    pub enum Frame {
27        /// Protobuf message bytes.
28        Message(Bytes),
29        /// `EndStream` trailer.
30        EndStream(EndStream),
31    }
32
33    /// Encode one Connect envelope frame.
34    ///
35    /// # Panics
36    ///
37    /// Panics when `message` is larger than the Connect envelope's `u32`
38    /// length field can represent.
39    #[must_use]
40    pub fn encode(flags: u8, message: &[u8]) -> Bytes {
41        let mut framed = BytesMut::with_capacity(5 + message.len());
42        framed.extend_from_slice(&[flags]);
43        framed.extend_from_slice(
44            &u32::try_from(message.len())
45                .expect("message length fits u32")
46                .to_be_bytes(),
47        );
48        framed.extend_from_slice(message);
49        framed.freeze()
50    }
51
52    /// Decode one complete Connect envelope frame.
53    ///
54    /// # Errors
55    ///
56    /// Returns an error for incomplete, malformed, compressed, or unsupported
57    /// Connect frames.
58    pub fn decode_one(input: &[u8]) -> Result<Frame, ConnectClientError> {
59        let Some((frame, _consumed)) = decode_next(input)? else {
60            return Err(ConnectClientError::PartialFrame { needed: 1 });
61        };
62        Ok(frame)
63    }
64
65    /// Decode the first frame from a buffer, returning `None` when more bytes are needed.
66    ///
67    /// # Errors
68    ///
69    /// Returns an error for malformed, compressed, or unsupported Connect
70    /// frames.
71    pub fn decode_next(input: &[u8]) -> Result<Option<(Frame, usize)>, ConnectClientError> {
72        if input.len() < 5 {
73            return Ok(None);
74        }
75        let flags = input[0];
76        let len = u32::from_be_bytes([input[1], input[2], input[3], input[4]]);
77        let len = usize::try_from(len).unwrap_or(usize::MAX);
78        let Some(frame_len) = 5usize.checked_add(len) else {
79            return Ok(None);
80        };
81        if input.len() < frame_len {
82            return Ok(None);
83        }
84        let body = Bytes::copy_from_slice(&input[5..frame_len]);
85        let frame = match flags {
86            MESSAGE_FLAGS => Ok(Frame::Message(body)),
87            END_STREAM_FLAGS => {
88                let trailer = if body.is_empty() {
89                    EndStream { error: None }
90                } else {
91                    let parsed = serde_json::from_slice::<TrailerBody>(&body)?;
92                    EndStream {
93                        error: parsed.error,
94                    }
95                };
96                Ok(Frame::EndStream(trailer))
97            }
98            flags if flags & 0x01 == 0x01 => Err(ConnectClientError::CompressedFrame),
99            other => Err(ConnectClientError::UnknownFrameFlags(other)),
100        }?;
101        Ok(Some((frame, frame_len)))
102    }
103}
104
105/// Connect `EndStream` frame.
106#[derive(Debug, Clone, PartialEq, Eq)]
107pub struct EndStream {
108    /// Optional Connect error from the trailer.
109    pub error: Option<ConnectTrailerError>,
110}
111
112/// Connect trailer error body.
113#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
114pub struct ConnectTrailerError {
115    /// Connect code string.
116    pub code: String,
117    /// Error message.
118    pub message: String,
119}
120
121#[derive(Debug, Deserialize)]
122pub(crate) struct TrailerBody {
123    pub(crate) error: Option<ConnectTrailerError>,
124}
125
126/// Minimal h2c Connect client.
127#[derive(Debug, Clone)]
128pub struct ConnectClient {
129    endpoint: String,
130    bearer: Option<String>,
131}
132
133impl ConnectClient {
134    /// Create a client for an HTTP endpoint.
135    #[must_use]
136    pub fn new(endpoint: impl Into<String>, bearer: Option<String>) -> Self {
137        Self {
138            endpoint: endpoint.into(),
139            bearer,
140        }
141    }
142
143    /// Execute a unary Connect request.
144    ///
145    /// # Errors
146    ///
147    /// Returns an error when the endpoint, connection, HTTP exchange, or
148    /// protobuf encoding/decoding fails, or when Connect returns an error.
149    pub async fn unary<Req, Resp>(
150        &self,
151        path: &str,
152        request: &Req,
153    ) -> Result<Resp, ConnectClientError>
154    where
155        Req: prost::Message,
156        Resp: prost::Message + Default,
157    {
158        let uri = self.uri_for(path)?;
159        let authority = authority_for(&uri)?;
160        let stream = tokio::net::TcpStream::connect(&authority).await?;
161        let io = TokioIo::new(stream);
162        let (mut sender, connection) = http2::handshake(TokioExecutor::new(), io).await?;
163        tokio::spawn(async move {
164            let _ = Box::pin(connection).await;
165        });
166
167        let mut encoded = Vec::with_capacity(prost::Message::encoded_len(request));
168        prost::Message::encode(request, &mut encoded)?;
169        let mut builder = Request::builder()
170            .method(Method::POST)
171            .uri(uri)
172            .header(CONTENT_TYPE, "application/proto");
173        if let Some(bearer) = &self.bearer {
174            builder = builder.header("authorization", format!("Bearer {bearer}"));
175        }
176        let response = sender
177            .send_request(builder.body(Full::new(Bytes::from(encoded)))?)
178            .await?;
179        if !response.status().is_success() {
180            let status = response.status();
181            let body = response.into_body().collect().await?.to_bytes();
182            return Err(response_error(status.as_u16(), &body));
183        }
184        let body = response.into_body().collect().await?.to_bytes();
185        Ok(<Resp as prost::Message>::decode(body)?)
186    }
187
188    /// Execute a Connect client-streaming/server-streaming request with one start message.
189    ///
190    /// # Errors
191    ///
192    /// Returns an error when the endpoint, connection, HTTP exchange, or
193    /// protobuf encoding fails, or when Connect rejects the stream.
194    pub async fn streaming<Req, Resp>(
195        &self,
196        path: &str,
197        request: &Req,
198    ) -> Result<mpsc::Receiver<Result<Resp, ConnectClientError>>, ConnectClientError>
199    where
200        Req: prost::Message,
201        Resp: prost::Message + Default + Send + 'static,
202    {
203        let uri = self.uri_for(path)?;
204        let authority = authority_for(&uri)?;
205        let stream = tokio::net::TcpStream::connect(&authority).await?;
206        let io = TokioIo::new(stream);
207        let (mut sender, connection) = http2::handshake(TokioExecutor::new(), io).await?;
208        tokio::spawn(async move {
209            let _ = Box::pin(connection).await;
210        });
211        let mut encoded = Vec::with_capacity(prost::Message::encoded_len(request));
212        prost::Message::encode(request, &mut encoded)?;
213        let frame = envelope::encode(envelope::MESSAGE_FLAGS, &encoded);
214        let body_stream = stream::iter([Ok::<_, std::convert::Infallible>(BodyFrame::data(frame))])
215            .chain(stream::pending());
216        let mut builder = Request::builder()
217            .method(Method::POST)
218            .uri(uri)
219            .header(CONTENT_TYPE, "application/connect+proto")
220            .header("connect-protocol-version", "1");
221        if let Some(bearer) = &self.bearer {
222            builder = builder.header("authorization", format!("Bearer {bearer}"));
223        }
224        let response = sender
225            .send_request(builder.body(StreamBody::new(body_stream))?)
226            .await?;
227        if !response.status().is_success() {
228            let status = response.status();
229            let body = response.into_body().collect().await?.to_bytes();
230            return Err(response_error(status.as_u16(), &body));
231        }
232
233        let (tx, rx) = mpsc::channel(STREAM_CHANNEL_CAPACITY);
234        tokio::spawn(read_streaming_response::<Resp>(response.into_body(), tx));
235        Ok(rx)
236    }
237
238    fn uri_for(&self, path: &str) -> Result<Uri, ConnectClientError> {
239        let endpoint = self.endpoint.trim_end_matches('/');
240        let uri = format!("{endpoint}{path}")
241            .parse::<Uri>()
242            .map_err(|e| ConnectClientError::InvalidEndpoint(e.to_string()))?;
243        if uri.scheme_str() != Some("http") {
244            return Err(ConnectClientError::InvalidEndpoint(
245                "only plaintext http endpoints are supported".into(),
246            ));
247        }
248        if uri.authority().is_none() {
249            return Err(ConnectClientError::InvalidEndpoint(
250                "missing authority".into(),
251            ));
252        }
253        Ok(uri)
254    }
255}
256
257fn authority_for(uri: &Uri) -> Result<String, ConnectClientError> {
258    uri.authority()
259        .ok_or_else(|| ConnectClientError::InvalidEndpoint("missing authority".into()))
260        .map(ToString::to_string)
261}
262
263fn response_error(status: u16, body: &[u8]) -> ConnectClientError {
264    if let Ok(error) = serde_json::from_slice::<ConnectTrailerError>(body) {
265        return error.into();
266    }
267    ConnectClientError::HttpStatus {
268        status,
269        message: String::from_utf8_lossy(body).into_owned(),
270    }
271}
272
273async fn read_streaming_response<Resp>(
274    mut body: hyper::body::Incoming,
275    tx: mpsc::Sender<Result<Resp, ConnectClientError>>,
276) where
277    Resp: prost::Message + Default,
278{
279    let mut buffer = BytesMut::new();
280    while let Some(frame) = body.frame().await {
281        let frame = match frame {
282            Ok(frame) => frame,
283            Err(error) => {
284                let _ = tx.send(Err(error.into())).await;
285                return;
286            }
287        };
288        let Ok(data) = frame.into_data() else {
289            continue;
290        };
291        buffer.extend_from_slice(&data);
292        loop {
293            let decoded = envelope::decode_next(&buffer);
294            let Some((frame, consumed)) = (match decoded {
295                Ok(frame) => frame,
296                Err(error) => {
297                    let _ = tx.send(Err(error)).await;
298                    return;
299                }
300            }) else {
301                break;
302            };
303            buffer.advance(consumed);
304            match frame {
305                envelope::Frame::Message(message) => {
306                    let decoded = Resp::decode(message).map_err(ConnectClientError::from);
307                    if tx.send(decoded).await.is_err() {
308                        return;
309                    }
310                }
311                envelope::Frame::EndStream(end) => {
312                    if let Some(error) = end.error {
313                        let _ = tx.send(Err(error.into())).await;
314                    }
315                    return;
316                }
317            }
318        }
319    }
320    if !buffer.is_empty() {
321        let _ = tx
322            .send(Err(ConnectClientError::PartialFrame { needed: 1 }))
323            .await;
324    }
325}
326
327/// Errors returned by the minimal Connect client.
328#[derive(Debug, thiserror::Error)]
329#[non_exhaustive]
330pub enum ConnectClientError {
331    /// Endpoint URL was invalid.
332    #[error("invalid endpoint: {0}")]
333    InvalidEndpoint(String),
334    /// TCP or process I/O failed.
335    #[error("io: {0}")]
336    Io(#[from] std::io::Error),
337    /// HTTP/2 connection failed.
338    #[error("hyper: {0}")]
339    Hyper(#[from] hyper::Error),
340    /// HTTP request construction failed.
341    #[error("http: {0}")]
342    Http(#[from] hyper::http::Error),
343    /// Protobuf encode or decode failed.
344    #[error("protobuf: {0}")]
345    Prost(#[from] prost::DecodeError),
346    /// Protobuf encode failed.
347    #[error("protobuf encode: {0}")]
348    ProstEncode(#[from] prost::EncodeError),
349    /// JSON trailer body failed to decode.
350    #[error("json: {0}")]
351    Json(#[from] serde_json::Error),
352    /// HTTP status was not successful.
353    #[error("http status {status}: {message}")]
354    HttpStatus {
355        /// HTTP status code.
356        status: u16,
357        /// Response body text.
358        message: String,
359    },
360    /// Connect trailer contained an error.
361    #[error("connect {code}: {message}")]
362    Connect {
363        /// Connect code string.
364        code: String,
365        /// Connect error message.
366        message: String,
367    },
368    /// Input did not contain a complete envelope frame.
369    #[error("partial frame; need {needed} more bytes")]
370    PartialFrame {
371        /// Number of additional bytes needed.
372        needed: usize,
373    },
374    /// Compressed frames are not supported by this minimal client.
375    #[error("compressed connect frames are unsupported")]
376    CompressedFrame,
377    /// Unknown frame flags were encountered.
378    #[error("unknown connect frame flags {0:#x}")]
379    UnknownFrameFlags(u8),
380}
381
382impl From<ConnectTrailerError> for ConnectClientError {
383    fn from(value: ConnectTrailerError) -> Self {
384        Self::Connect {
385            code: value.code,
386            message: value.message,
387        }
388    }
389}
390
391#[cfg(test)]
392mod tests {
393    use std::{convert::Infallible, time::Duration};
394
395    use bytes::Bytes;
396    use http_body_util::Full;
397    use hyper::{Response, StatusCode, service::service_fn};
398    use hyper_util::rt::{TokioExecutor, TokioIo};
399
400    use super::{
401        ConnectClient, ConnectClientError, STREAM_CHANNEL_CAPACITY,
402        envelope::{self, Frame},
403        response_error,
404    };
405    use crate::error::KrabkaError;
406
407    #[derive(Clone, PartialEq, prost::Message)]
408    struct Empty {}
409
410    async fn serve_once(status: StatusCode, body: Bytes) -> (String, tokio::task::JoinHandle<()>) {
411        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
412            .await
413            .expect("test listener binds");
414        let addr = listener.local_addr().expect("test listener has address");
415        let task = tokio::spawn(async move {
416            let (socket, _) = listener.accept().await.expect("test client connects");
417            let service = service_fn(move |_| {
418                let body = body.clone();
419                async move {
420                    Ok::<_, Infallible>(
421                        Response::builder()
422                            .status(status)
423                            .body(Full::new(body))
424                            .expect("test response builds"),
425                    )
426                }
427            });
428            let _ = hyper::server::conn::http2::Builder::new(TokioExecutor::new())
429                .serve_connection(TokioIo::new(socket), service)
430                .await;
431        });
432        (format!("http://{addr}"), task)
433    }
434
435    #[test]
436    fn envelope_round_trips() {
437        let msg = b"payload";
438        let framed = envelope::encode(0x00, msg);
439        assert_eq!(framed[0], 0);
440        assert_eq!(u32::from_be_bytes(framed[1..5].try_into().unwrap()), 7);
441        assert2::assert!(let Frame::Message(m) = envelope::decode_one(&framed).unwrap());
442        assert_eq!(m.as_ref(), msg);
443    }
444
445    #[test]
446    fn end_stream_trailer_decodes_connect_error() {
447        let framed = envelope::encode(0x02, br#"{"error":{"code":"not_found","message":"x"}}"#);
448        assert2::assert!(let Frame::EndStream(t) = envelope::decode_one(&framed).unwrap());
449        let error = t.error.unwrap();
450        assert_eq!(error.code, "not_found");
451        assert_eq!(error.message, "x");
452    }
453
454    #[test]
455    fn non_success_connect_body_decodes_connect_error() {
456        let error = response_error(
457            400,
458            br#"{"code":"invalid_argument","message":"bad request"}"#,
459        );
460
461        assert2::assert!(let ConnectClientError::Connect { code, message } = error);
462        assert_eq!(code, "invalid_argument");
463        assert_eq!(message, "bad request");
464    }
465
466    #[test]
467    fn non_connect_body_preserves_http_status() {
468        let error = response_error(503, b"upstream unavailable");
469
470        assert2::assert!(let ConnectClientError::HttpStatus { status, message } = error);
471        assert_eq!(status, 503);
472        assert_eq!(message, "upstream unavailable");
473    }
474
475    #[tokio::test]
476    async fn live_http_and_end_stream_errors_reach_the_sdk_taxonomy() {
477        let (endpoint, server) = serve_once(
478            StatusCode::BAD_REQUEST,
479            Bytes::from_static(br#"{"code":"invalid_argument","message":"bad request"}"#),
480        )
481        .await;
482        let error = ConnectClient::new(endpoint, None)
483            .unary::<_, Empty>("/test.Service/Unary", &Empty {})
484            .await
485            .expect_err("non-success Connect response fails");
486        server.abort();
487
488        assert2::assert!(let KrabkaError::InvalidArgument(message) = KrabkaError::from(error));
489        assert_eq!(message, "bad request");
490
491        let trailer = envelope::encode(
492            envelope::END_STREAM_FLAGS,
493            br#"{"error":{"code":"not_found","message":"missing"}}"#,
494        );
495        let (endpoint, server) = serve_once(StatusCode::OK, trailer).await;
496        let mut responses = ConnectClient::new(endpoint, None)
497            .streaming::<_, Empty>("/test.Service/Streaming", &Empty {})
498            .await
499            .expect("HTTP streaming response opens");
500        let error = tokio::time::timeout(Duration::from_secs(2), responses.recv())
501            .await
502            .expect("stream response arrives")
503            .expect("stream emits one error")
504            .expect_err("EndStream carries an error");
505        server.abort();
506
507        assert2::assert!(let KrabkaError::NotFound(message) = KrabkaError::from(error));
508        assert_eq!(message, "missing");
509    }
510
511    #[tokio::test]
512    async fn streaming_receiver_is_bounded_and_truncated_frames_are_errors() {
513        for body in [
514            Bytes::from_static(&[0x00, 0x00]),
515            Bytes::from_static(&[0x00, 0x00, 0x00, 0x00, 0x04, 0x01, 0x02]),
516        ] {
517            let (endpoint, server) = serve_once(StatusCode::OK, body).await;
518            let mut responses = ConnectClient::new(endpoint, None)
519                .streaming::<_, Empty>("/test.Service/Streaming", &Empty {})
520                .await
521                .expect("HTTP streaming response opens");
522
523            assert_eq!(responses.max_capacity(), STREAM_CHANNEL_CAPACITY);
524            let error = tokio::time::timeout(Duration::from_secs(2), responses.recv())
525                .await
526                .expect("stream response arrives")
527                .expect("truncated stream emits an error")
528                .expect_err("truncated stream must not complete cleanly");
529            assert2::assert!(matches!(error, ConnectClientError::PartialFrame { .. }));
530            server.abort();
531        }
532    }
533}
krabka_app_sdk/connect_client.rs

krabka_app_sdk/
connect_client.rs
