codetics-ocpp-service 0.1.0

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

//! One task per station connection.
//!
//! The task owns the socket and multiplexes, with a single `select!`:
//! inbound frames, the outbound command queue, the timeout of the CALL in
//! flight, keep-alive pings and service shutdown. This keeps the per-connection
//! footprint to one task, one small read buffer and one bounded channel.
//!
//! RPC rules applied here (OCPP-J §4.1.1): at most one CSMS-initiated CALL is
//! in flight per connection; further commands wait in the queue. Inbound CALLs
//! are processed in arrival order (the station also sends one at a time).

use std::time::Duration;

use ocpp_core::rpc::{self, Frame, FrameError};
use ocpp_core::{ErrorCode, OcppVersion, RpcError};
use ocpp_ws::{DeflateConfig, DeflateParams, Message, WsConfig, WsReader, WsWriter, close_code};
use tokio::io::{AsyncRead, AsyncWrite, ReadHalf, WriteHalf};
use tokio::sync::{mpsc, oneshot};
use tokio::time::{Instant, MissedTickBehavior};

use super::{Registration, StationCtx, v16, v21, v201};
use crate::db::{self, stations::Station};
use crate::hub::{CallFailure, CallOutcome, Command, ConnHandle};
use crate::state::App;

pub struct SessionParams {
    pub app: App,
    pub station: Station,
    /// `None` when no subprotocol could be agreed: the handshake is completed
    /// without `Sec-WebSocket-Protocol` and the connection closed (Part 4 §3.2).
    pub version: Option<OcppVersion>,
    pub deflate: Option<(DeflateParams, DeflateConfig)>,
    pub remote_addr: Option<String>,
}

struct Pending {
    id: String,
    action: String,
    reply: oneshot::Sender<CallOutcome>,
    deadline: Instant,
}

struct Session<S> {
    ctx: StationCtx,
    writer: WsWriter<WriteHalf<S>>,
    pending: Option<Pending>,
    last_rx: Instant,
}

pub async fn run<S>(p: SessionParams, io: S)
where
    S: AsyncRead + AsyncWrite + Send + 'static,
{
    let app = p.app;
    let ws_cfg = WsConfig {
        max_message_size: app.cfg.max_message_size,
        ..Default::default()
    };
    let (reader, mut writer) = ocpp_ws::server(io, ws_cfg, p.deflate);
    let Some(version) = p.version else {
        tracing::info!(station = %p.station.ocpp_identity, "no common OCPP version, closing");
        let _ = writer
            .send_close(
                close_code::PROTOCOL_ERROR,
                "no supported OCPP version offered",
            )
            .await;
        return;
    };

    let identity = p.station.ocpp_identity.clone();
    let station_id = p.station.id;
    let conn_id = app.hub.next_conn_id();
    let connection_id = format!("{}:{}", app.cfg.node_id, app.ids.next_id());
    let (cmd_tx, cmd_rx) = mpsc::channel(app.cfg.command_queue_size.max(1));
    let handle = ConnHandle {
        conn_id,
        station_id,
        version,
        connected_at: ocpp_core::DateTime::now(),
        remote_addr: p.remote_addr.clone(),
        tx: cmd_tx.clone(),
    };
    if let Some(old) = app.hub.register(&identity, handle) {
        tracing::info!(station = %identity, "new connection replaces the previous one");
        let _ = old.tx.try_send(Command::Close {
            code: close_code::POLICY_VIOLATION,
            reason: "replaced by a newer connection",
        });
    }
    if let Err(e) = db::stations::mark_connected(
        &app.db,
        station_id,
        &app.cfg.node_id,
        &connection_id,
        version.as_str(),
    )
    .await
    {
        tracing::error!(station = %identity, error = %e, "failed to record connection");
    }
    tracing::info!(station = %identity, version = %version, remote = ?p.remote_addr, deflate = p.deflate.is_some(), "station connected");

    let registration = Registration::from_db(&p.station.registration_status);
    if registration == Registration::Unknown {
        // Never booted (or booted on another CSMS): ask for a BootNotification.
        // The payload is identical in 1.6, 2.0.1 and 2.1.
        let (reply, _) = oneshot::channel();
        let _ = cmd_tx.try_send(Command::Call {
            action: "TriggerMessage".into(),
            payload: serde_json::json!({ "requestedMessage": "BootNotification" }),
            reply,
        });
    }
    drop(cmd_tx);

    let ctx = StationCtx {
        app: app.clone(),
        station_id,
        identity: identity.clone(),
        station_evse_id: p.station.evse_id.clone(),
        version,
        registration,
    };
    let mut session = Session {
        ctx,
        writer,
        pending: None,
        last_rx: Instant::now(),
    };
    let reason = session.run_loop(reader, cmd_rx).await;
    tracing::info!(station = %identity, reason, "station disconnected");

    if let Some(p) = session.pending.take() {
        let _ = p.reply.send(Err(CallFailure::Disconnected));
    }
    app.hub.unregister(&identity, conn_id);
    if let Err(e) = db::stations::mark_disconnected(&app.db, station_id, &connection_id).await {
        tracing::error!(station = %identity, error = %e, "failed to record disconnection");
    }
}

