pingora-core 0.9.1

// Copyright 2026 Cloudflare, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use bytes::Bytes;
use pingora_error::Result;
use pingora_http::{RequestHeader, ResponseHeader};
use std::time::Duration;

use super::v2::client::Http2Session;
use super::{custom::client::Session, v1::client::HttpSession as Http1Session};
use crate::protocols::{Digest, SocketAddr, Stream};

/// A type for Http client session. It can be either an Http1 connection or an Http2 stream.
pub enum HttpSession<S = ()> {
    H1(Http1Session),
    H2(Http2Session),
    Custom(S),
}

impl<S: Session> HttpSession<S> {
    pub fn as_http1(&self) -> Option<&Http1Session> {
        match self {
            Self::H1(s) => Some(s),
            Self::H2(_) => None,
            Self::Custom(_) => None,
        }
    }

    pub fn as_http2(&self) -> Option<&Http2Session> {
        match self {
            Self::H1(_) => None,
            Self::H2(s) => Some(s),
            Self::Custom(_) => None,
        }
    }

    pub fn as_custom(&self) -> Option<&S> {
        match self {
            Self::H1(_) => None,
            Self::H2(_) => None,
            Self::Custom(c) => Some(c),
        }
    }

    pub fn as_custom_mut(&mut self) -> Option<&mut S> {
        match self {
            Self::H1(_) => None,
            Self::H2(_) => None,
            Self::Custom(c) => Some(c),
        }
    }

    /// Write the request header to the server
    /// After the request header is sent. The caller can either start reading the response or
    /// sending request body if any.
    pub async fn write_request_header(&mut self, req: Box<RequestHeader>) -> Result<()> {
        match self {
            HttpSession::H1(h1) => {
                h1.write_request_header(req).await?;
                Ok(())
            }
            HttpSession::H2(h2) => h2.write_request_header(req, false),
            HttpSession::Custom(c) => c.write_request_header(req, false).await,
        }
    }

    /// Write a chunk of the request body.
    pub async fn write_request_body(&mut self, data: Bytes, end: bool) -> Result<()> {
        match self {
            HttpSession::H1(h1) => {
                // TODO: maybe h1 should also have the concept of `end`
                h1.write_body(&data).await?;
                Ok(())
            }
            HttpSession::H2(h2) => h2.write_request_body(data, end).await,
            HttpSession::Custom(c) => c.write_request_body(data, end).await,
        }
    }

    /// Signal that the request body has ended
    pub async fn finish_request_body(&mut self) -> Result<()> {
        match self {
            HttpSession::H1(h1) => {
                h1.finish_body().await?;
                Ok(())
            }
            HttpSession::H2(h2) => h2.finish_request_body(),
            HttpSession::Custom(c) => c.finish_request_body().await,
        }
    }

    /// Set the read timeout for reading header and body.
    ///
    /// The timeout is per read operation, not on the overall time reading the entire response
    pub fn set_read_timeout(&mut self, timeout: Option<Duration>) {
        match self {
            HttpSession::H1(h1) => h1.read_timeout = timeout,
            HttpSession::H2(h2) => h2.read_timeout = timeout,
            HttpSession::Custom(c) => c.set_read_timeout(timeout),
        }
    }

    /// Set the write timeout for writing header and body.
    ///
    /// The timeout is per write operation, not on the overall time writing the entire request.
    pub fn set_write_timeout(&mut self, timeout: Option<Duration>) {
        match self {
            HttpSession::H1(h1) => h1.write_timeout = timeout,
            HttpSession::H2(h2) => h2.write_timeout = timeout,
            HttpSession::Custom(c) => c.set_write_timeout(timeout),
        }
    }

    /// Read the response header from the server
    /// For http1, this function can be called multiple times, if the headers received are just
    /// informational headers.
    pub async fn read_response_header(&mut self) -> Result<()> {
        match self {
            HttpSession::H1(h1) => {
                h1.read_response().await?;
                Ok(())
            }
            HttpSession::H2(h2) => h2.read_response_header().await,
            HttpSession::Custom(c) => c.read_response_header().await,
        }
    }

