sozu-lib 2.3.0

//! Helpers shared between the H1 and H2 multiplexers.
//!
//! Both protocol state machines have converged on a handful of small
//! routines that operate on SocketHandler and StreamState in the same way.
//! Keeping them here prevents drift between the two write paths.
//!
//! That co-location is load-bearing for TLS close_notify ordering.
//! `ConnectionH1::close` and `ConnectionH2::close` both reach
//! `drain_tls_close_notify`, which emits close_notify once and then keeps
//! flushing rustls until the socket stops wanting a write, before the
//! underlying TCP socket is shut down. A per-protocol copy is how one of
//! the two drifts back to a single write attempt, which leaves a partial
//! TLS record in flight whenever the kernel send buffer is full — the
//! client reports that as a TLS decode error or an unexpected eof. The
//! drain loop and its round cap are documented on the helper itself.

use std::time::Instant;

use kawa::AsBuffer;

use crate::{
    protocol::http::editor::{HeaderEditMode, HeaderEditSnapshot},
    socket::{SocketHandler, SocketResult},
};

use super::Stream;

/// Most request bytes a lingering close reads and discards before it closes
/// anyway (RFC 9112 §9.6).
///
/// The drain exists so that the client's in-flight body does not make the
/// kernel answer the close with a reset, which destroys any part of the
/// response the client has not read yet. 4 MiB is the default ceiling of
/// Linux's send-buffer autotuning (`net.ipv4.tcp_wmem`), so it covers what a
/// client can already have queued when it sees the FIN, while bounding the
/// work one connection can make sozu do after its response.
pub(super) const LINGER_MAX_BYTES: usize = 4 * 1024 * 1024;

/// Wall-clock bound of the lingering close that follows a final GOAWAY
/// carrying an error code (a protocol violation, a flood). Never longer than
/// the graceful bound: the deadline is the smaller of this and the listener's
/// `request_timeout`, fixed when the linger is decided and not moved by what
/// the client sends.
///
/// That linger exists only so the GOAWAY reaches the client: a close with the
/// client's frames unread makes the kernel answer with a reset, which drops
/// the GOAWAY still queued in the socket (sozu-proxy/sozu#1861). By the time
/// the drain starts the GOAWAY and `close_notify` are already handed to the
/// kernel and the write side is shut down, so the drain only has to keep the
/// receive queue empty while the client reads them: a few round trips, which
/// one second covers even on an intercontinental path (~300 ms RTT), while
/// costing a misbehaving peer's slot and buffer at most a tenth of the
/// default 10 s `request_timeout`.
pub(super) const ERROR_GOAWAY_LINGER_TIMEOUT: std::time::Duration =
    std::time::Duration::from_secs(1);

/// Most bytes the lingering close after an error GOAWAY reads and discards
/// before it closes anyway.
///
/// Sōzu sets no `SO_RCVBUF` on a frontend socket, so a connection starts with
/// Linux's default receive buffer, `net.ipv4.tcp_rmem[1]` = 128 KiB, and
/// autotuning only grows it while the reader keeps up — which stops at the
/// final GOAWAY, when reading does. 256 KiB is that default buffer, what
/// can already be queued when the drain starts, plus as much again arriving
/// before the client reads the GOAWAY. A peer that sends more is still
/// sending regardless of the GOAWAY, and the reset it then gets is its own.
pub(super) const ERROR_GOAWAY_LINGER_MAX_BYTES: usize = 256 * 1024;

const _: () = assert!(
    ERROR_GOAWAY_LINGER_MAX_BYTES <= LINGER_MAX_BYTES,
    "the error-GOAWAY linger reads no more than a graceful one"
);

/// Size of one lingering read. A stack buffer: the drain allocates nothing.
pub(super) const LINGER_READ_CHUNK: usize = 16 * 1024;

