1use crate::WorkerError;
2use crate::api;
3use crate::decode_base64;
4use base64::Engine as _;
5use serde::{Deserialize, Serialize};
6
7#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
8#[serde(deny_unknown_fields)]
9pub struct ResourceLimits {
10 #[serde(default, skip_serializing_if = "Option::is_none")]
11 pub output_buffer_bytes: Option<u32>,
12 #[serde(default, skip_serializing_if = "Option::is_none")]
13 pub max_runtime_ms: Option<u32>,
14}
15
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(deny_unknown_fields)]
18pub struct StartRequest {
19 pub program: String,
20 #[serde(default, skip_serializing_if = "Vec::is_empty")]
21 pub argv: Vec<String>,
22 #[serde(default, skip_serializing_if = "Option::is_none")]
23 pub cwd: Option<String>,
24 #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
25 pub environment: std::collections::BTreeMap<String, String>,
26 #[serde(default, skip_serializing_if = "Vec::is_empty")]
27 pub environment_removals: Vec<String>,
28 #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
29 pub secret_references: std::collections::BTreeMap<String, String>,
30 pub client_key: String,
31 #[serde(default, skip_serializing_if = "Option::is_none")]
32 pub resource_limits: Option<ResourceLimits>,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
36#[serde(deny_unknown_fields)]
37pub struct Status {
38 pub handle: String,
39 pub client_key: String,
40 pub state: String,
41 #[serde(default)]
42 pub exit_code: Option<i32>,
43 #[serde(default)]
44 pub termination_reason: Option<String>,
45}
46
47#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
48#[serde(deny_unknown_fields)]
49pub struct WriteResponse {
50 pub written: u32,
51}
52
53#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
54#[serde(deny_unknown_fields)]
55pub struct ReadRequest {
56 pub handle: String,
57 #[serde(default)]
58 pub cursor: u64,
59 pub max_bytes: u32,
60}
61
62#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
63#[serde(deny_unknown_fields)]
64pub struct ReadResponse {
65 pub data_base64: String,
66 pub cursor: u64,
67 pub eof: bool,
68 pub process_exit: bool,
69 pub stream_gap: bool,
70 #[serde(default)]
71 pub dropped_bytes: u64,
72}
73
74impl ReadResponse {
75 pub fn data(&self) -> Result<Vec<u8>, WorkerError> {
76 decode_base64(&self.data_base64)
77 }
78}
79
80#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
81#[serde(deny_unknown_fields)]
82pub struct ExitResponse {
83 #[serde(default)]
84 pub exit_code: Option<i32>,
85 pub termination_reason: String,
86}
87
88#[derive(Serialize)]
89struct WriteRequest<'a> {
90 handle: &'a str,
91 data_base64: String,
92}
93
94#[derive(Serialize)]
95struct HandleRequest<'a> {
96 handle: &'a str,
97}
98
99#[derive(Serialize)]
100struct ClientKeyRequest<'a> {
101 client_key: &'a str,
102}
103
104pub fn start(request: &StartRequest) -> Result<Status, WorkerError> {
105 call("process.start", request)
106}
107
108pub fn attach(client_key: &str) -> Result<Status, WorkerError> {
109 call("process.attach", &ClientKeyRequest { client_key })
110}
111
112pub fn status(handle: &str) -> Result<Status, WorkerError> {
113 call("process.status", &HandleRequest { handle })
114}
115
116pub fn write_stdin(handle: &str, data: &[u8]) -> Result<WriteResponse, WorkerError> {
117 call(
118 "process.write_stdin",
119 &WriteRequest {
120 handle,
121 data_base64: base64::engine::general_purpose::STANDARD.encode(data),
122 },
123 )
124}
125
126pub fn close_stdin(handle: &str) -> Result<(), WorkerError> {
127 call::<_, serde_json::Value>("process.close_stdin", &HandleRequest { handle }).map(|_| ())
128}
129
130pub fn read_stdout(request: &ReadRequest) -> Result<ReadResponse, WorkerError> {
131 call("process.read_stdout", request)
132}
133
134pub fn read_stderr(request: &ReadRequest) -> Result<ReadResponse, WorkerError> {
135 call("process.read_stderr", request)
136}
137
138pub fn wait(handle: &str) -> Result<ExitResponse, WorkerError> {
139 call("process.wait", &HandleRequest { handle })
140}
141
142pub fn terminate(handle: &str) -> Result<(), WorkerError> {
143 call::<_, serde_json::Value>("process.terminate", &HandleRequest { handle }).map(|_| ())
144}
145
146pub fn kill(handle: &str) -> Result<(), WorkerError> {
147 call::<_, serde_json::Value>("process.kill", &HandleRequest { handle }).map(|_| ())
148}
149
150pub fn close(handle: &str) -> Result<(), WorkerError> {
151 call::<_, serde_json::Value>("process.close", &HandleRequest { handle }).map(|_| ())
152}
153
154fn call<RequestValue, ResponseValue>(
155 operation: &str,
156 request: &RequestValue,
157) -> Result<ResponseValue, WorkerError>
158where
159 RequestValue: Serialize,
160 ResponseValue: for<'de> Deserialize<'de>,
161{
162 api::call(operation, request).map_err(Into::into)
163}
164
165#[cfg(test)]
166mod tests {
167 use super::*;
168
169 #[test]
170 fn process_read_result_decodes_binary_data_without_text_assumptions() {
171 let result: ReadResponse = serde_json::from_value(serde_json::json!({
172 "data_base64": "AP8B",
173 "cursor": 3,
174 "eof": false,
175 "process_exit": false,
176 "stream_gap": false
177 }))
178 .unwrap();
179 assert_eq!(result.data().unwrap(), vec![0, 255, 1]);
180 }
181
182 #[test]
183 fn process_start_request_contains_secret_references_not_secret_values() {
184 let request = StartRequest {
185 program: "tool".into(),
186 argv: vec!["--json".into()],
187 cwd: None,
188 environment: std::collections::BTreeMap::new(),
189 environment_removals: Vec::new(),
190 secret_references: std::collections::BTreeMap::from([(
191 "TOKEN".into(),
192 "secret.token".into(),
193 )]),
194 client_key: "client".into(),
195 resource_limits: None,
196 };
197 let encoded = serde_json::to_value(request).unwrap();
198 assert_eq!(encoded["secret_references"]["TOKEN"], "secret.token");
199 assert!(encoded.get("TOKEN").is_none());
200 }
201}