brokerd: consult posts the allowlisted body

Co-Authored-By: ornith-1.5-35b-a3b <agent@wntrmute.dev>

crates/brokerd/src/consult.rs

@@ -4,10 +4,13 @@
 use std::ffi::OsString;
 use std::time::Duration;
 
-use bxcore::secrets::{Loaded, SecretError, SecretSpec};
+use bxcore::secrets::{Loaded, Secret, SecretError, SecretSpec};
 use serde::{Deserialize, Deserializer, Serialize};
 
+use crate::args::ToolArgs;
 use crate::config::Config;
+use crate::policy::Label;
+use crate::runner::{RunError, RunOutput, RunSpec, Runtime};
 
 pub const CONSULT: &str = "consult";
 /// `question` and the answer are cut to this many bytes, on a character boundary.
@@ -241,3 +244,99 @@ pub fn load_secret(
         .map_err(SecretStartError::Load)?;
     Ok(Some(loaded))
 }
+
+/// One POST of the allowlisted body. Tests inject a fake; production uses loopback HTTP/1.1.
+pub trait ConsultTransport: Send + Sync {
+    fn post(
+        &self,
+        url: &str,
+        bearer: &str,
+        body: &[u8],
+        timeout: Duration,
+    ) -> Result<Vec<u8>, String>;
+}
+
+/// Runs an allowed `consult` through `transport`.
+pub struct ConsultRuntime<T> {
+    consult: Consult,
+    secret: Secret,
+    transport: T,
+}
+
+impl<T: ConsultTransport> ConsultRuntime<T> {
+    pub fn new(consult: Consult, secret: Secret, transport: T) -> ConsultRuntime<T> {
+        ConsultRuntime {
+            consult,
+            secret,
+            transport,
+        }
+    }
+}
+
+impl<T: ConsultTransport> Runtime for ConsultRuntime<T> {
+    fn run(&self, spec: &RunSpec) -> Result<RunOutput, RunError> {
+        self.run_labelled(spec).map(|(output, _)| output)
+    }
+
+    fn run_labelled(&self, spec: &RunSpec) -> Result<(RunOutput, Option<Label>), RunError> {
+        let question = match spec.arguments() {
+            ToolArgs::Consult { question } => question.as_str(),
+            _ => {
+                return Err(RunError::Unavailable(
+                    "consult was routed a different tool".to_string(),
+                ));
+            }
+        };
+        let body = wire_body_parts(&self.consult.model, self.consult.temperature(), question);
+        let bytes = self
+            .transport
+            .post(
+                &self.consult.url,
+                self.secret.expose(),
+                body.as_bytes(),
+                self.consult.timeout(),
+            )
+            .map_err(RunError::Unavailable)?;
+        let content = parse_message_content(&bytes).map_err(RunError::Unavailable)?;
+        let neutralised = bxcore::sanitize::neutralise(&content);
+        let (content, truncated) = cap_bytes(&neutralised, QUESTION_CAP);
+        Ok((RunOutput { content, truncated }, None))
+    }
+}
+
+/// `choices[0].message.content`, ignoring unknown fields.
+pub fn parse_message_content(bytes: &[u8]) -> Result<String, String> {
+    let value: serde_json::Value =
+        serde_json::from_slice(bytes).map_err(|e| format!("consult response is not JSON: {e}"))?;
+    match value
+        .get("choices")
+        .and_then(|c| c.as_array())
+        .and_then(|a| a.first())
+        .and_then(|c| c.get("message"))
+        .and_then(|m| m.get("content"))
+        .and_then(|c| c.as_str())
+    {
+        Some(content) => Ok(content.to_string()),
+        None => Err("the response has no content".to_string()),
+    }
+}
+
+/// `consult` first, everything else to `inner`.
+pub struct ConsultStack {
+    pub consult: Box<dyn Runtime>,
+    pub inner: Box<dyn Runtime>,
+}
+
+impl Runtime for ConsultStack {
+    fn run(&self, spec: &RunSpec) -> Result<RunOutput, RunError> {
+        self.run_labelled(spec).map(|(output, _)| output)
+    }
+
+    fn run_labelled(&self, spec: &RunSpec) -> Result<(RunOutput, Option<Label>), RunError> {
+        if spec.tool() == crate::args::ToolName::Consult {
+            self.consult.run_labelled(spec)
+        } else {
+            self.inner.run_labelled(spec)
+        }
+    }
+}

crates/brokerd/src/consult_http.rs