/// Lingering close of a frontend that closes while its client may still be
/// sending: an H1 connection whose request was not received whole when its
/// response ended (RFC 9112 §9.6), or an H2 connection after its final
/// GOAWAY (RFC 9113 §6.8 leaves the peer free to send until it reads it) —
/// under [`ERROR_GOAWAY_LINGER_TIMEOUT`] and [`ERROR_GOAWAY_LINGER_MAX_BYTES`]
/// when that GOAWAY carries an error code.
///
/// Closing a socket whose receive queue still holds data, or that receives
/// data after the close, makes the kernel send a reset, and a reset discards
/// whatever part of the response is still queued. So the close is staged:
/// the write side is shut down once the response and any TLS `close_notify`
/// are flushed, then the rest of the request is read and dropped until the
/// client closes, [`LINGER_MAX_BYTES`] are read, or the deadline passes.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Linger {
    /// The connection closes without draining.
    Off,
    /// The close is decided but something is still to flush (the H1
    /// response's TLS `close_notify`, or the H2 final GOAWAY's): once it is
    /// out, shut the write side down and start draining.
    Pending { deadline: Instant },
    /// The write side is shut down; the rest of the request is being read
    /// and dropped.
    Draining { deadline: Instant, remaining: usize },
}

/// What one [`drain_discard`] pass decided.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum LingerRead {
    /// The socket has nothing more for now: wait for the next READABLE.
    Wait,
    /// The client closed, the read failed, or the byte budget is spent.
    Close,
}

/// Read and drop what a lingering frontend's client still sends, from the raw
/// socket — also under TLS: after `close_notify` nothing the client sends is
/// decrypted, only dropped. Reads into a stack buffer until `WouldBlock`
/// ([`LingerRead::Wait`]), the client's EOF, an error, or `remaining` reaches
/// zero ([`LingerRead::Close`]); `remaining` is decremented by every byte
/// read, and `on_read` is told each read's size.
pub(super) fn drain_discard(
    socket: &mio::net::TcpStream,
    remaining: &mut usize,
    mut on_read: impl FnMut(usize),
) -> LingerRead {
    use std::io::{ErrorKind, Read};
    let mut buf = [0u8; LINGER_READ_CHUNK];
    let mut socket = socket;
    loop {
        if *remaining == 0 {
            return LingerRead::Close;
        }
        let len = (*remaining).min(buf.len());
        match socket.read(&mut buf[..len]) {
            Ok(0) => return LingerRead::Close,
            Ok(size) => {
                debug_assert!(size <= *remaining, "a read never exceeds its buffer");
                *remaining -= size;
                on_read(size);
            }
            Err(e) if e.kind() == ErrorKind::WouldBlock => return LingerRead::Wait,
            Err(e) if e.kind() == ErrorKind::Interrupted => {}
            Err(_) => return LingerRead::Close,
        }
    }
}

/// Decision returned by [`end_stream_decision`] for the server-side end-of-stream path.
///
/// Both H1 and H2 must take the same action when a server-side stream ends, based
/// on whether the request was already partially consumed and whether the backend
/// has produced a forwardable response. Centralising the decision here keeps the
/// two protocols from drifting (e.g. H1 used to send 503 where H2 sent 502 for
/// "backend closed before any response").
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum EndStreamAction {
    /// Backend already produced a fully terminated response — finish forwarding
    /// it and mark the stream `Unlinked`.
    ForwardTerminated,
    /// Backend produced a partial response that the caller cannot finish (no
    /// keep-alive AND no Content-Length); mark the response terminated so the
    /// converter emits a final DATA frame with END_STREAM.
    CloseDelimited,
    /// Backend produced a partial response and went away mid-response: the
    /// connection is keep-alive (or otherwise expected to terminate cleanly),
    /// or the response already failed (Error phase, e.g. a body truncated by
    /// the backend close) after some of it reached the client. The caller must
    /// forcefully terminate with an internal error: RST_STREAM on H2, a
    /// connection close on H1.
    ForwardUnterminated,
    /// No response is available and the request was already partially consumed,
    /// so retrying is unsafe — send the given default status (502 Bad Gateway).
    /// Also chosen for a request an H2 backend encoded and did not send
    /// ([`Stream::front_bound_to_backend`]): it cannot be sent elsewhere.
    SendDefault(u16),
    /// No response is available and the request is untouched, so the caller may
    /// link the stream to a fresh backend and retry.
    Reconnect,
    /// No response is available and the request WAS written, but it went onto a
    /// reused keep-alive upstream whose peer had already closed it, and the
    /// whole request is still buffered: the caller may queue the captured bytes
    /// back onto the front kawa and link the stream to a fresh backend.
    ///
    /// Distinct from [`Self::Reconnect`] because the caller has extra work to
    /// do — `Stream::queue_upstream_replay` — before the stream can be linked.
    /// See `Stream::can_replay_on_fresh_upstream` for the boundary and its
    /// sources (sozu-proxy/sozu#1442).
    ReplayOnFreshBackend,
}

