codetics-ocpp-ws 0.1.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Kevin Vavelin, Codetics (https://codetics.fr)

//! Minimal WebSocket client, used by integration tests and the
//! charging-station simulator. [`connect`] opens plain TCP; [`handshake`]
//! runs over any established stream (e.g. a TLS stream).

use std::io;

use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf};
use tokio::net::TcpStream;

use crate::conn::{Role, WsConfig, WsReader, WsWriter};
use crate::deflate::{self, DeflateConfig};
use crate::{base64, handshake};

pub struct ClientOptions<'a> {
    /// `host:port`
    pub addr: &'a str,
    /// Request path, e.g. `/ocpp/CS001` (already percent-encoded).
    pub path: &'a str,
    pub subprotocols: &'a [&'a str],
    /// Offer permessage-deflate.
    pub deflate: bool,
    /// Value for `Authorization: Basic …`, as (username, password bytes).
    pub basic_auth: Option<(&'a str, &'a [u8])>,
}

pub struct ClientConnection<S = TcpStream> {
    pub reader: WsReader<ReadHalf<S>>,
    pub writer: WsWriter<WriteHalf<S>>,
    /// Subprotocol selected by the server, if any.
    pub subprotocol: Option<String>,
    pub deflate: bool,
}

#[derive(Debug)]
pub enum ConnectError {
    Io(io::Error),
    /// The server answered with a non-101 status.
    Status(u16),
    BadResponse(&'static str),
}

impl std::fmt::Display for ConnectError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ConnectError::Io(e) => write!(f, "{e}"),
            ConnectError::Status(s) => write!(f, "handshake rejected with HTTP {s}"),
            ConnectError::BadResponse(m) => write!(f, "bad handshake response: {m}"),
        }
    }
}

impl std::error::Error for ConnectError {}

impl From<io::Error> for ConnectError {
    fn from(e: io::Error) -> Self {
        ConnectError::Io(e)
    }
}

/// Opens a TCP connection to `opts.addr` and performs the WebSocket handshake.
pub async fn connect(opts: &ClientOptions<'_>) -> Result<ClientConnection, ConnectError> {
    let stream = TcpStream::connect(opts.addr).await?;
    stream.set_nodelay(true)?;
    handshake(stream, opts).await
}

/// Performs the WebSocket opening handshake over an established stream.
pub async fn handshake<S>(
    mut stream: S,
    opts: &ClientOptions<'_>,
) -> Result<ClientConnection<S>, ConnectError>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let key = {
        use std::hash::{BuildHasher, Hasher};
        let a = std::collections::hash_map::RandomState::new()
            .build_hasher()
            .finish();
        let b = std::collections::hash_map::RandomState::new()
            .build_hasher()
            .finish();
        let mut k = a.to_le_bytes().to_vec();
        k.extend_from_slice(&b.to_le_bytes());
        base64::encode(&k)
    };
    let mut req = format!(
        "GET {} HTTP/1.1\r\nHost: {}\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n",
        opts.path, opts.addr
    );
    if !opts.subprotocols.is_empty() {
        req.push_str(&format!(
            "Sec-WebSocket-Protocol: {}\r\n",
            opts.subprotocols.join(", ")
        ));
    }
    if opts.deflate {
        req.push_str("Sec-WebSocket-Extensions: permessage-deflate; client_max_window_bits\r\n");
    }
    if let Some((user, pass)) = opts.basic_auth {
        let mut creds = format!("{user}:").into_bytes();
        creds.extend_from_slice(pass);
        req.push_str(&format!(
            "Authorization: Basic {}\r\n",
            base64::encode(&creds)
        ));
    }
    req.push_str("\r\n");
    stream.write_all(req.as_bytes()).await?;

    let mut buf = Vec::with_capacity(1024);
    let header_end = loop {
        if let Some(i) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
            break i + 4;
        }
        if buf.len() > 16 * 1024 {
            return Err(ConnectError::BadResponse("response headers too large"));
        }
        if stream.read_buf(&mut buf).await? == 0 {
            return Err(ConnectError::BadResponse(
                "connection closed during handshake",
            ));
        }
    };
    let head = std::str::from_utf8(&buf[..header_end])
        .map_err(|_| ConnectError::BadResponse("non UTF-8 headers"))?;
    let mut lines = head.split("\r\n");
    let status: u16 = lines
        .next()
        .and_then(|l| l.split_whitespace().nth(1))
        .and_then(|s| s.parse().ok())
        .ok_or(ConnectError::BadResponse("invalid status line"))?;
    if status != 101 {
        return Err(ConnectError::Status(status));
    }
    let mut accept = None;
    let mut subprotocol = None;
    let mut extensions = None;
    for line in lines {
        if let Some((name, value)) = line.split_once(':') {
            let value = value.trim().to_owned();
            match name.trim().to_ascii_lowercase().as_str() {
                "sec-websocket-accept" => accept = Some(value),
                "sec-websocket-protocol" => subprotocol = Some(value),
                "sec-websocket-extensions" => extensions = Some(value),
                _ => {}
            }
        }
    }
    if accept.as_deref() != Some(handshake::accept_key(&key).as_str()) {
        return Err(ConnectError::BadResponse("invalid Sec-WebSocket-Accept"));
    }
    let params = extensions
        .as_deref()
        .and_then(deflate::accept_client_response);
    if params.is_some() && !opts.deflate {
        return Err(ConnectError::BadResponse(
            "server enabled an extension that was not offered",
        ));
    }
    let leftover = buf[header_end..].to_vec();
    let (r, w) = tokio::io::split(stream);
    let cfg = DeflateConfig::default();
    Ok(ClientConnection {
        reader: WsReader::new(r, Role::Client, WsConfig::default(), params, leftover),
        writer: WsWriter::new(w, Role::Client, params.map(|p| (p, cfg))),
        subprotocol,
        deflate: params.is_some(),
    })
}