bssh/ssh/control/
protocol.rs

1// Copyright 2025 Lablup Inc. and Jeongkyu Shin
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8
9use std::io;
10
11use serde::{Deserialize, Serialize};
12use thiserror::Error;
13use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
14
15use crate::forwarding::ForwardingDirective;
16use crate::ssh::SessionPolicy;
17use crate::ssh::tokio_client::AddressFamily;
18
19use super::ControlCommand;
20
21/// Version of bssh's length-prefixed JSON multiplexing protocol.
22pub const CONTROL_PROTOCOL_VERSION: u32 = 1;
23/// Largest accepted serialized control message.
24pub const MAX_CONTROL_FRAME_BYTES: usize = 1024 * 1024;
25/// Largest binary chunk carried in one data message.
26pub const MAX_CONTROL_DATA_BYTES: usize = 64 * 1024;
27const MAX_CONTROL_ENVIRONMENT_ENTRIES: usize = 4_096;
28
29/// One framed message on a bssh control socket.
30#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
31#[non_exhaustive]
32#[serde(tag = "type", content = "body", rename_all = "snake_case")]
33pub enum ControlMessage {
34    Request(ControlRequest),
35    Response(ControlResponse),
36    Data(ControlData),
37}
38
39impl ControlMessage {
40    fn validate(&self) -> Result<(), ControlProtocolError> {
41        match self {
42            Self::Request(request) => request.validate(),
43            Self::Response(_) => Ok(()),
44            Self::Data(data) => data.validate(),
45        }
46    }
47}
48
49/// Correlated client-to-master request.
50#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
51pub struct ControlRequest {
52    pub request_id: u64,
53    pub operation: ControlOperation,
54}
55
56impl ControlRequest {
57    /// Validate semantic limits that cannot be expressed by the JSON frame
58    /// length alone.
59    pub fn validate(&self) -> Result<(), ControlProtocolError> {
60        match &self.operation {
61            ControlOperation::Hello { version } if *version != CONTROL_PROTOCOL_VERSION => {
62                Err(ControlProtocolError::UnsupportedVersion {
63                    received: *version,
64                    supported: CONTROL_PROTOCOL_VERSION,
65                })
66            }
67            ControlOperation::Hello { .. } => Ok(()),
68            ControlOperation::OpenSession(session) => session.validate(),
69            ControlOperation::Command {
70                command, forwards, ..
71            } => match command {
72                ControlCommand::Forward | ControlCommand::Cancel if forwards.is_empty() => {
73                    Err(ControlProtocolError::MissingForwarding(*command))
74                }
75                ControlCommand::Check | ControlCommand::Exit | ControlCommand::Stop
76                    if !forwards.is_empty() =>
77                {
78                    Err(ControlProtocolError::UnexpectedForwarding(*command))
79                }
80                _ => Ok(()),
81            },
82        }
83    }
84}
85
86/// Request operation understood by a bssh control master.
87#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
88#[non_exhaustive]
89#[serde(tag = "operation", rename_all = "snake_case")]
90pub enum ControlOperation {
91    Hello {
92        version: u32,
93    },
94    OpenSession(SessionOpenRequest),
95    Command {
96        command: ControlCommand,
97        /// Raw directives in original CLI/config order.
98        forwards: Vec<ForwardingDirective>,
99        /// Address family used when the master parses raw forwarding specs.
100        address_family: AddressFamily,
101    },
102}
103
104/// Session policy transferred to the authenticated control master.
105#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
106pub struct SessionOpenRequest {
107    pub policy: SessionPolicy,
108    /// Client-side terminal name used when `policy.request_pty` is true.
109    pub terminal: Option<String>,
110}
111
112impl SessionOpenRequest {
113    /// Build a wire request. `LocalCommand` must execute in the invoking client
114    /// process, never in the long-lived master.
115    pub fn new(
116        policy: SessionPolicy,
117        terminal: Option<String>,
118    ) -> Result<Self, ControlProtocolError> {
119        let request = Self { policy, terminal };
120        request.validate()?;
121        Ok(request)
122    }
123
124    /// Enforce limits before sending or accepting a session request.
125    pub fn validate(&self) -> Result<(), ControlProtocolError> {
126        if self.policy.local_command.is_some() {
127            return Err(ControlProtocolError::LocalCommandNotAllowed);
128        }
129        if self.policy.environment.len() > MAX_CONTROL_ENVIRONMENT_ENTRIES {
130            return Err(ControlProtocolError::TooManyEnvironmentEntries {
131                count: self.policy.environment.len(),
132                maximum: MAX_CONTROL_ENVIRONMENT_ENTRIES,
133            });
134        }
135        Ok(())
136    }
137}
138
139/// Correlated master-to-client response.
140#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
141pub struct ControlResponse {
142    pub request_id: u64,
143    pub kind: ControlResponseKind,
144}
145
146/// Result variants needed by #286's session and five control commands.
147#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
148#[non_exhaustive]
149#[serde(tag = "result", rename_all = "snake_case")]
150pub enum ControlResponseKind {
151    Ok,
152    Alive { pid: u32 },
153    SessionOpened { session_id: u64 },
154    ExitStatus { session_id: u64, status: u32 },
155    Error { code: String, message: String },
156}
157
158/// Byte stream represented by a data message.
159#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
160#[non_exhaustive]
161#[serde(rename_all = "snake_case")]
162pub enum ControlDataStream {
163    Stdin,
164    Stdout,
165    Stderr,
166}
167
168/// Ordered session data or EOF indication.
169#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
170pub struct ControlData {
171    pub session_id: u64,
172    pub stream: ControlDataStream,
173    pub sequence: u64,
174    pub payload: Vec<u8>,
175    pub eof: bool,
176}
177
178impl ControlData {
179    /// Validate the per-message binary payload limit.
180    pub fn validate(&self) -> Result<(), ControlProtocolError> {
181        if self.payload.len() > MAX_CONTROL_DATA_BYTES {
182            return Err(ControlProtocolError::DataTooLarge {
183                length: self.payload.len(),
184                maximum: MAX_CONTROL_DATA_BYTES,
185            });
186        }
187        Ok(())
188    }
189}
190
191/// Serialize and write one bounded big-endian-length-prefixed JSON message.
192pub async fn write_control_message<W>(
193    writer: &mut W,
194    message: &ControlMessage,
195) -> Result<(), ControlProtocolError>
196where
197    W: AsyncWrite + Unpin,
198{
199    message.validate()?;
200    let encoded = serde_json::to_vec(message)?;
201    if encoded.len() > MAX_CONTROL_FRAME_BYTES {
202        return Err(ControlProtocolError::FrameTooLarge {
203            length: encoded.len(),
204            maximum: MAX_CONTROL_FRAME_BYTES,
205        });
206    }
207    let length = u32::try_from(encoded.len()).map_err(|_| ControlProtocolError::FrameTooLarge {
208        length: encoded.len(),
209        maximum: MAX_CONTROL_FRAME_BYTES,
210    })?;
211    writer.write_all(&length.to_be_bytes()).await?;
212    writer.write_all(&encoded).await?;
213    writer.flush().await?;
214    Ok(())
215}
216
217/// Read, decode, and validate one bounded big-endian-length-prefixed JSON
218/// message without allocating an attacker-declared oversized frame.
219pub async fn read_control_message<R>(reader: &mut R) -> Result<ControlMessage, ControlProtocolError>
220where
221    R: AsyncRead + Unpin,
222{
223    let mut header = [0u8; 4];
224    reader.read_exact(&mut header).await?;
225    let length = u32::from_be_bytes(header) as usize;
226    if length == 0 {
227        return Err(ControlProtocolError::EmptyFrame);
228    }
229    if length > MAX_CONTROL_FRAME_BYTES {
230        return Err(ControlProtocolError::FrameTooLarge {
231            length,
232            maximum: MAX_CONTROL_FRAME_BYTES,
233        });
234    }
235
236    let mut encoded = vec![0u8; length];
237    reader.read_exact(&mut encoded).await?;
238    let message = serde_json::from_slice::<ControlMessage>(&encoded)?;
239    message.validate()?;
240    Ok(message)
241}
242
243/// Framing, serialization, version, or semantic-limit failure.
244#[derive(Debug, Error)]
245#[non_exhaustive]
246pub enum ControlProtocolError {
247    #[error("control socket I/O failed: {0}")]
248    Io(#[from] io::Error),
249    #[error("control message JSON is invalid: {0}")]
250    Json(#[from] serde_json::Error),
251    #[error("control protocol frame must not be empty")]
252    EmptyFrame,
253    #[error("control protocol frame is {length} bytes, exceeding the {maximum}-byte limit")]
254    FrameTooLarge { length: usize, maximum: usize },
255    #[error("control data chunk is {length} bytes, exceeding the {maximum}-byte limit")]
256    DataTooLarge { length: usize, maximum: usize },
257    #[error("control protocol version {received} is unsupported; this build supports {supported}")]
258    UnsupportedVersion { received: u32, supported: u32 },
259    #[error(
260        "LocalCommand must execute in the invoking process and cannot be sent to a control master"
261    )]
262    LocalCommandNotAllowed,
263    #[error("session request contains {count} environment entries; maximum is {maximum}")]
264    TooManyEnvironmentEntries { count: usize, maximum: usize },
265    #[error("control command '{0}' requires at least one forwarding directive")]
266    MissingForwarding(ControlCommand),
267    #[error("control command '{0}' does not accept forwarding directives")]
268    UnexpectedForwarding(ControlCommand),
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274    use crate::ssh::SessionRequest;
275
276    fn session_policy() -> SessionPolicy {
277        SessionPolicy {
278            environment: vec![("LANG".into(), "C.UTF-8".into())],
279            local_command: None,
280            forward_agent: true,
281            request_pty: false,
282            stdin_null: false,
283            request: SessionRequest::Exec("printf test".into()),
284        }
285    }
286
287    #[tokio::test]
288    async fn request_round_trip_preserves_session_policy() {
289        let request = ControlMessage::Request(ControlRequest {
290            request_id: 7,
291            operation: ControlOperation::OpenSession(
292                SessionOpenRequest::new(session_policy(), Some("xterm-256color".into()))
293                    .expect("wire-safe session"),
294            ),
295        });
296        let (mut client, mut server) = tokio::io::duplex(16 * 1024);
297        let write = write_control_message(&mut client, &request);
298        let read = read_control_message(&mut server);
299        let (written, decoded) = tokio::join!(write, read);
300        written.expect("write succeeds");
301        assert_eq!(decoded.expect("read succeeds"), request);
302    }
303
304    #[tokio::test]
305    async fn data_round_trip_is_binary_safe() {
306        let message = ControlMessage::Data(ControlData {
307            session_id: 9,
308            stream: ControlDataStream::Stdout,
309            sequence: 3,
310            payload: vec![0, 255, b'\n', 0],
311            eof: true,
312        });
313        let (mut client, mut server) = tokio::io::duplex(16 * 1024);
314        let write = write_control_message(&mut client, &message);
315        let read = read_control_message(&mut server);
316        let (written, decoded) = tokio::join!(write, read);
317        written.expect("write succeeds");
318        assert_eq!(decoded.expect("read succeeds"), message);
319    }
320
321    #[tokio::test]
322    async fn oversized_declared_frame_is_rejected_before_payload_read() {
323        let (mut client, mut server) = tokio::io::duplex(16);
324        client
325            .write_all(&((MAX_CONTROL_FRAME_BYTES as u32) + 1).to_be_bytes())
326            .await
327            .expect("header write");
328        let error = read_control_message(&mut server)
329            .await
330            .expect_err("oversized frame");
331        assert!(matches!(error, ControlProtocolError::FrameTooLarge { .. }));
332    }
333
334    #[tokio::test]
335    async fn oversized_data_is_rejected_before_serialization() {
336        let message = ControlMessage::Data(ControlData {
337            session_id: 1,
338            stream: ControlDataStream::Stdin,
339            sequence: 0,
340            payload: vec![0; MAX_CONTROL_DATA_BYTES + 1],
341            eof: false,
342        });
343        let mut sink = tokio::io::sink();
344        let error = write_control_message(&mut sink, &message)
345            .await
346            .expect_err("oversized data");
347        assert!(matches!(error, ControlProtocolError::DataTooLarge { .. }));
348    }
349
350    #[test]
351    fn local_command_cannot_cross_the_master_boundary() {
352        let mut policy = session_policy();
353        policy.local_command = Some("touch /tmp/client-only".into());
354        assert!(matches!(
355            SessionOpenRequest::new(policy, None),
356            Err(ControlProtocolError::LocalCommandNotAllowed)
357        ));
358    }
359
360    #[test]
361    fn command_forwarding_shape_is_validated() {
362        let missing = ControlRequest {
363            request_id: 1,
364            operation: ControlOperation::Command {
365                command: ControlCommand::Forward,
366                forwards: Vec::new(),
367                address_family: AddressFamily::Any,
368            },
369        };
370        assert!(matches!(
371            missing.validate(),
372            Err(ControlProtocolError::MissingForwarding(
373                ControlCommand::Forward
374            ))
375        ));
376
377        let unexpected = ControlRequest {
378            request_id: 2,
379            operation: ControlOperation::Command {
380                command: ControlCommand::Check,
381                forwards: vec![ForwardingDirective::Dynamic("1080".into())],
382                address_family: AddressFamily::Any,
383            },
384        };
385        assert!(matches!(
386            unexpected.validate(),
387            Err(ControlProtocolError::UnexpectedForwarding(
388                ControlCommand::Check
389            ))
390        ));
391    }
392}
bssh/ssh/control/protocol.rs

bssh/ssh/control/
protocol.rs
