fips_core/transport/
websocket.rs

1//! Binary WebSocket physical transport.
2//!
3//! One binary WebSocket message carries one FIPS physical record. A tiny
4//! nonce-bound key-hint exchange precedes Noise IK when a client has only a
5//! seed URL; the hint is untrusted routing metadata and never bypasses FIPS
6//! identity authentication or ACLs.
7
8use super::tcp::stream::validate_stream_record;
9use super::{
10    ConnectionState, DiscoveredPeer, PacketBuffer, PacketTx, ReceivedPacket, Transport,
11    TransportAddr, TransportError, TransportId, TransportState, TransportType,
12};
13use crate::Identity;
14use crate::config::WebSocketConfig;
15use crate::dataplane::validate_direct_fsp_transport_fragment;
16use crate::discovery::local_udp::LocalKeyHint;
17use futures::{SinkExt, StreamExt};
18use rand::RngExt;
19use secp256k1::XOnlyPublicKey;
20use serde::Serialize;
21use std::collections::{HashMap, VecDeque};
22use std::net::SocketAddr;
23use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
24use std::sync::{Arc, Mutex as StdMutex};
25use std::time::{Duration, SystemTime, UNIX_EPOCH};
26use tokio::net::{TcpListener, TcpStream};
27use tokio::sync::{Mutex, Semaphore, mpsc};
28use tokio::task::JoinHandle;
29use tokio_tungstenite::tungstenite::handshake::server::{ErrorResponse, Request, Response};
30use tokio_tungstenite::tungstenite::protocol::WebSocketConfig as TungsteniteConfig;
31use tokio_tungstenite::tungstenite::{Bytes, Message};
32use tokio_tungstenite::{WebSocketStream, accept_hdr_async_with_config, connect_async_with_config};
33use tracing::{debug, info};
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36enum Direction {
37    Inbound,
38    Outbound,
39}
40
41struct Connection {
42    generation: u64,
43    tx: mpsc::Sender<Vec<u8>>,
44}
45
46type ConnectionPool = Arc<Mutex<HashMap<TransportAddr, Connection>>>;
47type ConnectionStates = Arc<StdMutex<HashMap<TransportAddr, ConnectionState>>>;
48type DiscoveryQueue = Arc<StdMutex<VecDeque<DiscoveredPeer>>>;
49
50#[derive(Debug, Default)]
51struct WebSocketStats {
52    connections_opened: AtomicU64,
53    connections_closed: AtomicU64,
54    connections_rejected: AtomicU64,
55    reconnect_attempts: AtomicU64,
56    frames_sent: AtomicU64,
57    frames_received: AtomicU64,
58    bytes_sent: AtomicU64,
59    bytes_received: AtomicU64,
60    invalid_frames: AtomicU64,
61    send_queue_full: AtomicU64,
62}
63
64#[derive(Debug, Serialize)]
65pub struct WebSocketStatsSnapshot {
66    pub connections_opened: u64,
67    pub connections_closed: u64,
68    pub connections_rejected: u64,
69    pub reconnect_attempts: u64,
70    pub frames_sent: u64,
71    pub frames_received: u64,
72    pub bytes_sent: u64,
73    pub bytes_received: u64,
74    pub invalid_frames: u64,
75    pub send_queue_full: u64,
76}
77
78impl WebSocketStats {
79    fn snapshot(&self) -> WebSocketStatsSnapshot {
80        let load = |value: &AtomicU64| value.load(Ordering::Relaxed);
81        WebSocketStatsSnapshot {
82            connections_opened: load(&self.connections_opened),
83            connections_closed: load(&self.connections_closed),
84            connections_rejected: load(&self.connections_rejected),
85            reconnect_attempts: load(&self.reconnect_attempts),
86            frames_sent: load(&self.frames_sent),
87            frames_received: load(&self.frames_received),
88            bytes_sent: load(&self.bytes_sent),
89            bytes_received: load(&self.bytes_received),
90            invalid_frames: load(&self.invalid_frames),
91            send_queue_full: load(&self.send_queue_full),
92        }
93    }
94}
95
96#[derive(Clone)]
97struct Runtime {
98    transport_id: TransportId,
99    config: WebSocketConfig,
100    local_pubkey: [u8; 32],
101    packet_tx: PacketTx,
102    pool: ConnectionPool,
103    states: ConnectionStates,
104    discoveries: DiscoveryQueue,
105    running: Arc<AtomicBool>,
106    total_slots: Arc<Semaphore>,
107    inbound_slots: Arc<Semaphore>,
108    generation: Arc<AtomicU64>,
109    stats: Arc<WebSocketStats>,
110}
111
112impl Runtime {
113    fn websocket_config(&self) -> TungsteniteConfig {
114        let mut config = TungsteniteConfig::default();
115        config.max_message_size = Some(self.config.max_frame_bytes());
116        config.max_frame_size = Some(self.config.max_frame_bytes());
117        config.max_write_buffer_size = self.config.max_frame_bytes().saturating_mul(2);
118        config.write_buffer_size = 0;
119        config
120    }
121
122    fn next_generation(&self) -> u64 {
123        self.generation.fetch_add(1, Ordering::Relaxed)
124    }
125
126    fn set_state(&self, addr: &TransportAddr, state: ConnectionState) {
127        self.states
128            .lock()
129            .unwrap_or_else(|error| error.into_inner())
130            .insert(addr.clone(), state);
131    }
132
133    fn clear_state_if(&self, addr: &TransportAddr, generation: u64) {
134        let connected_generation = self
135            .pool
136            .try_lock()
137            .ok()
138            .and_then(|pool| pool.get(addr).map(|connection| connection.generation));
139        if connected_generation != Some(generation) {
140            self.states
141                .lock()
142                .unwrap_or_else(|error| error.into_inner())
143                .remove(addr);
144        }
145    }
146}
147
148/// Generic WebSocket physical transport.
149pub struct WebSocketTransport {
150    transport_id: TransportId,
151    name: Option<String>,
152    config: WebSocketConfig,
153    state: TransportState,
154    local_addr: Option<SocketAddr>,
155    runtime: Runtime,
156    tasks: Arc<StdMutex<Vec<JoinHandle<()>>>>,
157}
158
159impl WebSocketTransport {
160    pub fn new(
161        transport_id: TransportId,
162        name: Option<String>,
163        config: WebSocketConfig,
164        packet_tx: PacketTx,
165        identity: &Identity,
166    ) -> Self {
167        let max_connections = config.max_connections();
168        let max_inbound = config.max_inbound_connections();
169        let runtime = Runtime {
170            transport_id,
171            config: config.clone(),
172            local_pubkey: identity.pubkey().serialize(),
173            packet_tx,
174            pool: Arc::new(Mutex::new(HashMap::new())),
175            states: Arc::new(StdMutex::new(HashMap::new())),
176            discoveries: Arc::new(StdMutex::new(VecDeque::new())),
177            running: Arc::new(AtomicBool::new(false)),
178            total_slots: Arc::new(Semaphore::new(max_connections)),
179            inbound_slots: Arc::new(Semaphore::new(max_inbound)),
180            generation: Arc::new(AtomicU64::new(1)),
181            stats: Arc::new(WebSocketStats::default()),
182        };
183        Self {
184            transport_id,
185            name,
186            config,
187            state: TransportState::Configured,
188            local_addr: None,
189            runtime,
190            tasks: Arc::new(StdMutex::new(Vec::new())),
191        }
192    }
193
194    pub fn name(&self) -> Option<&str> {
195        self.name.as_deref()
196    }
197
198    pub fn local_addr(&self) -> Option<SocketAddr> {
199        self.local_addr
200    }
201
202    pub fn public_url(&self) -> Option<&str> {
203        self.config.public_url.as_deref()
204    }
205
206    pub fn stats(&self) -> WebSocketStatsSnapshot {
207        self.runtime.stats.snapshot()
208    }
209
210    fn push_task(&self, task: JoinHandle<()>) {
211        self.tasks
212            .lock()
213            .unwrap_or_else(|error| error.into_inner())
214            .push(task);
215    }
216
217    pub async fn start_async(&mut self) -> Result<(), TransportError> {
218        if !self.state.can_start() {
219            return Err(TransportError::AlreadyStarted);
220        }
221        self.config
222            .validate()
223            .map_err(TransportError::StartFailed)?;
224        self.state = TransportState::Starting;
225        self.runtime.running.store(true, Ordering::Release);
226
227        if let Some(bind_addr) = self.config.bind_addr.as_deref() {
228            let bind_addr = bind_addr
229                .parse::<SocketAddr>()
230                .map_err(|error| TransportError::StartFailed(error.to_string()))?;
231            let listener = TcpListener::bind(bind_addr)
232                .await
233                .map_err(|error| TransportError::bind_failed(bind_addr, error))?;
234            self.local_addr = Some(
235                listener
236                    .local_addr()
237                    .map_err(|error| TransportError::StartFailed(error.to_string()))?,
238            );
239            self.push_task(tokio::spawn(run_accept_loop(
240                self.runtime.clone(),
241                listener,
242            )));
243        }
244
245        for seed_url in self.config.seed_urls.clone() {
246            let addr = TransportAddr::from_string(&seed_url);
247            self.runtime.set_state(&addr, ConnectionState::Connecting);
248            self.push_task(tokio::spawn(run_seed_dialer(self.runtime.clone(), addr)));
249        }
250
251        self.state = TransportState::Up;
252        info!(
253            transport_id = %self.transport_id,
254            local_addr = ?self.local_addr,
255            seeds = self.config.seed_urls.len(),
256            "WebSocket transport started"
257        );
258        Ok(())
259    }
260
261    pub async fn stop_async(&mut self) -> Result<(), TransportError> {
262        if !self.state.is_operational() {
263            return Err(TransportError::NotStarted);
264        }
265        self.runtime.running.store(false, Ordering::Release);
266        let tasks = {
267            let mut tasks = self.tasks.lock().unwrap_or_else(|error| error.into_inner());
268            std::mem::take(&mut *tasks)
269        };
270        for task in &tasks {
271            task.abort();
272        }
273        for task in tasks {
274            let _ = task.await;
275        }
276        self.runtime.pool.lock().await.clear();
277        self.runtime
278            .states
279            .lock()
280            .unwrap_or_else(|error| error.into_inner())
281            .clear();
282        self.local_addr = None;
283        self.state = TransportState::Down;
284        Ok(())
285    }
286
287    pub async fn send_async(
288        &self,
289        addr: &TransportAddr,
290        data: &[u8],
291    ) -> Result<usize, TransportError> {
292        if !self.state.is_operational() {
293            return Err(TransportError::NotStarted);
294        }
295        if data.len() > self.config.max_frame_bytes() {
296            return Err(TransportError::MtuExceeded {
297                packet_size: data.len(),
298                mtu: self.config.max_frame_bytes().min(u16::MAX as usize) as u16,
299            });
300        }
301        validate_websocket_record(data).map_err(TransportError::SendFailed)?;
302        let tx = self
303            .runtime
304            .pool
305            .lock()
306            .await
307            .get(addr)
308            .map(|connection| connection.tx.clone())
309            .ok_or(TransportError::NotStarted)?;
310        tx.try_send(data.to_vec()).map_err(|error| match error {
311            mpsc::error::TrySendError::Full(_) => {
312                self.runtime
313                    .stats
314                    .send_queue_full
315                    .fetch_add(1, Ordering::Relaxed);
316                TransportError::SendFailed("WebSocket send queue full".into())
317            }
318            mpsc::error::TrySendError::Closed(_) => {
319                TransportError::SendFailed("WebSocket connection closed".into())
320            }
321        })?;
322        Ok(data.len())
323    }
324
325    pub async fn connect_async(&self, addr: &TransportAddr) -> Result<(), TransportError> {
326        if !self.state.is_operational() {
327            return Err(TransportError::NotStarted);
328        }
329        let raw = addr
330            .as_str()
331            .ok_or_else(|| TransportError::InvalidAddress(addr.to_string()))?;
332        let candidate = WebSocketConfig {
333            seed_urls: vec![raw.to_owned()],
334            ..self.config.clone()
335        };
336        candidate
337            .validate()
338            .map_err(TransportError::InvalidAddress)?;
339        if self.connection_state_sync(addr) == ConnectionState::Connected {
340            return Ok(());
341        }
342        {
343            let mut states = self
344                .runtime
345                .states
346                .lock()
347                .unwrap_or_else(|error| error.into_inner());
348            if matches!(states.get(addr), Some(ConnectionState::Connecting)) {
349                return Ok(());
350            }
351            states.insert(addr.clone(), ConnectionState::Connecting);
352        }
353        self.push_task(tokio::spawn(run_one_shot_dial(
354            self.runtime.clone(),
355            addr.clone(),
356        )));
357        Ok(())
358    }
359
360    pub fn connection_state_sync(&self, addr: &TransportAddr) -> ConnectionState {
361        if let Ok(pool) = self.runtime.pool.try_lock()
362            && pool.contains_key(addr)
363        {
364            return ConnectionState::Connected;
365        }
366        self.runtime
367            .states
368            .lock()
369            .unwrap_or_else(|error| error.into_inner())
370            .get(addr)
371            .cloned()
372            .unwrap_or(ConnectionState::None)
373    }
374
375    pub async fn close_connection_async(&self, addr: &TransportAddr) {
376        self.runtime.pool.lock().await.remove(addr);
377        self.runtime
378            .states
379            .lock()
380            .unwrap_or_else(|error| error.into_inner())
381            .remove(addr);
382    }
383}
384
385impl Transport for WebSocketTransport {
386    fn transport_id(&self) -> TransportId {
387        self.transport_id
388    }
389
390    fn transport_type(&self) -> &TransportType {
391        &TransportType::WEBSOCKET
392    }
393
394    fn state(&self) -> TransportState {
395        self.state
396    }
397
398    fn mtu(&self) -> u16 {
399        self.config.mtu()
400    }
401
402    fn start(&mut self) -> Result<(), TransportError> {
403        Err(TransportError::NotSupported(
404            "use start_async() for WebSocket transport".into(),
405        ))
406    }
407
408    fn stop(&mut self) -> Result<(), TransportError> {
409        Err(TransportError::NotSupported(
410            "use stop_async() for WebSocket transport".into(),
411        ))
412    }
413
414    fn send(&self, _addr: &TransportAddr, _data: &[u8]) -> Result<(), TransportError> {
415        Err(TransportError::NotSupported(
416            "use send_async() for WebSocket transport".into(),
417        ))
418    }
419
420    fn discover(&self) -> Result<Vec<DiscoveredPeer>, TransportError> {
421        Ok(self
422            .runtime
423            .discoveries
424            .lock()
425            .unwrap_or_else(|error| error.into_inner())
426            .drain(..)
427            .collect())
428    }
429
430    fn auto_connect(&self) -> bool {
431        true
432    }
433
434    fn accept_connections(&self) -> bool {
435        self.config.accept_connections()
436    }
437}
438
439async fn run_accept_loop(runtime: Runtime, listener: TcpListener) {
440    while runtime.running.load(Ordering::Acquire) {
441        let Ok((stream, peer_addr)) = listener.accept().await else {
442            continue;
443        };
444        let Ok(inbound_permit) = runtime.inbound_slots.clone().try_acquire_owned() else {
445            runtime
446                .stats
447                .connections_rejected
448                .fetch_add(1, Ordering::Relaxed);
449            continue;
450        };
451        let Ok(total_permit) = runtime.total_slots.clone().try_acquire_owned() else {
452            runtime
453                .stats
454                .connections_rejected
455                .fetch_add(1, Ordering::Relaxed);
456            continue;
457        };
458        let accepted_runtime = runtime.clone();
459        tokio::spawn(async move {
460            let _inbound_permit = inbound_permit;
461            let _total_permit = total_permit;
462            if let Err(error) = accept_connection(accepted_runtime, stream, peer_addr).await {
463                debug!(%peer_addr, %error, "WebSocket inbound connection ended");
464            }
465        });
466    }
467}
468
469#[allow(clippy::result_large_err)]
470async fn accept_connection(
471    runtime: Runtime,
472    stream: TcpStream,
473    peer_addr: SocketAddr,
474) -> Result<(), TransportError> {
475    let path = runtime.config.path().to_owned();
476    let callback = move |request: &Request, response: Response| {
477        if request.uri().path() == path {
478            Ok(response)
479        } else {
480            let mut error = ErrorResponse::new(Some("not found".into()));
481            *error.status_mut() = tokio_tungstenite::tungstenite::http::StatusCode::NOT_FOUND;
482            Err(error)
483        }
484    };
485    let websocket =
486        accept_hdr_async_with_config(stream, callback, Some(runtime.websocket_config()))
487            .await
488            .map_err(|_| TransportError::ConnectionRefused)?;
489    let generation = runtime.next_generation();
490    let addr = TransportAddr::from_string(&format!("ws-peer://{peer_addr}/{generation}"));
491    run_connection(
492        runtime,
493        addr,
494        websocket,
495        generation,
496        Direction::Inbound,
497        false,
498    )
499    .await
500}
501
502async fn run_seed_dialer(runtime: Runtime, addr: TransportAddr) {
503    let mut delay_ms = runtime.config.reconnect_initial_ms();
504    while runtime.running.load(Ordering::Acquire) {
505        runtime.set_state(&addr, ConnectionState::Connecting);
506        runtime
507            .stats
508            .reconnect_attempts
509            .fetch_add(1, Ordering::Relaxed);
510        let result = dial_and_run(runtime.clone(), addr.clone()).await;
511        if !runtime.running.load(Ordering::Acquire) {
512            break;
513        }
514        match result {
515            Ok(()) => {
516                delay_ms = runtime.config.reconnect_initial_ms();
517            }
518            Err(error) => {
519                runtime.set_state(&addr, ConnectionState::Failed(error.to_string()));
520                debug!(remote_addr = %addr, %error, "WebSocket seed connection failed");
521                delay_ms = delay_ms
522                    .saturating_mul(2)
523                    .min(runtime.config.reconnect_max_ms());
524            }
525        }
526        tokio::time::sleep(Duration::from_millis(delay_ms)).await;
527    }
528}
529
530async fn run_one_shot_dial(runtime: Runtime, addr: TransportAddr) {
531    let result = dial_and_run(runtime.clone(), addr.clone()).await;
532    if let Err(error) = result {
533        runtime.set_state(&addr, ConnectionState::Failed(error.to_string()));
534    }
535}
536
537async fn dial_and_run(runtime: Runtime, addr: TransportAddr) -> Result<(), TransportError> {
538    let _slot = runtime
539        .total_slots
540        .clone()
541        .try_acquire_owned()
542        .map_err(|_| TransportError::ConnectionRefused)?;
543    let url = addr
544        .as_str()
545        .ok_or_else(|| TransportError::InvalidAddress(addr.to_string()))?;
546    let connect = connect_async_with_config(url, Some(runtime.websocket_config()), false);
547    let (websocket, _) = tokio::time::timeout(
548        Duration::from_millis(runtime.config.connect_timeout_ms()),
549        connect,
550    )
551    .await
552    .map_err(|_| TransportError::Timeout)?
553    .map_err(|error| TransportError::StartFailed(error.to_string()))?;
554    let generation = runtime.next_generation();
555    run_connection(
556        runtime,
557        addr,
558        websocket,
559        generation,
560        Direction::Outbound,
561        true,
562    )
563    .await
564}
565
566async fn run_connection<S>(
567    runtime: Runtime,
568    addr: TransportAddr,
569    websocket: WebSocketStream<S>,
570    generation: u64,
571    direction: Direction,
572    request_key_hint: bool,
573) -> Result<(), TransportError>
574where
575    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
576{
577    let (mut sink, mut stream) = websocket.split();
578    let (tx, mut rx) = mpsc::channel::<Vec<u8>>(runtime.config.max_send_queue());
579    {
580        let mut pool = runtime.pool.lock().await;
581        if pool.contains_key(&addr) {
582            return Err(TransportError::AlreadyStarted);
583        }
584        pool.insert(addr.clone(), Connection { generation, tx });
585    }
586    runtime.set_state(&addr, ConnectionState::Connected);
587    runtime
588        .stats
589        .connections_opened
590        .fetch_add(1, Ordering::Relaxed);
591
592    let mut pending_nonce = request_key_hint.then(|| rand::rng().random::<u64>());
593    if let Some(nonce) = pending_nonce {
594        sink.send(Message::Binary(
595            LocalKeyHint::Request { nonce }.encode().into(),
596        ))
597        .await
598        .map_err(|error| TransportError::SendFailed(error.to_string()))?;
599    }
600
601    let started = tokio::time::Instant::now();
602    let mut last_received = started;
603    let ping_secs = runtime.config.ping_interval_secs();
604    let idle_secs = runtime.config.idle_timeout_secs();
605    let mut ping = tokio::time::interval(if ping_secs == 0 {
606        Duration::from_secs(24 * 60 * 60)
607    } else {
608        Duration::from_secs(ping_secs)
609    });
610    ping.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
611    let mut check = tokio::time::interval(Duration::from_secs(1));
612    check.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
613
614    enum Event {
615        Outbound(Option<Vec<u8>>),
616        Inbound(Option<Result<Message, tokio_tungstenite::tungstenite::Error>>),
617        Ping,
618        Check,
619    }
620
621    let outcome = loop {
622        let event = tokio::select! {
623            outbound = rx.recv() => Event::Outbound(outbound),
624            inbound = stream.next() => Event::Inbound(inbound),
625            _ = ping.tick(), if ping_secs > 0 => Event::Ping,
626            _ = check.tick() => Event::Check,
627        };
628        match event {
629            Event::Outbound(Some(data)) => {
630                let len = data.len();
631                if let Err(error) = sink.send(Message::Binary(data.into())).await {
632                    break Err(TransportError::SendFailed(error.to_string()));
633                }
634                runtime.stats.frames_sent.fetch_add(1, Ordering::Relaxed);
635                runtime
636                    .stats
637                    .bytes_sent
638                    .fetch_add(len as u64, Ordering::Relaxed);
639            }
640            Event::Outbound(None) => break Ok(()),
641            Event::Inbound(Some(Ok(Message::Binary(data)))) => {
642                last_received = tokio::time::Instant::now();
643                if let Some(hint) = LocalKeyHint::decode(&data) {
644                    match hint {
645                        LocalKeyHint::Request { nonce } => {
646                            let reply = LocalKeyHint::Response {
647                                nonce,
648                                pubkey: runtime.local_pubkey,
649                            };
650                            if let Err(error) =
651                                sink.send(Message::Binary(reply.encode().into())).await
652                            {
653                                break Err(TransportError::SendFailed(error.to_string()));
654                            }
655                        }
656                        LocalKeyHint::Response { nonce, pubkey }
657                            if pending_nonce == Some(nonce) =>
658                        {
659                            pending_nonce = None;
660                            if pubkey != runtime.local_pubkey
661                                && let Ok(pubkey) = XOnlyPublicKey::from_slice(&pubkey)
662                            {
663                                runtime
664                                    .discoveries
665                                    .lock()
666                                    .unwrap_or_else(|error| error.into_inner())
667                                    .push_back(DiscoveredPeer::with_hint(
668                                        runtime.transport_id,
669                                        addr.clone(),
670                                        pubkey,
671                                    ));
672                            }
673                        }
674                        LocalKeyHint::Response { .. } => {}
675                    }
676                    continue;
677                }
678                if data.len() > runtime.config.max_frame_bytes()
679                    || validate_websocket_record(&data).is_err()
680                {
681                    runtime.stats.invalid_frames.fetch_add(1, Ordering::Relaxed);
682                    break Err(TransportError::RecvFailed(
683                        "invalid WebSocket FIPS physical record".into(),
684                    ));
685                }
686                let len = data.len();
687                let packet = ReceivedPacket::with_timestamp(
688                    runtime.transport_id,
689                    addr.clone(),
690                    PacketBuffer::new(data.to_vec()),
691                    now_ms(),
692                );
693                if runtime.packet_tx.send(packet).is_err() {
694                    break Err(TransportError::RecvFailed(
695                        "node packet channel closed".into(),
696                    ));
697                }
698                runtime
699                    .stats
700                    .frames_received
701                    .fetch_add(1, Ordering::Relaxed);
702                runtime
703                    .stats
704                    .bytes_received
705                    .fetch_add(len as u64, Ordering::Relaxed);
706            }
707            Event::Inbound(Some(Ok(Message::Ping(payload)))) => {
708                last_received = tokio::time::Instant::now();
709                if let Err(error) = sink.send(Message::Pong(payload)).await {
710                    break Err(TransportError::SendFailed(error.to_string()));
711                }
712            }
713            Event::Inbound(Some(Ok(Message::Pong(_)))) => {
714                last_received = tokio::time::Instant::now();
715            }
716            Event::Inbound(Some(Ok(Message::Close(_))) | None) => break Ok(()),
717            Event::Inbound(Some(Ok(Message::Text(_) | Message::Frame(_)))) => {
718                runtime.stats.invalid_frames.fetch_add(1, Ordering::Relaxed);
719                break Err(TransportError::RecvFailed(
720                    "WebSocket transport accepts binary messages only".into(),
721                ));
722            }
723            Event::Inbound(Some(Err(error))) => {
724                break Err(TransportError::RecvFailed(error.to_string()));
725            }
726            Event::Ping => {
727                if let Err(error) = sink.send(Message::Ping(Bytes::new())).await {
728                    break Err(TransportError::SendFailed(error.to_string()));
729                }
730            }
731            Event::Check => {
732                if pending_nonce.is_some()
733                    && started.elapsed()
734                        >= Duration::from_millis(runtime.config.key_hint_timeout_ms())
735                {
736                    break Err(TransportError::Timeout);
737                }
738                if idle_secs > 0 && last_received.elapsed() >= Duration::from_secs(idle_secs) {
739                    break Err(TransportError::Timeout);
740                }
741            }
742        }
743    };
744
745    {
746        let mut pool = runtime.pool.lock().await;
747        if pool
748            .get(&addr)
749            .is_some_and(|connection| connection.generation == generation)
750        {
751            pool.remove(&addr);
752        }
753    }
754    runtime.clear_state_if(&addr, generation);
755    runtime
756        .stats
757        .connections_closed
758        .fetch_add(1, Ordering::Relaxed);
759    debug!(
760        transport_id = %runtime.transport_id,
761        remote_addr = %addr,
762        ?direction,
763        "WebSocket physical connection closed"
764    );
765    outcome
766}
767
768fn now_ms() -> u64 {
769    SystemTime::now()
770        .duration_since(UNIX_EPOCH)
771        .unwrap_or_default()
772        .as_millis() as u64
773}
774
775fn validate_websocket_record(data: &[u8]) -> Result<(), String> {
776    if validate_direct_fsp_transport_fragment(data) {
777        return Ok(());
778    }
779    validate_stream_record(data)
780        .map_err(|error| format!("invalid FIPS physical record: {error}"))
781}
782
783#[cfg(test)]
784mod tests;
fips_core/transport/websocket.rs

fips_core/transport/
websocket.rs