/// Compute the canonical end-of-stream decision for a server-side stream.
///
/// Mirrors the H1 and H2 server end-of-stream logic so both protocols agree on
/// the outcome (in particular, "backend closed without response" is normalised
/// to **502 Bad Gateway**).
pub(super) fn end_stream_decision(stream: &Stream) -> EndStreamAction {
    if stream.back.is_main_phase() {
        if stream.back.is_terminated() {
            EndStreamAction::ForwardTerminated
        } else if !stream.context.keep_alive_backend {
            EndStreamAction::CloseDelimited
        } else {
            EndStreamAction::ForwardUnterminated
        }
    } else if stream.back.is_error() && stream.back.consumed {
        // The response failed after part of it was already handed to the
        // client (`Kawa::consume` sets `consumed` from the H1 write path and
        // the H2 converter). A default answer would clear the kawa and follow
        // those bytes as a second response: an H1 client would read it as body
        // bytes, an H2 client as DATA of the stream it already has HEADERS
        // for. End the stream in error instead (sozu-proxy/sozu#1633).
        EndStreamAction::ForwardUnterminated
    } else if stream.front.consumed {
        // The request reached the upstream. That is only replayable in the
        // stale-keep-alive case: no response byte was received, the method is
        // idempotent so re-issuing is permitted (RFC 9110 §9.2.2), and the
        // bytes were captured for exactly this. Every other shape stays 502.
        if stream.can_replay_on_fresh_upstream() {
            EndStreamAction::ReplayOnFreshBackend
        } else {
            EndStreamAction::SendDefault(502)
        }
    } else if stream.front_bound_to_backend {
        // Nothing reached the backend, but an H2 backend encoded the request:
        // `front.out` holds frames bound to that connection's stream id and
        // HPACK encoder, and the blocks they came from are gone. Linking it
        // to another backend would send those frames there, where they decode
        // against another table (sozu-proxy/sozu#1632); it cannot be encoded
        // again either. An H1 backend's output never sets the flag, so its
        // untouched requests keep `Reconnect`.
        debug_assert!(
            !stream.front.out.is_empty(),
            "a request bound to a backend holds that backend's unsent frames"
        );
        EndStreamAction::SendDefault(502)
    } else {
        EndStreamAction::Reconnect
    }
}

/// Drain rustls's pending TLS output before the underlying TCP socket is
/// shut down. A single write attempt is insufficient when the kernel send
/// buffer is full (common during large-response transfers): without this
/// drain loop, a partial TLS record can be left in-flight, producing
/// "TLS decode error / unexpected eof" on the client side.
///
/// Emits close_notify the first time it is called for a given session
/// (tracked by the caller's `close_notify_sent` flag) and then attempts
/// up to `MAX_DRAIN_ROUNDS` empty-vectored writes to flush rustls.
///
/// Returns `(tls_still_pending, drain_rounds)` so the caller can log or
/// react to an incomplete drain.
pub(super) fn drain_tls_close_notify<S: SocketHandler>(
    socket: &mut S,
    close_notify_sent: &mut bool,
) -> (bool, u32) {
    const MAX_DRAIN_ROUNDS: u32 = 16;
    if !*close_notify_sent {
        socket.socket_close();
        *close_notify_sent = true;
    }
    let mut drain_rounds = 0;
    while socket.socket_wants_write() && drain_rounds < MAX_DRAIN_ROUNDS {
        let (_size, status) = socket.socket_write_vectored(&[]);
        drain_rounds += 1;
        match status {
            SocketResult::WouldBlock | SocketResult::Error | SocketResult::Closed => break,
            SocketResult::Continue => {}
        }
    }
    (socket.socket_wants_write(), drain_rounds)
}

