pdk_websockets_lib/

lib.rs

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! PDK WebSockets Library
//!
//! Library for decoding and encoding WebSocket frames in Flex Gateway custom policies.
//! It wraps [`websocket-sans-io`] to provide ergonomic frame-level access for policies
//! that operate on WebSocket upgrade connections.
//!
//! ## Primary types
//!
//! - [`Frame`]: a single WebSocket frame with its payload and metadata
//! - [`FrameType`]: the kind of frame (Text, Binary, Ping, Pong, etc.)
//! - [`Decoder`]: incrementally decodes raw bytes into [`Frame`]s
//! - [`Encoder`]: re-encodes a collection of [`Frame`]s into bytes
//! - [`SinkResult`]: outcome of feeding bytes to the [`Decoder`]
//!
//! ## Example
//!
//! ```ignore
//! use pdk_websockets::{Decoder, Encoder, Frame, FrameType, SinkResult};
//!
//! let mut decoder = Decoder::default();
//! match decoder.sink(raw_bytes) {
//!     SinkResult::MidFrame => { /* pause and wait for more bytes */ }
//!     SinkResult::Complete(mut frames) => {
//!         // inspect frames
//!         if let Some(frame) = frames.first() {
//!             if let FrameType::Text = frame.frame_type() {
//!                 let text = String::from_utf8_lossy(frame.data());
//!             }
//!         }
//!         // modify frames, then re-encode
//!         frames.push(Frame::ping());
//!         let bytes = Encoder::default().encode(frames);
//!     }
//! }
//! ```

use websocket_sans_io::{
    FrameInfo, Opcode, WebsocketFrameDecoder, WebsocketFrameEncoder, WebsocketFrameEvent,
};

/// A single WebSocket frame.
#[derive(Clone, Debug)]
pub struct Frame {
    info: FrameInfo,
    data: Vec<u8>,
}

/// The kind of a WebSocket frame, derived from its opcode.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum FrameType {
    /// A UTF-8 text data frame.
    Text,
    /// A binary data frame.
    Binary,
    /// A continuation frame for a fragmented message.
    Continuation,
    /// A ping control frame.
    Ping,
    /// A pong control frame.
    Pong,
    /// A connection-close control frame.
    ConnectionClose,
    /// A frame with a reserved or unknown opcode.
    Reserved,
}

impl Frame {
    /// Creates a ping control frame.
    pub fn ping() -> Self {
        Self::control_frame(Opcode::Ping)
    }

    /// Creates a pong control frame.
    pub fn pong() -> Self {
        Self::control_frame(Opcode::Pong)
    }

    /// Creates a connection-close control frame.
    pub fn connection_close() -> Self {
        Self::control_frame(Opcode::ConnectionClose)
    }

    /// Creates a text data frame.
    ///
    /// Set `fin` to `true` for an unfragmented message or the final fragment,
    /// `false` for intermediate fragments.
    pub fn text<T: Into<Vec<u8>>>(text: T, fin: bool) -> Frame {
        Self::data_frame(Opcode::Text, text, fin)
    }

    /// Creates a binary data frame.
    ///
    /// Set `fin` to `true` for an unfragmented message or the final fragment,
    /// `false` for intermediate fragments.
    pub fn binary<T: Into<Vec<u8>>>(text: T, fin: bool) -> Frame {
        Self::data_frame(Opcode::Binary, text, fin)
    }

    /// Creates a continuation data frame for a fragmented message.
    ///
    /// Set `fin` to `true` for the final fragment, `false` for intermediate ones.
    pub fn continuation<T: Into<Vec<u8>>>(text: T, fin: bool) -> Frame {
        Self::data_frame(Opcode::Continuation, text, fin)
    }
}

impl Frame {
    /// Returns the [`FrameType`] of this frame.
    pub fn frame_type(&self) -> FrameType {
        match self.info.opcode {
            Opcode::Continuation => FrameType::Continuation,
            Opcode::Text => FrameType::Text,
            Opcode::Binary => FrameType::Binary,
            Opcode::ConnectionClose => FrameType::ConnectionClose,
            Opcode::Ping => FrameType::Ping,
            Opcode::Pong => FrameType::Pong,
            _ => FrameType::Reserved,
        }
    }

    /// Returns a reference to the frame's payload bytes.
    pub fn data(&self) -> &[u8] {
        &self.data
    }

    /// Consumes the frame and returns its payload bytes.
    pub fn take(self) -> Vec<u8> {
        self.data
    }

    /// Replaces the frame's payload with `data`.
    pub fn update<U: Into<Vec<u8>>>(&mut self, data: U) {
        let data = data.into();
        self.info.payload_length = data.len() as u64;
        self.data = data;
    }

    /// Returns `true` if this is the final fragment of a message (FIN bit is set).
    pub fn fin(&self) -> bool {
        self.info.fin
    }
}

/// Internal functions
impl Frame {
    fn control_frame(opcode: Opcode) -> Frame {
        Frame {
            info: Self::info(opcode, &[], true),
            data: Vec::new(),
        }
    }

