1use crate::crypto::encryption::restrict_to_owner;
34use serde::{Deserialize, Serialize};
35use std::collections::HashMap;
36use std::io::{BufRead, BufReader, Read, Write};
37use std::net::{Ipv4Addr, SocketAddr, TcpListener, TcpStream};
38use std::path::PathBuf;
39use std::sync::{Arc, Mutex};
40use std::time::{Duration, Instant};
41use subtle::ConstantTimeEq;
42use zeroize::Zeroizing;
43
44pub const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 900;
46
47const MAX_REQUEST_BYTES: u64 = 64 * 1024;
49
50#[derive(Serialize, Deserialize)]
54#[serde(tag = "op", rename_all = "snake_case")]
55pub enum Request {
56 Put {
58 token: String,
59 repo: String,
60 passphrase: String,
61 },
62 Get { token: String, repo: String },
64 Drop { token: String, repo: Option<String> },
66 Status { token: String },
68 Shutdown { token: String },
70}
71
72#[derive(Serialize, Deserialize, Debug)]
74#[serde(tag = "result", rename_all = "snake_case")]
75pub enum Response {
76 Passphrase {
77 passphrase: String,
78 },
79 Missing,
81 Ok,
82 Status {
83 entries: usize,
84 idle_timeout_secs: u64,
85 },
86 Denied,
87 Malformed {
88 message: String,
89 },
90}
91
92struct Entry {
93 passphrase: Zeroizing<String>,
94 last_used: Instant,
95}
96
97pub struct Store {
102 entries: HashMap<String, Entry>,
103 idle_timeout: Duration,
104}
105
106impl Store {
107 pub fn new(idle_timeout: Duration) -> Self {
108 Store {
109 entries: HashMap::new(),
110 idle_timeout,
111 }
112 }
113
114 pub fn put(&mut self, repo: String, passphrase: String) {
115 self.entries.insert(
116 repo,
117 Entry {
118 passphrase: Zeroizing::new(passphrase),
119 last_used: Instant::now(),
120 },
121 );
122 }
123
124 pub fn get(&mut self, repo: &str) -> Option<Zeroizing<String>> {
126 self.expire();
127 let entry = self.entries.get_mut(repo)?;
128 entry.last_used = Instant::now();
129 Some(entry.passphrase.clone())
130 }
131
132 pub fn drop_one(&mut self, repo: &str) {
133 self.entries.remove(repo);
134 }
135
136 pub fn drop_all(&mut self) {
137 self.entries.clear();
138 }
139
140 pub fn len(&mut self) -> usize {
141 self.expire();
142 self.entries.len()
143 }
144
145 pub fn is_empty(&mut self) -> bool {
146 self.len() == 0
147 }
148
149 fn expire(&mut self) {
150 let timeout = self.idle_timeout;
151 self.entries.retain(|_, e| e.last_used.elapsed() < timeout);
152 }
153}
154
155#[derive(Serialize, Deserialize)]
159pub struct Endpoint {
160 pub port: u16,
161 pub token: String,
162 pub idle_timeout_secs: u64,
163}
164
165pub fn endpoint_path() -> Result<PathBuf, String> {
166 let home = dirs::home_dir().ok_or("Could not determine home directory")?;
167 Ok(home.join(".lit").join("agent.json"))
168}
169
170impl Endpoint {
171 pub fn load() -> Result<Endpoint, String> {
172 let path = endpoint_path()?;
173 let raw = std::fs::read(&path)
174 .map_err(|_| "No agent is running (start one with `lit agent start`)".to_string())?;
175 serde_json::from_slice(&raw).map_err(|e| format!("Agent endpoint file is unreadable: {e}"))
176 }
177
178 fn save(&self) -> Result<(), String> {
179 let path = endpoint_path()?;
180 if let Some(parent) = path.parent() {
181 std::fs::create_dir_all(parent)
182 .map_err(|e| format!("Failed to create agent directory: {e}"))?;
183 }
184 let raw =
185 serde_json::to_vec(self).map_err(|e| format!("Failed to encode endpoint: {e}"))?;
186 std::fs::write(&path, raw).map_err(|e| format!("Failed to write endpoint: {e}"))?;
187
188 restrict_to_owner(&path)?;
192 Ok(())
193 }
194
195 fn remove() {
196 if let Ok(path) = endpoint_path() {
197 let _ = std::fs::remove_file(path);
198 }
199 }
200}
201
202fn generate_token() -> String {
204 use aes_gcm::aead::rand_core::RngCore;
205 use aes_gcm::aead::OsRng;
206
207 let mut bytes = [0u8; 32];
208 OsRng.fill_bytes(&mut bytes);
209 hex::encode(bytes)
210}
211
212fn token_matches(presented: &str, expected: &str) -> bool {
215 let a = presented.as_bytes();
216 let b = expected.as_bytes();
217 if a.len() != b.len() {
218 return false;
219 }
220 a.ct_eq(b).into()
221}
222
223fn token_of(req: &Request) -> &str {
224 match req {
225 Request::Put { token, .. }
226 | Request::Get { token, .. }
227 | Request::Drop { token, .. }
228 | Request::Status { token }
229 | Request::Shutdown { token } => token,
230 }
231}
232
233fn apply(req: Request, store: &Arc<Mutex<Store>>) -> (Response, bool) {
237 let mut store = match store.lock() {
238 Ok(s) => s,
239 Err(_) => {
240 return (
241 Response::Malformed {
242 message: "agent state is poisoned".to_string(),
243 },
244 false,
245 )
246 }
247 };
248
249 match req {
250 Request::Put {
251 repo, passphrase, ..
252 } => {
253 store.put(repo, passphrase);
254 (Response::Ok, false)
255 }
256 Request::Get { repo, .. } => match store.get(&repo) {
257 Some(p) => (
258 Response::Passphrase {
259 passphrase: p.to_string(),
260 },
261 false,
262 ),
263 None => (Response::Missing, false),
264 },
265 Request::Drop { repo, .. } => {
266 match repo {
267 Some(r) => store.drop_one(&r),
268 None => store.drop_all(),
269 }
270 (Response::Ok, false)
271 }
272 Request::Status { .. } => (
273 Response::Status {
274 entries: store.len(),
275 idle_timeout_secs: store.idle_timeout.as_secs(),
276 },
277 false,
278 ),
279 Request::Shutdown { .. } => {
280 store.drop_all();
281 (Response::Ok, true)
282 }
283 }
284}
285
286fn handle_connection(
290 stream: &mut TcpStream,
291 expected_token: &str,
292 store: &Arc<Mutex<Store>>,
293) -> bool {
294 let _ = stream.set_read_timeout(Some(Duration::from_secs(5)));
296 let _ = stream.set_write_timeout(Some(Duration::from_secs(5)));
297
298 let Ok(peer) = stream.try_clone() else {
299 return false;
300 };
301
302 let mut line = String::new();
305 if BufReader::new(peer.take(MAX_REQUEST_BYTES))
306 .read_line(&mut line)
307 .is_err()
308 {
309 return false;
310 }
311
312 let (response, shutdown) = match serde_json::from_str::<Request>(line.trim()) {
313 Ok(req) => {
314 if token_matches(token_of(&req), expected_token) {
315 apply(req, store)
316 } else {
317 (Response::Denied, false)
321 }
322 }
323 Err(e) => (
324 Response::Malformed {
325 message: e.to_string(),
326 },
327 false,
328 ),
329 };
330
331 if let Ok(mut body) = serde_json::to_vec(&response) {
332 body.push(b'\n');
333 let _ = stream.write_all(&body);
334 let _ = stream.flush();
335 }
336
337 shutdown
338}
339
340pub fn serve(idle_timeout: Duration) -> Result<(), String> {
342 if Endpoint::load().is_ok() && ping().is_ok() {
343 return Err("An agent is already running (`lit agent stop` to replace it)".to_string());
344 }
345
346 let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
349 .map_err(|e| format!("Failed to bind agent socket: {e}"))?;
350 let port = listener
351 .local_addr()
352 .map_err(|e| format!("Failed to read agent port: {e}"))?
353 .port();
354
355 let token = generate_token();
356 Endpoint {
357 port,
358 token: token.clone(),
359 idle_timeout_secs: idle_timeout.as_secs(),
360 }
361 .save()?;
362
363 let store = Arc::new(Mutex::new(Store::new(idle_timeout)));
364
365 for incoming in listener.incoming() {
366 let mut stream = match incoming {
367 Ok(s) => s,
368 Err(_) => continue,
369 };
370 if handle_connection(&mut stream, &token, &store) {
371 break;
372 }
373 }
374
375 if let Ok(mut s) = store.lock() {
376 s.drop_all();
377 }
378 Endpoint::remove();
379 Ok(())
380}
381
382fn request(req: &Request) -> Result<Response, String> {
384 let endpoint = Endpoint::load()?;
385 let mut stream = TcpStream::connect(SocketAddr::from((Ipv4Addr::LOCALHOST, endpoint.port)))
386 .map_err(|_| "No agent is running (start one with `lit agent start`)".to_string())?;
387 stream
388 .set_read_timeout(Some(Duration::from_secs(5)))
389 .map_err(|e| format!("Failed to configure agent socket: {e}"))?;
390
391 let mut body = serde_json::to_vec(req).map_err(|e| format!("Failed to encode request: {e}"))?;
392 body.push(b'\n');
393 stream
394 .write_all(&body)
395 .map_err(|e| format!("Failed to reach agent: {e}"))?;
396
397 let mut line = String::new();
398 BufReader::new(&stream)
399 .read_line(&mut line)
400 .map_err(|e| format!("Failed to read agent reply: {e}"))?;
401
402 serde_json::from_str(line.trim()).map_err(|e| format!("Agent sent an unreadable reply: {e}"))
403}
404
405fn token() -> Result<String, String> {
406 Ok(Endpoint::load()?.token)
407}
408
409pub fn ping() -> Result<(), String> {
411 match request(&Request::Status { token: token()? })? {
412 Response::Status { .. } => Ok(()),
413 _ => Err("Agent did not answer a status request".to_string()),
414 }
415}
416
417pub fn get(repo: &str) -> Option<Zeroizing<String>> {
421 let token = token().ok()?;
422 match request(&Request::Get {
423 token,
424 repo: repo.to_string(),
425 })
426 .ok()?
427 {
428 Response::Passphrase { passphrase } => Some(Zeroizing::new(passphrase)),
429 _ => None,
430 }
431}
432
433pub fn put(repo: &str, passphrase: &str) -> Result<(), String> {
434 match request(&Request::Put {
435 token: token()?,
436 repo: repo.to_string(),
437 passphrase: passphrase.to_string(),
438 })? {
439 Response::Ok => Ok(()),
440 other => Err(format!("Agent refused to store the passphrase: {other:?}")),
441 }
442}
443
444pub fn drop_entry(repo: Option<&str>) -> Result<(), String> {
445 match request(&Request::Drop {
446 token: token()?,
447 repo: repo.map(|r| r.to_string()),
448 })? {
449 Response::Ok => Ok(()),
450 other => Err(format!("Agent refused: {other:?}")),
451 }
452}
453
454pub fn status() -> Result<(usize, u64), String> {
455 match request(&Request::Status { token: token()? })? {
456 Response::Status {
457 entries,
458 idle_timeout_secs,
459 } => Ok((entries, idle_timeout_secs)),
460 other => Err(format!("Agent refused: {other:?}")),
461 }
462}
463
464pub fn shutdown() -> Result<(), String> {
465 let result = request(&Request::Shutdown { token: token()? });
466 Endpoint::remove();
469 match result? {
470 Response::Ok => Ok(()),
471 other => Err(format!("Agent refused to stop: {other:?}")),
472 }
473}
474
475#[cfg(test)]
476mod tests {
477 use super::*;
478
479 #[test]
480 fn test_entries_expire_when_idle() {
481 let mut store = Store::new(Duration::from_millis(50));
482 store.put("repo".to_string(), "hunter2".to_string());
483 assert!(store.get("repo").is_some());
484
485 std::thread::sleep(Duration::from_millis(80));
486 assert!(
487 store.get("repo").is_none(),
488 "an entry left alone past the timeout should be gone"
489 );
490 assert!(store.is_empty());
491 }
492
493 #[test]
494 fn test_use_refreshes_the_timeout() {
495 let mut store = Store::new(Duration::from_millis(120));
498 store.put("repo".to_string(), "hunter2".to_string());
499
500 for _ in 0..4 {
501 std::thread::sleep(Duration::from_millis(50));
502 assert!(store.get("repo").is_some(), "use should keep it alive");
503 }
504 }
505
506 #[test]
507 fn test_drop_all_forgets_everything() {
508 let mut store = Store::new(Duration::from_secs(60));
509 store.put("a".to_string(), "one".to_string());
510 store.put("b".to_string(), "two".to_string());
511 assert_eq!(store.len(), 2);
512
513 store.drop_all();
514 assert!(store.is_empty());
515 }
516
517 #[test]
518 fn test_token_comparison_rejects_wrong_and_short_tokens() {
519 let real = generate_token();
520 assert!(token_matches(&real, &real));
521 assert!(!token_matches("", &real));
522 assert!(!token_matches(&real[..real.len() - 1], &real));
523
524 let mut wrong = real.clone();
525 let last = if wrong.ends_with('a') { 'b' } else { 'a' };
527 wrong.pop();
528 wrong.push(last);
529 assert!(!token_matches(&wrong, &real));
530 }
531
532 #[test]
533 fn test_generated_tokens_differ() {
534 assert_ne!(generate_token(), generate_token());
535 }
536
537 #[test]
540 fn test_wrong_token_is_denied_and_changes_nothing() {
541 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
542 let real = generate_token();
543
544 let req = Request::Put {
545 token: "not-the-token".to_string(),
546 repo: "repo".to_string(),
547 passphrase: "hunter2".to_string(),
548 };
549 assert!(!token_matches(token_of(&req), &real));
550
551 assert!(store.lock().unwrap().is_empty());
553 }
554
555 #[test]
556 fn test_put_then_get_round_trips_through_apply() {
557 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
558
559 let (resp, stop) = apply(
560 Request::Put {
561 token: String::new(),
562 repo: "repo".to_string(),
563 passphrase: "hunter2".to_string(),
564 },
565 &store,
566 );
567 assert!(matches!(resp, Response::Ok));
568 assert!(!stop);
569
570 let (resp, _) = apply(
571 Request::Get {
572 token: String::new(),
573 repo: "repo".to_string(),
574 },
575 &store,
576 );
577 match resp {
578 Response::Passphrase { passphrase } => assert_eq!(passphrase, "hunter2"),
579 other => panic!("expected the passphrase back, got {other:?}"),
580 }
581
582 let (_, stop) = apply(
583 Request::Shutdown {
584 token: String::new(),
585 },
586 &store,
587 );
588 assert!(stop, "shutdown should stop the agent");
589 assert!(
590 store.lock().unwrap().is_empty(),
591 "shutdown should clear what it held"
592 );
593 }
594
595 #[test]
596 fn test_get_for_unknown_repo_is_missing_not_an_error() {
597 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
598 let (resp, _) = apply(
599 Request::Get {
600 token: String::new(),
601 repo: "never-stored".to_string(),
602 },
603 &store,
604 );
605 assert!(matches!(resp, Response::Missing));
606 }
607}