/// Apply per-frontend response-side header edits to a response kawa
/// just before [`kawa::Kawa::prepare`] runs. Mirrors the request-side
/// pass that `Router::route_from_request` runs against the front kawa
/// at routing time. Each edit's [`HeaderEditMode`] selects the action:
///
/// - [`HeaderEditMode::Append`] with a non-empty `val` → append the
///   header before the `end_header` flag block.
/// - [`HeaderEditMode::Append`] with an empty `val` → fall back to the
///   legacy delete encoding (HAProxy `del-header` parity).
/// - [`HeaderEditMode::SetIfAbsent`] → append the header only when no
///   non-elided header with the same name already exists on the
///   response (RFC 6797 §6.1 single-`Strict-Transport-Security`
///   requirement, generalised to any operator-supplied
///   set-only-when-missing edit). The presence check runs against the
///   pre-edit response state, so a `SetIfAbsent` does NOT race against
///   a sibling `Set` (or empty-`val` Append) issued in the same batch.
/// - [`HeaderEditMode::Set`] → remove every existing header with the
///   matching name from `kawa.blocks`, then append the new value.
///   Used by typed policies that want to override a backend-supplied
///   header unconditionally (e.g. HSTS with
///   `force_replace_backend = true`).
///
/// Both H1 and H2 emission paths funnel through here so the on-wire
/// header set stays consistent.
pub(super) fn apply_response_header_edits<T: AsBuffer>(
    kawa: &mut kawa::Kawa<T>,
    edits: &[HeaderEditSnapshot],
) {
    use kawa::{Block, Pair, Store};

    if edits.is_empty() {
        return;
    }

    // Snapshot the lowercased header names already present on the
    // response BEFORE any edit runs. Used by `SetIfAbsent` to honour
    // upstream-supplied headers (e.g. backend already sent its own
    // `Strict-Transport-Security`) without racing against a sibling
    // `Set` edit in the same batch.
    //
    // Skip the scan when the batch contains no `SetIfAbsent` entry —
    // `Append` and `Set` never consult `existing_keys`, so the
    // common operator-defined response-header path (Append-only)
    // pays zero allocation overhead.
    let needs_existing_snapshot = edits
        .iter()
        .any(|e| matches!(e.mode, HeaderEditMode::SetIfAbsent));
    let existing_keys: Vec<Vec<u8>> = if needs_existing_snapshot {
        let buf = kawa.storage.buffer();
        kawa.blocks
            .iter()
            .filter_map(|block| {
                if let Block::Header(Pair { key, val: _ }) = block {
                    // Skip elided headers — see `Store::Empty` note below
                    // in the retain pass for the panic rationale.
                    if matches!(key, Store::Empty) {
                        None
                    } else {
                        Some(key.data(buf).iter().map(u8::to_ascii_lowercase).collect())
                    }
                } else {
                    None
                }
            })
            .collect()
    } else {
        Vec::new()
    };

    let keys_to_drop: Vec<Vec<u8>> = edits
        .iter()
        .filter(|e| {
            // Explicit `Set` (delete-then-insert in one entry) OR legacy
            // empty-`val` Append (delete encoding preserved for callers
            // still using the implicit shape).
            matches!(e.mode, HeaderEditMode::Set)
                || (matches!(e.mode, HeaderEditMode::Append) && e.val.is_empty())
        })
        .map(|e| e.key.iter().map(u8::to_ascii_lowercase).collect())
        .collect();
    if !keys_to_drop.is_empty() {
        let buf = kawa.storage.buffer();
        kawa.blocks.retain(|block| {
            if let Block::Header(Pair { key, val: _ }) = block {
                // Skip elided headers — kawa's earlier passes (HPACK
                // decoder, H1 header parser) rewrite headers they want to
                // suppress to `Pair { key: Store::Empty, val: Store::Empty }`
                // rather than removing them in-place. Calling `.data()` on
                // `Store::Empty` panics in `kawa-0.6.8/src/storage/repr.rs`
                // (the impl unwraps a `None` slice). Pinning this guard
                // explicitly until kawa changes its policy.
                if matches!(key, Store::Empty) {
                    return true;
                }
                // Both `keys_to_drop` and `key_lower` are already
                // ASCII-lowercased, so a byte-equality compare is enough;
                // a second `compare_no_case` pass would just refold
                // bytes that are already canonical.
                let key_lower: Vec<u8> = key.data(buf).iter().map(u8::to_ascii_lowercase).collect();
                !keys_to_drop
                    .iter()
                    .any(|k| k.as_slice() == key_lower.as_slice())
            } else {
                true
            }
        });
    }

    let end_header_idx = end_of_headers_index(kawa);

    let mut to_insert: Vec<Block> = Vec::new();
    for edit in edits {
        match edit.mode {
            // Empty-`val` Append is the legacy delete shape — the
            // retain pass above already removed any matching headers.
            HeaderEditMode::Append if edit.val.is_empty() => continue,
            HeaderEditMode::Append => {}
            HeaderEditMode::SetIfAbsent => {
                let key_lower: Vec<u8> = edit.key.iter().map(u8::to_ascii_lowercase).collect();
                if existing_keys
                    .iter()
                    .any(|k| k.as_slice() == key_lower.as_slice())
                {
                    // Upstream already supplied this header — honour
                    // the operator's "set only when missing" intent and
                    // skip the insert (RFC 6797 §6.1).
                    continue;
                }
            }
            // `Set` is delete-then-insert: the retain pass above has
            // already dropped any pre-existing header with the same
            // name, and this branch falls through to the insert below
            // unconditionally (no presence check).
            HeaderEditMode::Set => {}
        }
        to_insert.push(Block::Header(Pair {
            key: Store::from_slice(&edit.key),
            val: Store::from_slice(&edit.val),
        }));
    }
    if !to_insert.is_empty() {
        // Insert each edit one position deeper than the previous so we
        // preserve operator order at the wire. `insert_at` points at the
        // existing `Block::Flags { end_header: true }` (or
        // `kawa.blocks.len()` when absent); each successive insert pushes
        // that anchor right by one — the `+ offset` keeps the flag block
        // (and any post-headers blocks) trailing the new entries even
        // when elided headers sit between them.
        let insert_at = end_header_idx.unwrap_or(kawa.blocks.len());
        for (offset, block) in to_insert.into_iter().enumerate() {
            kawa.blocks.insert(insert_at + offset, block);
        }
    }
}

