client.rs 72.7%

120 of 165 lines run, 95 of 190 functions · run #36 on main · run missed

//! A client for the daemon over its Unix socket (spec §8.5), signing every
//! request with the caller's key: the shims and the CLI use it, and so do
//! the tests.

use anyhow::{Context, Result};
use http_body_util::{BodyExt, Full};
use hyper::body::Bytes;
use spectra_core::keys::KeyPair;
use spectra_core::proof::Proof;
use std::path::{Path, PathBuf};

/// A response, its body read whole.
#[derive(Debug)]
pub struct Response {
    pub status: u16,
    pub headers: Vec<(String, String)>,
    pub body: Vec<u8>,
}

impl Response {
    pub fn json(&self) -> Result<serde_json::Value> {
        serde_json::from_slice(&self.body).with_context(|| format!("not JSON: {}", String::from_utf8_lossy(&self.body)))
    }
    pub fn header(&self, name: &str) -> Option<&str> {
        self.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str())
    }
}

/// Where the daemon is: its Unix socket, or a server over TLS with a
/// pinned certificate (spec §8.6).
#[derive(Clone)]
pub enum Endpoint {
    Unix(PathBuf),
    Tls { addr: String, config: std::sync::Arc<rustls::ClientConfig> },
}

impl Endpoint {
    /// `https://HOST:PORT` accepting exactly the certificate pinned by `pin`.
    pub fn tls(url: &str, pin: [u8; 32]) -> Result<Endpoint> {
        Ok(Endpoint::Tls { addr: host_port(url)?, config: crate::tls::Pinned::new(Some(pin)).client_config()? })
    }

    pub fn describe(&self) -> String {
        match self {
            Endpoint::Unix(p) => p.display().to_string(),
            Endpoint::Tls { addr, .. } => format!("https://{addr}"),
        }
    }
}

/// `host:port` of an `https://` URL.
pub fn host_port(url: &str) -> Result<String> {
    let rest = url.strip_prefix("https://").with_context(|| format!("{url} is not an https:// URL"))?;
    let hp = rest.split('/').next().unwrap_or("");
    if hp.is_empty() {
        anyhow::bail!("{url} names no host");
    }
    // The port defaults to Spectra's.
    let has_port = if hp.starts_with('[') { hp.contains("]:") } else { hp.contains(':') };
    Ok(if has_port { hp.to_string() } else { format!("{hp}:7350") })
}

trait Io: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send {}
impl<T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send> Io for T {}

pub struct Client {
    endpoint: Endpoint,
    signer: Option<KeyPair>,
}

impl Client {
    pub fn new(socket: &Path, signer: Option<KeyPair>) -> Client {
        Client { endpoint: Endpoint::Unix(socket.to_path_buf()), signer }
    }

    pub fn at(endpoint: Endpoint, signer: Option<KeyPair>) -> Client {
        Client { endpoint, signer }
    }

    pub fn endpoint(&self) -> &Endpoint {
        &self.endpoint
    }

    /// The key this client signs with, if any.
    pub fn signer(&self) -> Option<&KeyPair> {
        self.signer.as_ref()
    }

    async fn io(&self) -> Result<Box<dyn Io>> {
        Ok(match &self.endpoint {
            Endpoint::Unix(socket) => Box::new(
                tokio::net::UnixStream::connect(socket)
                    .await
                    .with_context(|| format!("connecting to {} (is `spectra serve` running?)", socket.display()))?,
            ),
            Endpoint::Tls { addr, config } => Box::new(crate::tls::connect(addr, config.clone()).await?),
        })
    }

    /// Whether the daemon answers at the endpoint.
    pub async fn reachable(&self) -> bool {
        self.io().await.is_ok()
    }

    /// Whether a daemon answers on the socket.
    pub async fn alive(socket: &Path) -> bool {
        tokio::net::UnixStream::connect(socket).await.is_ok()
    }