impl<S: AsyncRead + AsyncWrite + Send + 'static> Session<S> {
    async fn run_loop(
        &mut self,
        mut reader: WsReader<ReadHalf<S>>,
        mut cmd_rx: mpsc::Receiver<Command>,
    ) -> &'static str {
        let app = self.ctx.app.clone();
        let mut shutdown = app.shutdown.clone();
        let ping_every = app.cfg.ping_interval;
        let mut ping = tokio::time::interval(if ping_every.is_zero() {
            Duration::from_secs(3600)
        } else {
            ping_every
        });
        ping.set_missed_tick_behavior(MissedTickBehavior::Delay);
        ping.tick().await;
        let far = Instant::now() + Duration::from_secs(86_400 * 365);

        loop {
            let deadline = self.pending.as_ref().map(|p| p.deadline).unwrap_or(far);
            tokio::select! {
                _ = shutdown.changed() => {
                    let _ = self.writer.send_close(close_code::GOING_AWAY, "server shutting down").await;
                    return "server shutdown";
                }
                msg = reader.recv() => {
                    self.last_rx = Instant::now();
                    match msg {
                        Ok(Some(Message::Text(text))) => {
                            if let Err(e) = self.on_text(&text).await {
                                tracing::debug!(station = %self.ctx.identity, error = %e, "write failed");
                                return "write error";
                            }
                        }
                        Ok(Some(Message::Ping(data))) => {
                            if self.writer.send_pong(&data).await.is_err() {
                                return "write error";
                            }
                        }
                        Ok(Some(Message::Pong(_))) => {}
                        Ok(Some(Message::Binary(_))) => {
                            let _ = self.writer.send_close(close_code::UNSUPPORTED_DATA, "OCPP-J uses text frames").await;
                            return "binary frame";
                        }
                        Ok(Some(Message::Close(frame))) => {
                            let code = frame.as_ref().map(|f| f.0).unwrap_or(close_code::NORMAL);
                            let _ = self.writer.send_close(code, "").await;
                            return "closed by station";
                        }
                        Ok(None) => return "connection lost",
                        Err(e) => {
                            tracing::info!(station = %self.ctx.identity, error = %e, "WebSocket error");
                            let _ = self.writer.send_close(e.close_code(), "").await;
                            return "protocol error";
                        }
                    }
                }
                cmd = cmd_rx.recv(), if self.pending.is_none() => match cmd {
                    Some(Command::Call { action, payload, reply }) => {
                        if reply.is_closed() && action != "TriggerMessage" {
                            continue; // The API caller gave up while queued.
                        }
                        let id = app.ids.next_id();
                        let text = match rpc::encode_call(&id, &action, &payload) {
                            Ok(t) => t,
                            Err(e) => {
                                let _ = reply.send(Err(CallFailure::Error(RpcError::internal(e.to_string()))));
                                continue;
                            }
                        };
                        if self.writer.send_text(&text).await.is_err() {
                            let _ = reply.send(Err(CallFailure::Disconnected));
                            return "write error";
                        }
                        tracing::debug!(station = %self.ctx.identity, %action, %id, "CALL sent");
                        self.pending = Some(Pending { id, action, reply, deadline: Instant::now() + app.cfg.call_timeout });
                    }
                    Some(Command::Close { code, reason }) => {
                        let _ = self.writer.send_close(code, reason).await;
                        return reason;
                    }
                    None => {
                        let _ = self.writer.send_close(close_code::POLICY_VIOLATION, "connection replaced").await;
                        return "unregistered";
                    }
                },
                _ = tokio::time::sleep_until(deadline), if self.pending.is_some() => {
                    let p = self.pending.take().expect("pending checked");
                    tracing::warn!(station = %self.ctx.identity, action = %p.action, id = %p.id, "CALL timed out");
                    let _ = p.reply.send(Err(CallFailure::Timeout));
                }
                _ = ping.tick(), if !ping_every.is_zero() => {
                    let idle = self.last_rx.elapsed();
                    if idle >= ping_every * 3 {
                        let _ = self.writer.send_close(close_code::GOING_AWAY, "keep-alive timeout").await;
                        return "keep-alive timeout";
                    }
                    if idle >= ping_every && self.writer.send_ping(b"").await.is_err() {
                        return "write error";
                    }
                }
            }
        }
    }

    async fn send(&mut self, text: String) -> Result<(), ocpp_ws::WsError> {
        self.writer.send_text(&text).await
    }

    async fn send_error(
        &mut self,
        id: &str,
        err: &RpcError,
        result_error: bool,
    ) -> Result<(), ocpp_ws::WsError> {
        let text = rpc::encode_error(id, err, self.ctx.version.dialect(), result_error);
        self.send(text).await
    }

    async fn on_text(&mut self, text: &str) -> Result<(), ocpp_ws::WsError> {
        let version = self.ctx.version;
        let frame = match Frame::parse(text) {
            Ok(Frame::Send { .. } | Frame::CallResultError { .. })
                if version != OcppVersion::V21 =>
            {
                // Message types 5 and 6 only exist in OCPP 2.1.
                return self.unknown_message_type(text, None).await;
            }
            Ok(f) => f,
            Err(FrameError::UnknownMessageType { id, type_id }) => {
                return self
                    .unknown_message_type(text, id.map(|i| (i, type_id)))
                    .await;
            }
            Err(e @ FrameError::Malformed { .. }) => {
                tracing::warn!(station = %self.ctx.identity, error = %e, "malformed frame");
                let code = if e.reply_id() == rpc::UNKNOWN_MESSAGE_ID {
                    ErrorCode::RpcFrameworkError
                } else {
                    ErrorCode::FormatViolation
                };
                let id = e.reply_id().to_owned();
                return self
                    .send_error(&id, &RpcError::new(code, e.to_string()), false)
                    .await;
            }
        };

        match frame {
            Frame::Call {
                id,
                action,
                payload,
            } => {
                if !self.ctx.registration.allows(&action) {
                    let err = RpcError::new(
                        ErrorCode::SecurityError,
                        "Charging Station is not accepted by the CSMS; send BootNotification first",
                    );
                    return self.send_error(&id, &err, false).await;
                }
                let result = match version {
                    OcppVersion::V16 => v16::handle(&mut self.ctx, &action, payload.get()).await,
                    OcppVersion::V201 => v201::handle(&mut self.ctx, &action, payload.get()).await,
                    OcppVersion::V21 => v21::handle(&mut self.ctx, &action, payload.get()).await,
                };
                match result {
                    Ok(response) => match rpc::encode_result(&id, &response) {
                        Ok(t) => self.send(t).await,
                        Err(e) => {
                            self.send_error(&id, &RpcError::internal(e.to_string()), false)
                                .await
                        }
                    },
                    Err(err) => {
                        if err.code != ErrorCode::InternalError {
                            tracing::info!(station = %self.ctx.identity, %action, error = %err, "CALL rejected");
                        }
                        self.send_error(&id, &err, false).await
                    }
                }
            }
            Frame::Send {
                action, payload, ..
            } => {
                v21::handle_send(&mut self.ctx, &action, payload.get()).await;
                Ok(())
            }
            Frame::CallResult { id, payload } => {
                let Some(p) = self.take_pending(&id) else {
                    return Ok(());
                };
                let validated = match version {
                    OcppVersion::V16 => v16::validate_response(&p.action, payload.get()),
                    OcppVersion::V201 => v201::validate_response(&p.action, payload.get()),
                    OcppVersion::V21 => v21::validate_response(&p.action, payload.get()),
                };
                match validated {
                    Ok(v) => {
                        let _ = p.reply.send(Ok(v));
                    }
                    Err(violation) => {
                        tracing::warn!(station = %self.ctx.identity, action = %p.action, error = %violation, "invalid CALLRESULT");
                        let raw =
                            serde_json::from_str(payload.get()).unwrap_or(serde_json::Value::Null);
                        let err = RpcError::from(violation.clone());
                        let _ = p
                            .reply
                            .send(Err(CallFailure::InvalidResponse { violation, raw }));
                        if version == OcppVersion::V21 {
                            // 2.1: tell the station its response could not be processed.
                            return self.send_error(&id, &err, true).await;
                        }
                    }
                }
                Ok(())
            }
            Frame::CallError {
                id,
                code,
                description,
                details,
            } => {
                if let Some(p) = self.take_pending(&id) {
                    let _ = p.reply.send(Err(CallFailure::Error(rpc::error_from_frame(
                        &code,
                        description,
                        &details,
                    ))));
                } else {
                    tracing::warn!(station = %self.ctx.identity, %id, %code, %description, "CALLERROR for unknown message id");
                }
                Ok(())
            }
            Frame::CallResultError {
                id,
                code,
                description,
                ..
            } => {
                tracing::warn!(station = %self.ctx.identity, %id, %code, %description, "station could not process our CALLRESULT");
                Ok(())
            }
        }
    }

    fn take_pending(&mut self, id: &str) -> Option<Pending> {
        match &self.pending {
            Some(p) if p.id == id => self.pending.take(),
            _ => {
                tracing::debug!(station = %self.ctx.identity, %id, "response for a call that is not pending (late or duplicate)");
                None
            }
        }
    }

    /// 1.6 and 2.1: ignore. 2.0.1 Part 4 §4.4: answer MessageTypeNotSupported.
    async fn unknown_message_type(
        &mut self,
        text: &str,
        id: Option<(String, u64)>,
    ) -> Result<(), ocpp_ws::WsError> {
        tracing::warn!(station = %self.ctx.identity, frame = %text.chars().take(120).collect::<String>(), "unsupported message type");
        if self.ctx.version == OcppVersion::V201 {
            let id = id
                .map(|(i, _)| i)
                .or_else(|| Frame::parse(text).ok().map(|f| f.id().to_owned()))
                .unwrap_or_else(|| rpc::UNKNOWN_MESSAGE_ID.into());
            let err = RpcError::new(
                ErrorCode::MessageTypeNotSupported,
                "message type not supported",
            );
            return self.send_error(&id, &err, false).await;
        }
        Ok(())
    }
}