@@ -0,0 +1,111 @@
+//! Loopback HTTP/1.1 POST for `consult`. HTTPS is a later milestone.
+
+use std::io::{Read, Write};
+use std::net::TcpStream;
+use std::time::{Duration, Instant};
+
+use crate::consult::{ConsultTransport, parse_loopback_http};
+
+/// Largest response body accepted, in bytes.
+pub const MAX_RESPONSE: usize = 262_144;
+
+/// Production transport: one TCP connection to the configured loopback URL.
+pub struct HttpTransport;
+
+impl ConsultTransport for HttpTransport {
+    fn post(
+        &self,
+        url: &str,
+        bearer: &str,
+        body: &[u8],
+        timeout: Duration,
+    ) -> Result<Vec<u8>, String> {
+        post_http(url, bearer, body, timeout)
+    }
+}
+
+fn post_http(url: &str, bearer: &str, body: &[u8], timeout: Duration) -> Result<Vec<u8>, String> {
+    let (host, port, path) = parse_loopback_http(url)?;
+    let addr = format!("{host}:{port}");
+    let mut stream =
+        TcpStream::connect(&addr).map_err(|e| format!("cannot connect to {addr}: {e}"))?;
+    stream
+        .set_read_timeout(Some(timeout))
+        .map_err(|e| format!("cannot set read timeout: {e}"))?;
+    stream
+        .set_write_timeout(Some(timeout))
+        .map_err(|e| format!("cannot set write timeout: {e}"))?;
+    let host_header = if port == 80 {
+        host
+    } else {
+        format!("{host}:{port}")
+    };
+    let length = body.len();
+    let header = format!(
+        "POST {path} HTTP/1.1\r\nHost: {host_header}\r\nAuthorization: Bearer {bearer}\r\nContent-Type: application/json\r\nContent-Length: {length}\r\nConnection: close\r\n\r\n"
+    );
+    stream
+        .write_all(header.as_bytes())
+        .map_err(|e| format!("cannot write consult request: {e}"))?;
+    stream
+        .write_all(body)
+        .map_err(|e| format!("cannot write consult request: {e}"))?;
+    stream
+        .flush()
+        .map_err(|e| format!("cannot write consult request: {e}"))?;
+    read_response(&mut stream, timeout)
+}
+
+fn read_response(stream: &mut TcpStream, timeout: Duration) -> Result<Vec<u8>, String> {
+    let until = Instant::now()
+        .checked_add(timeout)
+        .ok_or_else(|| "the consult timeout is too large".to_string())?;
+    let mut buf = Vec::new();
+    let mut chunk = [0u8; 4096];
+    loop {
+        if Instant::now() >= until {
+            return Err("consult timed out".to_string());
+        }
+        let left = until.saturating_duration_since(Instant::now());
+        stream
+            .set_read_timeout(Some(left))
+            .map_err(|e| format!("cannot set read timeout: {e}"))?;
+        match stream.read(&mut chunk) {
+            Ok(0) => break,
+            Ok(n) => {
+                buf.extend_from_slice(&chunk[..n]);
+                if buf.len() > MAX_RESPONSE {
+                    return Err("consult response is too large".to_string());
+                }
+            }
+            Err(e) if e.kind() == std::io::ErrorKind::TimedOut => {
+                return Err("consult timed out".to_string());
+            }
+            Err(e) => return Err(format!("cannot read consult response: {e}")),
+        }
+    }
+    split_http_body(&buf)
+}
+
+fn split_http_body(buf: &[u8]) -> Result<Vec<u8>, String> {
+    const SEP: &[u8] = b"\r\n\r\n";
+    let sep = buf
+        .windows(SEP.len())
+        .position(|w| w == SEP)
+        .ok_or_else(|| "consult response has no header terminator".to_string())?;
+    let (head, rest) = match buf.split_at_checked(sep) {
+        Some(pair) => pair,
+        None => return Err("consult response is truncated".to_string()),
+    };
+    let body = match rest.get(SEP.len()..) {
+        Some(body) => body,
+        None => return Err("consult response is truncated".to_string()),
+    };
+    let head = std::str::from_utf8(head)
+        .map_err(|_| "consult response headers are not UTF-8".to_string())?;
+    let status_line = head.lines().next().unwrap_or("");
+    if !status_line.contains(" 200 ") && !status_line.ends_with(" 200") {
+        return Err(format!("consult HTTP status is not 200: {status_line}"));
+    }
+    Ok(body.to_vec())
+}

crates/brokerd/src/lib.rs

@@ -8,6 +8,7 @@ pub mod broker;
 pub mod builtin;
 pub mod config;
 pub mod consult;
+pub mod consult_http;
 pub mod container;
 pub mod courier;
 pub mod grants;

crates/brokerd/src/main.rs

@@ -7,6 +7,8 @@ use std::sync::Arc;
 
 use brokerd::audit::{AuditError, RECOVERED_NOTICE};
 use brokerd::config::Config;
+use brokerd::consult::{self, ConsultRuntime, ConsultStack};
+use brokerd::consult_http::HttpTransport;
 use brokerd::container::Podman;
 use brokerd::memory_port::{MEMORY_WAIT, MemoryPort, Routed};
 use brokerd::runner::{Refusing, Runtime};
@@ -43,6 +45,19 @@ fn main() -> ExitCode {
         }
     };
 
+    let consult_secret = match consult::load_secret(&cfg, &|k| std::env::var_os(k)) {
+        Ok(loaded) => loaded,
+        Err(e) => {
+            eprintln!("brokerd: {e}");
+            return ExitCode::from(1);
+        }
+    };
+    if let Some(loaded) = &consult_secret
+        && let Some(warning) = &loaded.warning
+    {
+        eprintln!("{warning}");
+    }
+
     // One brokerd per home (M5b spec, section 9a). The lock comes before anything is made, the
     // audit log's own lock included, so a second copy never takes the sockets of the one that
     // runs. Held until the process ends.
@@ -87,6 +102,14 @@ fn main() -> ExitCode {
         other: runtime,
     });
 
+    let runtime: Box<dyn Runtime> = match (cfg.consult.clone(), consult_secret) {
+        (Some(consult), Some(loaded)) => Box::new(ConsultStack {
+            consult: Box::new(ConsultRuntime::new(consult, loaded.secret, HttpTransport)),
+            inner: runtime,
+        }),
+        _ => runtime,
+    };
+
     let started = match serve::start(cfg, accept_break, runtime, log) {
         Ok(started) => started,
         Err(ServeError::Audit(e @ AuditError::NothingToAccept)) => {