/// Locate the `Block::Flags { end_header: true }` block in `kawa.blocks`.
///
/// Both the request-side rewrite helper (`mux::router`) and the
/// response-side edit helper above use this anchor as the insertion
/// point for new header blocks: anything after the end-of-headers flag
/// is body / chunks / trailers and must NOT receive a Header block. A
/// shared helper keeps the predicate canonical so a future change to
/// `kawa::Flags` (e.g. adding fields) only needs touching once.
pub(super) fn end_of_headers_index<T: AsBuffer>(kawa: &kawa::Kawa<T>) -> Option<usize> {
    kawa.blocks.iter().position(|b| {
        matches!(
            b,
            kawa::Block::Flags(kawa::Flags {
                end_header: true,
                ..
            })
        )
    })
}

#[cfg(test)]
mod tests {
    use super::apply_response_header_edits;
    use crate::protocol::http::editor::{HeaderEditMode, HeaderEditSnapshot};
    use kawa::{
        AsBuffer, Block, Buffer, Flags, Kawa, Kind, Pair, SliceBuffer, StatusLine, Store, Version,
    };

    fn make_kawa<'a>(buf: &'a mut [u8]) -> Kawa<SliceBuffer<'a>> {
        Kawa::new(Kind::Response, Buffer::new(SliceBuffer(buf)))
    }

    fn pretty_blocks<T: AsBuffer>(kawa: &Kawa<T>) -> Vec<(Vec<u8>, Vec<u8>)> {
        let buf = kawa.storage.buffer();
        kawa.blocks
            .iter()
            .filter_map(|b| {
                if let Block::Header(Pair { key, val }) = b {
                    if matches!(key, Store::Empty) {
                        Some((b"<elided>".to_vec(), b"".to_vec()))
                    } else {
                        Some((key.data(buf).to_vec(), val.data(buf).to_vec()))
                    }
                } else {
                    None
                }
            })
            .collect()
    }

    /// Insertion offset arithmetic: when the response already carries
    /// elided headers (kawa marks suppressed headers as `Store::Empty`
    /// instead of removing them), inserting two new headers must keep
    /// them contiguous AND keep the `end_header` flag block trailing.
    #[test]
    fn test_response_edit_offsets_with_elided_headers() {
        let mut buf = vec![0u8; 4096];
        let mut kawa = make_kawa(&mut buf);
        kawa.detached.status_line = StatusLine::Response {
            version: Version::V11,
            code: 200,
            status: Store::Static(b"200"),
            reason: Store::Static(b"OK"),
        };
        // Pre-populate with a real header, an elided header (suppressed
        // by an earlier kawa pass), another real header, and finally the
        // end-of-headers flag block.
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Static(b"server"),
            val: Store::Static(b"sozu"),
        }));
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Empty,
            val: Store::Empty,
        }));
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Static(b"content-length"),
            val: Store::Static(b"3"),
        }));
        kawa.blocks.push_back(Block::Flags(Flags {
            end_body: false,
            end_chunk: false,
            end_header: true,
            end_stream: false,
        }));

        apply_response_header_edits(
            &mut kawa,
            &[
                HeaderEditSnapshot {
                    key: b"x-frame-options".to_vec(),
                    val: b"DENY".to_vec(),
                    mode: HeaderEditMode::Append,
                },
                HeaderEditSnapshot {
                    key: b"x-content-type-options".to_vec(),
                    val: b"nosniff".to_vec(),
                    mode: HeaderEditMode::Append,
                },
            ],
        );

        let names: Vec<Vec<u8>> = pretty_blocks(&kawa).into_iter().map(|(k, _)| k).collect();
        assert_eq!(
            names,
            vec![
                b"server".to_vec(),
                b"<elided>".to_vec(),
                b"content-length".to_vec(),
                b"x-frame-options".to_vec(),
                b"x-content-type-options".to_vec(),
            ],
            "new headers must land in order, between the last real header and the end-of-headers flag"
        );

        // The end-of-headers flag must still be the final block.
        assert!(
            matches!(
                kawa.blocks.back(),
                Some(Block::Flags(Flags {
                    end_header: true,
                    ..
                }))
            ),
            "end_header flag must remain the last block"
        );
    }

    /// Delete-by-name with elided headers around the target: the elided
    /// blocks must survive the retain pass (their `key.data()` would
    /// panic) and the targeted header must be dropped.
    #[test]
    fn test_response_edit_delete_skips_elided_blocks() {
        let mut buf = vec![0u8; 4096];
        let mut kawa = make_kawa(&mut buf);
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Empty,
            val: Store::Empty,
        }));
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Static(b"x-internal"),
            val: Store::Static(b"secret"),
        }));
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Empty,
            val: Store::Empty,
        }));
        kawa.blocks.push_back(Block::Flags(Flags {
            end_body: false,
            end_chunk: false,
            end_header: true,
            end_stream: false,
        }));

        apply_response_header_edits(
            &mut kawa,
            &[HeaderEditSnapshot {
                key: b"X-Internal".to_vec(),
                val: Vec::new(),
                mode: HeaderEditMode::Append,
            }],
        );

        let names: Vec<Vec<u8>> = pretty_blocks(&kawa).into_iter().map(|(k, _)| k).collect();
        assert_eq!(
            names,
            vec![b"<elided>".to_vec(), b"<elided>".to_vec()],
            "x-internal must be dropped; elided blocks must survive the retain pass"
        );
    }

    /// Delete-then-set: the helper applies all deletes first, then
    /// inserts the non-empty edits. Two edits with the same key
    /// (one delete, one set) effectively replace the value.
    #[test]
    fn test_response_edit_delete_then_set_replaces() {
        let mut buf = vec![0u8; 4096];
        let mut kawa = make_kawa(&mut buf);
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Static(b"cache-control"),
            val: Store::Static(b"public"),
        }));
        kawa.blocks.push_back(Block::Flags(Flags {
            end_body: false,
            end_chunk: false,
            end_header: true,
            end_stream: false,
        }));

        apply_response_header_edits(
            &mut kawa,
            &[
                HeaderEditSnapshot {
                    key: b"Cache-Control".to_vec(),
                    val: Vec::new(),
                    mode: HeaderEditMode::Append,
                },
                HeaderEditSnapshot {
                    key: b"Cache-Control".to_vec(),
                    val: b"no-store".to_vec(),
                    mode: HeaderEditMode::Append,
                },
            ],
        );

        let pairs = pretty_blocks(&kawa);
        assert_eq!(pairs.len(), 1, "only the replacement header must remain");
        assert_eq!(pairs[0].0, b"Cache-Control");
        assert_eq!(pairs[0].1, b"no-store");
    }

    /// `SetIfAbsent` MUST NOT add a second header when the upstream
    /// response already carries one with the same name. This is the
    /// HSTS-on-listener-default contract: if the backend already
    /// emitted its own `Strict-Transport-Security`, the operator's
    /// listener-default policy steps aside (RFC 6797 §6.1).
    #[test]
    fn test_response_edit_set_if_absent_skips_when_present() {
        let mut buf = vec![0u8; 4096];
        let mut kawa = make_kawa(&mut buf);
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Static(b"strict-transport-security"),
            val: Store::Static(b"max-age=12345"),
        }));
        kawa.blocks.push_back(Block::Flags(Flags {
            end_body: false,
            end_chunk: false,
            end_header: true,
            end_stream: false,
        }));

        apply_response_header_edits(
            &mut kawa,
            &[HeaderEditSnapshot {
                key: b"strict-transport-security".to_vec(),
                val: b"max-age=31536000".to_vec(),
                mode: HeaderEditMode::SetIfAbsent,
            }],
        );

        let pairs = pretty_blocks(&kawa);
        assert_eq!(
            pairs.len(),
            1,
            "SetIfAbsent must not duplicate an existing header"
        );
        assert_eq!(pairs[0].0, b"strict-transport-security");
        assert_eq!(
            pairs[0].1, b"max-age=12345",
            "the upstream-supplied STS value must be preserved unchanged"
        );
    }

    /// `SetIfAbsent` MUST insert the new header when no upstream
    /// header with the same name is present. The new entry lands
    /// before the end-of-headers flag (same insertion point as a
    /// regular `Append`).
    #[test]
    fn test_response_edit_set_if_absent_inserts_when_absent() {
        let mut buf = vec![0u8; 4096];
        let mut kawa = make_kawa(&mut buf);
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Static(b"server"),
            val: Store::Static(b"sozu"),
        }));
        kawa.blocks.push_back(Block::Flags(Flags {
            end_body: false,
            end_chunk: false,
            end_header: true,
            end_stream: false,
        }));

        apply_response_header_edits(
            &mut kawa,
            &[HeaderEditSnapshot {
                key: b"strict-transport-security".to_vec(),
                val: b"max-age=31536000".to_vec(),
                mode: HeaderEditMode::SetIfAbsent,
            }],
        );

        let pairs = pretty_blocks(&kawa);
        assert_eq!(
            pairs.len(),
            2,
            "SetIfAbsent must insert exactly one new header when absent"
        );
        assert_eq!(pairs[0].0, b"server");
        assert_eq!(pairs[1].0, b"strict-transport-security");
        assert_eq!(pairs[1].1, b"max-age=31536000");

        // The end-of-headers flag must remain the final block.
        assert!(
            matches!(
                kawa.blocks.back(),
                Some(Block::Flags(Flags {
                    end_header: true,
                    ..
                }))
            ),
            "end_header flag must remain the last block"
        );
    }

    #[test]
    fn test_response_edit_set_replaces_existing_header() {
        // `HeaderEditMode::Set` is delete-then-insert: a pre-existing
        // header with the matching name (case-insensitive) must be
        // dropped, and the new value inserted exactly once.
        let mut buf = vec![0u8; 4096];
        let mut kawa = make_kawa(&mut buf);
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Static(b"Strict-Transport-Security"),
            val: Store::Static(b"max-age=300"),
        }));
        kawa.blocks.push_back(Block::Flags(Flags {
            end_body: false,
            end_chunk: false,
            end_header: true,
            end_stream: false,
        }));

        apply_response_header_edits(
            &mut kawa,
            &[HeaderEditSnapshot {
                key: b"strict-transport-security".to_vec(),
                val: b"max-age=31536000; includeSubDomains".to_vec(),
                mode: HeaderEditMode::Set,
            }],
        );

        let pairs = pretty_blocks(&kawa);
        assert_eq!(
            pairs.len(),
            1,
            "Set must drop the pre-existing header and insert exactly one"
        );
        assert_eq!(pairs[0].0, b"strict-transport-security");
        assert_eq!(pairs[0].1, b"max-age=31536000; includeSubDomains");
    }

    #[test]
    fn test_response_edit_set_inserts_when_absent() {
        // `Set` with no pre-existing header must still insert the new
        // value (the retain pass is a no-op when nothing matches).
        let mut buf = vec![0u8; 4096];
        let mut kawa = make_kawa(&mut buf);
        kawa.blocks.push_back(Block::Header(Pair {
            key: Store::Static(b"server"),
            val: Store::Static(b"sozu"),
        }));
        kawa.blocks.push_back(Block::Flags(Flags {
            end_body: false,
            end_chunk: false,
            end_header: true,
            end_stream: false,
        }));

        apply_response_header_edits(
            &mut kawa,
            &[HeaderEditSnapshot {
                key: b"strict-transport-security".to_vec(),
                val: b"max-age=31536000".to_vec(),
                mode: HeaderEditMode::Set,
            }],
        );

        let pairs = pretty_blocks(&kawa);
        assert_eq!(pairs.len(), 2);
        assert_eq!(pairs[0].0, b"server");
        assert_eq!(pairs[1].0, b"strict-transport-security");
        assert_eq!(pairs[1].1, b"max-age=31536000");
    }

    /// `end_stream_decision` on a response already in the Error phase: once
    /// part of it was handed to the client (`consumed`), a default answer
    /// would follow those bytes as a second response, so the stream ends in
    /// error. Before any byte left, a clean 502 is still the answer.
    ///
    /// Red on `d5161919`, which answered `SendDefault(502)` in both cases, so
    /// a 502 was written into the body of a truncated 200
    /// (sozu-proxy/sozu#1633).
    #[test]
    fn an_errored_response_already_handed_to_the_client_is_not_replaced() {
        use std::{cell::RefCell, rc::Rc};

        use rusty_ulid::Ulid;

        use super::{EndStreamAction, end_stream_decision};
        use crate::{pool::Pool, protocol::mux::test_support::test_context};

        let pool = Rc::new(RefCell::new(Pool::with_capacity(4, 8, 16384)));
        let mut context = test_context(&pool);
        let id = context
            .create_stream(Ulid::generate(), 1 << 16)
            .expect("the test pool must hand out stream buffers");
        let stream = &mut context.streams[id];
        stream.front.consumed = true;
        stream
            .back
            .parsing_phase
            .error(kawa::ParsingErrorKind::Processing {
                message: "INTERNAL_ERROR",
            });
        assert_eq!(
            end_stream_decision(stream),
            EndStreamAction::SendDefault(502),
            "nothing of the response left yet: a clean 502 replaces it"
        );
        stream.back.consumed = true;
        assert_eq!(
            end_stream_decision(stream),
            EndStreamAction::ForwardUnterminated,
            "part of the response left: it ends in error, never followed by a 502"
        );
    }
}