winx_code_agent/daemon/
control.rs

1use std::fmt::Write as FmtWrite;
2use std::io::ErrorKind;
3use std::os::unix::fs::PermissionsExt;
4use std::path::{Path, PathBuf};
5
6use sha2::{Digest, Sha256};
7use tokio::net::{UnixListener, UnixStream};
8
9use super::protocol::{
10    read_json_frame, write_json_frame, ConfigureSessionParams, HelloResult, JournalReadParams,
11    RpcError, RpcRequest, RpcResponse, RunActionParams, SessionInfo, SessionParams,
12    MAX_FRAME_BYTES, PROTOCOL_MAJOR, PROTOCOL_MINOR,
13};
14use crate::daemon::DaemonClient;
15use crate::errors::{Result, WinxError};
16use crate::runtime::ensure_daemon_at;
17use crate::types::normalize_thread_id;
18
19/// Stable control plane. Each logical session is owned by a separate guardian
20/// process, so restarting this process does not close any PTY master.
21pub struct ControlServer {
22    listener: UnixListener,
23    socket_path: PathBuf,
24    guardian_dir: PathBuf,
25    guardian_binary: PathBuf,
26    epoch: String,
27}
28
29impl ControlServer {
30    pub async fn bind(socket_path: impl AsRef<Path>) -> Result<Self> {
31        let socket_path = socket_path.as_ref().to_path_buf();
32        let parent = socket_path.parent().ok_or_else(|| {
33            WinxError::ConfigurationError("daemon socket must have a parent directory".to_string())
34        })?;
35        tokio::fs::create_dir_all(parent).await?;
36        tokio::fs::set_permissions(parent, std::fs::Permissions::from_mode(0o700)).await?;
37        let guardian_dir = parent.join("guardians");
38        tokio::fs::create_dir_all(&guardian_dir).await?;
39        tokio::fs::set_permissions(&guardian_dir, std::fs::Permissions::from_mode(0o700)).await?;
40
41        if tokio::fs::try_exists(&socket_path).await? {
42            if UnixStream::connect(&socket_path).await.is_ok() {
43                return Err(WinxError::ConfigurationError(format!(
44                    "a winxd instance is already listening at {}",
45                    socket_path.display()
46                )));
47            }
48            tokio::fs::remove_file(&socket_path).await?;
49        }
50        let listener = UnixListener::bind(&socket_path)?;
51        tokio::fs::set_permissions(&socket_path, std::fs::Permissions::from_mode(0o600)).await?;
52
53        Ok(Self {
54            listener,
55            socket_path,
56            guardian_dir,
57            guardian_binary: guardian_binary()?,
58            epoch: format!("{:016x}", rand::random::<u64>()),
59        })
60    }
61
62    pub async fn serve(self) -> Result<()> {
63        loop {
64            let (stream, _) = self.listener.accept().await?;
65            if !same_uid(&stream)? {
66                tracing::warn!("Rejected winxd control connection from a different uid");
67                continue;
68            }
69            let guardian_dir = self.guardian_dir.clone();
70            let guardian_binary = self.guardian_binary.clone();
71            let epoch = self.epoch.clone();
72            tokio::spawn(async move {
73                if let Err(error) =
74                    serve_connection(stream, guardian_dir, guardian_binary, epoch).await
75                {
76                    tracing::debug!("winxd control client disconnected: {error}");
77                }
78            });
79        }
80    }
81}
82
83impl Drop for ControlServer {
84    fn drop(&mut self) {
85        let _ = std::fs::remove_file(&self.socket_path);
86    }
87}
88
89async fn serve_connection(
90    mut stream: UnixStream,
91    guardian_dir: PathBuf,
92    guardian_binary: PathBuf,
93    epoch: String,
94) -> Result<()> {
95    loop {
96        let request: RpcRequest = match read_json_frame(&mut stream).await {
97            Ok(request) => request,
98            Err(error) if error.kind() == ErrorKind::UnexpectedEof => return Ok(()),
99            Err(error) => return Err(error.into()),
100        };
101        let response = dispatch(request, &guardian_dir, &guardian_binary, &epoch).await;
102        write_json_frame(&mut stream, &response).await?;
103    }
104}
105
106async fn dispatch(
107    request: RpcRequest,
108    guardian_dir: &Path,
109    guardian_binary: &Path,
110    epoch: &str,
111) -> RpcResponse {
112    if request.jsonrpc != "2.0" {
113        return rpc_error(request.id, -32600, "JSON-RPC version must be 2.0");
114    }
115    if request.method == "winx.hello" {
116        return rpc_result(
117            request.id,
118            &HelloResult {
119                protocol_major: PROTOCOL_MAJOR,
120                protocol_minor: PROTOCOL_MINOR,
121                capabilities: vec![
122                    "per_session_guardians".to_string(),
123                    "planned_control_restart".to_string(),
124                    "session.configure".to_string(),
125                    "shell.run_action".to_string(),
126                    "session.list".to_string(),
127                    "session.read_output".to_string(),
128                    "session.kill".to_string(),
129                    "multi_consumer_cursors".to_string(),
130                    "idempotency".to_string(),
131                ],
132                max_frame_bytes: MAX_FRAME_BYTES,
133                daemon_epoch: epoch.to_string(),
134                daemon_pid: std::process::id(),
135            },
136        );
137    }
138
139    if request.method == "session.list" {
140        return match list_sessions(guardian_dir).await {
141            Ok(sessions) => rpc_result(request.id, &sessions),
142            Err(error) => rpc_error(request.id, -32000, &error.to_string()),
143        };
144    }
145
146    let thread_id = match request_thread_id(&request) {
147        Ok(thread_id) if !thread_id.is_empty() => thread_id,
148        Ok(_) => return rpc_error(request.id, -32602, "explicit thread_id is required"),
149        Err(error) => return rpc_error(request.id, -32602, &error.to_string()),
150    };
151    let guardian_socket = guardian_socket(guardian_dir, &thread_id);
152    if matches!(request.method.as_str(), "session.configure" | "shell.run_action") {
153        if let Err(error) = ensure_daemon_at(&guardian_socket, guardian_binary).await {
154            return rpc_error(request.id, -32001, &error.to_string());
155        }
156    }
157
158    let kill_after = request.method == "session.kill";
159    let request_id = request.id;
160    let response = match relay(&guardian_socket, request).await {
161        Ok(response) => response,
162        Err(error) => return rpc_error(request_id, -32002, &error.to_string()),
163    };
164    if kill_after && response.error.is_none() {
165        terminate_guardian(&guardian_socket).await;
166    }
167    response
168}
169
170fn request_thread_id(request: &RpcRequest) -> Result<String> {
171    let thread_id = match request.method.as_str() {
172        "session.configure" => {
173            decode::<ConfigureSessionParams>(request.params.clone())?.snapshot.chat_id
174        }
175        "shell.run_action" => decode::<RunActionParams>(request.params.clone())?.command.thread_id,
176        "session.read_output" => decode::<JournalReadParams>(request.params.clone())?.thread_id,
177        "session.info" | "session.kill" | "session.interrupt" => {
178            decode::<SessionParams>(request.params.clone())?.thread_id
179        }
180        _ => {
181            return Err(WinxError::InvalidInput(format!(
182                "unsupported control method {}",
183                request.method
184            )))
185        }
186    };
187    Ok(normalize_thread_id(&thread_id))
188}
189
190async fn relay(socket: &Path, request: RpcRequest) -> Result<RpcResponse> {
191    let mut stream = UnixStream::connect(socket).await.map_err(|error| {
192        WinxError::ShellInitializationError(format!(
193            "guardian at {} is unavailable: {error}",
194            socket.display()
195        ))
196    })?;
197    let hello_id = rand::random::<u64>();
198    write_json_frame(
199        &mut stream,
200        &RpcRequest {
201            jsonrpc: "2.0".to_string(),
202            id: hello_id,
203            method: "winx.hello".to_string(),
204            params: serde_json::json!({ "protocol_major": PROTOCOL_MAJOR }),
205        },
206    )
207    .await?;
208    let hello: RpcResponse = read_json_frame(&mut stream).await?;
209    let hello: HelloResult = decode(hello.result.ok_or_else(|| {
210        WinxError::ConfigurationError("guardian handshake omitted result".to_string())
211    })?)?;
212    if hello.protocol_major != PROTOCOL_MAJOR {
213        return Err(WinxError::ConfigurationError(format!(
214            "guardian protocol major {} is incompatible with control {}",
215            hello.protocol_major, PROTOCOL_MAJOR
216        )));
217    }
218    write_json_frame(&mut stream, &request).await?;
219    Ok(read_json_frame(&mut stream).await?)
220}
221
222async fn list_sessions(guardian_dir: &Path) -> Result<Vec<SessionInfo>> {
223    let mut sessions = Vec::new();
224    let mut entries = tokio::fs::read_dir(guardian_dir).await?;
225    while let Some(entry) = entries.next_entry().await? {
226        let path = entry.path();
227        if path.extension().and_then(|extension| extension.to_str()) != Some("sock") {
228            continue;
229        }
230        let request = RpcRequest {
231            jsonrpc: "2.0".to_string(),
232            id: rand::random::<u64>(),
233            method: "session.list".to_string(),
234            params: serde_json::json!({}),
235        };
236        let Ok(response) = relay(&path, request).await else { continue };
237        let Some(result) = response.result else { continue };
238        if let Ok(mut guardian_sessions) = serde_json::from_value::<Vec<SessionInfo>>(result) {
239            sessions.append(&mut guardian_sessions);
240        }
241    }
242    sessions.sort_by(|left, right| left.thread_id.cmp(&right.thread_id));
243    Ok(sessions)
244}
245
246async fn terminate_guardian(socket: &Path) {
247    let Ok(hello) = DaemonClient::new(socket).hello().await else { return };
248    let Ok(pid) = i32::try_from(hello.daemon_pid) else { return };
249    // The handshake came from this exact same-uid guardian socket, so the pid is
250    // resolved rather than guessed. SIGTERM affects only that guardian process.
251    unsafe {
252        libc::kill(pid, libc::SIGTERM);
253    }
254}
255
256fn guardian_socket(guardian_dir: &Path, thread_id: &str) -> PathBuf {
257    let digest = Sha256::digest(thread_id.as_bytes());
258    let mut name = String::with_capacity(24);
259    for byte in &digest[..12] {
260        let _ = write!(name, "{byte:02x}");
261    }
262    guardian_dir.join(format!("{name}.sock"))
263}
264
265fn decode<T: serde::de::DeserializeOwned>(value: serde_json::Value) -> Result<T> {
266    serde_json::from_value(value)
267        .map_err(|error| WinxError::DeserializationError(error.to_string()))
268}
269
270fn guardian_binary() -> Result<PathBuf> {
271    if let Some(path) = std::env::var_os("WINX_GUARDIAN_BIN") {
272        return Ok(PathBuf::from(path));
273    }
274    let executable = std::env::current_exe()?;
275    let sibling = executable.with_file_name("winx-guardian");
276    if sibling.is_file() {
277        Ok(sibling)
278    } else {
279        Err(WinxError::ConfigurationError(format!(
280            "winx-guardian not found beside {} (set WINX_GUARDIAN_BIN)",
281            executable.display()
282        )))
283    }
284}
285
286fn rpc_result<T: serde::Serialize>(id: u64, result: &T) -> RpcResponse {
287    match serde_json::to_value(result) {
288        Ok(result) => {
289            RpcResponse { jsonrpc: "2.0".to_string(), id, result: Some(result), error: None }
290        }
291        Err(error) => rpc_error(id, -32603, &error.to_string()),
292    }
293}
294
295fn rpc_error(id: u64, code: i32, message: &str) -> RpcResponse {
296    RpcResponse {
297        jsonrpc: "2.0".to_string(),
298        id,
299        result: None,
300        error: Some(RpcError { code, message: message.to_string() }),
301    }
302}
303
304fn same_uid(stream: &UnixStream) -> Result<bool> {
305    let credentials = stream.peer_cred()?;
306    // SAFETY: geteuid has no preconditions and reads process credentials only.
307    Ok(credentials.uid() == unsafe { libc::geteuid() })
308}
winx_code_agent/daemon/control.rs

winx_code_agent/daemon/
control.rs