    /// Send `method path` with `body` and extra headers, signed.
    pub async fn send(&self, method: &str, path: &str, headers: &[(&str, String)], body: Vec<u8>) -> Result<Response> {
        let stream = self.io().await?;
        let (mut sender, conn) = hyper::client::conn::http1::handshake(hyper_util::rt::TokioIo::new(stream)).await?;
        tokio::spawn(async move {
            let _ = conn.await;
        });
        // One request a connection: the server closes it after answering, so
        // nothing lingers to hold up its shutdown.
        let mut req = hyper::Request::builder().method(method).uri(path).header("host", "spectra").header("connection", "close");
        // JSON unless the caller says otherwise (a form, say).
        if !headers.iter().any(|(k, _)| k.eq_ignore_ascii_case("content-type")) {
            req = req.header("content-type", "application/json");
        }
        for (k, v) in headers {
            req = req.header(*k, v);
        }
        if let Some(pair) = &self.signer {
            for (k, v) in sign_headers(pair, &format!("{method} {path}"), &body) {
                req = req.header(k, v);
            }
        }
        let resp = sender.send_request(req.body(Full::new(Bytes::from(body)))?).await?;
        let status = resp.status().as_u16();
        let headers = resp.headers().iter().map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string())).collect();
        let body = resp.into_body().collect().await?.to_bytes().to_vec();
        Ok(Response { status, headers, body })
    }

    /// A JSON-RPC request to `/mcp` with the protocol headers of `version`.
    pub async fn mcp(&self, version: &str, request: &serde_json::Value) -> Result<Response> {
        let method = request.get("method").and_then(|m| m.as_str()).unwrap_or("").to_string();
        let mut headers = vec![
            ("accept", "application/json, text/event-stream".to_string()),
            ("mcp-protocol-version", version.to_string()),
            ("mcp-method", method.clone()),
        ];
        if let Some(name) = request.pointer("/params/name").or_else(|| request.pointer("/params/uri")).and_then(|n| n.as_str()) {
            headers.push(("mcp-name", name.to_string()));
        }
        self.send("POST", "/mcp", &headers, serde_json::to_vec(request)?).await
    }
}

impl Client {
    /// Pass a request through as it is, unsigned, and hand back the
    /// response with its body still streaming (Radar's proxy, spec §8.6).
    pub async fn forward(&self, req: hyper::Request<Full<Bytes>>) -> Result<hyper::Response<hyper::body::Incoming>> {
        let stream = self.io().await?;
        let (mut sender, conn) = hyper::client::conn::http1::handshake(hyper_util::rt::TokioIo::new(stream)).await?;
        tokio::spawn(async move {
            let _ = conn.await;
        });
        Ok(sender.send_request(req).await?)
    }

    /// Send a signed GET and hand each line of the streamed body to `line`
    /// as it arrives (an SSE stream, say); stop when it returns false.
    pub async fn stream(&self, path: &str, mut line: impl FnMut(&str) -> bool) -> Result<u16> {
        let stream = self.io().await?;
        let (mut sender, conn) = hyper::client::conn::http1::handshake(hyper_util::rt::TokioIo::new(stream)).await?;
        tokio::spawn(async move {
            let _ = conn.await;
        });
        let mut req = hyper::Request::builder().method("GET").uri(path).header("host", "spectra").header("connection", "close");
        if let Some(pair) = &self.signer {
            for (k, v) in sign_headers(pair, &format!("GET {path}"), b"") {
                req = req.header(k, v);
            }
        }
        let resp = sender.send_request(req.body(Full::new(Bytes::new()))?).await?;
        let status = resp.status().as_u16();
        let mut body = resp.into_body();
        let mut buf = String::new();
        while let Some(frame) = body.frame().await {
            if let Ok(data) = frame?.into_data() {
                buf.push_str(&String::from_utf8_lossy(&data));
                while let Some(i) = buf.find('\n') {
                    let l: String = buf.drain(..=i).collect();
                    if !line(l.trim_end()) {
                        return Ok(status);
                    }
                }
            }
        }
        Ok(status)
    }
}