    /// Read response body
    ///
    /// `None` when no more body to read.
    pub async fn read_response_body(&mut self) -> Result<Option<Bytes>> {
        match self {
            HttpSession::H1(h1) => h1.read_body_bytes().await,
            HttpSession::H2(h2) => h2.read_response_body().await,
            HttpSession::Custom(c) => c.read_response_body().await,
        }
    }

    /// No (more) body to read
    pub fn response_done(&mut self) -> bool {
        match self {
            HttpSession::H1(h1) => h1.is_body_done(),
            HttpSession::H2(h2) => h2.response_finished(),
            HttpSession::Custom(c) => c.response_finished(),
        }
    }

    /// Give up the http session abruptly.
    ///
    /// This is a failure path: the request is abandoned mid-message, so each
    /// protocol signals it in whatever way lets the peer tell this apart from a
    /// request that was completed.
    /// For H1 this will close the underlying connection
    /// For H2 this will send RST_STREAM frame to end this stream if the stream has not ended at all
    pub async fn shutdown(&mut self) {
        match self {
            Self::H1(s) => s.shutdown().await,
            Self::H2(s) => s.shutdown(),
            Self::Custom(c) => c.abandon("shutdown").await,
        }
    }

    /// Get the response header of the server
    ///
    /// `None` if the response header is not read yet.
    pub fn response_header(&self) -> Option<&ResponseHeader> {
        match self {
            Self::H1(s) => s.resp_header(),
            Self::H2(s) => s.response_header(),
            Self::Custom(c) => c.response_header(),
        }
    }

    /// Return the [Digest] of the connection
    ///
    /// For reused connection, the timing in the digest will reflect its initial handshakes
    /// The caller should check if the connection is reused to avoid misuse of the timing field.
    pub fn digest(&self) -> Option<&Digest> {
        match self {
            Self::H1(s) => Some(s.digest()),
            Self::H2(s) => s.digest(),
            Self::Custom(c) => c.digest(),
        }
    }

    /// Return a mutable [Digest] reference for the connection.
    ///
    /// Will return `None` if this is an H2 session and multiple streams are open.
    pub fn digest_mut(&mut self) -> Option<&mut Digest> {
        match self {
            Self::H1(s) => Some(s.digest_mut()),
            Self::H2(s) => s.digest_mut(),
            Self::Custom(s) => s.digest_mut(),
        }
    }

    /// Return the server (peer) address of the connection.
    pub fn server_addr(&self) -> Option<&SocketAddr> {
        match self {
            Self::H1(s) => s.server_addr(),
            Self::H2(s) => s.server_addr(),
            Self::Custom(s) => s.server_addr(),
        }
    }

    /// Return the client (local) address of the connection.
    pub fn client_addr(&self) -> Option<&SocketAddr> {
        match self {
            Self::H1(s) => s.client_addr(),
            Self::H2(s) => s.client_addr(),
            Self::Custom(s) => s.client_addr(),
        }
    }

