taler-rust

sse.rs (7689B)

---

 /*
   This file is part of TALER
   Copyright (C) 2026 Taler Systems SA
 
   TALER is free software; you can redistribute it and/or modify it under the
   terms of the GNU Affero General Public License as published by the Free Software
   Foundation; either version 3, or (at your option) any later version.
 
   TALER is distributed in the hope that it will be useful, but WITHOUT ANY
   WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
   A PARTICULAR PURPOSE.  See the GNU Affero General Public License for more details.
 
   You should have received a copy of the GNU Affero General Public License along with
   TALER; see the file COPYING.  If not, see <http://www.gnu.org/licenses/>
 */
 
 use std::pin::Pin;
 
 use compact_str::CompactString;
 use futures_util::{Stream, StreamExt as _, stream};
 use tokio_util::{
     bytes::Bytes,
     codec::{FramedRead, LinesCodec, LinesCodecError},
     io::StreamReader,
 };
 use tracing::trace;
 
 #[derive(Debug, Default, PartialEq, Eq)]
 pub struct SseMessage {
     pub event: CompactString,
     pub data: String,
 }
 
 type SseStream = dyn Stream<Item = std::result::Result<String, LinesCodecError>> + Send;
 
 /// Server-sent event client
 pub struct SseClient {
     pub last_event_id: Option<CompactString>,
     pub reconnection_time: Option<u64>,
     stream: Pin<Box<SseStream>>,
 }
 
 impl SseClient {
     pub fn new() -> Self {
         Self {
             last_event_id: None,
             reconnection_time: None,
             stream: Box::pin(stream::empty()),
         }
     }
 
     pub fn connect<E: std::error::Error + Send + Sync + 'static>(
         &mut self,
         stream: impl Stream<Item = Result<Bytes, E>> + 'static + Send,
     ) {
         let stream = stream.map(|it| it.map_err(std::io::Error::other));
         let lines = FramedRead::new(StreamReader::new(stream), LinesCodec::new());
         self.stream = Box::pin(lines);
     }
 
     pub async fn next(&mut self) -> Option<Result<SseMessage, LinesCodecError>> {
         // TODO add tests
         let mut event = CompactString::new("message");
         let mut data = None::<String>;
         while let Some(res) = self.stream.next().await {
             let line = match res {
                 Ok(line) => line,
                 Err(e) => return Some(Err(e)),
             };
             // Parse line
             let (field, value): (&str, &str) = if line.is_empty() {
                 if let Some(data) = data.take() {
                     return Some(Ok(SseMessage { event, data }));
                 } else {
                     event = CompactString::new("message");
                     continue;
                 }
             } else if let Some(comment) = line.strip_prefix(':') {
                 trace!(target: "sse", "{comment}");
                 continue;
             } else if let Some((k, v)) = line.split_once(':') {
                 (k, v.strip_prefix(' ').unwrap_or(v))
             } else {
                 (&line, "")
             };
 
             // Process field
             match field {
                 "event" => event = CompactString::new(value),
                 "data" => match data.as_mut() {
                     Some(data) => {
                         data.push('\n');
                         data.push_str(value);
                     }
                     None => data = Some(value.to_string()),
                 },
                 "id" => {
                     if !value.contains('\0') {
                         self.last_event_id = Some(CompactString::new(value))
                     }
                 }
                 "retry" => {
                     if value.as_bytes().iter().all(|c| c.is_ascii_digit())
                         && let Ok(int) = value.parse::<u64>()
                     {
                         self.reconnection_time = Some(int)
                     }
                 }
                 _ => continue,
             }
         }
         None
     }
 }
 
 impl Default for SseClient {
     fn default() -> Self {
         Self::new()
     }
 }
 
 #[tokio::test]
 pub async fn protocol() {
     pub async fn test(
         stream: &'static str,
         result: &[(&str, &str)],
         last_event_id: Option<&str>,
         reconnection_time: Option<u64>,
     ) {
         let stream = stream::iter(
             stream
                 .as_bytes()
                 .chunks(12)
                 .map(|chunk| std::io::Result::Ok(Bytes::from_static(chunk))),
         );
         let mut client = SseClient::new();
         client.connect(stream);
         let mut res = Vec::new();
         while let Some(msg) = client.next().await {
             res.push(msg.unwrap());
         }
         assert_eq!(
             result,
             &res.iter()
                 .map(|m| (m.event.as_ref(), m.data.as_ref()))
                 .collect::<Vec<_>>()
         );
         assert_eq!(client.last_event_id.as_deref(), last_event_id);
         assert_eq!(client.reconnection_time, reconnection_time);
     }
 
     macro_rules! check {
         // Handle multiple tuples + optional id and retry
         ($stream:expr $(, ($e:expr, $d:expr))* $(, id: $id:expr)? $(, retry: $retry:expr)?) => {
             test(
                 $stream,
                 &[ $( ($e, $d) ),* ],
                 { let mut _id = None; $( _id = Some($id); )? _id },
                 { let mut _r = None; $( _r = Some($retry); )? _r }
             ).await
         };
     }
 
     check!("data\n\n", ("message", ""));
     check!("data:key:value\n\n", ("message", "key:value"));
     check!("data: value\n\n", ("message", "value"));
     check!("data:first\ndata:second\n\n", ("message", "first\nsecond"));
 
     check!(
         "event:first\nevent:second\ndata:test\n\n",
         ("second", "test")
     );
 
     check!("data:test\r\n\r\n", ("message", "test"));
     check!("data:test\r\r");
     check!("data:line1\r\ndata:line2\n\n", ("message", "line1\nline2"));
     check!("data:test\n");
     check!("data:test\n\n\n", ("message", "test"));
     check!("data:\ndata:\n\n", ("message", "\n"));
     check!("data:\ndata:content\ndata:\n\n", ("message", "\ncontent\n"));
     check!("data:Hello 世界 🌍\n\n", ("message", "Hello 世界 🌍"));
     check!("data:   \n\n", ("message", "  "));
     check!("id:123\ndata:test\n\n", ("message", "test"), id: "123");
     check!("id:first\nid:second\ndata:test\n\n", ("message", "test"), id: "second");
     check!("id:first\nid:second\nid:\ndata:test\n\n", ("message", "test"), id: "");
     check!("id:\ndata:test\n\n", ("message", "test"), id: "");
     check!(
         "id:test:123\ndata:test\n\n",
         ("message", "test"),
         id: "test:123"
     );
     check!("id:123\x00456\ndata:test\n\n", ("message", "test"));
     check!("id:123\n\n", id: "123");
     check!(
         "id:123\ndata:first\n\ndata:second\n\n",
         ("message", "first"), ("message", "second"),
         id: "123"
     );
     check!("event:customEvent\ndata:test\n\n", ("customEvent", "test"));
     check!("event:\ndata:test\n\n", ("", "test"));
     check!("event:my event\ndata:test\n\n", ("my event", "test"));
 
     check!("retry:3000\n\n", retry: 3000);
     check!("retry:0\n\n", retry: 0);
     check!("retry:abc\n\n");
     check!("retry:-1000\n\n");
     check!("retry:1000.5\n\n");
     check!("retry:1000\nretry:2000\n\n", retry: 2000);
 
     check!(":comment\n\n");
     check!(": comment\n\n");
     check!(":comment\ndata:test\n\n", ("message", "test"));
     check!("unknown:value\ndata:test\n\n", ("message", "test"));
     check!("datta:test\n\n");
     check!(" data:test\n\n");
     check!("data :test\n\n");
     check!("data:\tvalue\n\n", ("message", "\tvalue"));
     check!("id:123\n\n", id: "123");
     check!("event:test\n\n");
     check!("event:test\n\ndata:value\n\n", ("message", "value"));
 }