From d1f25b5ae6d1c22c652aff3401e55f1a78773d5f Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Thu, 29 Oct 2026 21:54:03 +0100 Subject: [PATCH] claude-self-run-lease: Self runs borrow the one Claude token over a unix socket bound to the peer's uid, its pidfd and its unit's InvocationID; the connector sets the reservation, settle bills reported per-model usage uncapped, no report bills the reservation (squash of eureka-body/claude-self-run-lease 480cf94..b5f26ef; Soul s1-s3) Co-Authored-By: Claude Opus 5.5 --- crates/claude-connector/src/daemon.rs | 171 ++- crates/claude-connector/src/lease.rs | 729 +++++++++++ crates/claude-connector/src/lib.rs | 16 +- crates/claude-connector/src/main.rs | 29 +- crates/claude-connector/src/unit.rs | 361 +++++ .../claude-connector/tests/config_document.rs | 219 +++- .../tests/fixtures/lease_v1/grant.bin | Bin 0 -> 109 bytes .../tests/fixtures/lease_v1/refusal.bin | Bin 0 -> 104 bytes .../tests/fixtures/lease_v1/request.bin | Bin 0 -> 66 bytes .../tests/fixtures/lease_v1/settle.bin | Bin 0 -> 336 bytes .../tests/fixtures/lease_v1/settled.bin | Bin 0 -> 81 bytes crates/claude-connector/tests/lease.rs | 1166 +++++++++++++++++ crates/model-connector/src/serve.rs | 38 +- crates/model-connector/src/service.rs | 29 +- crates/model-connector/src/service/lease.rs | 691 ++++++++++ 15 files changed, 3421 insertions(+), 28 deletions(-) create mode 100644 crates/claude-connector/src/lease.rs create mode 100644 crates/claude-connector/src/unit.rs create mode 100644 crates/claude-connector/tests/fixtures/lease_v1/grant.bin create mode 100644 crates/claude-connector/tests/fixtures/lease_v1/refusal.bin create mode 100644 crates/claude-connector/tests/fixtures/lease_v1/request.bin create mode 100644 crates/claude-connector/tests/fixtures/lease_v1/settle.bin create mode 100644 crates/claude-connector/tests/fixtures/lease_v1/settled.bin create mode 100644 crates/claude-connector/tests/lease.rs create mode 100644 crates/model-connector/src/service/lease.rs diff --git a/crates/claude-connector/src/daemon.rs b/crates/claude-connector/src/daemon.rs index 674975d..9f549fc 100644 --- a/crates/claude-connector/src/daemon.rs +++ b/crates/claude-connector/src/daemon.rs @@ -6,17 +6,22 @@ use std::collections::{BTreeMap, BTreeSet}; use std::net::{SocketAddr, TcpListener}; use std::path::{Path, PathBuf}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::Duration; use cultcache_rs::{CultCache, DatabaseEntry, SingleFileMessagePackBackingStore}; use model_connector::{ - CallerAdmission, CallerConfig, MIN_PAYLOAD_BYTES, ModelPrice, ServeError, ServeLimits, - ServiceError, load_admissions, serve, + CallerAdmission, CallerConfig, MIN_PAYLOAD_BYTES, ModelPrice, ProviderBackend, ServeError, + ServeLimits, ServiceError, TransportService, load_admissions, serve_service, }; +use serde::{Deserialize, Serialize}; use thiserror::Error; use crate::backend::readiness; +#[cfg(target_os = "linux")] +use crate::lease::{LeaseCaller, LeaseLimits, serve_leases, user_uid}; +#[cfg(target_os = "linux")] +use crate::unit::SystemdUnits; use crate::{ ClaudeCodeBackend, ClaudeCodeConfig, ClaudeCodeLimits, CredentialTokenSource, StderrSink, }; @@ -24,6 +29,18 @@ use crate::{ const CONFIG_KEY: &str = "runtime"; const MODEL_PREFIX: &str = "claude-"; +/// A unix account that may lease the subscription token for a Self run, the +/// Self unit (`mind-self@.service`) its process must run in, +/// the caller whose spend window its leases reserve in, and what each lease +/// reserves there. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct LeaseCallerConfig { + pub unix_user: String, + pub caller_id: String, + pub unit_instance: String, + pub reservation_micro_usd: u64, +} + #[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)] #[cultcache(type = "gamecult.model_connector.claude_connector_config.v1")] pub struct ClaudeConnectorConfig { @@ -72,6 +89,17 @@ pub struct ClaudeConnectorConfig { /// ids exactly as the CLI reports them in its init event. #[cultcache(key = 18)] pub prices: BTreeMap, + /// Where the Self-run lease socket is bound (a systemd RuntimeDirectory), + /// or none when this connector lends no token. + #[cultcache(key = 19)] + pub lease_socket: Option, + /// The accounts allowed to lease, each bound to a configured caller. + /// Empty exactly when `lease_socket` is none. + #[cultcache(key = 20)] + pub lease_callers: Vec, + /// The furthest a lease's expiry may lie from now. + #[cultcache(key = 21)] + pub lease_max_ms: u64, } /// A refusal. It names the field or stage that failed and nothing the @@ -205,9 +233,39 @@ impl ClaudeConnectorConfig { { return Err(ConfigError::Invalid("prices")); } + self.validate_leases(&caller.caller_runtime_id, caller.spend_ceiling_micro_usd)?; Ok(listen) } + fn validate_leases(&self, caller_id: &str, spend_ceiling: u64) -> Result<(), ConfigError> { + let Some(socket) = &self.lease_socket else { + return if self.lease_callers.is_empty() { + Ok(()) + } else { + Err(ConfigError::Invalid("lease_callers")) + }; + }; + if !socket.is_absolute() { + return Err(ConfigError::Invalid("lease_socket")); + } + if self.lease_max_ms == 0 { + return Err(ConfigError::Invalid("lease_max_ms")); + } + let mut users = BTreeSet::new(); + if self.lease_callers.is_empty() + || !self.lease_callers.iter().all(|lease| { + is_trimmed_nonempty(&lease.unix_user) + && lease.caller_id == caller_id + && is_unit_instance(&lease.unit_instance) + && (1..=spend_ceiling).contains(&lease.reservation_micro_usd) + && users.insert(lease.unix_user.as_str()) + }) + { + return Err(ConfigError::Invalid("lease_callers")); + } + Ok(()) + } + fn backend_config(&self) -> ClaudeCodeConfig { ClaudeCodeConfig { claude_binary: self.claude_binary.clone(), @@ -236,6 +294,14 @@ impl ClaudeConnectorConfig { } } +/// The `` of a unit name, as plain as a systemd instance can be. +fn is_unit_instance(text: &str) -> bool { + !text.is_empty() + && text + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_')) +} + fn is_trimmed_nonempty(text: &str) -> bool { !text.is_empty() && text.trim() == text } @@ -394,9 +460,85 @@ fn readiness_of(config: &ClaudeConnectorConfig) -> Result, &CredentialTokenSource::from_environment(backend.token_credential.clone()), ) .map_err(ConfigError::Unusable)?; + if config.lease_socket.is_some() { + lease_callers_of(config)?; + } Ok(admissions) } +/// The accounts allowed to lease, resolved to the uids the kernel will report +/// for their connections. An account the system does not have is not ready. +#[cfg(target_os = "linux")] +fn lease_callers_of(config: &ClaudeConnectorConfig) -> Result, ConfigError> { + config + .lease_callers + .iter() + .map(|lease| { + user_uid(&lease.unix_user) + .map(|uid| LeaseCaller { + uid, + caller_id: lease.caller_id.clone(), + unit_instance: lease.unit_instance.clone(), + reservation_micro_usd: lease.reservation_micro_usd, + }) + .ok_or(ConfigError::Unusable("lease_callers")) + }) + .collect() +} + +/// Peer credentials identify a lease caller only on Linux. +#[cfg(not(target_os = "linux"))] +fn lease_callers_of(_: &ClaudeConnectorConfig) -> Result<(), ConfigError> { + Err(ConfigError::Unusable("lease_socket")) +} + +/// Binds the lease socket (mode 0660; the unit's RuntimeDirectory and group +/// decide who can reach it, and the kernel's peer uid decides who may lease) +/// and serves it beside the RUDP loop. The lease listener is half of the +/// connector: if it stops, so does the process. +#[cfg(target_os = "linux")] +fn start_leases( + config: &ClaudeConnectorConfig, + socket: &Path, + service: Arc>, +) -> Result<(), ConfigError> { + use std::os::unix::fs::{FileTypeExt, PermissionsExt}; + use std::os::unix::net::UnixListener; + + let callers = lease_callers_of(config)?; + // A socket a previous run left is replaced; anything else at the path is not ours. + if std::fs::symlink_metadata(socket).is_ok_and(|found| found.file_type().is_socket()) { + let _ = std::fs::remove_file(socket); + } + let listener = UnixListener::bind(socket).map_err(|_| ConfigError::Listen)?; + std::fs::set_permissions(socket, std::fs::Permissions::from_mode(0o660)) + .map_err(|_| ConfigError::Listen)?; + let limits = LeaseLimits { + socket_timeout: Duration::from_millis(config.socket_timeout_ms), + max_lease_ms: config.lease_max_ms, + max_connections: config.max_connections, + }; + let tokens = Arc::new(CredentialTokenSource::from_environment( + config.token_credential.clone(), + )); + std::thread::Builder::new() + .name("claude-lease-listener".to_string()) + .spawn(move || { + let _ = serve_leases( + &listener, + callers, + limits, + service, + tokens, + Arc::new(SystemdUnits::system()), + ); + eprintln!("claude-connector: the lease listener stopped"); + std::process::exit(1); + }) + .map_err(|_| ConfigError::Serving("lease listener"))?; + Ok(()) +} + /// Whether `serve` would start with this document and this environment. pub fn check_ready(config: &ClaudeConnectorConfig) -> Result<(), ConfigError> { readiness_of(config).map(drop) @@ -413,14 +555,21 @@ pub fn serve_config(path: &Path) -> Result<(), ConfigError> { Box::new(LogSink(std::sync::Mutex::new(std::io::stderr()))), )); let listener = TcpListener::bind(listen).map_err(|_| ConfigError::Listen)?; - serve( - listener, - config.serve_limits(), - admissions, - config.prices.clone(), - &config.replay_store, - backend, - )?; + let service = Arc::new(Mutex::new( + TransportService::open( + &config.replay_store, + backend.provider(), + config.prices.clone(), + admissions, + config.max_expiry_skew_ms, + ) + .map_err(|error| ConfigError::from(ServeError::Service(error)))?, + )); + #[cfg(target_os = "linux")] + if let Some(socket) = &config.lease_socket { + start_leases(&config, socket, service.clone())?; + } + serve_service(listener, config.serve_limits(), service, backend)?; Ok(()) } diff --git a/crates/claude-connector/src/lease.rs b/crates/claude-connector/src/lease.rs new file mode 100644 index 0000000..d8badf3 --- /dev/null +++ b/crates/claude-connector/src/lease.rs @@ -0,0 +1,729 @@ +//! The Self-run credential lease. The connector is the only holder of the +//! subscription token and the one owner of Self-run admission, reservation and +//! the spend ledger (rulings self-runs-through-connector, +//! self-run-credential-lease and self-run-lease-billing-and-binding). A Self +//! run asks over a unix socket for a lease. The connector decides who may ask +//! (the peer's uid, and the systemd unit invocation its process runs in), +//! reserves its own configured amount in the caller's spend window, and only +//! then hands back one pipe fd that holds the token. The run settles with the +//! per-model usage it reports, and the connector bills that at true cost, +//! uncapped; a run that vanishes or lets its lease expire is billed the whole +//! reservation. Nothing the run says about itself decides who it is, what it +//! reserves, or what the plan allows. +//! +//! The token exists in this process only as the value the token source returns +//! and in the pipe. It is not in an environment, an argument, a file, a log or +//! a reply frame. The connector spawns nothing here: the Self run's `claude` +//! stays in the Self run's own unit. + +use std::ffi::CString; +use std::io::{self, Read, Write}; +use std::mem::{size_of, zeroed}; +use std::os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, OwnedFd, RawFd}; +use std::os::unix::net::{UnixListener, UnixStream}; +use std::ptr; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::thread; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use model_connector::{ + LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseSettlement, LeaseUse, ModelUsage, + TransportService, +}; +use serde::{Deserialize, Serialize}; + +use crate::TokenSource; + +/// The largest frame either side sends or accepts, in bytes. It is wire +/// contract: a Huginn client sizes against it, so it is pinned as a literal. +pub const MAX_LEASE_FRAME_BYTES: usize = 65_536; +const _: () = assert!(MAX_LEASE_FRAME_BYTES == 65_536); + +/// One message of the lease protocol: a u32 big-endian length, then this as a +/// MessagePack named map `{kind, value}`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", content = "value", rename_all = "snake_case")] +pub enum LeaseMessage { + ClaudeLeaseRequest(LeaseRequest), + /// Travels with exactly one SCM_RIGHTS fd: the read end of a pipe that + /// holds the token, its write end closed. + ClaudeLeaseGrant(LeaseGranted), + ClaudeLeaseRefusal(LeaseRefused), + ClaudeLeaseSettle(LeaseSettle), + ClaudeLeaseSettled(LeaseSettled), +} + +/// A request names only when the lease ends. Who asks, and what is reserved, +/// are the connector's to know. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct LeaseRequest { + pub expires_at_unix_ms: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct LeaseGranted { + pub lease_id: String, + pub reserved_micro_usd: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum LeaseRefusalKind { + SpendCeiling, + CostBound, + UnknownCaller, + Expired, + Busy, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct LeaseRefused { + pub kind: LeaseRefusalKind, + pub detail: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct LeaseSettle { + pub lease_id: String, + /// What the run used of each model, as the run reports it. + pub usage: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct LeaseSettled { + pub billed_micro_usd: u64, + /// How far the bill passed the reservation; zero when it did not. + pub overrun_micro_usd: u64, +} + +/// A frame: the length, then the named-map MessagePack of `message`. +pub fn encode_frame(message: &LeaseMessage) -> Vec { + // Every field of the message is plain data, so encoding cannot fail. + let body = rmp_serde::to_vec_named(message).unwrap_or_default(); + let mut frame = Vec::with_capacity(4 + body.len()); + frame.extend_from_slice(&(body.len() as u32).to_be_bytes()); + frame.extend_from_slice(&body); + frame +} + +/// The message in one whole frame, or `None` when the frame is not exactly one +/// well-formed message within the bound. +pub fn decode_frame(frame: &[u8]) -> Option { + let (length, body) = frame.split_first_chunk::<4>()?; + let length = u32::from_be_bytes(*length) as usize; + if length == 0 || length > MAX_LEASE_FRAME_BYTES || body.len() != length { + return None; + } + rmp_serde::from_slice(body).ok() +} + +/// A reader whose waiting can be bounded, so a frame has one deadline across +/// all of its reads. +trait Bounded: Read { + fn within(&mut self, left: Duration) -> io::Result<()>; +} + +impl Bounded for UnixStream { + fn within(&mut self, left: Duration) -> io::Result<()> { + self.set_read_timeout(Some(left)) + } +} + +impl Bounded for &[u8] { + fn within(&mut self, _: Duration) -> io::Result<()> { + Ok(()) + } +} + +/// Fills `buffer` or fails: when `deadline` passes, however the bytes trickle. +fn read_exact_by( + stream: &mut impl Bounded, + buffer: &mut [u8], + deadline: Instant, +) -> io::Result<()> { + let mut filled = 0; + while filled < buffer.len() { + let left = deadline + .checked_duration_since(Instant::now()) + .filter(|left| !left.is_zero()) + .ok_or_else(|| io::Error::from(io::ErrorKind::TimedOut))?; + stream.within(left)?; + match stream.read(&mut buffer[filled..]) { + Ok(0) => return Err(io::Error::from(io::ErrorKind::UnexpectedEof)), + Ok(read) => filled += read, + Err(error) + if matches!( + error.kind(), + io::ErrorKind::Interrupted | io::ErrorKind::WouldBlock + ) => {} + Err(error) => return Err(error), + } + } + Ok(()) +} + +/// One whole frame, complete by `deadline`. +fn read_message(stream: &mut impl Bounded, deadline: Instant) -> io::Result { + let invalid = || io::Error::from(io::ErrorKind::InvalidData); + let mut length = [0_u8; 4]; + read_exact_by(stream, &mut length, deadline)?; + let length = u32::from_be_bytes(length) as usize; + if length == 0 || length > MAX_LEASE_FRAME_BYTES { + return Err(invalid()); + } + let mut body = vec![0_u8; length]; + read_exact_by(stream, &mut body, deadline)?; + rmp_serde::from_slice(&body).map_err(|_| invalid()) +} + +/// The most a pipe takes in one atomic write. +const PIPE_BUF: usize = 4096; + +/// A unix account allowed to lease, the Self unit its process must run in, +/// and the caller whose spend window its leases reserve in. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LeaseCaller { + pub uid: u32, + pub caller_id: String, + /// The `` of `mind-self@.service`. + pub unit_instance: String, + /// What each of this caller's leases reserves, whatever the run asks. + pub reservation_micro_usd: u64, +} + +impl LeaseCaller { + /// The one unit a lease to this caller may be requested from. + pub fn unit(&self) -> String { + format!("mind-self@{}.service", self.unit_instance) + } +} + +/// The systemd unit invocation a process runs in. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PeerUnit { + pub unit: String, + pub invocation_id: String, +} + +/// Which invocation of `unit` a process is running, as the system records it. +/// The lease listener asks this port; `SystemdUnits` answers it in production. +/// `None` when the process is not in `unit` or the system cannot say. The +/// process is named by the handle the kernel gave for the peer, never by a pid. +pub trait PeerUnits: Send + Sync { + fn invocation_in(&self, process: BorrowedFd<'_>, unit: &str) -> Option; +} + +/// The bounds of the lease listener. +#[derive(Debug, Clone, Copy)] +pub struct LeaseLimits { + /// How long a peer has to send its request frame, whole. + pub socket_timeout: Duration, + /// The furthest an expiry may lie from now. + pub max_lease_ms: u64, + pub max_connections: usize, +} + +/// Who is on the other end of a unix stream, as the kernel recorded them at +/// connect. +#[derive(Debug)] +pub struct PeerCredentials { + pub uid: u32, + /// A handle on the peer's process (SO_PEERPIDFD, Linux 6.5 and later). It + /// pins the process's identity: a recycled pid cannot be read through it. + /// `None` when the kernel cannot give one, and then no unit is vouched + /// for. + pub process: Option, +} + +/// The uid and process handle of the process on the other end of `stream`. +pub fn peer_credentials(stream: &UnixStream) -> io::Result { + // SAFETY: `credentials` and `length` are valid for the call, and `length` + // is the size of the buffer getsockopt writes into. + let (result, credentials) = unsafe { + let mut credentials: libc::ucred = zeroed(); + let mut length = size_of::() as libc::socklen_t; + let result = libc::getsockopt( + stream.as_raw_fd(), + libc::SOL_SOCKET, + libc::SO_PEERCRED, + ptr::from_mut(&mut credentials).cast(), + &mut length, + ); + (result, credentials) + }; + if result != 0 { + return Err(io::Error::last_os_error()); + } + // SAFETY: as above, with a `c_int` the kernel fills with a new descriptor + // that this function then owns. + let (result, pidfd) = unsafe { + let mut pidfd: libc::c_int = -1; + let mut length = size_of::() as libc::socklen_t; + let result = libc::getsockopt( + stream.as_raw_fd(), + libc::SOL_SOCKET, + libc::SO_PEERPIDFD, + ptr::from_mut(&mut pidfd).cast(), + &mut length, + ); + (result, pidfd) + }; + Ok(PeerCredentials { + uid: credentials.uid, + // SAFETY: on success `pidfd` is a descriptor nothing else owns. + process: (result == 0 && pidfd >= 0).then(|| unsafe { OwnedFd::from_raw_fd(pidfd) }), + }) +} + +/// The uid of the account `name`, if the system has one. +pub fn user_uid(name: &str) -> Option { + let name = CString::new(name).ok()?; + let mut buffer = vec![0 as libc::c_char; 16 * 1024]; + // SAFETY: every pointer is valid for the call and `buffer` is as long as + // the length passed; `found` is read only when the call reports success. + unsafe { + let mut entry: libc::passwd = zeroed(); + let mut found: *mut libc::passwd = ptr::null_mut(); + let result = libc::getpwnam_r( + name.as_ptr(), + &mut entry, + buffer.as_mut_ptr(), + buffer.len(), + &mut found, + ); + (result == 0 && !found.is_null()).then_some(entry.pw_uid) + } +} + +/// A pipe whose read end holds `token` and whose write end is closed, so a +/// reader gets the token and then end of file. The token is smaller than +/// PIPE_BUF, so the write is whole and cannot block. +fn token_pipe(token: &[u8]) -> io::Result { + if token.len() > PIPE_BUF { + return Err(io::Error::from(io::ErrorKind::InvalidInput)); + } + let mut ends = [0 as RawFd; 2]; + // SAFETY: `ends` has room for the two descriptors pipe2 writes. + if unsafe { libc::pipe2(ends.as_mut_ptr(), libc::O_CLOEXEC) } != 0 { + return Err(io::Error::last_os_error()); + } + // SAFETY: pipe2 just returned these two descriptors and nothing else owns them. + let (read, write) = unsafe { (OwnedFd::from_raw_fd(ends[0]), OwnedFd::from_raw_fd(ends[1])) }; + // SAFETY: the pointer and length describe `token`. + let written = unsafe { libc::write(write.as_raw_fd(), token.as_ptr().cast(), token.len()) }; + if usize::try_from(written).ok() != Some(token.len()) { + return Err(io::Error::last_os_error()); + } + drop(write); + Ok(read) +} + +/// Writes `frame` to `stream` with `fd` attached as SCM_RIGHTS. +fn send_with_fd(stream: &UnixStream, frame: &[u8], fd: RawFd) -> io::Result<()> { + let fd_size = size_of::() as u32; + // SAFETY: CMSG_SPACE is a pure size computation. + let space = unsafe { libc::CMSG_SPACE(fd_size) } as usize; + // u64 elements align the control buffer for a cmsghdr. + let mut control = vec![0_u64; space.div_ceil(size_of::())]; + let mut vector = libc::iovec { + iov_base: frame.as_ptr().cast_mut().cast(), + iov_len: frame.len(), + }; + // SAFETY: msghdr is plain data for which all-zero is valid; the pointers + // set below outlive the sendmsg call, and the cmsg header written lies + // inside `control`, which CMSG_SPACE sized for one descriptor. + let sent = unsafe { + let mut message: libc::msghdr = zeroed(); + message.msg_iov = &mut vector; + message.msg_iovlen = 1; + message.msg_control = control.as_mut_ptr().cast(); + message.msg_controllen = space as _; + let header = libc::CMSG_FIRSTHDR(&message); + (*header).cmsg_level = libc::SOL_SOCKET; + (*header).cmsg_type = libc::SCM_RIGHTS; + (*header).cmsg_len = libc::CMSG_LEN(fd_size) as _; + ptr::copy_nonoverlapping( + ptr::from_ref(&fd).cast::(), + libc::CMSG_DATA(header), + size_of::(), + ); + loop { + let sent = libc::sendmsg(stream.as_raw_fd(), &message, libc::MSG_NOSIGNAL); + if sent >= 0 || io::Error::last_os_error().kind() != io::ErrorKind::Interrupted { + break sent; + } + } + }; + let Ok(sent) = usize::try_from(sent) else { + return Err(io::Error::last_os_error()); + }; + // The descriptor travels with the first byte; any rest is plain data. + let mut stream = stream; + stream.write_all(&frame[sent..]) +} + +fn unix_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |elapsed| { + u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX) + }) +} + +struct Permit(Arc); + +impl Permit { + fn acquire(active: &Arc, limit: usize) -> Option { + active + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |held| { + (held < limit).then_some(held + 1) + }) + .ok() + .map(|_| Self(active.clone())) + } +} + +impl Drop for Permit { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } +} + +struct Lender { + callers: Vec, + limits: LeaseLimits, + service: Arc>, + tokens: Arc, + units: Arc, +} + +/// Serves leases on `listener` until it fails. Each connection is one lease: +/// request, then grant or refusal, then (after a grant) settle. A caller is +/// the account whose uid the kernel reports for the connection, and only a +/// process in that caller's Self unit may lease, once per unit invocation; no +/// field of any message names either. +pub fn serve_leases( + listener: &UnixListener, + callers: Vec, + limits: LeaseLimits, + service: Arc>, + tokens: Arc, + units: Arc, +) -> io::Result<()> { + let lender = Arc::new(Lender { + callers, + limits, + service, + tokens, + units, + }); + let active = Arc::new(AtomicUsize::new(0)); + loop { + let (stream, _) = listener.accept()?; + let Some(permit) = Permit::acquire(&active, limits.max_connections) else { + continue; + }; + let lender = lender.clone(); + thread::Builder::new() + .name("claude-lease".to_string()) + .spawn(move || { + let _permit = permit; + // A connection that fails is dropped; the lease it may have + // held is settled inside `lend`. + let _ = lender.lend(stream); + })?; + } +} + +impl Lender { + fn lend(&self, mut stream: UnixStream) -> io::Result<()> { + stream.set_write_timeout(Some(self.limits.socket_timeout))?; + let peer = peer_credentials(&stream)?; + let request_deadline = Instant::now() + self.limits.socket_timeout; + let LeaseMessage::ClaudeLeaseRequest(request) = + read_message(&mut stream, request_deadline)? + else { + return Err(io::Error::from(io::ErrorKind::InvalidData)); + }; + let Some(caller) = self.callers.iter().find(|caller| caller.uid == peer.uid) else { + return refuse( + &mut stream, + LeaseRefusalKind::UnknownCaller, + "this account may not lease", + ); + }; + // The unit and its invocation come from the system's record of the + // peer's process, never from the peer. + let Some(invocation) = peer + .process + .as_ref() + .and_then(|process| self.units.invocation_in(process.as_fd(), &caller.unit())) + .filter(|found| found.unit == caller.unit()) + else { + return refuse( + &mut stream, + LeaseRefusalKind::UnknownCaller, + "this process is not in the Self unit", + ); + }; + let now = unix_ms(); + if request.expires_at_unix_ms <= now + || request.expires_at_unix_ms > now.saturating_add(self.limits.max_lease_ms) + { + return refuse( + &mut stream, + LeaseRefusalKind::Expired, + "the expiry is not in the future or is past the lease bound", + ); + } + // The reservation is claimed in the ledger before the token is read, + // let alone lent. A ledger that fails closes the connection with no + // answer and no lease. + let admission = self + .service + .lock() + .map_err(|_| io::Error::other("service lock"))? + .begin_lease( + &caller.caller_id, + &invocation.invocation_id, + caller.reservation_micro_usd, + request.expires_at_unix_ms, + now, + ) + .map_err(|_| io::Error::other("ledger"))?; + let grant = match admission { + LeaseAdmission::Granted(grant) => grant, + LeaseAdmission::Refused(why) => { + let (kind, detail) = match why { + LeaseRefusal::SpendCeiling => ( + LeaseRefusalKind::SpendCeiling, + "the window cannot hold this reservation", + ), + LeaseRefusal::CostBound => ( + LeaseRefusalKind::CostBound, + "the reservation must be above zero", + ), + LeaseRefusal::UnknownCaller => ( + LeaseRefusalKind::UnknownCaller, + "this account is bound to no admitted caller", + ), + LeaseRefusal::Busy => ( + LeaseRefusalKind::Busy, + "a lease is open for this caller or this unit invocation leased", + ), + }; + return refuse(&mut stream, kind, detail); + } + }; + self.hold(stream, caller, &grant, request.expires_at_unix_ms) + } + + /// Everything after a reservation exists: lend the token, then settle the + /// lease exactly once whatever happens. + fn hold( + &self, + mut stream: UnixStream, + caller: &LeaseCaller, + grant: &LeaseGrant, + expires_at_unix_ms: u64, + ) -> io::Result<()> { + // A token that cannot be read or piped was never lent: nothing was + // used, so nothing is billed. + let Some(read_end) = self + .tokens + .read() + .and_then(|token| token_pipe(token.as_bytes()).ok()) + else { + self.settle(caller, grant, LeaseUse::NeverLent)?; + return Ok(()); + }; + let frame = encode_frame(&LeaseMessage::ClaudeLeaseGrant(LeaseGranted { + lease_id: grant.lease_id.clone(), + reserved_micro_usd: grant.reserved_micro_usd, + })); + let sent = send_with_fd(&stream, &frame, read_end.as_raw_fd()); + drop(read_end); + if sent.is_err() { + // The descriptor may have been delivered: the reservation is billed. + self.settle(caller, grant, LeaseUse::Reservation)?; + return Ok(()); + } + let outcome = self.await_settle(&mut stream, grant, expires_at_unix_ms); + match outcome { + Some(settle) => { + let settlement = self.settle(caller, grant, LeaseUse::Reported(&settle.usage))?; + let _ = stream.write_all(&encode_frame(&LeaseMessage::ClaudeLeaseSettled( + LeaseSettled { + billed_micro_usd: settlement.billed_micro_usd, + overrun_micro_usd: settlement.overrun_micro_usd, + }, + ))); + } + None => { + self.settle(caller, grant, LeaseUse::Reservation)?; + } + } + Ok(()) + } + + /// The run's settle, if its whole frame arrives for this lease before the + /// lease expires. Anything else (end of file, silence, a frame still + /// trickling at the expiry, a malformed frame, the wrong lease) is no + /// report. + fn await_settle( + &self, + stream: &mut UnixStream, + grant: &LeaseGrant, + expires_at_unix_ms: u64, + ) -> Option { + let remaining = expires_at_unix_ms.saturating_sub(unix_ms()); + let deadline = Instant::now() + Duration::from_millis(remaining); + match read_message(stream, deadline).ok()? { + LeaseMessage::ClaudeLeaseSettle(settle) if settle.lease_id == grant.lease_id => { + Some(settle) + } + _ => None, + } + } + + fn settle( + &self, + caller: &LeaseCaller, + grant: &LeaseGrant, + usage: LeaseUse<'_>, + ) -> io::Result { + self.service + .lock() + .map_err(|_| io::Error::other("service lock"))? + .settle_lease(&caller.caller_id, &grant.lease_id, usage) + .map_err(|_| io::Error::other("ledger")) + } +} + +fn refuse(stream: &mut UnixStream, kind: LeaseRefusalKind, detail: &str) -> io::Result<()> { + stream.write_all(&encode_frame(&LeaseMessage::ClaudeLeaseRefusal( + LeaseRefused { + kind, + detail: detail.to_string(), + }, + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A settle frame whose body is exactly `body_bytes` long. + fn settle_frame_of(body_bytes: usize) -> Vec { + let frame_with = |pad: usize| { + encode_frame(&LeaseMessage::ClaudeLeaseSettle(LeaseSettle { + lease_id: "x".repeat(pad), + usage: Vec::new(), + })) + }; + // The string header is the same width between 256 and 65535 bytes. + let base = frame_with(300).len() - 4 - 300; + frame_with(body_bytes - base) + } + + fn within(frame: &[u8]) -> io::Result { + read_message(&mut &frame[..], Instant::now() + Duration::from_secs(5)) + } + + #[test] + fn the_frame_bound_is_sixty_five_thousand_five_hundred_thirty_six() { + assert_eq!(MAX_LEASE_FRAME_BYTES, 65_536); + } + + #[test] + fn a_frame_is_one_whole_message_within_the_bound() { + let at_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES); + assert_eq!(at_bound.len(), 4 + MAX_LEASE_FRAME_BYTES); + assert!(decode_frame(&at_bound).is_some()); + let past_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES + 1); + assert!(decode_frame(&past_bound).is_none()); + let mut trailing = at_bound.clone(); + trailing.push(0); + assert!(decode_frame(&trailing).is_none()); + assert!(decode_frame(&at_bound[..at_bound.len() - 1]).is_none()); + assert!(decode_frame(&[0, 0, 0, 0]).is_none()); + assert!(decode_frame(&[0, 0]).is_none()); + } + + #[test] + fn a_stream_frame_over_the_bound_is_refused_before_it_is_read() { + let past_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES + 1); + assert!(within(&past_bound).is_err()); + let at_bound = settle_frame_of(MAX_LEASE_FRAME_BYTES); + assert!(within(&at_bound).is_ok()); + assert!(within(&[0_u8, 0, 0, 0]).is_err()); + // Short of the length it announced is end of file, not a message. + assert!(within(&at_bound[..at_bound.len() - 1]).is_err()); + } + + #[test] + fn a_frame_that_trickles_past_its_deadline_is_refused_however_small_each_read() { + let (mut reader, mut writer) = UnixStream::pair().unwrap(); + let frame = encode_frame(&LeaseMessage::ClaudeLeaseSettle(LeaseSettle { + lease_id: "lease:x".to_string(), + usage: Vec::new(), + })); + let trickle = thread::spawn(move || { + for byte in frame { + if writer.write_all(&[byte]).is_err() { + return; + } + thread::sleep(Duration::from_millis(60)); + } + }); + // Every read lands well inside any per-read timeout; the frame as a + // whole does not. + let started = Instant::now(); + let outcome = read_message(&mut reader, started + Duration::from_millis(400)); + assert_eq!(outcome.unwrap_err().kind(), io::ErrorKind::TimedOut); + assert!(started.elapsed() < Duration::from_millis(1_000)); + drop(reader); + trickle.join().unwrap(); + } + + /// A reader that fails the way a reset connection does. + struct Reset; + + impl Read for Reset { + fn read(&mut self, _: &mut [u8]) -> io::Result { + Err(io::Error::from(io::ErrorKind::ConnectionReset)) + } + } + + impl Bounded for Reset { + fn within(&mut self, _: Duration) -> io::Result<()> { + Ok(()) + } + } + + #[test] + fn a_read_that_fails_ends_the_frame_at_once_not_at_the_deadline() { + let started = Instant::now(); + let outcome = read_message(&mut Reset, started + Duration::from_secs(30)); + assert_eq!(outcome.unwrap_err().kind(), io::ErrorKind::ConnectionReset); + assert!(started.elapsed() < Duration::from_secs(1)); + } + + #[test] + fn a_frame_that_arrives_in_pieces_within_its_deadline_is_read() { + let (mut reader, mut writer) = UnixStream::pair().unwrap(); + let frame = encode_frame(&LeaseMessage::ClaudeLeaseRequest(LeaseRequest { + expires_at_unix_ms: 7, + })); + let expected = decode_frame(&frame); + let pieces = thread::spawn(move || { + for part in frame.chunks(3) { + writer.write_all(part).unwrap(); + thread::sleep(Duration::from_millis(10)); + } + }); + let read = read_message(&mut reader, Instant::now() + Duration::from_secs(5)).unwrap(); + assert_eq!(Some(read), expected); + pieces.join().unwrap(); + } +} diff --git a/crates/claude-connector/src/lib.rs b/crates/claude-connector/src/lib.rs index 327a078..22ff46f 100644 --- a/crates/claude-connector/src/lib.rs +++ b/crates/claude-connector/src/lib.rs @@ -8,18 +8,32 @@ mod config; mod daemon; mod handle; mod launch; +#[cfg(target_os = "linux")] +mod lease; mod process; mod quota; mod stream; +#[cfg(target_os = "linux")] +mod unit; pub use backend::{ClaudeCodeBackend, CredentialTokenSource, StderrSink, TokenSource}; pub use config::{ClaudeCodeConfig, ClaudeCodeLimits}; -pub use daemon::{ClaudeConnectorConfig, ConfigError, check_ready, load_config, write_config}; +pub use daemon::{ + ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, check_ready, load_config, write_config, +}; #[cfg(unix)] pub use daemon::{check_prompt_dir, serve_config}; pub use handle::{FsSessionStore, SessionStore}; pub use launch::ChildEnv; +#[cfg(target_os = "linux")] +pub use lease::{ + LeaseCaller, LeaseGranted, LeaseLimits, LeaseMessage, LeaseRefusalKind, LeaseRefused, + LeaseRequest, LeaseSettle, LeaseSettled, MAX_LEASE_FRAME_BYTES, PeerCredentials, PeerUnit, + PeerUnits, decode_frame, encode_frame, peer_credentials, serve_leases, user_uid, +}; #[cfg(unix)] pub use process::StdProcessPort; pub use process::{ProcessPort, RunningChild, StdoutRead}; pub use quota::QuotaCell; +#[cfg(target_os = "linux")] +pub use unit::SystemdUnits; diff --git a/crates/claude-connector/src/main.rs b/crates/claude-connector/src/main.rs index c58e7f5..0ea296f 100644 --- a/crates/claude-connector/src/main.rs +++ b/crates/claude-connector/src/main.rs @@ -13,7 +13,7 @@ use std::path::PathBuf; use std::process::ExitCode; use claude_connector::{ - ClaudeConnectorConfig, ConfigError, check_ready, load_config, write_config, + ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, check_ready, load_config, write_config, }; use model_connector::{CallerConfig, ModelPrice, Rates, frame_bytes_for_payload}; @@ -76,9 +76,10 @@ struct Flags { single: BTreeMap<&'static str, String>, models: Vec, prices: Vec, + lease_users: Vec, } -const SINGLE: [&str; 22] = [ +const SINGLE: [&str; 26] = [ "--listen", "--replay-store", "--claude-binary", @@ -101,6 +102,10 @@ const SINGLE: [&str; 22] = [ "--max-payload-bytes", "--max-output-tokens", "--socket-timeout-ms", + "--lease-socket", + "--lease-max-ms", + "--lease-unit-instance", + "--lease-reservation-micro-usd", ]; impl Flags { @@ -109,6 +114,7 @@ impl Flags { single: BTreeMap::new(), models: Vec::new(), prices: Vec::new(), + lease_users: Vec::new(), }; let mut args = args.iter(); while let Some(&name) = args.next() { @@ -118,6 +124,7 @@ impl Flags { match name { "--allowed-model" => flags.models.push(value.to_string()), "--price" => flags.prices.push(value.to_string()), + "--lease-user" => flags.lease_users.push(value.to_string()), _ => { let Some(known) = SINGLE.iter().find(|known| **known == name) else { return Err(Failure::Flags); @@ -206,6 +213,19 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> { return Err(Failure::Malformed("--price")); } } + let caller_id = flags.text("--caller-id")?; + let lease_callers = flags + .lease_users + .iter() + .map(|unix_user| { + Ok(LeaseCallerConfig { + unix_user: unix_user.clone(), + caller_id: caller_id.clone(), + unit_instance: flags.text("--lease-unit-instance")?, + reservation_micro_usd: flags.number("--lease-reservation-micro-usd")?, + }) + }) + .collect::, Failure>>()?; let config = ClaudeConnectorConfig { listen: flags.text("--listen")?, replay_store: flags.path("--replay-store")?, @@ -226,7 +246,7 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> { socket_timeout_ms: flags.number_or("--socket-timeout-ms", 10_000)?, max_expiry_skew_ms: 300_000, callers: vec![CallerConfig { - caller_runtime_id: flags.text("--caller-id")?, + caller_runtime_id: caller_id.clone(), connection_key_file: flags.path("--connection-key-file")?, connection_key_epoch: flags.number("--connection-key-epoch")?, allowed_models: flags.models.clone(), @@ -237,6 +257,9 @@ fn init(path: &str, args: &[&str]) -> Result<(), Failure> { spend_window_ms: flags.number("--spend-window-ms")?, }], prices, + lease_socket: flags.single.get("--lease-socket").map(PathBuf::from), + lease_callers, + lease_max_ms: flags.number_or("--lease-max-ms", 3_600_000)?, }; write_config(path.as_ref(), &config)?; Ok(()) diff --git a/crates/claude-connector/src/unit.rs b/crates/claude-connector/src/unit.rs new file mode 100644 index 0000000..1ab8dc9 --- /dev/null +++ b/crates/claude-connector/src/unit.rs @@ -0,0 +1,361 @@ +//! Which systemd unit invocation a process runs in, from the system's own +//! record: the process's cgroup says the unit, and the service manager says +//! that unit's current InvocationID. The lease listener asks this through the +//! `PeerUnits` port; the answer never comes from the peer. +//! +//! This is the one place the connector asks systemd anything. It runs +//! `systemctl show` for one read-only property of one unit and takes nothing +//! from the output but a 32-digit hex id. + +use std::fs; +use std::io::Read; +use std::os::fd::{AsRawFd, BorrowedFd}; +use std::path::PathBuf; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +use crate::lease::{PeerUnit, PeerUnits}; + +/// Answers `PeerUnits` from `//cgroup` and `systemctl show`. +pub struct SystemdUnits { + proc_root: PathBuf, + systemctl: PathBuf, + /// How long the manager has to answer before the lease is refused. + deadline: Duration, +} + +impl SystemdUnits { + /// The running system's `/proc` and `systemctl`. + pub fn system() -> Self { + Self { + proc_root: PathBuf::from("/proc"), + systemctl: PathBuf::from("/usr/bin/systemctl"), + deadline: Duration::from_secs(5), + } + } + + /// The pid a process handle names now, from the handle's own record. The + /// handle holds the process's identity, so the number cannot have been + /// recycled; a process that has exited reads as `-1`. + fn pid_of(&self, process: BorrowedFd<'_>) -> Option { + let info = fs::read_to_string( + self.proc_root + .join("self/fdinfo") + .join(process.as_raw_fd().to_string()), + ) + .ok()?; + info.lines() + .find_map(|line| line.strip_prefix("Pid:"))? + .trim() + .parse() + .ok() + } + + fn unit_in_cgroup(&self, pid: u32) -> Option { + let cgroup = fs::read_to_string(self.proc_root.join(pid.to_string()).join("cgroup")); + unit_of_cgroup(&cgroup.ok()?) + } + + /// `systemctl show` of `unit`'s InvocationID, killed and refused if it has + /// not exited by the deadline. + fn invocation_of(&self, unit: &str) -> Option { + let mut child = Command::new(&self.systemctl) + .env_clear() + .args(["show", "--property=InvocationID", "--value", "--", unit]) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .spawn() + .ok()?; + let give_up = Instant::now() + self.deadline; + let status = loop { + match child.try_wait() { + Ok(Some(status)) => break status, + Ok(None) if Instant::now() < give_up => { + std::thread::sleep(Duration::from_millis(10)); + } + _ => { + let _ = child.kill(); + let _ = child.wait(); + return None; + } + } + }; + if !status.success() { + return None; + } + // The id is 33 bytes; anything past a small bound is not an answer. + let mut output = Vec::new(); + child + .stdout + .take()? + .take(256) + .read_to_end(&mut output) + .ok()?; + invocation_id_of(std::str::from_utf8(&output).ok()?) + } +} + +/// The service unit a cgroup file puts its process in: the last component of +/// the unified-hierarchy path, if it names a `.service`. A process in a +/// sub-cgroup, a scope, a slice or the root is in no service unit this +/// recognises. +pub fn unit_of_cgroup(cgroup: &str) -> Option { + let path = cgroup.lines().find_map(|line| line.strip_prefix("0::"))?; + let unit = path.rsplit('/').next()?; + (unit.len() > ".service".len() && unit.ends_with(".service")).then(|| unit.to_string()) +} + +/// An InvocationID as `systemctl show --value` prints it: 32 lowercase hex +/// digits and a newline. Anything else (including the empty answer of a unit +/// that is not running) is not an invocation. +pub fn invocation_id_of(output: &str) -> Option { + let id = output.strip_suffix('\n').unwrap_or(output); + (id.len() == 32 && id.bytes().all(|b| matches!(b, b'0'..=b'9' | b'a'..=b'f'))) + .then(|| id.to_string()) +} + +impl PeerUnits for SystemdUnits { + fn invocation_in(&self, process: BorrowedFd<'_>, unit: &str) -> Option { + let pid = self.pid_of(process)?; + // A process in any other unit is not asked about. + if self.unit_in_cgroup(pid).as_deref() != Some(unit) { + return None; + } + let invocation_id = self.invocation_of(unit)?; + // The same process must still be in that unit now: one that left (or + // died) while the manager was asked is not vouched for. + (self.pid_of(process) == Some(pid) && self.unit_in_cgroup(pid).as_deref() == Some(unit)) + .then(|| PeerUnit { + unit: unit.to_string(), + invocation_id, + }) + } +} + +#[cfg(test)] +mod tests { + use std::os::fd::{AsFd, AsRawFd, OwnedFd}; + + use super::*; + + #[test] + fn a_service_unit_is_the_last_cgroup_component() { + let service = "0::/system.slice/system-mind\\x2dself.slice/mind-self@huginn.service\n"; + assert_eq!( + unit_of_cgroup(service).as_deref(), + Some("mind-self@huginn.service") + ); + // Several hierarchies listed: the unified one decides. + let hybrid = "12:cpu:/foo\n0::/system.slice/mind-self@a.service\n"; + assert_eq!( + unit_of_cgroup(hybrid).as_deref(), + Some("mind-self@a.service") + ); + } + + #[test] + fn a_process_outside_a_service_unit_has_none() { + for cgroup in [ + "0::/\n", + "0::/user.slice/user-1000.slice/session-3.scope\n", + "0::/user.slice/user-1000.slice/user@1000.service/app.slice/x.service/child\n", + "0::/system.slice/mind-self@huginn.service/sub\n", + "0::/system.slice/.service\n", + "1:name=systemd:/system.slice/a.service\n", + "", + ] { + assert_eq!(unit_of_cgroup(cgroup), None, "{cgroup:?}"); + } + } + + #[test] + fn an_invocation_id_is_exactly_32_lowercase_hex_digits() { + let id = "0123456789abcdef0123456789abcdef"; + assert_eq!(invocation_id_of(&format!("{id}\n")).as_deref(), Some(id)); + assert_eq!(invocation_id_of(id).as_deref(), Some(id)); + for bad in [ + "", + "\n", + "0123456789ABCDEF0123456789ABCDEF\n", + "0123456789abcdef0123456789abcde\n", + "0123456789abcdef0123456789abcdef0\n", + "0123456789abcdef0123456789abcdeg\n", + "0123456789abcdef0123456789abcdef\n\n", + ] { + assert_eq!(invocation_id_of(bad), None, "{bad:?}"); + } + } + + const ID: &str = "0123456789abcdef0123456789abcdef"; + const UNIT: &str = "mind-self@eureka.service"; + + /// A fake `/proc` with one process in `cgroup`, and a fake `systemctl` + /// that records its arguments, runs `side_effect`, and exits `status` + /// after printing `output`. + struct Fake { + dir: PathBuf, + units: SystemdUnits, + // A script being written while another test forks is "text file + // busy" when it is run: fakes take turns. + _turn: std::sync::MutexGuard<'static, ()>, + } + + static TURN: std::sync::Mutex<()> = std::sync::Mutex::new(()); + + impl Fake { + fn new(name: &str, cgroup: &str, output: &str, status: i32, side_effect: &str) -> Self { + use std::os::unix::fs::PermissionsExt; + let turn = TURN + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let dir = std::env::temp_dir().join(format!("unit-{}-{name}", std::process::id())); + let _ = fs::remove_dir_all(&dir); + let side_effect = side_effect.replace("{d}", &dir.display().to_string()); + fs::create_dir_all(dir.join("proc/42")).unwrap(); + fs::write(dir.join("proc/42/cgroup"), cgroup).unwrap(); + let script = dir.join("systemctl"); + fs::write( + &script, + format!( + "#!/bin/sh\nprintf '%s' \"$*\" > {d}/args\n{side_effect}\nprintf '%s' '{output}'\nexit {status}\n", + d = dir.display() + ), + ) + .unwrap(); + fs::set_permissions(&script, fs::Permissions::from_mode(0o755)).unwrap(); + let units = SystemdUnits { + proc_root: dir.join("proc"), + systemctl: script, + deadline: Duration::from_secs(5), + }; + Self { + dir, + units, + _turn: turn, + } + } + + /// A process handle that names process `pid`, as the kernel's own + /// record of the handle says (`-1` once the process has exited). + fn process(&self, pid: i64) -> OwnedFd { + let handle: OwnedFd = fs::File::open("/dev/null").unwrap().into(); + let info = self.dir.join("proc/self/fdinfo"); + fs::create_dir_all(&info).unwrap(); + fs::write( + info.join(handle.as_raw_fd().to_string()), + format!("pos:\t0\nflags:\t02000002\nPid:\t{pid}\nNSpid:\t{pid}\n"), + ) + .unwrap(); + handle + } + + fn invocation_in(&self, pid: i64, unit: &str) -> Option { + self.units.invocation_in(self.process(pid).as_fd(), unit) + } + + fn args(&self) -> String { + fs::read_to_string(self.dir.join("args")).unwrap() + } + } + + impl Drop for Fake { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.dir); + } + } + + fn eureka_cgroup() -> String { + format!("0::/system.slice/{UNIT}\n") + } + + #[test] + fn a_process_in_the_unit_is_that_unit_at_the_invocation_the_manager_reports() { + let fake = Fake::new("ok", &eureka_cgroup(), &format!("{ID}\n"), 0, ""); + assert_eq!( + fake.invocation_in(42, UNIT), + Some(PeerUnit { + unit: UNIT.to_string(), + invocation_id: ID.to_string() + }) + ); + // The manager is asked for exactly the invocation of that unit. + assert_eq!( + fake.args(), + format!("show --property=InvocationID --value -- {UNIT}") + ); + // A handle on a process that has exited, or on one with no cgroup + // record, is in no unit. + assert_eq!(fake.invocation_in(-1, UNIT), None); + assert_eq!(fake.invocation_in(43, UNIT), None); + } + + #[test] + fn a_process_outside_the_unit_is_not_asked_about() { + let scope = Fake::new("scope", "0::/user.slice/session-3.scope\n", ID, 0, ""); + assert_eq!(scope.invocation_in(42, UNIT), None); + assert!(!scope.dir.join("args").exists()); + drop(scope); + // Another service, even a real one, is not the caller's unit: the + // manager is never spawned for it. + let other = Fake::new("other", "0::/system.slice/ssh.service\n", ID, 0, ""); + assert_eq!(other.invocation_in(42, UNIT), None); + assert!(!other.dir.join("args").exists()); + } + + #[test] + fn a_manager_that_fails_or_names_no_invocation_vouches_for_nothing() { + let failed = Fake::new("fail", &eureka_cgroup(), ID, 1, ""); + assert_eq!(failed.invocation_in(42, UNIT), None); + drop(failed); + let not_running = Fake::new("idle", &eureka_cgroup(), "\n", 0, ""); + assert_eq!(not_running.invocation_in(42, UNIT), None); + let missing = SystemdUnits { + proc_root: not_running.dir.join("proc"), + systemctl: not_running.dir.join("no-such-systemctl"), + deadline: Duration::from_secs(5), + }; + assert_eq!( + missing.invocation_in(not_running.process(42).as_fd(), UNIT), + None + ); + } + + #[test] + fn a_manager_that_does_not_answer_in_time_is_killed_and_the_lease_refused() { + let mut fake = Fake::new("hung", &eureka_cgroup(), ID, 0, "exec sleep 30"); + fake.units.deadline = Duration::from_millis(200); + let started = Instant::now(); + assert_eq!(fake.invocation_in(42, UNIT), None); + assert!(started.elapsed() < Duration::from_secs(5)); + } + + #[test] + fn a_process_that_left_its_unit_while_the_manager_was_asked_is_not_vouched_for() { + let fake = Fake::new( + "moves", + &eureka_cgroup(), + ID, + 0, + "printf '0::/system.slice/other.service\n' > {d}/proc/42/cgroup", + ); + assert_eq!(fake.invocation_in(42, UNIT), None); + } + + #[test] + fn a_process_that_exited_while_the_manager_was_asked_is_not_vouched_for() { + let fake = Fake::new("dies", &eureka_cgroup(), ID, 0, ""); + let handle = fake.process(42); + let side = fake + .dir + .join("proc/self/fdinfo") + .join(handle.as_raw_fd().to_string()); + let mut script = fs::read_to_string(fake.dir.join("systemctl")).unwrap(); + script = script.replace( + "exit 0", + &format!("printf 'Pid:\t-1\n' > {}\nexit 0", side.display()), + ); + fs::write(fake.dir.join("systemctl"), script).unwrap(); + assert_eq!(fake.units.invocation_in(handle.as_fd(), UNIT), None); + } +} diff --git a/crates/claude-connector/tests/config_document.rs b/crates/claude-connector/tests/config_document.rs index 141f471..80285a2 100644 --- a/crates/claude-connector/tests/config_document.rs +++ b/crates/claude-connector/tests/config_document.rs @@ -8,7 +8,9 @@ use std::process::{Command, Output}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; -use claude_connector::{ClaudeConnectorConfig, ConfigError, load_config, write_config}; +use claude_connector::{ + ClaudeConnectorConfig, ConfigError, LeaseCallerConfig, load_config, write_config, +}; use cultcache_rs::{ CacheBackingStore, CultCache, CultCacheEnvelope, DatabaseEntry, SingleFileMessagePackBackingStore, @@ -92,6 +94,14 @@ fn valid() -> ClaudeConnectorConfig { spend_window_ms: 3_600_000, }], prices: BTreeMap::from([(HAIKU.to_string(), price()), (OPUS.to_string(), price())]), + lease_socket: Some("/run/claude-connector/lease.sock".into()), + lease_callers: vec![LeaseCallerConfig { + unix_user: "mind-eureka-self".to_string(), + caller_id: "persona-a-caller".to_string(), + unit_instance: "eureka".to_string(), + reservation_micro_usd: 5_000_000, + }], + lease_max_ms: 3_600_000, } } @@ -989,3 +999,210 @@ fn a_private_regular_file_is_not_a_prompt_dir() { assert!(claude_connector::check_prompt_dir(&file).is_err()); assert!(claude_connector::check_prompt_dir(&scratch.path("absent")).is_err()); } + +#[test] +fn leases_are_configured_together_or_not_at_all() { + let cases: Vec<(&str, Change)> = vec![ + ( + "lease_socket", + Box::new(|c| c.lease_socket = Some("lease.sock".into())), + ), + ("lease_callers", Box::new(|c| c.lease_callers.clear())), + ("lease_callers", Box::new(|c| c.lease_socket = None)), + ( + "lease_callers", + Box::new(|c| c.lease_callers[0].caller_id = "another-caller".to_string()), + ), + ( + "lease_callers", + Box::new(|c| c.lease_callers[0].unix_user = " ".to_string()), + ), + ( + "lease_callers", + Box::new(|c| { + let again = c.lease_callers[0].clone(); + c.lease_callers.push(again); + }), + ), + ("lease_max_ms", Box::new(|c| c.lease_max_ms = 0)), + ( + "lease_callers", + Box::new(|c| c.lease_callers[0].reservation_micro_usd = 0), + ), + ( + "lease_callers", + Box::new(|c| c.lease_callers[0].reservation_micro_usd = 50_000_001), + ), + ( + "lease_callers", + Box::new(|c| c.lease_callers[0].unit_instance = String::new()), + ), + ( + "lease_callers", + Box::new(|c| c.lease_callers[0].unit_instance = "a/b".to_string()), + ), + ]; + for (field, change) in cases { + let mut config = valid(); + change(&mut config); + assert_eq!(refused(&config), field); + } + // The whole ceiling is a reservation the caller can hold. + let mut whole = valid(); + whole.lease_callers[0].reservation_micro_usd = 50_000_000; + whole.validate().unwrap(); + let mut none = valid(); + none.lease_socket = None; + none.lease_callers.clear(); + none.validate().unwrap(); +} + +#[test] +fn config_init_binds_each_lease_user_to_the_configured_caller() { + let scratch = Scratch::new(); + let made = init( + &scratch, + &init_flags( + &scratch, + &[ + "--lease-socket", + "/run/claude-connector/lease.sock", + "--lease-user", + "mind-eureka-self", + "--lease-max-ms", + "120000", + "--lease-unit-instance", + "eureka", + "--lease-reservation-micro-usd", + "5000000", + ], + ), + ); + assert!(made.status.success(), "{}", shown(&made)); + let config = load_config(&scratch.path("connector.cc")).unwrap(); + assert_eq!( + config.lease_callers, + [LeaseCallerConfig { + unix_user: "mind-eureka-self".to_string(), + caller_id: "persona-a-caller".to_string(), + unit_instance: "eureka".to_string(), + reservation_micro_usd: 5_000_000, + }] + ); + assert_eq!(config.lease_max_ms, 120_000); +} + +#[cfg(target_os = "linux")] +#[test] +fn config_check_refuses_a_lease_account_the_system_does_not_have() { + let scratch = Scratch::new(); + let (mut flags, credentials) = prepare(&scratch); + let socket = scratch.path("lease.sock"); + flags.extend( + [ + "--lease-socket", + socket.to_str().unwrap(), + "--lease-user", + "no-such-account-4f1c", + "--lease-unit-instance", + "eureka", + "--lease-reservation-micro-usd", + "5000000", + ] + .map(String::from), + ); + let made = init(&scratch, &flags); + assert!(made.status.success(), "{}", shown(&made)); + let config = scratch.path("connector.cc"); + let checked = run( + &["config", "check", config.to_str().unwrap()], + &[("CREDENTIALS_DIRECTORY", &credentials)], + ); + assert!(!checked.status.success()); + assert!( + shown(&checked).contains("lease_callers"), + "{}", + shown(&checked) + ); + assert!(!shown(&checked).contains("no-such-account")); +} + +#[cfg(target_os = "linux")] +#[test] +fn serve_binds_the_lease_socket_and_answers_a_stranger_with_a_refusal() { + use std::io::{Read, Write}; + use std::os::unix::fs::PermissionsExt; + use std::os::unix::net::UnixStream; + use std::process::Stdio; + use std::time::Instant; + + use claude_connector::{ + LeaseMessage, LeaseRefusalKind, LeaseRefused, LeaseRequest, decode_frame, encode_frame, + }; + + let scratch = Scratch::new(); + let (mut flags, credentials) = prepare(&scratch); + let socket = scratch.path("lease.sock"); + // The test runs as someone other than `nobody`, so it is a stranger. + flags.extend( + [ + "--lease-socket", + socket.to_str().unwrap(), + "--lease-user", + "nobody", + "--lease-unit-instance", + "eureka", + "--lease-reservation-micro-usd", + "5000000", + ] + .map(String::from), + ); + let made = init(&scratch, &flags); + assert!(made.status.success(), "{}", shown(&made)); + let config = scratch.path("connector.cc"); + let mut server = Command::new(env!("CARGO_BIN_EXE_claude-connector")) + .args(["serve", "--config", config.to_str().unwrap()]) + .env_clear() + .env("CREDENTIALS_DIRECTORY", &credentials) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + let started = Instant::now(); + while !socket.exists() && started.elapsed() < Duration::from_secs(10) { + std::thread::sleep(Duration::from_millis(50)); + } + let outcome = std::panic::catch_unwind(|| { + assert_eq!( + std::fs::metadata(&socket).unwrap().permissions().mode() & 0o777, + 0o660 + ); + let mut stream = UnixStream::connect(&socket).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + stream + .write_all(&encode_frame(&LeaseMessage::ClaudeLeaseRequest( + LeaseRequest { + expires_at_unix_ms: u64::MAX / 2, + }, + ))) + .unwrap(); + let mut reply = vec![0_u8; 4096]; + let read = stream.read(&mut reply).unwrap(); + assert!(matches!( + decode_frame(&reply[..read]), + Some(LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { + kind: LeaseRefusalKind::UnknownCaller, + .. + })) + )); + }); + // The server is still serving: it did not return when it started. + let running = server.try_wait().unwrap().is_none(); + server.kill().unwrap(); + server.wait().unwrap(); + outcome.unwrap(); + assert!(running); +} diff --git a/crates/claude-connector/tests/fixtures/lease_v1/grant.bin b/crates/claude-connector/tests/fixtures/lease_v1/grant.bin new file mode 100644 index 0000000000000000000000000000000000000000..554e3a64cb35edbf81585cd63ecbe95744d580df GIT binary patch literal 109 zcmZQzV90D*lAW2CvMD(yu{0$$J|{J?I5j@KC^4^OX<1@UX=>97h)`zAO*Jsf%D~Xb r*u>P#+`=+3DLExI4V}NKD784Xs4O)lJ~uPDC_lcmIOQAzr}_Z^$PFl0 literal 0 HcmV?d00001 diff --git a/crates/claude-connector/tests/fixtures/lease_v1/refusal.bin b/crates/claude-connector/tests/fixtures/lease_v1/refusal.bin new file mode 100644 index 0000000000000000000000000000000000000000..22a4a25e3ceba32fc73bbfaf38dd854feff2550d GIT binary patch literal 104 zcmXBL!3~2j3KM{AsLOPL()t?<|E=k6otQfW-a1l#;wZPj&fC#$7NP yxM8fAdt=)cwvG#YhZkJ#P*Z>Gf?&1(Ei!05ds%5>_Y#zg2}RaM6%BRa(BcQ7qADT) literal 0 HcmV?d00001 diff --git a/crates/claude-connector/tests/fixtures/lease_v1/request.bin b/crates/claude-connector/tests/fixtures/lease_v1/request.bin new file mode 100644 index 0000000000000000000000000000000000000000..fddaa3785fe2e36c3d51f2ec7bb63649d2decda3 GIT binary patch literal 66 zcmZQzV6bailAW2CvL!huu{0$$J|{J?I5oZ~wXig`xMXQrVoqsl3W!){%1t#e%gVse z$k@cx%-q51Y}lBbn4FOsUzD1d0uwAX*aQ(QFUl-Qg^4nLS({&4f~1#OVFOTCaeiKA lUOH4!3e@fmP`fh{GqX!^TFy8D({ctZRx_YlY=dF#0RUJliN*i` literal 0 HcmV?d00001 diff --git a/crates/claude-connector/tests/fixtures/lease_v1/settled.bin b/crates/claude-connector/tests/fixtures/lease_v1/settled.bin new file mode 100644 index 0000000000000000000000000000000000000000..710d4b817e73b57863b3e0adaf6b682b017db04c GIT binary patch literal 81 zcmZQzVDN2PlAW2CvL!huu{0$$J|{J?I5oaFwWK5`HDzg8Voqsl(}twX93U?~H#4~? bKfbg$ u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64 +} + +fn own_uid() -> u32 { + // SAFETY: geteuid has no preconditions and cannot fail. + unsafe { libc::geteuid() } +} + +/// The 32-digit invocation id of the test's `n`th unit invocation. +fn invocation(n: u64) -> String { + format!("{n:032x}") +} + +type Probe = Box; + +/// A token source that counts its reads and can run a probe at the moment of +/// each one. +struct FakeTokens { + token: Option<&'static str>, + reads: AtomicUsize, + probe: Mutex>, +} + +impl TokenSource for FakeTokens { + fn read(&self) -> Option> { + self.reads.fetch_add(1, Ordering::SeqCst); + if let Some(probe) = self.probe.lock().unwrap().as_ref() { + probe(); + } + self.token.map(|token| Zeroizing::new(token.to_string())) + } +} + +/// The system's record of where every peer runs, as the test sets it, and the +/// pid each handle it was asked about names (by the kernel's own record). +struct FakeUnits { + peer: Mutex>, + asked_about: Mutex>, +} + +impl PeerUnits for FakeUnits { + fn invocation_in(&self, process: BorrowedFd<'_>, unit: &str) -> Option { + let info = + std::fs::read_to_string(format!("/proc/self/fdinfo/{}", process.as_raw_fd())).unwrap(); + let pid = info + .lines() + .find_map(|line| line.strip_prefix("Pid:")) + .unwrap() + .trim() + .to_string(); + self.asked_about + .lock() + .unwrap() + .push((pid, unit.to_string())); + self.peer.lock().unwrap().clone() + } +} + +struct Harness { + dir: PathBuf, + socket: PathBuf, + service: Arc>, + tokens: Arc, + /// The fake behind the listener, when it has one. + units: Option>, +} + +impl Harness { + /// The peers connect from invocation `n` of the Self unit. + fn enter(&self, n: u64) { + self.run_in(UNIT, n); + } + + fn run_in(&self, unit: &str, n: u64) { + *self.units.as_ref().unwrap().peer.lock().unwrap() = Some(PeerUnit { + unit: unit.to_string(), + invocation_id: invocation(n), + }); + } +} + +impl Drop for Harness { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.dir); + } +} + +fn price() -> ModelPrice { + // One micro-USD per token in every class. + ModelPrice { + rates: Rates { + input: 1_000_000, + cache_write_short: 1_000_000, + cache_write_long: 1_000_000, + cache_read: 1_000_000, + output: 1_000_000, + }, + long_context: None, + overhead_tokens: 0, + } +} + +/// What a harness is built from. +struct Setup { + ceiling: u64, + reservation: u64, + uid: u32, + token: Option<&'static str>, + max_connections: usize, +} + +fn setup(ceiling: u64, reservation: u64) -> Setup { + Setup { + ceiling, + reservation, + uid: own_uid(), + token: Some(TOKEN), + max_connections: 8, + } +} + +/// A listener serving `uid` as the caller `eureka`, each of whose leases +/// reserves `reservation`, with a spend window so long that it never rolls +/// over during a test. Peers start in invocation 1 of the Self unit. +fn start(setup: Setup) -> Harness { + let units = Arc::new(FakeUnits { + peer: Mutex::new(None), + asked_about: Mutex::new(Vec::new()), + }); + let harness = start_with(setup, units.clone(), Some(units), INSTANCE); + harness.enter(1); + harness +} + +fn start_with( + setup: Setup, + units: Arc, + fake: Option>, + instance: &str, +) -> Harness { + let dir = std::env::temp_dir().join(format!( + "claude-lease-{}-{}", + std::process::id(), + SEQUENCE.fetch_add(1, Ordering::SeqCst) + )); + std::fs::create_dir_all(&dir).unwrap(); + let socket = dir.join("lease.sock"); + let service = Arc::new(Mutex::new( + TransportService::open( + &dir.join("replay.cc"), + "claude-code", + BTreeMap::from([(MODEL.to_string(), price())]), + [CallerAdmission::new( + CALLER, + "lease-test-connection-key", + 1, + [MODEL.to_string()], + 1, + 65_536, + 32_768, + SpendCeiling(setup.ceiling), + SpendWindowMs(1 << 50), + ) + .unwrap()], + 300_000, + ) + .unwrap(), + )); + let tokens = Arc::new(FakeTokens { + token: setup.token, + reads: AtomicUsize::new(0), + probe: Mutex::new(None), + }); + let listener = UnixListener::bind(&socket).unwrap(); + let serving = (service.clone(), tokens.clone(), units); + let caller = LeaseCaller { + uid: setup.uid, + caller_id: CALLER.to_string(), + unit_instance: instance.to_string(), + reservation_micro_usd: setup.reservation, + }; + std::thread::spawn(move || { + let _ = serve_leases( + &listener, + vec![caller], + LeaseLimits { + socket_timeout: Duration::from_secs(2), + max_lease_ms: 3_600_000, + max_connections: setup.max_connections, + }, + serving.0, + serving.1, + serving.2, + ); + }); + Harness { + dir, + socket, + service, + tokens, + units: fake, + } +} + +/// One frame, and the descriptors that arrived with it. An empty frame is end +/// of file. +fn receive(stream: &UnixStream) -> (Vec, Vec) { + use std::os::fd::AsRawFd; + let mut buffer = vec![0_u8; 70_000]; + let mut control = vec![0_u64; 32]; + let mut vector = libc::iovec { + iov_base: buffer.as_mut_ptr().cast(), + iov_len: buffer.len(), + }; + // SAFETY: the buffers outlive the call and the lengths describe them; the + // descriptors read out of the control messages were just installed in this + // process by the kernel and are owned by nothing else. + unsafe { + let mut message: libc::msghdr = std::mem::zeroed(); + message.msg_iov = &mut vector; + message.msg_iovlen = 1; + message.msg_control = control.as_mut_ptr().cast(); + message.msg_controllen = (control.len() * 8) as _; + let received = libc::recvmsg(stream.as_raw_fd(), &mut message, libc::MSG_CMSG_CLOEXEC); + if received < 0 { + // A listener that drops a connection with our request unread + // resets it: that is end of file for this test's purposes. + let error = std::io::Error::last_os_error(); + assert_eq!(error.kind(), std::io::ErrorKind::ConnectionReset, "{error}"); + return (Vec::new(), Vec::new()); + } + let mut fds = Vec::new(); + let mut header = libc::CMSG_FIRSTHDR(&message); + while !header.is_null() { + if (*header).cmsg_level == libc::SOL_SOCKET && (*header).cmsg_type == libc::SCM_RIGHTS { + let bytes = (*header).cmsg_len as usize - libc::CMSG_LEN(0) as usize; + let data = libc::CMSG_DATA(header).cast::(); + for index in 0..bytes / std::mem::size_of::() { + fds.push(OwnedFd::from_raw_fd(data.add(index).read_unaligned())); + } + } + header = libc::CMSG_NXTHDR(&message, header); + } + buffer.truncate(received as usize); + (buffer, fds) + } +} + +fn connect(harness: &Harness) -> UnixStream { + let stream = UnixStream::connect(&harness.socket).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + stream +} + +fn send(stream: &mut UnixStream, message: &LeaseMessage) { + stream.write_all(&encode_frame(message)).unwrap(); +} + +fn request(ttl_ms: u64) -> LeaseMessage { + LeaseMessage::ClaudeLeaseRequest(LeaseRequest { + expires_at_unix_ms: now_ms() + ttl_ms, + }) +} + +/// What one request to the listener came back with. +struct Answer { + stream: UnixStream, + raw: Vec, + message: Option, + fds: Vec, +} + +fn ask(harness: &Harness, ttl_ms: u64) -> Answer { + ask_at(&harness.socket, ttl_ms) +} + +fn ask_at(socket: &std::path::Path, ttl_ms: u64) -> Answer { + let mut stream = UnixStream::connect(socket).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(10))) + .unwrap(); + // A listener at its limit has already dropped us: the write may fail. + let _ = stream.write_all(&encode_frame(&request(ttl_ms))); + let (raw, fds) = receive(&stream); + let message = decode_frame(&raw); + Answer { + stream, + raw, + message, + fds, + } +} + +fn refusal_of(answer: &Answer) -> Option { + match &answer.message { + Some(LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { kind, .. })) => Some(*kind), + _ => None, + } +} + +fn granted(answer: &Answer) -> LeaseGranted { + match &answer.message { + Some(LeaseMessage::ClaudeLeaseGrant(grant)) => grant.clone(), + other => panic!("expected a grant, got {other:?}"), + } +} + +fn granted_ok(answer: &Answer) -> bool { + matches!(answer.message, Some(LeaseMessage::ClaudeLeaseGrant(_))) && answer.fds.len() == 1 +} + +/// Asks, as a fresh unit invocation each time, until the caller's previous +/// lease has been settled by the listener. +fn ask_after_busy(harness: &Harness, first_invocation: u64) -> Answer { + let deadline = Instant::now() + Duration::from_secs(5); + for attempt in 0.. { + harness.enter(first_invocation + attempt); + let answer = ask(harness, 60_000); + if refusal_of(&answer) != Some(LeaseRefusalKind::Busy) || Instant::now() > deadline { + return answer; + } + std::thread::sleep(Duration::from_millis(20)); + } + unreachable!() +} + +fn used(model: &str, tokens: u64) -> ModelUsage { + ModelUsage { + model: model.to_string(), + usage: Usage { + input_tokens: 0, + cache_read_tokens: 0, + cache_write_tokens: 0, + output_tokens: tokens, + reasoning_tokens: None, + }, + } +} + +fn settle_frame(answer: &Answer, usage: Vec) -> Vec { + encode_frame(&LeaseMessage::ClaudeLeaseSettle(LeaseSettle { + lease_id: granted(answer).lease_id, + usage, + })) +} + +fn settle(answer: &mut Answer, usage: Vec) -> LeaseSettled { + let frame = settle_frame(answer, usage); + answer.stream.write_all(&frame).unwrap(); + let (raw, _) = receive(&answer.stream); + match decode_frame(&raw) { + Some(LeaseMessage::ClaudeLeaseSettled(settled)) => settled, + other => panic!("expected a settled reply, got {other:?}"), + } +} + +fn contains(haystack: &[u8], needle: &str) -> bool { + haystack + .windows(needle.len()) + .any(|window| window == needle.as_bytes()) +} + +/// Asks the ledger directly for a reservation of `cost`, as an invocation +/// nobody else uses. A grant is left open. +fn probe_reserve(harness: &Harness, name: &str, cost: u64) -> LeaseAdmission { + harness + .service + .lock() + .unwrap() + .begin_lease(CALLER, name, cost, FAR, now_ms()) + .unwrap() +} + +fn ceiling_refused() -> LeaseAdmission { + LeaseAdmission::Refused(LeaseRefusal::SpendCeiling) +} + +#[test] +fn a_grant_passes_one_fd_holding_the_token_then_end_of_file() { + let harness = start(setup(1_000, 600)); + let mut answer = ask(&harness, 60_000); + assert_eq!(granted(&answer).reserved_micro_usd, 600); + assert_eq!(answer.fds.len(), 1); + // Reading to end of file proves the write end is closed: nobody holds + // the pipe open for another write. + let mut read = Vec::new(); + File::from(answer.fds.remove(0)) + .read_to_end(&mut read) + .unwrap(); + assert_eq!(read, TOKEN.as_bytes()); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1); +} + +#[test] +fn the_lease_is_named_for_the_unit_invocation_not_for_anything_the_peer_says() { + let harness = start(setup(1_000, 600)); + harness.enter(0xabc); + let answer = ask(&harness, 60_000); + assert_eq!( + granted(&answer).lease_id, + format!("lease:{}", invocation(0xabc)) + ); +} + +#[test] +fn the_token_is_in_the_pipe_and_nowhere_else() { + let harness = start(setup(1_000, 600)); + let mut answer = ask(&harness, 60_000); + let mut frames = vec![answer.raw.clone()]; + let settled = settle(&mut answer, vec![used(MODEL, 10)]); + frames.push(encode_frame(&LeaseMessage::ClaudeLeaseSettled(settled))); + harness.enter(2); + let refused = ask(&harness, 1); + frames.push(refused.raw.clone()); + for frame in &frames { + assert!(!contains(frame, TOKEN)); + } + assert!(std::env::vars().all(|(name, value)| !name.contains(TOKEN) && !value.contains(TOKEN))); + assert!(std::env::args().all(|argument| !argument.contains(TOKEN))); + // The one place it is: the pipe. + let mut read = String::new(); + File::from(answer.fds.remove(0)) + .read_to_string(&mut read) + .unwrap(); + assert_eq!(read, TOKEN); +} + +#[test] +fn a_peer_uid_outside_the_lease_callers_is_refused_and_nothing_is_reserved_or_read() { + let harness = start(Setup { + uid: own_uid() + 1, + ..setup(1_000, 600) + }); + let answer = ask(&harness, 60_000); + assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::UnknownCaller)); + assert!(answer.fds.is_empty()); + assert!(!contains(&answer.raw, TOKEN)); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); + // The ledger was not touched: the caller's whole ceiling is still free. + assert!(matches!( + probe_reserve(&harness, "probe", 1_000), + LeaseAdmission::Granted(_) + )); +} + +#[test] +fn the_configured_uid_outside_the_self_unit_is_refused_and_nothing_is_reserved_or_read() { + let harness = start(setup(1_000, 600)); + for (label, peer) in [ + ("another instance", Some(("mind-self@other.service", 1))), + ("a user service", Some(("huginn.service", 1))), + ("no unit", None), + ] { + *harness.units.as_ref().unwrap().peer.lock().unwrap() = peer.map(|(unit, n)| PeerUnit { + unit: unit.to_string(), + invocation_id: invocation(n), + }); + let answer = ask(&harness, 60_000); + assert_eq!( + refusal_of(&answer), + Some(LeaseRefusalKind::UnknownCaller), + "{label}" + ); + assert!(answer.fds.is_empty(), "{label}"); + } + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); + assert!(matches!( + probe_reserve(&harness, "probe", 1_000), + LeaseAdmission::Granted(_) + )); +} + +#[test] +fn one_invocation_of_the_self_unit_leases_once_and_the_next_invocation_leases_again() { + let harness = start(setup(1_000_000, 600)); + let mut first = ask(&harness, 60_000); + assert!(granted_ok(&first)); + settle(&mut first, vec![used(MODEL, 10)]); + // Settled, the caller holds nothing; the invocation has still leased. + let again = ask(&harness, 60_000); + assert_eq!(refusal_of(&again), Some(LeaseRefusalKind::Busy)); + assert!(again.fds.is_empty()); + harness.enter(2); + assert!(granted_ok(&ask(&harness, 60_000))); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 2); +} + +#[test] +fn the_reservation_is_claimed_before_the_token_is_read_or_lent() { + let harness = start(setup(1_000, 600)); + let seen = Arc::new(Mutex::new(None)); + let (service, sink) = (harness.service.clone(), seen.clone()); + *harness.tokens.probe.lock().unwrap() = Some(Box::new(move || { + // From inside the token read: the caller already holds a lease. + let admission = service + .lock() + .unwrap() + .begin_lease(CALLER, "probe", 1, FAR, now_ms()) + .unwrap(); + *sink.lock().unwrap() = Some(admission); + })); + let answer = ask(&harness, 60_000); + assert_eq!(answer.fds.len(), 1); + assert_eq!( + *seen.lock().unwrap(), + Some(LeaseAdmission::Refused(LeaseRefusal::Busy)) + ); +} + +#[test] +fn a_reservation_the_window_cannot_hold_is_refused_and_lends_nothing() { + let harness = start(setup(1_000, 1_001)); + let answer = ask(&harness, 60_000); + assert_eq!(refusal_of(&answer), Some(LeaseRefusalKind::SpendCeiling)); + assert!(answer.fds.is_empty()); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); + // Exactly the ceiling fits. + let harness = start(setup(1_000, 1_000)); + assert!(granted_ok(&ask(&harness, 60_000))); +} + +#[test] +fn a_free_lease_and_a_bad_expiry_are_refused() { + let harness = start(setup(1_000, 0)); + assert_eq!( + refusal_of(&ask(&harness, 60_000)), + Some(LeaseRefusalKind::CostBound) + ); + let harness = start(setup(1_000, 100)); + for ttl in [0, 3_600_001 + 1_000] { + let mut stream = connect(&harness); + let expires_at_unix_ms = if ttl == 0 { 1 } else { now_ms() + ttl }; + send( + &mut stream, + &LeaseMessage::ClaudeLeaseRequest(LeaseRequest { expires_at_unix_ms }), + ); + let (raw, fds) = receive(&stream); + assert!(fds.is_empty()); + assert!(matches!( + decode_frame(&raw), + Some(LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { + kind: LeaseRefusalKind::Expired, + .. + })) + )); + } + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); +} + +#[test] +fn a_second_open_lease_for_the_caller_is_busy_even_from_another_invocation() { + let harness = start(setup(1_000_000, 100)); + let first = ask(&harness, 60_000); + assert!(granted_ok(&first)); + harness.enter(2); + let second = ask(&harness, 60_000); + assert_eq!(refusal_of(&second), Some(LeaseRefusalKind::Busy)); + assert!(second.fds.is_empty()); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1); +} + +#[test] +fn settlement_bills_the_true_cost_of_each_models_use_and_the_window_holds_it() { + let harness = start(setup(1_000, 600)); + let mut first = ask(&harness, 60_000); + // One micro-USD a token: 250 of the priced model, and 100 of a model the + // table has never heard of, which bills at the dearest known price. + let settled = settle( + &mut first, + vec![used(MODEL, 250), used("claude-new-9", 100)], + ); + assert_eq!( + (settled.billed_micro_usd, settled.overrun_micro_usd), + (350, 0) + ); + // 350 spent: 650 remain, not 651. + assert_eq!(probe_reserve(&harness, "over", 651), ceiling_refused()); + assert!(matches!( + probe_reserve(&harness, "fits", 650), + LeaseAdmission::Granted(_) + )); +} + +#[test] +fn usage_above_the_reservation_is_billed_in_full_and_the_overrun_is_reported_and_counted() { + let harness = start(setup(1_000, 100)); + let mut first = ask(&harness, 60_000); + let settled = settle(&mut first, vec![used(MODEL, 700), used(MODEL, 5_000)]); + assert_eq!(settled.billed_micro_usd, 5_700); + assert_eq!(settled.overrun_micro_usd, 5_600); + // The window is past its ceiling: nothing more reserves in it. + assert_eq!(probe_reserve(&harness, "next", 1), ceiling_refused()); +} + +#[test] +fn a_plan_quota_the_run_reports_is_not_part_of_the_protocol() { + // A settle in the shape the first protocol used, with a plan_quota the + // run vouches for: the extra key is ignored and the usage is billed. + let harness = start(setup(1_000, 600)); + let mut first = ask(&harness, 60_000); + let legacy = serde_json::json!({ + "kind": "claude_lease_settle", + "value": { + "lease_id": granted(&first).lease_id, + "usage": [{"model": MODEL, "usage": { + "input_tokens": 0, "cache_read_tokens": 0, "cache_write_tokens": 0, + "output_tokens": 40, "reasoning_tokens": null}}], + "plan_quota": [{"window": "five_hour", "utilization_permille": 0, + "resets_at_unix_ms": 1, "status": "allowed", "observed_at_unix_ms": 1}] + } + }); + let body = rmp_serde::to_vec_named(&legacy).unwrap(); + let mut frame = (body.len() as u32).to_be_bytes().to_vec(); + frame.extend_from_slice(&body); + first.stream.write_all(&frame).unwrap(); + let (raw, _) = receive(&first.stream); + assert!(matches!( + decode_frame(&raw), + Some(LeaseMessage::ClaudeLeaseSettled(LeaseSettled { + billed_micro_usd: 40, + overrun_micro_usd: 0 + })) + )); +} + +#[test] +fn a_connection_dropped_without_settling_is_billed_the_reservation() { + let harness = start(setup(1_000, 600)); + let first = ask(&harness, 60_000); + assert!(granted_ok(&first)); + drop(first); + let over = ask_after_busy(&harness, 2); + assert_eq!(refusal_of(&over), Some(LeaseRefusalKind::SpendCeiling)); + assert_eq!(probe_reserve(&harness, "over", 401), ceiling_refused()); + assert!(matches!( + probe_reserve(&harness, "fits", 400), + LeaseAdmission::Granted(_) + )); +} + +#[test] +fn a_lease_that_expires_unsettled_is_closed_and_billed_the_reservation() { + let harness = start(setup(1_000, 600)); + let first = ask(&harness, 400); + assert!(granted_ok(&first)); + let started = Instant::now(); + // The listener, not the client, ends the lease at its expiry. + let (raw, fds) = receive(&first.stream); + assert!(raw.is_empty() && fds.is_empty()); + assert!(started.elapsed() >= Duration::from_millis(250)); + harness.enter(2); + let over = ask(&harness, 60_000); + assert_eq!(refusal_of(&over), Some(LeaseRefusalKind::SpendCeiling)); +} + +#[test] +fn a_settle_trickled_past_the_expiry_is_no_report_and_bills_the_reservation() { + let harness = start(setup(1_000, 600)); + let mut answer = ask(&harness, 1_500); + let frame = settle_frame(&answer, vec![used(MODEL, 10)]); + let chunk = frame.len() / 4 + 1; + // Each write lands inside any per-read wait; the frame does not land + // inside the lease. + for (i, part) in frame.chunks(chunk).enumerate() { + if i > 0 { + std::thread::sleep(Duration::from_millis(1_200)); + } + let _ = answer.stream.write_all(part); + } + let (raw, _) = receive(&answer.stream); + assert!(raw.is_empty(), "{:?}", decode_frame(&raw)); + assert_eq!(probe_reserve(&harness, "over", 401), ceiling_refused()); +} + +#[test] +fn a_request_trickled_past_the_socket_timeout_is_dropped_without_a_reservation() { + let harness = start(setup(1_000, 600)); + let mut stream = connect(&harness); + let started = Instant::now(); + // Three bytes of a length, 0.9 s apart, then silence: every wait is inside + // the 2 s socket timeout, the frame is not. + for byte in [0_u8, 0, 0] { + stream.write_all(&[byte]).unwrap(); + std::thread::sleep(Duration::from_millis(900)); + } + let (raw, fds) = receive(&stream); + assert!(raw.is_empty() && fds.is_empty()); + assert!( + started.elapsed() < Duration::from_millis(3_300), + "{:?}", + started.elapsed() + ); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); + assert!(matches!( + probe_reserve(&harness, "probe", 1_000), + LeaseAdmission::Granted(_) + )); +} + +#[test] +fn a_settle_for_another_lease_is_no_report_and_bills_the_reservation() { + let harness = start(setup(1_000, 600)); + let mut first = ask(&harness, 60_000); + send( + &mut first.stream, + &LeaseMessage::ClaudeLeaseSettle(LeaseSettle { + lease_id: "lease:someone-else".to_string(), + usage: vec![used(MODEL, 1)], + }), + ); + let (raw, _) = receive(&first.stream); + assert!(raw.is_empty()); + harness.enter(2); + let over = ask(&harness, 60_000); + assert_eq!(refusal_of(&over), Some(LeaseRefusalKind::SpendCeiling)); +} + +#[test] +fn a_settle_that_reports_nothing_is_billed_the_reservation() { + let harness = start(setup(1_000, 600)); + let mut first = ask(&harness, 60_000); + let settled = settle(&mut first, vec![]); + assert_eq!( + (settled.billed_micro_usd, settled.overrun_micro_usd), + (600, 0) + ); + // 600 spent: the whole ceiling no longer reserves. + assert_eq!(probe_reserve(&harness, "all", 1_000), ceiling_refused()); +} + +#[test] +fn a_token_that_cannot_be_read_is_never_lent_and_bills_nothing() { + let harness = start(Setup { + token: None, + ..setup(1_000, 1_000) + }); + let first = ask(&harness, 60_000); + assert!(first.raw.is_empty() && first.fds.is_empty()); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1); + // Nothing was billed: the whole ceiling reserves again and gets as far + // as the token read, not a refusal. + let again = ask_after_busy(&harness, 2); + assert!(again.message.is_none() && again.fds.is_empty()); + assert!(harness.tokens.reads.load(Ordering::SeqCst) >= 2); +} + +#[test] +fn a_malformed_or_oversized_frame_is_dropped_without_a_reservation() { + let harness = start(setup(1_000, 1_000)); + for bytes in [ + vec![0xff, 0xff, 0xff, 0xff], + vec![0, 0, 0, 0], + vec![0, 0, 0, 3, 0xc1, 0xc1, 0xc1], + ] { + let mut stream = connect(&harness); + stream.write_all(&bytes).unwrap(); + let started = Instant::now(); + let (raw, fds) = receive(&stream); + assert!(raw.is_empty() && fds.is_empty()); + // Refused at once, not after the socket timeout of a read that was + // never going to complete. + assert!(started.elapsed() < Duration::from_millis(1_000)); + } + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); + assert!(granted_ok(&ask(&harness, 60_000))); +} + +#[test] +fn every_message_kind_is_the_pinned_bytes_of_its_fixture() { + let expires = 1_793_307_000_000; + let lease = format!("lease:{}", "0123456789abcdef0123456789abcdef"); + let messages = [ + ( + &include_bytes!("fixtures/lease_v1/request.bin")[..], + LeaseMessage::ClaudeLeaseRequest(LeaseRequest { + expires_at_unix_ms: expires, + }), + ), + ( + &include_bytes!("fixtures/lease_v1/grant.bin")[..], + LeaseMessage::ClaudeLeaseGrant(LeaseGranted { + lease_id: lease.clone(), + reserved_micro_usd: 600_000, + }), + ), + ( + &include_bytes!("fixtures/lease_v1/refusal.bin")[..], + LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { + kind: LeaseRefusalKind::SpendCeiling, + detail: "the window cannot hold this reservation".to_string(), + }), + ), + ( + &include_bytes!("fixtures/lease_v1/settle.bin")[..], + LeaseMessage::ClaudeLeaseSettle(LeaseSettle { + lease_id: lease, + usage: vec![ + ModelUsage { + model: "claude-opus-5-1".to_string(), + usage: Usage { + input_tokens: 1_200, + cache_read_tokens: 30_000, + cache_write_tokens: 500, + output_tokens: 800, + reasoning_tokens: Some(100), + }, + }, + ModelUsage { + model: "claude-haiku-5-1".to_string(), + usage: Usage { + input_tokens: 400, + cache_read_tokens: 0, + cache_write_tokens: 0, + output_tokens: 60, + reasoning_tokens: None, + }, + }, + ], + }), + ), + ( + &include_bytes!("fixtures/lease_v1/settled.bin")[..], + LeaseMessage::ClaudeLeaseSettled(LeaseSettled { + billed_micro_usd: 123_456, + overrun_micro_usd: 23_456, + }), + ), + ]; + for (bytes, message) in messages { + assert_eq!(encode_frame(&message), bytes, "{message:?}"); + assert_eq!(decode_frame(bytes), Some(message)); + } +} + +#[test] +fn connections_past_the_limit_are_dropped_and_finished_ones_free_their_place() { + let harness = start(Setup { + max_connections: 2, + ..setup(1_000, 0) + }); + let held: Vec = (0..2).map(|_| connect(&harness)).collect(); + // Let the listener take both before the third arrives. + std::thread::sleep(Duration::from_millis(300)); + let third = ask(&harness, 60_000); + assert!(third.raw.is_empty() && third.fds.is_empty()); + drop(held); + // Finished connections give their places back: more refusals in a row + // than the limit holds, each answered. + let first = ask_until_answered(&harness); + assert_eq!(refusal_of(&first), Some(LeaseRefusalKind::CostBound)); + for _ in 0..4 { + let again = ask(&harness, 60_000); + assert_eq!(refusal_of(&again), Some(LeaseRefusalKind::CostBound)); + } +} + +/// A free lease, asked until the listener answers instead of dropping. +fn ask_until_answered(harness: &Harness) -> Answer { + let deadline = Instant::now() + Duration::from_secs(5); + loop { + let answer = ask(harness, 60_000); + if !answer.raw.is_empty() || Instant::now() > deadline { + return answer; + } + std::thread::sleep(Duration::from_millis(20)); + } +} + +#[test] +fn an_account_resolves_to_its_uid_and_a_missing_one_to_none() { + use claude_connector::user_uid; + assert_eq!(user_uid("root"), Some(0)); + assert_eq!(user_uid("nobody"), Some(65534)); + assert_eq!(user_uid("no-such-account-4f1c"), None); + assert_eq!(user_uid("bad name"), None); +} + +#[test] +fn a_token_up_to_one_atomic_pipe_write_is_lent_and_a_longer_one_is_not() { + let fits: &'static str = Box::leak("t".repeat(4096).into_boxed_str()); + let harness = start(Setup { + token: Some(fits), + ..setup(1_000, 100) + }); + let mut answer = ask(&harness, 60_000); + let mut read = Vec::new(); + File::from(answer.fds.remove(0)) + .read_to_end(&mut read) + .unwrap(); + assert_eq!(read.len(), 4096); + + let too_long: &'static str = Box::leak("t".repeat(4097).into_boxed_str()); + let harness = start(Setup { + token: Some(too_long), + ..setup(1_000, 100) + }); + let answer = ask(&harness, 60_000); + assert!(answer.raw.is_empty() && answer.fds.is_empty()); +} + +// ---- a second real uid (Soul s1, kept as a regression) ---- + +fn describe(answer: &Answer) -> String { + let kind = match &answer.message { + Some(LeaseMessage::ClaudeLeaseGrant(_)) => "grant".to_string(), + Some(LeaseMessage::ClaudeLeaseRefusal(r)) => format!("refusal:{:?}", r.kind), + Some(other) => format!("other:{other:?}"), + None => format!("none(raw={})", answer.raw.len()), + }; + format!("kind={kind} fds={}", answer.fds.len()) +} + +/// The child half of the second-uid test: connects to `LEASE_SOCK` as whoever +/// it was started as and prints one line saying what came back. The fake +/// canary token is never printed, only whether it matched. +#[test] +#[ignore = "run as a child of a_second_real_uid_cannot_lease"] +fn lease_probe_client() { + let socket = PathBuf::from(std::env::var("LEASE_SOCK").unwrap()); + let mut answer = ask_at(&socket, 60_000); + let shown = describe(&answer); + let mut token_matches = false; + if !answer.fds.is_empty() { + let mut read = Vec::new(); + File::from(answer.fds.remove(0)) + .read_to_end(&mut read) + .unwrap(); + token_matches = read == TOKEN.as_bytes(); + } + println!( + "LEASEPROBE uid={} {} token_matches_canary={token_matches}", + own_uid(), + shown + ); +} + +/// Runs the probe client as `reuid`, with primary group `regid` and +/// supplementary `groups`, and returns its LEASEPROBE line. The marker is +/// found anywhere on a line: the test harness prints it after its own prefix. +fn probe_as(socket: &std::path::Path, reuid: u32, regid: u32, groups: &str) -> String { + let exe = std::env::current_exe().unwrap(); + let out = std::process::Command::new("setpriv") + .arg(format!("--reuid={reuid}")) + .arg(format!("--regid={regid}")) + .arg(format!("--groups={groups}")) + .arg(&exe) + .args([ + "--exact", + "lease_probe_client", + "--ignored", + "--nocapture", + "--test-threads=1", + ]) + .env("LEASE_SOCK", socket) + .output() + .unwrap(); + let stdout = String::from_utf8_lossy(&out.stdout).to_string(); + stdout + .lines() + .find_map(|line| line.find("LEASEPROBE").map(|at| line[at..].to_string())) + .unwrap_or_else(|| { + format!( + "NO-LINE status={:?} stdout={stdout:?} stderr={}", + out.status, + String::from_utf8_lossy(&out.stderr) + ) + }) +} + +fn have(tool: &str) -> bool { + std::process::Command::new(tool) + .arg("--version") + .output() + .is_ok() +} + +#[test] +fn a_second_real_uid_cannot_lease() { + use std::os::unix::fs::PermissionsExt; + if own_uid() != 0 || !have("setpriv") { + eprintln!("SKIPPED a_second_real_uid_cannot_lease: needs root and setpriv"); + return; + } + // Numeric uids need no account: 4242 is the configured Self user, 4243 a + // second account that shares its groups. + let harness = start(Setup { + uid: 4242, + ..setup(1_000_000, 600) + }); + std::fs::set_permissions(&harness.dir, std::fs::Permissions::from_mode(0o755)).unwrap(); + std::fs::set_permissions(&harness.socket, std::fs::Permissions::from_mode(0o666)).unwrap(); + // Another uid, in the caller's group and the lease group. + let a = probe_as(&harness.socket, 4243, 4243, "4242,4244"); + // Another uid whose primary gid is the caller's. + let b = probe_as(&harness.socket, 4243, 4242, "4242"); + // Root, which is not configured. + let root = describe(&ask(&harness, 60_000)); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 0); + // The configured uid does lease. + let c = probe_as(&harness.socket, 4242, 4242, "4242"); + assert!(a.contains("kind=refusal:UnknownCaller fds=0"), "{a}"); + assert!(b.contains("kind=refusal:UnknownCaller fds=0"), "{b}"); + assert!(root.contains("kind=refusal:UnknownCaller fds=0"), "{root}"); + assert!( + c.contains("kind=grant fds=1") && c.contains("token_matches_canary=true"), + "{c}" + ); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 1); +} + +// ---- a real systemd, where there is one ---- + +/// The child half of the systemd test, run inside a transient service: leases, +/// settles, and leases again from the same invocation. +#[test] +#[ignore = "run inside a transient unit by a_real_systemd_unit_invocation_leases_once"] +fn lease_in_unit_client() { + let socket = PathBuf::from(std::env::var("LEASE_SOCK").unwrap()); + let mut first = ask_at(&socket, 60_000); + let first_kind = describe(&first); + if granted_ok(&first) { + settle(&mut first, Vec::new()); + } + let second = ask_at(&socket, 60_000); + println!( + "LEASEPROBE first[{first_kind}] second[{}]", + describe(&second) + ); +} + +fn run_in_unit(socket: &std::path::Path) -> String { + let exe = std::env::current_exe().unwrap(); + let out = std::process::Command::new("systemd-run") + .args([ + "--unit=mind-self@leasetest.service", + "--pipe", + "--wait", + "--collect", + "--quiet", + ]) + .arg(format!("--setenv=LEASE_SOCK={}", socket.display())) + .arg(&exe) + .args([ + "--exact", + "lease_in_unit_client", + "--ignored", + "--nocapture", + "--test-threads=1", + ]) + .output() + .unwrap(); + let stdout = String::from_utf8_lossy(&out.stdout).to_string(); + stdout + .lines() + .find_map(|line| line.find("LEASEPROBE").map(|at| line[at..].to_string())) + .unwrap_or_else(|| { + format!( + "NO-LINE status={:?} stdout={stdout:?} stderr={}", + out.status, + String::from_utf8_lossy(&out.stderr) + ) + }) +} + +#[test] +fn a_real_systemd_unit_invocation_leases_once() { + if own_uid() != 0 + || !std::path::Path::new("/run/systemd/system").exists() + || !have("systemd-run") + { + eprintln!( + "SKIPPED a_real_systemd_unit_invocation_leases_once: needs root, systemd as init and systemd-run" + ); + return; + } + // The listener asks the production resolver, not the fake. + let harness = start_with( + setup(1_000_000, 600), + Arc::new(SystemdUnits::system()), + None, + "leasetest", + ); + let leased_then_busy = "first[kind=grant fds=1] second[kind=refusal:Busy fds=0]"; + // Two invocations of the same unit: each leases once, as its own. + let one = run_in_unit(&harness.socket); + let two = run_in_unit(&harness.socket); + assert!(one.contains(leased_then_busy), "{one}"); + assert!(two.contains(leased_then_busy), "{two}"); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 2); + // A process outside any such unit is not a Self run, though its uid is + // the configured one. + let outside = ask(&harness, 60_000); + assert_eq!(refusal_of(&outside), Some(LeaseRefusalKind::UnknownCaller)); + assert_eq!(harness.tokens.reads.load(Ordering::SeqCst), 2); +} + +#[test] +fn the_peer_is_named_by_a_process_handle_not_a_pid_number() { + let (one, _other) = UnixStream::pair().unwrap(); + let peer = claude_connector::peer_credentials(&one).unwrap(); + assert_eq!(peer.uid, own_uid()); + let process = peer.process.expect("SO_PEERPIDFD on this kernel"); + let info = + std::fs::read_to_string(format!("/proc/self/fdinfo/{}", process.as_raw_fd())).unwrap(); + assert!( + info.lines() + .any(|line| line.strip_prefix("Pid:").map(str::trim) + == Some(&std::process::id().to_string())), + "{info}" + ); +} + +#[test] +fn the_unit_is_asked_of_the_callers_unit_through_the_peers_handle() { + let harness = start(setup(1_000, 600)); + let answer = ask(&harness, 60_000); + assert!(refusal_of(&answer).is_none() && answer.fds.len() == 1); + assert_eq!( + *harness.units.as_ref().unwrap().asked_about.lock().unwrap(), + vec![(std::process::id().to_string(), UNIT.to_string())] + ); +} diff --git a/crates/model-connector/src/serve.rs b/crates/model-connector/src/serve.rs index 26f5e2d..cf76afc 100644 --- a/crates/model-connector/src/serve.rs +++ b/crates/model-connector/src/serve.rs @@ -171,15 +171,7 @@ pub fn serve( backend: Arc, ) -> Result<(), ServeError> { limits.validate(&admissions)?; - if !listener - .local_addr() - .map_err(ServeError::Listen)? - .ip() - .is_loopback() - { - return Err(ServeError::InvalidConfig("listener address")); - } - backend.ready().map_err(ServeError::BackendUnready)?; + admit_listener(&listener, backend.as_ref())?; let service = Arc::new(Mutex::new(TransportService::open( replay_store, backend.provider(), @@ -190,6 +182,34 @@ pub fn serve( accept_loop(|| listener.accept(), limits, service, backend) } +/// `serve` over a service the caller opened and keeps a handle to, so another +/// front of the same ledger (the Self-run credential lease) reserves in the +/// same windows. The limits must already have been validated against the +/// service's admissions. +pub fn serve_service( + listener: TcpListener, + limits: ServeLimits, + service: Arc>, + backend: Arc, +) -> Result<(), ServeError> { + admit_listener(&listener, backend.as_ref())?; + accept_loop(|| listener.accept(), limits, service, backend) +} + +/// What must be true before a serve loop touches the ledger: the listener is +/// on loopback and the backend can take work. +fn admit_listener(listener: &TcpListener, backend: &dyn ProviderBackend) -> Result<(), ServeError> { + if !listener + .local_addr() + .map_err(ServeError::Listen)? + .ip() + .is_loopback() + { + return Err(ServeError::InvalidConfig("listener address")); + } + backend.ready().map_err(ServeError::BackendUnready) +} + /// The accept loop over any source of connections, so a test can present a /// peer address that a loopback listener never reports. fn accept_loop( diff --git a/crates/model-connector/src/service.rs b/crates/model-connector/src/service.rs index c269916..16853dc 100644 --- a/crates/model-connector/src/service.rs +++ b/crates/model-connector/src/service.rs @@ -317,6 +317,9 @@ fn prune_to_latest( struct LiveRequest { window: u64, reserved_micro_usd: u64, + /// A Self run's credential lease, not a provider call: it holds its + /// reservation in the window like a call does but takes no request slot. + lease: bool, } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -338,6 +341,12 @@ enum ReplayState { charged_micro_usd: u64, overrun_micro_usd: u64, }, + /// A credential lease that settled: it has no answer to replay, only the + /// charge it left in its window and how far that passed the reservation. + LeaseSettled { + charged_micro_usd: u64, + overrun_micro_usd: u64, + }, } #[derive(Clone, Debug, PartialEq, Eq, DatabaseEntry)] @@ -586,6 +595,9 @@ impl TransportService { ReplayState::Active { reserved_micro_usd } => *reserved_micro_usd, ReplayState::Completed { charged_micro_usd, .. + } + | ReplayState::LeaseSettled { + charged_micro_usd, .. } => *charged_micro_usd, }; add_spend( @@ -787,6 +799,9 @@ impl TransportService { .map_err(replay_store_error)?; Ok(TransportAdmission::Reply(stand_in)) } + ReplayState::LeaseSettled { .. } => { + self.reply_refusal(&invocation, Refusal::ReplayConflict, &security) + } ReplayState::Active { .. } => { let refusal = if self.live.contains_key(&identity) { Refusal::InFlight @@ -810,8 +825,10 @@ impl TransportService { }; let active_for_caller = self .live - .keys() - .filter(|key| key.caller_runtime_id == invocation.caller_runtime_id) + .iter() + .filter(|(key, request)| { + !request.lease && key.caller_runtime_id == invocation.caller_runtime_id + }) .count(); if active_for_caller >= max_concurrent_requests { return self.reply_refusal(&invocation, Refusal::Capacity, &security); @@ -873,6 +890,7 @@ impl TransportService { LiveRequest { window, reserved_micro_usd, + lease: false, }, ); Ok(TransportAdmission::Execute(Box::new(ExecutionClaim { @@ -1234,6 +1252,9 @@ fn replay_store_error(error: impl std::fmt::Display) -> ServiceError { ServiceError::ReplayStore(error.to_string()) } +mod lease; +pub use lease::{LeaseAdmission, LeaseGrant, LeaseRefusal, LeaseSettlement, LeaseUse, ModelUsage}; + fn contract_refusal(error: &ContractError) -> Refusal { match error { ContractError::Expiry => Refusal::Expired, @@ -2525,7 +2546,9 @@ mod tests { ReplayState::Completed { charged_micro_usd, .. } => charged_micro_usd, - ReplayState::Active { .. } => panic!("{request_id} is not completed"), + ReplayState::Active { .. } | ReplayState::LeaseSettled { .. } => { + panic!("{request_id} is not completed") + } } } diff --git a/crates/model-connector/src/service/lease.rs b/crates/model-connector/src/service/lease.rs new file mode 100644 index 0000000..6e9be71 --- /dev/null +++ b/crates/model-connector/src/service/lease.rs @@ -0,0 +1,691 @@ +//! A credential lease: a Self run's reservation in the caller's spend window. +//! The lease is one record of the same ledger a provider call reserves in, not +//! a second ledger: it is claimed before anything is lent, counts against the +//! window's ceiling while open, and settles to the true cost of what the run +//! reports it used, per model and uncapped: a charge above the reservation is +//! kept as the lease's overrun, and still counts against the window. A lease +//! nothing settled is billed the whole reservation, as a call a crash left +//! behind is. + +use super::*; + +/// The request-id namespace of leases in a caller's replay records. +const LEASE_PREFIX: &str = "lease:"; + +/// What an admitted lease holds. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LeaseGrant { + pub lease_id: String, + pub reserved_micro_usd: u64, +} + +/// What a run used of one model. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ModelUsage { + pub model: String, + pub usage: Usage, +} + +/// What a settled lease was billed. `overrun_micro_usd` is how far the charge +/// passed the reservation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct LeaseSettlement { + pub billed_micro_usd: u64, + pub overrun_micro_usd: u64, +} + +/// What a lease's settlement is billed on. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LeaseUse<'a> { + /// No report: the whole reservation. A run that cannot say what it used + /// pays what was set aside for it. + Reservation, + /// The run's report. A report that names no use (an empty list) is no + /// report, and bills the reservation. + Reported(&'a [ModelUsage]), + /// The token was never lent, so nothing was used and nothing is billed. + NeverLent, +} + +/// What `usage` costs at the table's prices: each entry at its model's price, +/// and a model with no price at the dearest price the table holds, so an +/// unknown name never bills less than any known one. `None` when the table +/// prices nothing. +fn lease_cost(prices: &PriceTable, usage: &[ModelUsage]) -> Option { + let dearest = |usage: &Usage| { + prices + .values() + .map(|price| cost_micro_usd(price, usage, None)) + .max() + }; + usage + .iter() + .map(|entry| match prices.get(&entry.model) { + Some(price) => Some(cost_micro_usd(price, &entry.usage, None)), + None => dearest(&entry.usage), + }) + .try_fold(0_u64, |total, cost| Some(total.saturating_add(cost?))) +} + +/// Why no lease was reserved. Nothing was written to the ledger. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LeaseRefusal { + /// The window's spend plus open reservations plus this one passes the + /// caller's ceiling. + SpendCeiling, + /// The reservation asked for is zero: a lease that bills nothing is a + /// free pass. + CostBound, + UnknownCaller, + /// The caller holds a lease already, or this run id was leased before. + Busy, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum LeaseAdmission { + Granted(LeaseGrant), + Refused(LeaseRefusal), +} + +impl TransportService { + /// Reserves `max_cost_micro_usd` in `caller_runtime_id`'s current window + /// for the unit invocation `invocation_id`, which leases once. A ledger + /// that cannot be written answers `Err` and reserves nothing. A lease takes no request slot. + pub fn begin_lease( + &mut self, + caller_runtime_id: &str, + invocation_id: &str, + max_cost_micro_usd: u64, + expires_at_unix_ms: u64, + now_unix_ms: u64, + ) -> Result { + let Some(caller) = self.callers.get(caller_runtime_id) else { + return Ok(LeaseAdmission::Refused(LeaseRefusal::UnknownCaller)); + }; + let spend_ceiling = caller.spend_ceiling_micro_usd; + let spend_window_ms = caller.spend_window_ms; + let connection_key_epoch = caller.connection_key_epoch; + require_id(invocation_id, "invocation_id")?; + if invocation_id.len() + LEASE_PREFIX.len() > MAX_IDENTITY_BYTES { + return Err(ContractError::Invalid("invocation_id").into()); + } + if max_cost_micro_usd == 0 { + return Ok(LeaseAdmission::Refused(LeaseRefusal::CostBound)); + } + let identity = RequestIdentity { + caller_runtime_id: caller_runtime_id.to_string(), + request_id: format!("{LEASE_PREFIX}{invocation_id}"), + }; + let key = replay_key(&identity); + let holding = self + .live + .iter() + .any(|(held, request)| request.lease && held.caller_runtime_id == caller_runtime_id); + if holding + || self + .replay_store + .get::(&key) + .map_err(replay_store_error)? + .is_some() + { + return Ok(LeaseAdmission::Refused(LeaseRefusal::Busy)); + } + let admitted_at_unix_ms = admission_clock( + &self.window_spend, + caller_runtime_id, + now_unix_ms, + spend_window_ms, + ); + let window = admitted_at_unix_ms / spend_window_ms; + prune_windows( + &mut self.window_spend, + &self.live, + caller_runtime_id, + window, + ); + if self + .window_spend(caller_runtime_id, window) + .saturating_add(max_cost_micro_usd) + > spend_ceiling + { + return Ok(LeaseAdmission::Refused(LeaseRefusal::SpendCeiling)); + } + let mut digest = Sha256::new(); + digest.update(b"claude_lease_request"); + digest.update((invocation_id.len() as u64).to_be_bytes()); + digest.update(invocation_id.as_bytes()); + digest.update(max_cost_micro_usd.to_be_bytes()); + digest.update(expires_at_unix_ms.to_be_bytes()); + let record = ReplayRecord { + caller_runtime_id: caller_runtime_id.to_string(), + request_id: identity.request_id.clone(), + connection_key_epoch, + invocation_sha256: digest.finalize().into(), + expires_at_unix_ms, + state: ReplayState::Active { + reserved_micro_usd: max_cost_micro_usd, + }, + admitted_at_unix_ms, + }; + // Reserving and claiming are one write, as for a call: a ledger that + // cannot be written lends nothing. + self.replay_store + .put(key, &record) + .map_err(replay_store_error)?; + add_spend( + &mut self.window_spend, + caller_runtime_id, + window, + 0, + max_cost_micro_usd, + ); + let lease_id = identity.request_id.clone(); + self.live.insert( + identity, + LiveRequest { + window, + reserved_micro_usd: max_cost_micro_usd, + lease: true, + }, + ); + Ok(LeaseAdmission::Granted(LeaseGrant { + lease_id, + reserved_micro_usd: max_cost_micro_usd, + })) + } + + /// Settles an open lease and returns what it was billed. A report bills + /// the true cost of every model's reported use, uncapped (the overrun is + /// recorded); no report, or an empty one, bills the whole reservation; + /// only a token never lent bills nothing. A failed ledger + /// write frees the caller's lease and keeps the reservation counted, as a + /// call's does. + pub fn settle_lease( + &mut self, + caller_runtime_id: &str, + lease_id: &str, + usage: LeaseUse<'_>, + ) -> Result { + let identity = RequestIdentity { + caller_runtime_id: caller_runtime_id.to_string(), + request_id: lease_id.to_string(), + }; + if !self + .live + .get(&identity) + .is_some_and(|request| request.lease) + { + return Err(ServiceError::NoActiveRequest); + } + let key = replay_key(&identity); + let active = self + .replay_store + .get::(&key) + .map_err(replay_store_error)? + .ok_or(ServiceError::NoActiveRequest)?; + let ReplayState::Active { reserved_micro_usd } = active.state else { + return Err(ServiceError::ActiveRequestMismatch); + }; + let caller = self + .callers + .get(caller_runtime_id) + .ok_or(ServiceError::CallerNotAdmitted)?; + let spend_window_ms = caller.spend_window_ms; + let charged_micro_usd = match usage { + LeaseUse::Reservation | LeaseUse::Reported([]) => reserved_micro_usd, + LeaseUse::NeverLent => 0, + LeaseUse::Reported(usage) => { + lease_cost(&self.prices, usage).unwrap_or(reserved_micro_usd) + } + }; + let overrun_micro_usd = charged_micro_usd.saturating_sub(reserved_micro_usd); + let settled = ReplayRecord { + state: ReplayState::LeaseSettled { + charged_micro_usd, + overrun_micro_usd, + }, + ..active.clone() + }; + self.replay_store + .put(key, &settled) + .map_err(replay_store_error) + .inspect_err(|_| { + self.live.remove(&identity); + })?; + add_spend( + &mut self.window_spend, + caller_runtime_id, + active.admitted_at_unix_ms / spend_window_ms, + reserved_micro_usd, + charged_micro_usd, + ); + self.live.remove(&identity); + prune_to_latest(&mut self.window_spend, &self.live, caller_runtime_id); + Ok(LeaseSettlement { + billed_micro_usd: charged_micro_usd, + overrun_micro_usd, + }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::Ordering; + + use super::*; + use crate::fixtures::*; + + const CALLER: &str = "eureka"; + const KEY: &str = "lease-distinct-test-key"; + const NOW: u64 = 7_200_500; + + fn open_with(store: impl CacheBackingStore + 'static, ceiling: u64) -> TransportService { + TransportService::with_store( + store, + "fixture-sse", + fixture_prices(), + [CallerAdmission::new( + CALLER, + KEY, + 1, + ["gpt-5.4".to_string()], + 1, + 65_536, + 32_768, + SpendCeiling(ceiling), + SpendWindowMs(WINDOW_MS), + ) + .unwrap()], + 5_000, + ) + .unwrap() + } + + fn open(store: &TestStore, ceiling: u64) -> TransportService { + TransportService::open( + &store.path, + "fixture-sse", + fixture_prices(), + [CallerAdmission::new( + CALLER, + KEY, + 1, + ["gpt-5.4".to_string()], + 1, + 65_536, + 32_768, + SpendCeiling(ceiling), + SpendWindowMs(WINDOW_MS), + ) + .unwrap()], + 5_000, + ) + .unwrap() + } + + fn lease(service: &mut TransportService, run: &str, cost: u64) -> LeaseAdmission { + service + .begin_lease(CALLER, run, cost, NOW + 60_000, NOW) + .unwrap() + } + + fn granted(admission: LeaseAdmission) -> LeaseGrant { + match admission { + LeaseAdmission::Granted(grant) => grant, + LeaseAdmission::Refused(why) => panic!("refused {why:?}"), + } + } + + fn held(service: &TransportService) -> u64 { + service.window_spend(CALLER, NOW / WINDOW_MS) + } + + fn output(tokens: u64) -> Vec { + used("gpt-5.4", tokens) + } + + fn used(model: &str, tokens: u64) -> Vec { + vec![ModelUsage { + model: model.to_string(), + usage: Usage { + input_tokens: 0, + cache_read_tokens: 0, + cache_write_tokens: 0, + output_tokens: tokens, + reasoning_tokens: None, + }, + }] + } + + fn billed(settlement: LeaseSettlement) -> u64 { + settlement.billed_micro_usd + } + + #[test] + fn a_lease_reserves_in_the_callers_window_and_names_what_it_holds() { + let store = TestStore::new(); + let mut service = open(&store, 1_000); + let grant = granted(lease(&mut service, "run-1", 600)); + assert_eq!(grant.reserved_micro_usd, 600); + assert_eq!(held(&service), 600); + assert_eq!( + billed( + service + .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation) + .unwrap() + ), + 600 + ); + } + + #[test] + fn the_window_ceiling_counts_open_and_settled_reservations_inclusively() { + let store = TestStore::new(); + let mut service = open(&store, 1_000); + let first = granted(lease(&mut service, "run-1", 600)); + service + .settle_lease(CALLER, &first.lease_id, LeaseUse::Reservation) + .unwrap(); + assert_eq!( + lease(&mut service, "run-2", 401), + LeaseAdmission::Refused(LeaseRefusal::SpendCeiling) + ); + assert_eq!(held(&service), 600); + granted(lease(&mut service, "run-3", 400)); + assert_eq!(held(&service), 1_000); + } + + #[test] + fn a_caller_holds_one_lease_at_a_time_and_an_invocation_id_is_leased_once() { + let store = TestStore::new(); + let mut service = open(&store, 1_000_000); + let first = granted(lease(&mut service, "run-1", 100)); + assert_eq!( + lease(&mut service, "run-2", 100), + LeaseAdmission::Refused(LeaseRefusal::Busy) + ); + service + .settle_lease(CALLER, &first.lease_id, LeaseUse::Reservation) + .unwrap(); + assert_eq!( + lease(&mut service, "run-1", 100), + LeaseAdmission::Refused(LeaseRefusal::Busy) + ); + granted(lease(&mut service, "run-2", 100)); + } + + #[test] + fn a_stranger_and_a_free_lease_are_refused_without_a_ledger_write() { + let store = TestStore::new(); + let mut service = open(&store, 1_000); + assert_eq!( + service + .begin_lease("stranger", "run-1", 100, NOW + 1, NOW) + .unwrap(), + LeaseAdmission::Refused(LeaseRefusal::UnknownCaller) + ); + assert_eq!( + lease(&mut service, "run-1", 0), + LeaseAdmission::Refused(LeaseRefusal::CostBound) + ); + assert_eq!(held(&service), 0); + // Nothing was claimed: the same run id leases once it asks sensibly. + granted(lease(&mut service, "run-1", 100)); + } + + #[test] + fn settlement_bills_the_true_cost_uncapped_and_records_the_overrun() { + let store = TestStore::new(); + let mut service = open(&store, 1_000_000); + // Output at 15 USD per million tokens: 10 tokens bill 150 micro-USD. + let small = granted(lease(&mut service, "run-1", 10_000)); + let settled = service + .settle_lease(CALLER, &small.lease_id, LeaseUse::Reported(&output(10)[..])) + .unwrap(); + assert_eq!( + (settled.billed_micro_usd, settled.overrun_micro_usd), + (150, 0) + ); + assert_eq!(held(&service), 150); + // 1000 tokens bill 15_000: the reservation does not cap it, and the + // 5_000 past it is recorded as the overrun. + let big = granted(lease(&mut service, "run-2", 10_000)); + let settled = service + .settle_lease( + CALLER, + &big.lease_id, + LeaseUse::Reported(&output(1_000)[..]), + ) + .unwrap(); + assert_eq!( + (settled.billed_micro_usd, settled.overrun_micro_usd), + (15_000, 5_000) + ); + assert_eq!(held(&service), 15_150); + } + + #[test] + fn settlement_prices_each_model_at_its_own_rate_and_an_unknown_model_at_the_dearest() { + let store = TestStore::new(); + let cheap = ModelPrice { + rates: Rates { + input: 1_000_000, + cache_write_short: 1_000_000, + cache_write_long: 1_000_000, + cache_read: 1_000_000, + output: 2_000_000, + }, + ..fixture_price() + }; + let prices = PriceTable::from([ + ("gpt-5.4".to_string(), fixture_price()), + ("cheap".to_string(), cheap), + ]); + let mut service = TransportService::open( + &store.path, + "fixture-sse", + prices, + [CallerAdmission::new( + CALLER, + KEY, + 1, + ["gpt-5.4".to_string()], + 1, + 65_536, + 32_768, + SpendCeiling(1_000_000), + SpendWindowMs(WINDOW_MS), + ) + .unwrap()], + 5_000, + ) + .unwrap(); + let mixed: Vec = [("gpt-5.4", 10), ("cheap", 10), ("unheard-of", 10)] + .into_iter() + .flat_map(|(model, tokens)| used(model, tokens)) + .collect(); + let grant = granted(lease(&mut service, "run-1", 10_000)); + // 10 tokens at 15, 10 at 2, and 10 at the dearest (15) per million. + assert_eq!( + billed( + service + .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&mixed[..])) + .unwrap() + ), + 150 + 20 + 150 + ); + // A run that reports no use pays the reservation, not nothing. + let grant = granted(lease(&mut service, "run-2", 10_000)); + assert_eq!( + billed( + service + .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&[])) + .unwrap() + ), + 10_000 + ); + } + + #[test] + fn settling_an_empty_report_cannot_free_the_reservation() { + let store = TestStore::new(); + let mut service = open(&store, 1_000); + let grant = granted(lease(&mut service, "run-1", 600)); + assert_eq!( + billed( + service + .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&[])) + .unwrap() + ), + 600 + ); + assert_eq!(held(&service), 600); + assert!(matches!( + lease(&mut service, "run-2", 1_000), + LeaseAdmission::Refused(LeaseRefusal::SpendCeiling) + )); + } + + #[test] + fn a_token_never_lent_bills_nothing_and_frees_the_window() { + let store = TestStore::new(); + let mut service = open(&store, 1_000); + let grant = granted(lease(&mut service, "run-1", 600)); + let settled = service + .settle_lease(CALLER, &grant.lease_id, LeaseUse::NeverLent) + .unwrap(); + assert_eq!( + (settled.billed_micro_usd, settled.overrun_micro_usd), + (0, 0) + ); + assert_eq!(held(&service), 0); + granted(lease(&mut service, "run-2", 1_000)); + } + + #[test] + fn a_lease_settles_once() { + let store = TestStore::new(); + let mut service = open(&store, 1_000_000); + let grant = granted(lease(&mut service, "run-1", 100)); + service + .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation) + .unwrap(); + assert!(matches!( + service.settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation), + Err(ServiceError::NoActiveRequest) + )); + assert_eq!(held(&service), 100); + } + + #[test] + fn a_lease_left_open_by_a_crash_counts_at_its_reservation_after_restart() { + let store = TestStore::new(); + let mut service = open(&store, 1_000_000); + granted(lease(&mut service, "run-1", 700)); + drop(service); + let service = open(&store, 1_000_000); + assert_eq!(held(&service), 700); + } + + #[test] + fn a_ledger_that_cannot_be_written_lends_nothing() { + let flaky = FlakyStore::new(1_000); + let allowance = flaky.allowance(); + let mut service = open_with(flaky, 1_000_000); + allowance.store(0, Ordering::SeqCst); + assert!( + service + .begin_lease(CALLER, "run-1", 100, NOW + 1, NOW) + .is_err() + ); + assert_eq!(held(&service), 0); + allowance.store(1_000, Ordering::SeqCst); + granted(lease(&mut service, "run-1", 100)); + } + + #[test] + fn a_failed_settle_write_frees_the_lease_and_keeps_the_reservation_counted() { + let flaky = FlakyStore::new(1_000); + let allowance = flaky.allowance(); + let mut service = open_with(flaky, 1_000_000); + let grant = granted(lease(&mut service, "run-1", 100)); + allowance.store(0, Ordering::SeqCst); + assert!( + service + .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&output(1)[..])) + .is_err() + ); + assert_eq!(held(&service), 100); + allowance.store(1_000, Ordering::SeqCst); + granted(lease(&mut service, "run-2", 100)); + } + + fn begin_call(service: &mut TransportService, id: &str) -> TransportAdmission { + let mut call = request(); + call.request_id = id.to_string(); + let invocation = + TransportInvocation::new(CALLER, NOW + 1_000, [7; 32], MAX_COST, call).unwrap(); + let sealed = crate::envelope::encrypt_message( + CALLER, + id, + crate::envelope::EnvelopeKind::Invocation, + &crate::contract::wire_encode(&invocation).unwrap(), + &security(KEY), + ) + .unwrap(); + service.begin(&sealed, NOW, 1 << 20).unwrap() + } + + #[test] + fn a_lease_takes_no_request_slot_from_the_callers_calls() { + // The caller admits one concurrent call; an open lease must not use it. + let store = TestStore::new(); + let mut service = open(&store, 1 << 40); + granted(lease(&mut service, "run-1", 100)); + match begin_call(&mut service, "call-1") { + TransportAdmission::Execute(_) => {} + TransportAdmission::Reply(_) => panic!("the lease held the call's slot"), + } + } + + #[test] + fn a_call_in_flight_does_not_make_the_caller_busy_for_a_lease() { + let store = TestStore::new(); + let mut service = open(&store, 1 << 40); + assert!(matches!( + begin_call(&mut service, "call-1"), + TransportAdmission::Execute(_) + )); + granted(lease(&mut service, "run-1", 100)); + } + + #[test] + fn an_invocation_id_is_bounded_so_its_lease_id_fits_an_identity() { + let store = TestStore::new(); + let mut service = open(&store, 1 << 40); + let longest = "r".repeat(MAX_IDENTITY_BYTES - LEASE_PREFIX.len()); + let grant = granted(lease(&mut service, &longest, 100)); + assert_eq!(grant.lease_id.len(), MAX_IDENTITY_BYTES); + service + .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reservation) + .unwrap(); + assert!( + service + .begin_lease(CALLER, &format!("{longest}r"), 100, NOW + 1, NOW) + .is_err() + ); + } + + #[test] + fn a_settled_lease_leaves_its_charge_not_its_reservation_across_a_restart() { + let store = TestStore::new(); + let mut service = open(&store, 1_000_000); + let grant = granted(lease(&mut service, "run-1", 10_000)); + service + .settle_lease(CALLER, &grant.lease_id, LeaseUse::Reported(&output(10)[..])) + .unwrap(); + drop(service); + let service = open(&store, 1_000_000); + assert_eq!(held(&service), 150); + } +}