    /// Get the reference of the [Stream] that this HTTP/1 session is operating upon.
    /// None if the HTTP session is over H2
    pub fn stream(&self) -> Option<&Stream> {
        match self {
            Self::H1(s) => Some(s.stream()),
            Self::H2(_) => None,
            Self::Custom(_) => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{HttpSession, Session};
    use crate::protocols::http::custom::{BodyWrite, CustomMessageWrite};
    use crate::protocols::{Digest, SocketAddr, UniqueIDType};
    use async_trait::async_trait;
    use bytes::Bytes;
    use futures::Stream;
    use http::HeaderMap;
    use pingora_error::Result;
    use pingora_http::{RequestHeader, ResponseHeader};
    use std::sync::{Arc, Mutex};
    use std::time::Duration;

    /// Records which shutdown entry point the session dispatcher reached. Every
    /// other method is `unreachable!()`: this exists to pin the dispatch, not to
    /// model a protocol.
    struct ShutdownRecordingCustom {
        shutdown_calls: Arc<Mutex<Vec<String>>>,
    }

    #[async_trait]
    impl Session for ShutdownRecordingCustom {
        async fn shutdown(&mut self, code: u32, ctx: &str) {
            self.shutdown_calls
                .lock()
                .unwrap()
                .push(format!("shutdown({code}, {ctx})"));
        }

        async fn abandon(&mut self, ctx: &str) {
            self.shutdown_calls
                .lock()
                .unwrap()
                .push(format!("abandon({ctx})"));
        }

        async fn write_request_header(
            &mut self,
            _req: Box<RequestHeader>,
            _end: bool,
        ) -> Result<()> {
            unreachable!("not used by the shutdown dispatch test")
        }

        async fn write_request_body(&mut self, _data: Bytes, _end: bool) -> Result<()> {
            unreachable!("not used by the shutdown dispatch test")
        }

        async fn finish_request_body(&mut self) -> Result<()> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn set_read_timeout(&mut self, _timeout: Option<Duration>) {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn set_write_timeout(&mut self, _timeout: Option<Duration>) {
            unreachable!("not used by the shutdown dispatch test")
        }

        async fn read_response_header(&mut self) -> Result<()> {
            unreachable!("not used by the shutdown dispatch test")
        }

        async fn read_response_body(&mut self) -> Result<Option<Bytes>> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn response_finished(&self) -> bool {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn response_header(&self) -> Option<&ResponseHeader> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn was_upgraded(&self) -> bool {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn digest(&self) -> Option<&Digest> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn digest_mut(&mut self) -> Option<&mut Digest> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn server_addr(&self) -> Option<&SocketAddr> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn client_addr(&self) -> Option<&SocketAddr> {
            unreachable!("not used by the shutdown dispatch test")
        }

        async fn read_trailers(&mut self) -> Result<Option<HeaderMap>> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn fd(&self) -> UniqueIDType {
            unreachable!("not used by the shutdown dispatch test")
        }

        async fn check_response_end_or_error(&mut self, _headers: bool) -> Result<bool> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn take_request_body_writer(&mut self) -> Option<Box<dyn BodyWrite>> {
            unreachable!("not used by the shutdown dispatch test")
        }

        async fn finish_custom(&mut self) -> Result<()> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn take_custom_message_reader(
            &mut self,
        ) -> Option<Box<dyn Stream<Item = Result<Bytes>> + Unpin + Send + Sync + 'static>> {
            unreachable!("not used by the shutdown dispatch test")
        }

        async fn drain_custom_messages(&mut self) -> Result<()> {
            unreachable!("not used by the shutdown dispatch test")
        }

        fn take_custom_message_writer(&mut self) -> Option<Box<dyn CustomMessageWrite>> {
            unreachable!("not used by the shutdown dispatch test")
        }
    }

    /// `HttpSession::shutdown` abandons a request mid-message, so it must take the
    /// entry point that lets a custom protocol convey exactly that. Mirrors the
    /// server-side `custom_session_shutdown_signals_an_incomplete_message`.
    ///
    /// This method has no callers in-tree today, so the test guards the routing
    /// rather than a live path — which is precisely when a silent regression back
    /// to the benign `shutdown(0, ..)` would go unnoticed.
    #[tokio::test]
    async fn custom_session_shutdown_signals_an_incomplete_message() {
        let shutdown_calls = Arc::new(Mutex::new(Vec::new()));
        let mut session = HttpSession::Custom(ShutdownRecordingCustom {
            shutdown_calls: shutdown_calls.clone(),
        });

        session.shutdown().await;

        assert_eq!(*shutdown_calls.lock().unwrap(), ["abandon(shutdown)"]);
    }
}