1use crate::{PrayError, PrayResult};
2use base64::{engine::general_purpose::STANDARD, Engine as _};
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use std::io::{Read, Write};
6
7pub const SSH_RPC_SPEC: &str = "pray-ssh-rpc-v1";
8pub const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
11pub struct RpcRequest {
12 pub spec: String,
13 pub id: String,
14 pub method: String,
15 #[serde(default)]
16 pub params: Value,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
20pub struct RpcResponse {
21 pub spec: String,
22 pub id: String,
23 pub status: u16,
24 pub content_type: String,
25 #[serde(skip_serializing_if = "Option::is_none")]
26 pub body_encoding: Option<String>,
27 pub body: Value,
28}
29
30impl RpcRequest {
31 pub fn new(id: impl Into<String>, method: impl Into<String>, params: Value) -> Self {
32 Self {
33 spec: SSH_RPC_SPEC.to_string(),
34 id: id.into(),
35 method: method.into(),
36 params,
37 }
38 }
39}
40
41impl RpcResponse {
42 pub fn json_ok(id: impl Into<String>, body: Value) -> Self {
43 Self {
44 spec: SSH_RPC_SPEC.to_string(),
45 id: id.into(),
46 status: 200,
47 content_type: "application/json".to_string(),
48 body_encoding: None,
49 body,
50 }
51 }
52
53 pub fn binary_ok(id: impl Into<String>, bytes: &[u8]) -> Self {
54 Self {
55 spec: SSH_RPC_SPEC.to_string(),
56 id: id.into(),
57 status: 200,
58 content_type: "application/octet-stream".to_string(),
59 body_encoding: Some("base64".to_string()),
60 body: Value::String(STANDARD.encode(bytes)),
61 }
62 }
63
64 pub fn error(id: impl Into<String>, status: u16, message: impl Into<String>) -> Self {
65 Self {
66 spec: SSH_RPC_SPEC.to_string(),
67 id: id.into(),
68 status,
69 content_type: "application/json".to_string(),
70 body_encoding: None,
71 body: serde_json::json!({ "error": message.into() }),
72 }
73 }
74
75 pub fn decode_body_bytes(&self) -> PrayResult<Vec<u8>> {
76 if self.content_type == "application/octet-stream"
77 && self.body_encoding.as_deref() == Some("base64")
78 {
79 let encoded = self.body.as_str().ok_or_else(|| {
80 PrayError::Resolution("rpc binary body must be a base64 string".to_string())
81 })?;
82 STANDARD.decode(encoded).map_err(|error| {
83 PrayError::Resolution(format!("rpc binary body base64 decode failed: {error}"))
84 })
85 } else if self.body.is_string() {
86 Ok(self.body.as_str().unwrap_or_default().as_bytes().to_vec())
87 } else {
88 serde_json::to_vec(&self.body).map_err(|error| PrayError::Manifest(error.to_string()))
89 }
90 }
91
92 pub fn decode_json_body<T: for<'de> Deserialize<'de>>(&self) -> PrayResult<T> {
93 if self.status / 100 != 2 {
94 return Err(PrayError::Resolution(format!(
95 "rpc {} failed with status {}",
96 self.id, self.status
97 )));
98 }
99 serde_json::from_value(self.body.clone()).map_err(|error| PrayError::Parse {
100 kind: "ssh rpc response",
101 message: error.to_string(),
102 })
103 }
104}
105
106pub fn write_frame(writer: &mut impl Write, payload: &[u8]) -> PrayResult<()> {
107 if payload.len() > MAX_FRAME_BYTES {
108 return Err(PrayError::Unsupported(format!(
109 "rpc frame exceeds maximum size of {MAX_FRAME_BYTES} bytes"
110 )));
111 }
112 let length = u32::try_from(payload.len())
113 .map_err(|_| PrayError::Unsupported("rpc frame length overflow".to_string()))?;
114 writer
115 .write_all(&length.to_be_bytes())
116 .map_err(PrayError::from)?;
117 writer.write_all(payload).map_err(PrayError::from)?;
118 writer.flush().map_err(PrayError::from)?;
119 Ok(())
120}
121
122pub fn read_frame(reader: &mut impl Read) -> PrayResult<Vec<u8>> {
123 let mut length_bytes = [0u8; 4];
124 match reader.read_exact(&mut length_bytes) {
125 Ok(()) => {}
126 Err(error) if error.kind() == std::io::ErrorKind::UnexpectedEof => {
127 return Err(PrayError::Resolution("rpc stream closed".to_string()));
128 }
129 Err(error) => return Err(error.into()),
130 }
131 let length = u32::from_be_bytes(length_bytes) as usize;
132 if length > MAX_FRAME_BYTES {
133 return Err(PrayError::Unsupported(format!(
134 "rpc frame exceeds maximum size of {MAX_FRAME_BYTES} bytes"
135 )));
136 }
137 let mut payload = vec![0u8; length];
138 reader.read_exact(&mut payload).map_err(PrayError::from)?;
139 Ok(payload)
140}
141
142pub fn call_stdio(
143 reader: &mut impl Read,
144 writer: &mut impl Write,
145 request: &RpcRequest,
146) -> PrayResult<RpcResponse> {
147 let payload =
148 serde_json::to_vec(request).map_err(|error| PrayError::Manifest(error.to_string()))?;
149 write_frame(writer, &payload)?;
150 let response_bytes = read_frame(reader)?;
151 serde_json::from_slice(&response_bytes).map_err(|error| PrayError::Parse {
152 kind: "ssh rpc response",
153 message: error.to_string(),
154 })
155}
156
157#[cfg(test)]
158mod tests {
159 use super::*;
160
161 #[test]
162 fn frame_round_trip_preserves_payload() {
163 let payload = br#"{"spec":"pray-ssh-rpc-v1"}"#;
164 let mut buffer = Vec::new();
165 write_frame(&mut buffer, payload).expect("write frame");
166 let mut cursor = std::io::Cursor::new(buffer);
167 let decoded = read_frame(&mut cursor).expect("read frame");
168 assert_eq!(decoded, payload);
169 }
170
171 #[test]
172 fn binary_response_round_trips_base64() {
173 let response = RpcResponse::binary_ok("1", b"artifact-bytes");
174 let bytes = response.decode_body_bytes().expect("decode bytes");
175 assert_eq!(bytes, b"artifact-bytes");
176 }
177}