use crate::crypto::encryption::restrict_to_owner;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::io::{BufRead, BufReader, Read, Write};
use std::net::{Ipv4Addr, SocketAddr, TcpListener, TcpStream};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use subtle::ConstantTimeEq;
use zeroize::Zeroizing;
pub const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 900;
const MAX_REQUEST_BYTES: u64 = 64 * 1024;
#[derive(Serialize, Deserialize)]
#[serde(tag = "op", rename_all = "snake_case")]
pub enum Request {
Hello { nonce: String },
Put {
token: String,
repo: String,
passphrase: String,
},
Get { token: String, repo: String },
Drop { token: String, repo: Option<String> },
Status { token: String },
Shutdown { token: String },
}
#[derive(Serialize, Deserialize, Debug)]
#[serde(tag = "result", rename_all = "snake_case")]
pub enum Response {
Hello {
proof: String,
},
Passphrase {
passphrase: String,
},
Missing,
Ok,
Status {
entries: usize,
idle_timeout_secs: u64,
},
Denied,
Malformed {
message: String,
},
}
struct Entry {
passphrase: Zeroizing<String>,
last_used: Instant,
}
pub struct Store {
entries: HashMap<String, Entry>,
idle_timeout: Duration,
}
impl Store {
pub fn new(idle_timeout: Duration) -> Self {
Store {
entries: HashMap::new(),
idle_timeout,
}
}
pub fn put(&mut self, repo: String, passphrase: String) {
self.entries.insert(
repo,
Entry {
passphrase: Zeroizing::new(passphrase),
last_used: Instant::now(),
},
);
}
pub fn get(&mut self, repo: &str) -> Option<Zeroizing<String>> {
self.expire();
let entry = self.entries.get_mut(repo)?;
entry.last_used = Instant::now();
Some(entry.passphrase.clone())
}
pub fn drop_one(&mut self, repo: &str) {
self.entries.remove(repo);
}
pub fn drop_all(&mut self) {
self.entries.clear();
}
pub fn len(&mut self) -> usize {
self.expire();
self.entries.len()
}
pub fn is_empty(&mut self) -> bool {
self.len() == 0
}
fn expire(&mut self) {
let timeout = self.idle_timeout;
self.entries.retain(|_, e| e.last_used.elapsed() < timeout);
}
}
#[derive(Serialize, Deserialize)]
pub struct Endpoint {
pub port: u16,
pub token: String,
pub idle_timeout_secs: u64,
}
pub fn endpoint_path() -> Result<PathBuf, String> {
let home = dirs::home_dir().ok_or("Could not determine home directory")?;
Ok(home.join(".lit").join("agent.json"))
}
impl Endpoint {
pub fn load() -> Result<Endpoint, String> {
Self::load_from(&endpoint_path()?)
}
pub(crate) fn load_from(path: &std::path::Path) -> Result<Endpoint, String> {
let raw = std::fs::read(path)
.map_err(|_| "No agent is running (start one with `lit agent start`)".to_string())?;
serde_json::from_slice(&raw).map_err(|e| format!("Agent endpoint file is unreadable: {e}"))
}
fn save(&self) -> Result<(), String> {
self.save_to(&endpoint_path()?)
}
pub(crate) fn save_to(&self, path: &std::path::Path) -> Result<(), String> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| format!("Failed to create agent directory: {e}"))?;
let _ = crate::crypto::encryption::restrict_dir_to_owner(parent);
}
let raw =
serde_json::to_vec(self).map_err(|e| format!("Failed to encode endpoint: {e}"))?;
let temp = path.with_extension("tmp");
std::fs::write(&temp, raw).map_err(|e| format!("Failed to write endpoint: {e}"))?;
restrict_to_owner(&temp)?;
std::fs::rename(&temp, path).map_err(|e| {
let _ = std::fs::remove_file(&temp);
format!("Failed to write endpoint: {e}")
})?;
Ok(())
}
fn remove() {
if let Ok(path) = endpoint_path() {
let _ = std::fs::remove_file(path);
}
}
}
fn generate_token() -> String {
use aes_gcm::aead::rand_core::RngCore;
use aes_gcm::aead::OsRng;
let mut bytes = [0u8; 32];
OsRng.fill_bytes(&mut bytes);
hex::encode(bytes)
}
fn token_matches(presented: &str, expected: &str) -> bool {
let a = presented.as_bytes();
let b = expected.as_bytes();
if a.len() != b.len() {
return false;
}
a.ct_eq(b).into()
}
fn token_of(req: &Request) -> Option<&str> {
match req {
Request::Put { token, .. }
| Request::Get { token, .. }
| Request::Drop { token, .. }
| Request::Status { token }
| Request::Shutdown { token } => Some(token),
Request::Hello { .. } => None,
}
}
fn proof_for(token: &str, nonce: &str) -> String {
use hmac::{Hmac, Mac};
use sha2::Sha256;
let mut mac = <Hmac<Sha256> as Mac>::new_from_slice(token.as_bytes())
.expect("HMAC accepts keys of any length");
mac.update(nonce.as_bytes());
hex::encode(mac.finalize().into_bytes())
}
fn apply(req: Request, store: &Arc<Mutex<Store>>) -> (Response, bool) {
let mut store = match store.lock() {
Ok(s) => s,
Err(_) => {
return (
Response::Malformed {
message: "agent state is poisoned".to_string(),
},
false,
)
}
};
match req {
Request::Hello { .. } => (Response::Denied, false),
Request::Put {
repo, passphrase, ..
} => {
store.put(repo, passphrase);
(Response::Ok, false)
}
Request::Get { repo, .. } => match store.get(&repo) {
Some(p) => (
Response::Passphrase {
passphrase: p.to_string(),
},
false,
),
None => (Response::Missing, false),
},
Request::Drop { repo, .. } => {
match repo {
Some(r) => store.drop_one(&r),
None => store.drop_all(),
}
(Response::Ok, false)
}
Request::Status { .. } => (
Response::Status {
entries: store.len(),
idle_timeout_secs: store.idle_timeout.as_secs(),
},
false,
),
Request::Shutdown { .. } => {
store.drop_all();
(Response::Ok, true)
}
}
}
fn respond_to(req: Request, expected_token: &str, store: &Arc<Mutex<Store>>) -> (Response, bool) {
match req {
Request::Hello { nonce } => (
Response::Hello {
proof: proof_for(expected_token, &nonce),
},
false,
),
other => match token_of(&other) {
Some(t) if token_matches(t, expected_token) => apply(other, store),
_ => (Response::Denied, false),
},
}
}
fn handle_connection(
stream: &mut TcpStream,
expected_token: &str,
store: &Arc<Mutex<Store>>,
) -> bool {
let _ = stream.set_read_timeout(Some(Duration::from_secs(5)));
let _ = stream.set_write_timeout(Some(Duration::from_secs(5)));
let Ok(peer) = stream.try_clone() else {
return false;
};
let mut reader = BufReader::new(peer.take(MAX_REQUEST_BYTES));
for _ in 0..2 {
let mut line = String::new();
match reader.read_line(&mut line) {
Ok(0) | Err(_) => return false,
Ok(_) => {}
}
let (response, shutdown) = match serde_json::from_str::<Request>(line.trim()) {
Ok(req) => respond_to(req, expected_token, store),
Err(e) => (
Response::Malformed {
message: e.to_string(),
},
false,
),
};
let closing = !matches!(response, Response::Hello { .. });
if let Ok(mut body) = serde_json::to_vec(&response) {
body.push(b'\n');
if stream.write_all(&body).is_err() {
return false;
}
let _ = stream.flush();
}
if closing {
return shutdown;
}
}
false
}
pub fn serve(idle_timeout: Duration) -> Result<(), String> {
if Endpoint::load().is_ok() && ping().is_ok() {
return Err("An agent is already running (`lit agent stop` to replace it)".to_string());
}
let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
.map_err(|e| format!("Failed to bind agent socket: {e}"))?;
let port = listener
.local_addr()
.map_err(|e| format!("Failed to read agent port: {e}"))?
.port();
let token = generate_token();
Endpoint {
port,
token: token.clone(),
idle_timeout_secs: idle_timeout.as_secs(),
}
.save()?;
let store = Arc::new(Mutex::new(Store::new(idle_timeout)));
for incoming in listener.incoming() {
let mut stream = match incoming {
Ok(s) => s,
Err(_) => continue,
};
if handle_connection(&mut stream, &token, &store) {
break;
}
}
if let Ok(mut s) = store.lock() {
s.drop_all();
}
Endpoint::remove();
Ok(())
}
fn write_line(stream: &mut TcpStream, req: &Request) -> Result<(), String> {
let mut body = serde_json::to_vec(req).map_err(|e| format!("Failed to encode request: {e}"))?;
body.push(b'\n');
stream
.write_all(&body)
.map_err(|e| format!("Failed to reach agent: {e}"))
}
fn read_response(reader: &mut impl BufRead) -> Result<Response, String> {
let mut line = String::new();
reader
.read_line(&mut line)
.map_err(|e| format!("Failed to read agent reply: {e}"))?;
serde_json::from_str(line.trim()).map_err(|e| format!("Agent sent an unreadable reply: {e}"))
}
fn request(req: &Request) -> Result<Response, String> {
let endpoint = Endpoint::load()?;
let mut stream = TcpStream::connect(SocketAddr::from((Ipv4Addr::LOCALHOST, endpoint.port)))
.map_err(|_| "No agent is running (start one with `lit agent start`)".to_string())?;
stream
.set_read_timeout(Some(Duration::from_secs(5)))
.map_err(|e| format!("Failed to configure agent socket: {e}"))?;
stream
.set_write_timeout(Some(Duration::from_secs(5)))
.map_err(|e| format!("Failed to configure agent socket: {e}"))?;
let peer = stream
.try_clone()
.map_err(|e| format!("Failed to read from agent: {e}"))?;
let mut reader = BufReader::new(peer.take(MAX_REQUEST_BYTES));
let nonce = generate_token();
write_line(
&mut stream,
&Request::Hello {
nonce: nonce.clone(),
},
)?;
let proved = matches!(
read_response(&mut reader),
Ok(Response::Hello { ref proof }) if token_matches(proof, &proof_for(&endpoint.token, &nonce))
);
if !proved {
return Err(
"Whatever is listening on the agent's port could not prove it is the agent; \
refusing to send anything to it. Run `lit agent stop` and start a new one."
.to_string(),
);
}
write_line(&mut stream, req)?;
read_response(&mut reader)
}
fn token() -> Result<String, String> {
Ok(Endpoint::load()?.token)
}
pub fn ping() -> Result<(), String> {
match request(&Request::Status { token: token()? })? {
Response::Status { .. } => Ok(()),
_ => Err("Agent did not answer a status request".to_string()),
}
}
pub fn get(repo: &str) -> Option<Zeroizing<String>> {
let token = token().ok()?;
match request(&Request::Get {
token,
repo: repo.to_string(),
})
.ok()?
{
Response::Passphrase { passphrase } => Some(Zeroizing::new(passphrase)),
_ => None,
}
}
pub fn put(repo: &str, passphrase: &str) -> Result<(), String> {
match request(&Request::Put {
token: token()?,
repo: repo.to_string(),
passphrase: passphrase.to_string(),
})? {
Response::Ok => Ok(()),
other => Err(format!("Agent refused to store the passphrase: {other:?}")),
}
}
pub fn drop_entry(repo: Option<&str>) -> Result<(), String> {
match request(&Request::Drop {
token: token()?,
repo: repo.map(|r| r.to_string()),
})? {
Response::Ok => Ok(()),
other => Err(format!("Agent refused: {other:?}")),
}
}
pub fn status() -> Result<(usize, u64), String> {
match request(&Request::Status { token: token()? })? {
Response::Status {
entries,
idle_timeout_secs,
} => Ok((entries, idle_timeout_secs)),
other => Err(format!("Agent refused: {other:?}")),
}
}
pub fn shutdown() -> Result<(), String> {
match request(&Request::Shutdown { token: token()? }) {
Ok(Response::Ok) => {
Endpoint::remove();
Ok(())
}
Ok(other) => Err(format!("Agent refused to stop: {other:?}")),
Err(e) if e.contains("No agent is running") || e.contains("could not prove") => {
Endpoint::remove();
Err(e)
}
Err(e) => Err(format!(
"{e}. The agent may still be running and holding a passphrase; \
its endpoint file has been left in place so it can be reached again."
)),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_entries_expire_when_idle() {
let mut store = Store::new(Duration::from_millis(50));
store.put("repo".to_string(), "hunter2".to_string());
assert!(store.get("repo").is_some());
std::thread::sleep(Duration::from_millis(80));
assert!(
store.get("repo").is_none(),
"an entry left alone past the timeout should be gone"
);
assert!(store.is_empty());
}
#[test]
fn test_use_refreshes_the_timeout() {
let mut store = Store::new(Duration::from_millis(120));
store.put("repo".to_string(), "hunter2".to_string());
for _ in 0..4 {
std::thread::sleep(Duration::from_millis(50));
assert!(store.get("repo").is_some(), "use should keep it alive");
}
}
#[test]
fn test_drop_all_forgets_everything() {
let mut store = Store::new(Duration::from_secs(60));
store.put("a".to_string(), "one".to_string());
store.put("b".to_string(), "two".to_string());
assert_eq!(store.len(), 2);
store.drop_all();
assert!(store.is_empty());
}
#[test]
fn test_token_comparison_rejects_wrong_and_short_tokens() {
let real = generate_token();
assert!(token_matches(&real, &real));
assert!(!token_matches("", &real));
assert!(!token_matches(&real[..real.len() - 1], &real));
let mut wrong = real.clone();
let last = if wrong.ends_with('a') { 'b' } else { 'a' };
wrong.pop();
wrong.push(last);
assert!(!token_matches(&wrong, &real));
}
#[test]
fn test_generated_tokens_differ() {
assert_ne!(generate_token(), generate_token());
}
#[test]
fn test_wrong_token_is_denied_and_changes_nothing() {
let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
let real = generate_token();
let (resp, stop) = respond_to(
Request::Put {
token: "not-the-token".to_string(),
repo: "repo".to_string(),
passphrase: "hunter2".to_string(),
},
&real,
&store,
);
assert!(matches!(resp, Response::Denied));
assert!(!stop);
assert!(
store.lock().unwrap().is_empty(),
"an unauthenticated Put must store nothing"
);
}
#[test]
fn test_handshake_proves_the_peer_holds_the_token() {
let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
let real = generate_token();
let nonce = generate_token();
let (resp, _) = respond_to(
Request::Hello {
nonce: nonce.clone(),
},
&real,
&store,
);
let proof = match resp {
Response::Hello { proof } => proof,
other => panic!("expected a proof, got {other:?}"),
};
assert!(token_matches(&proof, &proof_for(&real, &nonce)));
let impostor = generate_token();
assert!(!token_matches(&proof, &proof_for(&impostor, &nonce)));
let other_nonce = generate_token();
assert!(!token_matches(&proof, &proof_for(&real, &other_nonce)));
}
#[test]
fn test_hello_carries_no_token_but_grants_nothing() {
assert!(token_of(&Request::Hello {
nonce: "n".to_string()
})
.is_none());
let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
let (resp, stop) = apply(
Request::Hello {
nonce: "n".to_string(),
},
&store,
);
assert!(matches!(resp, Response::Denied));
assert!(!stop);
}
#[test]
fn test_put_then_get_round_trips_through_apply() {
let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
let (resp, stop) = apply(
Request::Put {
token: String::new(),
repo: "repo".to_string(),
passphrase: "hunter2".to_string(),
},
&store,
);
assert!(matches!(resp, Response::Ok));
assert!(!stop);
let (resp, _) = apply(
Request::Get {
token: String::new(),
repo: "repo".to_string(),
},
&store,
);
match resp {
Response::Passphrase { passphrase } => assert_eq!(passphrase, "hunter2"),
other => panic!("expected the passphrase back, got {other:?}"),
}
let (_, stop) = apply(
Request::Shutdown {
token: String::new(),
},
&store,
);
assert!(stop, "shutdown should stop the agent");
assert!(
store.lock().unwrap().is_empty(),
"shutdown should clear what it held"
);
}
#[test]
fn test_endpoint_is_written_restricted_and_by_rename() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("nested").join("agent.json");
let endpoint = Endpoint {
port: 4242,
token: generate_token(),
idle_timeout_secs: 900,
};
endpoint.save_to(&path).unwrap();
let read_back = Endpoint::load_from(&path).unwrap();
assert_eq!(read_back.port, endpoint.port);
assert_eq!(read_back.token, endpoint.token);
assert!(
!path.with_extension("tmp").exists(),
"the temporary endpoint file was left behind"
);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(mode & 0o777, 0o600, "the token file is readable by others");
let dir_mode = std::fs::metadata(path.parent().unwrap())
.unwrap()
.permissions()
.mode();
assert_eq!(dir_mode & 0o777, 0o700, "the agent directory is listable");
}
}
#[test]
fn test_endpoint_can_be_replaced() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("agent.json");
for port in [1111u16, 2222] {
Endpoint {
port,
token: generate_token(),
idle_timeout_secs: 60,
}
.save_to(&path)
.unwrap();
assert_eq!(Endpoint::load_from(&path).unwrap().port, port);
}
}
#[test]
fn test_get_for_unknown_repo_is_missing_not_an_error() {
let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
let (resp, _) = apply(
Request::Get {
token: String::new(),
repo: "never-stored".to_string(),
},
&store,
);
assert!(matches!(resp, Response::Missing));
}
}