fix: Adding a event size cap and stream timeout to prevent compromised BMC…

[rohanharkin]

#138
Adds a configurable per-event byte cap that aborts the SSE stream if a single event exceeds the limit without terminating, preventing a compromised BMC from exhausting memory by sending an endless unterminated event.
Adds a configurable idle timeout that aborts the stream if no event arrives within the window, catching connections that have stalled

This was tested manually using a mock BMC that opens a text/event-stream response and then streams an SSE event that never terminates, against a client that consumes the stream through EventService/Bmc::stream in a memory-capped container.

Two tests added as well
sse_aborts_when_event_exceeds_byte_limit()
sse_aborts_on_idle_timeout()

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yoks]

SSEStream implementation discrards comments, this means for healthy stream which sends heartbeats via comments would not:

• Count towards byte-limit (same issue with unbounded growth)
• Would terminate idle stream even if it sends heartbeats

I do not know how to properly fix it with SSEStream as it not supports it. So either commnet and contribution to upstream lib to fix it. Or rewrite SSE stream ourselves to properly fix it.

[yoks]

It does seems like unrelated change.

[rohanharkin]

before adding this CI was failing as the additional fields went past Clippys struct_field_name_threshold. It was then failing because of the field client within Client

[yoks]

Well, this is valid lint then. As your sse prefixed params can be moved to SSE options struct

[rohanharkin]

Sorry I was confused - I added this in the wrong place

[yoks]

I think we still can address it by making SSL params their separate struct and reuse it in Client and ClientParams.

[rohanharkin]

Ah okay yes, will do that now

[rohanharkin]

Clippy will still fail the CI on the client: ReqwestClient - I can keep the allow in or rename it what do you think is better

[yoks]

You can (and arguably should), reuse it for Client struct as well, as you copy same param there

[yoks]

After playing with Claude with this change, it looks like this would be a better shape here:

#[derive(Clone, Copy)]
enum DelimiterState {
    /// Inside a line.
    Data,
    /// Saw `\r`; a following `\n` belongs to the same terminator.
    Cr,
    /// Saw one complete line terminator.
    Eol,
}

// This should fix the issue with heartbeat and byte counting, as we look at raw data
fn cap_event_bytes(
    bytes: impl Stream<Item = Result<Bytes, reqwest::Error>>,
    limit: usize,
) -> impl Stream<Item = Result<Bytes, SseByteError>> {
    use DelimiterState::{Cr, Data, Eol};
    bytes.scan((0_usize, Data), move |(count, state), chunk| {
        let item = chunk.map_err(SseByteError::Upstream).and_then(|bytes| {
            for &b in bytes.iter() {
                *count += 1;
                *state = match (*state, b) {
                    (Data, b'\r') => Cr,
                    (Data | Cr, b'\n') => Eol,
                    (Cr | Eol, b'\r') => {
                        *count = 0; // blank line: event delimited
                        Cr
                    }
                    (Eol, b'\n') => {
                        *count = 0; // blank line: event delimited
                        Eol
                    }
                    _ => Data,
                };
            }
            if *count > limit {
                Err(SseByteError::EventTooLarge { limit })
            } else {
                Ok(bytes)
            }
        });
        future::ready(Some(item))
    })
}

fn event_to_item<T: DeserializeOwned>(
    event: Result<sse_stream::Sse, sse_stream::Error>,
) -> Option<Result<T, BmcError>> {
    match event {
        Err(err) => Some(Err(map_sse_error(err))),
        Ok(sse) => sse.data.map(|data| {
            serde_path_to_error::deserialize(&mut serde_json::Deserializer::from_str(&data))
                .map_err(BmcError::JsonError)
        }),
    }
}

fn with_idle_timeout<T>(
    events: impl Stream<Item = Result<T, BmcError>>,
    idle: Option<Duration>,
) -> impl Stream<Item = Result<T, BmcError>> {
    match idle {
        // tokio_stream is additional dependency, but it should be fine 
        Some(d) => tokio_stream::StreamExt::timeout(events, d)
            .map(move |item| item.unwrap_or(Err(BmcError::SseIdleTimeout { idle: d })))
            .left_stream(),
        None => events.right_stream(),
    }
}

fn fuse_on_fatal<T>(
    events: impl Stream<Item = Result<T, BmcError>>,
) -> impl Stream<Item = Result<T, BmcError>> {
    events.scan(false, |errored, item| {
        // is_stream_fatal is impl on BmcError which checks if it JSON error or not, so we would
        // terminate on non fata errors
        let stop = std::mem::replace(errored, matches!(&item, Err(e) if e.is_stream_fatal()));
        future::ready((!stop).then_some(item))
    })
}

Ans SSE function can be compacted to this:

async fn sse<T: Send + Sized + for<'de> serde::Deserialize<'de>>(
    &self,
    url: Url,
    credentials: &BmcCredentials,
    custom_headers: &HeaderMap,
) -> Result<BoxTryStream<T, Self::Error>, Self::Error> {
    let request = auth_headers(self.client.get(url), credentials)
        .headers(custom_headers.clone())
        .header(header::ACCEPT, "text/event-stream")
        .timeout(self.sse_total_timeout.unwrap_or(Duration::MAX)); // note what it does default to MAX here

    let response = self.send(request.build()?).await?;

    if response.status() != StatusCode::OK {
        return Err(BmcError::HttpError {
            status: response.status(),
            error_response: None,
        });
    }

    let capped = cap_event_bytes(response.bytes_stream(), self.sse_max_event_bytes);
    let events = sse_stream::SseStream::from_bytes_stream(capped)
        .filter_map(|event| future::ready(event_to_item(event)));

    Ok(Box::pin(fuse_on_fatal(with_idle_timeout(
        events,
        self.sse_idle_timeout,
    ))))
}

[rohanharkin]

SSEStream implementation discrards comments, this means for healthy stream which sends heartbeats via comments would not:

• Count towards byte-limit (same issue with unbounded growth)
• Would terminate idle stream even if it sends heartbeats

I do not know how to properly fix it with SSEStream as it not supports it. So either commnet and contribution to upstream lib to fix it. Or rewrite SSE stream ourselves to properly fix it.

It seems like a lot to rewrite SSE for this change - Probably makes more sense as you suggested to comment it and contribute to SSE

[yoks]

You have SseOptions::default(), why not reuse it here?

[yoks]

This is no longer the case, it is SseOptions now

[yoks]

Ok i see, lint acts on client. Lets rename it client to inner. This is because it is wrapper on the ReqwestClient anyway.

[yoks]

Looks good, but can you address few last nits? After that we can merge it.