/// The signature headers of spec §8.5 for `method_path` and `body`.
pub fn sign_headers(pair: &KeyPair, method_path: &str, body: &[u8]) -> Vec<(&'static str, String)> {
    let p = Proof::sign(pair, method_path, body);
    vec![
        ("spectra-key", p.sig.key.to_hex()),
        ("spectra-time", p.time_ms.to_string()),
        ("spectra-nonce", hex::encode(p.nonce)),
        ("spectra-signature", hex::encode(spectra_core::cbor::encode(&p.sig.to_value()))),
    ]
}

/// The JSON-RPC result inside a response, whether it came as JSON or as
/// the last `data:` of an SSE stream.
pub fn rpc_result(r: &Response) -> Result<serde_json::Value> {
    let is_sse = r.header("content-type").is_some_and(|c| c.starts_with("text/event-stream"));
    if !is_sse {
        return r.json();
    }
    let text = String::from_utf8_lossy(&r.body);
    let last = text
        .lines()
        .rev()
        .filter_map(|l| l.strip_prefix("data:"))
        .map(str::trim)
        .find(|l| !l.is_empty())
        .context("an SSE response with no data")?;
    Ok(serde_json::from_str(last)?)
}

`/`	Search
`g` `c`	Code
`g` `h`	History
`g` `b`	Bookmarks
`g` `t`	Trust
`y`	Permalink to this state (files and trees)
`?`	This sheet