    fn data_frame<D: Into<Vec<u8>>>(opcode: Opcode, data: D, fin: bool) -> Frame {
        let data = data.into();
        Frame {
            info: Self::info(opcode, &data, fin),
            data,
        }
    }

    fn info(opcode: Opcode, data: &[u8], fin: bool) -> FrameInfo {
        FrameInfo {
            opcode,
            payload_length: data.len() as u64,
            mask: Some(rand::random()),
            fin,
            reserved: 0,
        }
    }

    fn encode(mut self, encoder: &mut WebsocketFrameEncoder, result: &mut Vec<u8>) {
        result.extend(encoder.start_frame(&self.info));
        if self.info.payload_length != 0 {
            encoder.transform_frame_payload(&mut self.data);
            result.extend(self.data);
        }
    }
}

/// Incrementally decodes raw bytes into [`Frame`]s.
///
/// Feed chunks of bytes via [`Decoder::sink`]. If the last frame in a chunk is
/// incomplete, `sink` returns [`SinkResult::MidFrame`] and buffers the partial
/// state internally. Call `sink` again with the next chunk to continue.
#[derive(Default)]
pub struct Decoder {
    decoder: WebsocketFrameDecoder,
    started: bool,
    ongoing: Vec<u8>,
    parsed: Vec<Frame>,
}

/// The result of feeding bytes to [`Decoder::sink`].
pub enum SinkResult {
    /// All bytes were consumed and every frame found is complete.
    /// Contains the decoded frames.
    Complete(Vec<Frame>),
    /// The last frame in the supplied bytes is incomplete.
    /// The decoder has buffered the partial state; supply more bytes to finish it.
    MidFrame,
}

impl Decoder {
    /// Drains and returns all frames that have been fully decoded so far,
    /// without consuming partially-decoded state.
    pub fn take_complete(&mut self) -> Vec<Frame> {
        self.parsed.split_off(0)
    }

    /// Feeds a chunk of raw bytes into the decoder.
    ///
    /// Returns [`SinkResult::Complete`] with all fully parsed frames when every
    /// byte in `body` has been consumed, or [`SinkResult::MidFrame`] when the
    /// final frame in the chunk is still incomplete.
    pub fn sink(&mut self, mut body: Vec<u8>) -> SinkResult {
        let mut position = 0;
        while position < body.len() {
            // this should never happen as infallible due to feature selection
            let frame = self.decoder.add_data(&mut body[position..]).unwrap();

            if let Some(event) = frame.event {
                match event {
                    WebsocketFrameEvent::Start { .. } => {
                        self.started = true;
                    }
                    WebsocketFrameEvent::PayloadChunk { .. } => {
                        self.ongoing
                            .extend_from_slice(&body[position..position + frame.consumed_bytes]);
                    }
                    WebsocketFrameEvent::End { frame_info, .. } => {
                        self.parsed.push(Frame {
                            info: frame_info,
                            data: self.ongoing.split_off(0),
                        });
                        self.started = false;
                    }
                }
                position += frame.consumed_bytes;
            }
        }

        if self.started {
            // call with empty data to see if end event is there
            let frame = self.decoder.add_data(&mut []).unwrap();
            if let Some(WebsocketFrameEvent::End { frame_info, .. }) = frame.event {
                self.parsed.push(Frame {
                    info: frame_info,
                    data: self.ongoing.split_off(0),
                });
                self.started = false;
            }
        }

        if self.started {
            SinkResult::MidFrame
        } else {
            SinkResult::Complete(self.parsed.split_off(0))
        }
    }
}

/// Encodes a collection of [`Frame`]s back into raw bytes suitable for
/// writing to a WebSocket connection.
#[derive(Default)]
pub struct Encoder {}

impl Encoder {
    /// Encodes each frame in `frame` and returns the concatenated byte representation for the flow client to server direction.
    pub fn encode_client(&mut self, frame: Vec<Frame>) -> Vec<u8> {
        let mut encoder = WebsocketFrameEncoder::new();

        let mut result = Vec::new();
        for frame in frame {
            frame.encode(&mut encoder, &mut result);
        }

        result
    }

    /// Encodes each frame in `frame` and returns the concatenated byte representation for the flow server to client direction.
    pub fn encode_server(&mut self, mut frame: Vec<Frame>) -> Vec<u8> {
        for frame in &mut frame {
            frame.info.mask = None;
        }

        self.encode_client(frame)
    }
}

#[derive(Default, PartialEq)]
enum State {
    #[default]
    Http,
    Websocket,
}

#[derive(Default)]
pub struct UpgradeTracker {
    state: State,
}

impl UpgradeTracker {
    pub fn ready(&self) -> bool {
        self.state == State::Websocket
    }

    pub fn track_upgrade_headers(&mut self, status: Option<&str>, upgrade: Option<&str>) {
        if self.state == State::Http {
            let is_switching = status == Some("101");
            let is_websocket = upgrade == Some("websocket");
            if is_switching && is_websocket {
                self.state = State::Websocket;
            }
        }
    }
}