1		//! A client for the daemon over its Unix socket (spec §8.5), signing every
2		//! request with the caller's key: the shims and the CLI use it, and so do
3		//! the tests.
4
5		use anyhow::{Context, Result};
6		use http_body_util::{BodyExt, Full};
7		use hyper::body::Bytes;
8		use spectra_core::keys::KeyPair;
9		use spectra_core::proof::Proof;
10		use std::path::{Path, PathBuf};
11
12		/// A response, its body read whole.
13		#[derive(Debug)]
14		pub struct Response {
15		pub status: u16,
16		pub headers: Vec<(String, String)>,
17		pub body: Vec<u8>,
18		}
19
20		impl Response {
21	118×	pub fn json(&self) -> Result<serde_json::Value> {
22	118×	serde_json::from_slice(&self.body).with_context(\|\| format!("not JSON: {}", String::from_utf8_lossy(&self.body)))
23	118×	}
24	46×	pub fn header(&self, name: &str) -> Option<&str> {
25	47×	self.headers.iter().find(\|(k, _)\| k.eq_ignore_ascii_case(name)).map(\|(_, v)\| v.as_str())
26	46×	}
27		}
28
29		/// Where the daemon is: its Unix socket, or a server over TLS with a
30		/// pinned certificate (spec §8.6).
31		#[derive(Clone)]
32		pub enum Endpoint {
33		Unix(PathBuf),
34		Tls { addr: String, config: std::sync::Arc<rustls::ClientConfig> },
35		}
36
37		impl Endpoint {
38		/// `https://HOST:PORT` accepting exactly the certificate pinned by `pin`.
39	23×	pub fn tls(url: &str, pin: [u8; 32]) -> Result<Endpoint> {
40	23×	Ok(Endpoint::Tls { addr: host_port(url)?, config: crate::tls::Pinned::new(Some(pin)).client_config()? })
41	23×	}
42
43	0×	pub fn describe(&self) -> String {
44	0×	match self {
45	0×	Endpoint::Unix(p) => p.display().to_string(),
46	0×	Endpoint::Tls { addr, .. } => format!("https://{addr}"),
47		}
48	0×	}
49		}
50
51		/// `host:port` of an `https://` URL.
52	38×	pub fn host_port(url: &str) -> Result<String> {
53	38×	let rest = url.strip_prefix("https://").with_context(\|\| format!("{url} is not an https:// URL"))?;
54	38×	let hp = rest.split('/').next().unwrap_or("");
55	38×	if hp.is_empty() {
56	0×	anyhow::bail!("{url} names no host");
57	38×	}
58		// The port defaults to Spectra's.
59	38×	let has_port = if hp.starts_with('[') { hp.contains("]:") } else { hp.contains(':') };
60	38×	Ok(if has_port { hp.to_string() } else { format!("{hp}:7350") })
61	38×	}
62
63		trait Io: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send {}
64		impl<T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send> Io for T {}
65
66		pub struct Client {
67		endpoint: Endpoint,
68		signer: Option<KeyPair>,
69		}
70
71		impl Client {
72	47×	pub fn new(socket: &Path, signer: Option<KeyPair>) -> Client {
73	47×	Client { endpoint: Endpoint::Unix(socket.to_path_buf()), signer }
74	47×	}
75
76	47×	pub fn at(endpoint: Endpoint, signer: Option<KeyPair>) -> Client {
77	47×	Client { endpoint, signer }
78	47×	}
79
80	0×	pub fn endpoint(&self) -> &Endpoint {
81	0×	&self.endpoint
82	0×	}
83
84		/// The key this client signs with, if any.
85	2×	pub fn signer(&self) -> Option<&KeyPair> {
86	2×	self.signer.as_ref()
87	2×	}
88
89	178×	async fn io(&self) -> Result<Box<dyn Io>> {
90	178×	Ok(match &self.endpoint {
91	118×	Endpoint::Unix(socket) => Box::new(
92	118×	tokio::net::UnixStream::connect(socket)
93	118×	.await
94	118×	.with_context(\|\| format!("connecting to {} (is `spectra serve` running?)", socket.display()))?,
95		),
96	60×	Endpoint::Tls { addr, config } => Box::new(crate::tls::connect(addr, config.clone()).await?),
97		})
98	178×	}
99
100		/// Whether the daemon answers at the endpoint.
101	0×	pub async fn reachable(&self) -> bool {
102	0×	self.io().await.is_ok()
103	0×	}
104
105		/// Whether a daemon answers on the socket.
106	119×	pub async fn alive(socket: &Path) -> bool {
107	119×	tokio::net::UnixStream::connect(socket).await.is_ok()
108	119×	}
109
110		/// Send `method path` with `body` and extra headers, signed.
111	176×	pub async fn send(&self, method: &str, path: &str, headers: &[(&str, String)], body: Vec<u8>) -> Result<Response> {
112	176×	let stream = self.io().await?;
113	174×	let (mut sender, conn) = hyper::client::conn::http1::handshake(hyper_util::rt::TokioIo::new(stream)).await?;
114	174×	tokio::spawn(async move {
115	174×	let _ = conn.await;
116	174×	});
117		// One request a connection: the server closes it after answering, so
118		// nothing lingers to hold up its shutdown.
119	174×	let mut req = hyper::Request::builder().method(method).uri(path).header("host", "spectra").header("connection", "close");
120		// JSON unless the caller says otherwise (a form, say).
121	224×	if !headers.iter().any(\|(k, _)\| k.eq_ignore_ascii_case("content-type")) {
122	174×	req = req.header("content-type", "application/json");
123	174×	}
124	224×	for (k, v) in headers {
125	224×	req = req.header(*k, v);
126	224×	}
127	174×	if let Some(pair) = &self.signer {
128	544×	for (k, v) in sign_headers(pair, &format!("{method} {path}"), &body) {
129	544×	req = req.header(k, v);
130	544×	}
131	38×	}
132	174×	let resp = sender.send_request(req.body(Full::new(Bytes::from(body)))?).await?;
133	174×	let status = resp.status().as_u16();
134	746×	let headers = resp.headers().iter().map(\|(k, v)\| (k.to_string(), v.to_str().unwrap_or("").to_string())).collect();
135	174×	let body = resp.into_body().collect().await?.to_bytes().to_vec();
136	174×	Ok(Response { status, headers, body })
137	176×	}
138
139		/// A JSON-RPC request to `/mcp` with the protocol headers of `version`.
140	57×	pub async fn mcp(&self, version: &str, request: &serde_json::Value) -> Result<Response> {
141	57×	let method = request.get("method").and_then(\|m\| m.as_str()).unwrap_or("").to_string();
142	57×	let mut headers = vec![
143	57×	("accept", "application/json, text/event-stream".to_string()),
144	57×	("mcp-protocol-version", version.to_string()),
145	57×	("mcp-method", method.clone()),
146		];
147	57×	if let Some(name) = request.pointer("/params/name").or_else(\|\| request.pointer("/params/uri")).and_then(\|n\| n.as_str()) {
148	46×	headers.push(("mcp-name", name.to_string()));
149	46×	}
150	57×	self.send("POST", "/mcp", &headers, serde_json::to_vec(request)?).await
151	57×	}
152		}
153
154		impl Client {
155		/// Pass a request through as it is, unsigned, and hand back the
156		/// response with its body still streaming (Radar's proxy, spec §8.6).
157	2×	pub async fn forward(&self, req: hyper::Request<Full<Bytes>>) -> Result<hyper::Response<hyper::body::Incoming>> {
158	2×	let stream = self.io().await?;
159	2×	let (mut sender, conn) = hyper::client::conn::http1::handshake(hyper_util::rt::TokioIo::new(stream)).await?;
160	2×	tokio::spawn(async move {
161	2×	let _ = conn.await;
162	2×	});
163	2×	Ok(sender.send_request(req).await?)
164	2×	}
165
166		/// Send a signed GET and hand each line of the streamed body to `line`
167		/// as it arrives (an SSE stream, say); stop when it returns false.
168	0×	pub async fn stream(&self, path: &str, mut line: impl FnMut(&str) -> bool) -> Result<u16> {
169	0×	let stream = self.io().await?;
170	0×	let (mut sender, conn) = hyper::client::conn::http1::handshake(hyper_util::rt::TokioIo::new(stream)).await?;
171	0×	tokio::spawn(async move {
172	0×	let _ = conn.await;
173	0×	});
174	0×	let mut req = hyper::Request::builder().method("GET").uri(path).header("host", "spectra").header("connection", "close");
175	0×	if let Some(pair) = &self.signer {
176	0×	for (k, v) in sign_headers(pair, &format!("GET {path}"), b"") {
177	0×	req = req.header(k, v);
178	0×	}
179	0×	}
180	0×	let resp = sender.send_request(req.body(Full::new(Bytes::new()))?).await?;
181	0×	let status = resp.status().as_u16();
182	0×	let mut body = resp.into_body();
183	0×	let mut buf = String::new();
184	0×	while let Some(frame) = body.frame().await {
185	0×	if let Ok(data) = frame?.into_data() {
186	0×	buf.push_str(&String::from_utf8_lossy(&data));
187	0×	while let Some(i) = buf.find('\n') {
188	0×	let l: String = buf.drain(..=i).collect();
189	0×	if !line(l.trim_end()) {
190	0×	return Ok(status);
191	0×	}
192		}
193	0×	}
194		}
195	0×	Ok(status)
196	0×	}
197		}
198
199		/// The signature headers of spec §8.5 for `method_path` and `body`.
200	138×	pub fn sign_headers(pair: &KeyPair, method_path: &str, body: &[u8]) -> Vec<(&'static str, String)> {
201	138×	let p = Proof::sign(pair, method_path, body);
202	138×	vec![
203	138×	("spectra-key", p.sig.key.to_hex()),
204	138×	("spectra-time", p.time_ms.to_string()),
205	138×	("spectra-nonce", hex::encode(p.nonce)),
206	138×	("spectra-signature", hex::encode(spectra_core::cbor::encode(&p.sig.to_value()))),
207		]
208	138×	}
209
210		/// The JSON-RPC result inside a response, whether it came as JSON or as
211		/// the last `data:` of an SSE stream.
212	44×	pub fn rpc_result(r: &Response) -> Result<serde_json::Value> {
213	44×	let is_sse = r.header("content-type").is_some_and(\|c\| c.starts_with("text/event-stream"));
214	44×	if !is_sse {
215	39×	return r.json();
216	5×	}
217	5×	let text = String::from_utf8_lossy(&r.body);
218	5×	let last = text
219	5×	.lines()
220	5×	.rev()
221	10×	.filter_map(\|l\| l.strip_prefix("data:"))
222	5×	.map(str::trim)
223	5×	.find(\|l\| !l.is_empty())
224	5×	.context("an SSE response with no data")?;
225	5×	Ok(serde_json::from_str(last)?)
226	44×	}

Coverage / crates/spectra-mcp/src/client.rs 72.7%

Keyboard