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)]
57#[serde(tag = "op", rename_all = "snake_case")]
58pub enum Request {
59 Hello { nonce: String },
62 Put {
64 token: String,
65 repo: String,
66 passphrase: String,
67 },
68 Get { token: String, repo: String },
70 Drop { token: String, repo: Option<String> },
72 Status { token: String },
74 Shutdown { token: String },
76}
77
78#[derive(Serialize, Deserialize, Debug)]
80#[serde(tag = "result", rename_all = "snake_case")]
81pub enum Response {
82 Hello {
84 proof: String,
85 },
86 Passphrase {
87 passphrase: String,
88 },
89 Missing,
91 Ok,
92 Status {
93 entries: usize,
94 idle_timeout_secs: u64,
95 },
96 Denied,
97 Malformed {
98 message: String,
99 },
100}
101
102struct Entry {
103 passphrase: Zeroizing<String>,
104 last_used: Instant,
105}
106
107pub struct Store {
112 entries: HashMap<String, Entry>,
113 idle_timeout: Duration,
114}
115
116impl Store {
117 pub fn new(idle_timeout: Duration) -> Self {
118 Store {
119 entries: HashMap::new(),
120 idle_timeout,
121 }
122 }
123
124 pub fn put(&mut self, repo: String, passphrase: String) {
125 self.entries.insert(
126 repo,
127 Entry {
128 passphrase: Zeroizing::new(passphrase),
129 last_used: Instant::now(),
130 },
131 );
132 }
133
134 pub fn get(&mut self, repo: &str) -> Option<Zeroizing<String>> {
136 self.expire();
137 let entry = self.entries.get_mut(repo)?;
138 entry.last_used = Instant::now();
139 Some(entry.passphrase.clone())
140 }
141
142 pub fn drop_one(&mut self, repo: &str) {
143 self.entries.remove(repo);
144 }
145
146 pub fn drop_all(&mut self) {
147 self.entries.clear();
148 }
149
150 pub fn len(&mut self) -> usize {
151 self.expire();
152 self.entries.len()
153 }
154
155 pub fn is_empty(&mut self) -> bool {
156 self.len() == 0
157 }
158
159 fn expire(&mut self) {
160 let timeout = self.idle_timeout;
161 self.entries.retain(|_, e| e.last_used.elapsed() < timeout);
162 }
163}
164
165#[derive(Serialize, Deserialize)]
169pub struct Endpoint {
170 pub port: u16,
171 pub token: String,
172 pub idle_timeout_secs: u64,
173}
174
175pub fn endpoint_path() -> Result<PathBuf, String> {
176 let home = dirs::home_dir().ok_or("Could not determine home directory")?;
177 Ok(home.join(".lit").join("agent.json"))
178}
179
180impl Endpoint {
181 pub fn load() -> Result<Endpoint, String> {
182 let path = endpoint_path()?;
183 let raw = std::fs::read(&path)
184 .map_err(|_| "No agent is running (start one with `lit agent start`)".to_string())?;
185 serde_json::from_slice(&raw).map_err(|e| format!("Agent endpoint file is unreadable: {e}"))
186 }
187
188 fn save(&self) -> Result<(), String> {
189 let path = endpoint_path()?;
190 if let Some(parent) = path.parent() {
191 std::fs::create_dir_all(parent)
192 .map_err(|e| format!("Failed to create agent directory: {e}"))?;
193 }
194 let raw =
195 serde_json::to_vec(self).map_err(|e| format!("Failed to encode endpoint: {e}"))?;
196 std::fs::write(&path, raw).map_err(|e| format!("Failed to write endpoint: {e}"))?;
197
198 restrict_to_owner(&path)?;
202 Ok(())
203 }
204
205 fn remove() {
206 if let Ok(path) = endpoint_path() {
207 let _ = std::fs::remove_file(path);
208 }
209 }
210}
211
212fn generate_token() -> String {
214 use aes_gcm::aead::rand_core::RngCore;
215 use aes_gcm::aead::OsRng;
216
217 let mut bytes = [0u8; 32];
218 OsRng.fill_bytes(&mut bytes);
219 hex::encode(bytes)
220}
221
222fn token_matches(presented: &str, expected: &str) -> bool {
225 let a = presented.as_bytes();
226 let b = expected.as_bytes();
227 if a.len() != b.len() {
228 return false;
229 }
230 a.ct_eq(b).into()
231}
232
233fn token_of(req: &Request) -> Option<&str> {
234 match req {
235 Request::Put { token, .. }
236 | Request::Get { token, .. }
237 | Request::Drop { token, .. }
238 | Request::Status { token }
239 | Request::Shutdown { token } => Some(token),
240 Request::Hello { .. } => None,
242 }
243}
244
245fn proof_for(token: &str, nonce: &str) -> String {
256 use hmac::{Hmac, Mac};
257 use sha2::Sha256;
258
259 let mut mac = <Hmac<Sha256> as Mac>::new_from_slice(token.as_bytes())
260 .expect("HMAC accepts keys of any length");
261 mac.update(nonce.as_bytes());
262 hex::encode(mac.finalize().into_bytes())
263}
264
265fn apply(req: Request, store: &Arc<Mutex<Store>>) -> (Response, bool) {
269 let mut store = match store.lock() {
270 Ok(s) => s,
271 Err(_) => {
272 return (
273 Response::Malformed {
274 message: "agent state is poisoned".to_string(),
275 },
276 false,
277 )
278 }
279 };
280
281 match req {
282 Request::Hello { .. } => (Response::Denied, false),
284 Request::Put {
285 repo, passphrase, ..
286 } => {
287 store.put(repo, passphrase);
288 (Response::Ok, false)
289 }
290 Request::Get { repo, .. } => match store.get(&repo) {
291 Some(p) => (
292 Response::Passphrase {
293 passphrase: p.to_string(),
294 },
295 false,
296 ),
297 None => (Response::Missing, false),
298 },
299 Request::Drop { repo, .. } => {
300 match repo {
301 Some(r) => store.drop_one(&r),
302 None => store.drop_all(),
303 }
304 (Response::Ok, false)
305 }
306 Request::Status { .. } => (
307 Response::Status {
308 entries: store.len(),
309 idle_timeout_secs: store.idle_timeout.as_secs(),
310 },
311 false,
312 ),
313 Request::Shutdown { .. } => {
314 store.drop_all();
315 (Response::Ok, true)
316 }
317 }
318}
319
320fn respond_to(req: Request, expected_token: &str, store: &Arc<Mutex<Store>>) -> (Response, bool) {
322 match req {
323 Request::Hello { nonce } => (
324 Response::Hello {
325 proof: proof_for(expected_token, &nonce),
326 },
327 false,
328 ),
329 other => match token_of(&other) {
330 Some(t) if token_matches(t, expected_token) => apply(other, store),
334 _ => (Response::Denied, false),
335 },
336 }
337}
338
339fn handle_connection(
343 stream: &mut TcpStream,
344 expected_token: &str,
345 store: &Arc<Mutex<Store>>,
346) -> bool {
347 let _ = stream.set_read_timeout(Some(Duration::from_secs(5)));
349 let _ = stream.set_write_timeout(Some(Duration::from_secs(5)));
350
351 let Ok(peer) = stream.try_clone() else {
352 return false;
353 };
354
355 let mut reader = BufReader::new(peer.take(MAX_REQUEST_BYTES));
358
359 for _ in 0..2 {
362 let mut line = String::new();
363 match reader.read_line(&mut line) {
364 Ok(0) | Err(_) => return false,
365 Ok(_) => {}
366 }
367
368 let (response, shutdown) = match serde_json::from_str::<Request>(line.trim()) {
369 Ok(req) => respond_to(req, expected_token, store),
370 Err(e) => (
371 Response::Malformed {
372 message: e.to_string(),
373 },
374 false,
375 ),
376 };
377
378 let closing = !matches!(response, Response::Hello { .. });
379
380 if let Ok(mut body) = serde_json::to_vec(&response) {
381 body.push(b'\n');
382 if stream.write_all(&body).is_err() {
383 return false;
384 }
385 let _ = stream.flush();
386 }
387
388 if closing {
390 return shutdown;
391 }
392 }
393
394 false
395}
396
397pub fn serve(idle_timeout: Duration) -> Result<(), String> {
399 if Endpoint::load().is_ok() && ping().is_ok() {
400 return Err("An agent is already running (`lit agent stop` to replace it)".to_string());
401 }
402
403 let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
406 .map_err(|e| format!("Failed to bind agent socket: {e}"))?;
407 let port = listener
408 .local_addr()
409 .map_err(|e| format!("Failed to read agent port: {e}"))?
410 .port();
411
412 let token = generate_token();
413 Endpoint {
414 port,
415 token: token.clone(),
416 idle_timeout_secs: idle_timeout.as_secs(),
417 }
418 .save()?;
419
420 let store = Arc::new(Mutex::new(Store::new(idle_timeout)));
421
422 for incoming in listener.incoming() {
423 let mut stream = match incoming {
424 Ok(s) => s,
425 Err(_) => continue,
426 };
427 if handle_connection(&mut stream, &token, &store) {
428 break;
429 }
430 }
431
432 if let Ok(mut s) = store.lock() {
433 s.drop_all();
434 }
435 Endpoint::remove();
436 Ok(())
437}
438
439fn write_line(stream: &mut TcpStream, req: &Request) -> Result<(), String> {
440 let mut body = serde_json::to_vec(req).map_err(|e| format!("Failed to encode request: {e}"))?;
441 body.push(b'\n');
442 stream
443 .write_all(&body)
444 .map_err(|e| format!("Failed to reach agent: {e}"))
445}
446
447fn read_response(reader: &mut impl BufRead) -> Result<Response, String> {
448 let mut line = String::new();
449 reader
450 .read_line(&mut line)
451 .map_err(|e| format!("Failed to read agent reply: {e}"))?;
452 serde_json::from_str(line.trim()).map_err(|e| format!("Agent sent an unreadable reply: {e}"))
453}
454
455fn request(req: &Request) -> Result<Response, String> {
463 let endpoint = Endpoint::load()?;
464 let mut stream = TcpStream::connect(SocketAddr::from((Ipv4Addr::LOCALHOST, endpoint.port)))
465 .map_err(|_| "No agent is running (start one with `lit agent start`)".to_string())?;
466 stream
467 .set_read_timeout(Some(Duration::from_secs(5)))
468 .map_err(|e| format!("Failed to configure agent socket: {e}"))?;
469 stream
470 .set_write_timeout(Some(Duration::from_secs(5)))
471 .map_err(|e| format!("Failed to configure agent socket: {e}"))?;
472
473 let peer = stream
474 .try_clone()
475 .map_err(|e| format!("Failed to read from agent: {e}"))?;
476 let mut reader = BufReader::new(peer.take(MAX_REQUEST_BYTES));
477
478 let nonce = generate_token();
479 write_line(
480 &mut stream,
481 &Request::Hello {
482 nonce: nonce.clone(),
483 },
484 )?;
485
486 let proved = matches!(
491 read_response(&mut reader),
492 Ok(Response::Hello { ref proof }) if token_matches(proof, &proof_for(&endpoint.token, &nonce))
493 );
494
495 if !proved {
496 return Err(
497 "Whatever is listening on the agent's port could not prove it is the agent; \
498 refusing to send anything to it. Run `lit agent stop` and start a new one."
499 .to_string(),
500 );
501 }
502
503 write_line(&mut stream, req)?;
506 read_response(&mut reader)
507}
508
509fn token() -> Result<String, String> {
510 Ok(Endpoint::load()?.token)
511}
512
513pub fn ping() -> Result<(), String> {
515 match request(&Request::Status { token: token()? })? {
516 Response::Status { .. } => Ok(()),
517 _ => Err("Agent did not answer a status request".to_string()),
518 }
519}
520
521pub fn get(repo: &str) -> Option<Zeroizing<String>> {
525 let token = token().ok()?;
526 match request(&Request::Get {
527 token,
528 repo: repo.to_string(),
529 })
530 .ok()?
531 {
532 Response::Passphrase { passphrase } => Some(Zeroizing::new(passphrase)),
533 _ => None,
534 }
535}
536
537pub fn put(repo: &str, passphrase: &str) -> Result<(), String> {
538 match request(&Request::Put {
539 token: token()?,
540 repo: repo.to_string(),
541 passphrase: passphrase.to_string(),
542 })? {
543 Response::Ok => Ok(()),
544 other => Err(format!("Agent refused to store the passphrase: {other:?}")),
545 }
546}
547
548pub fn drop_entry(repo: Option<&str>) -> Result<(), String> {
549 match request(&Request::Drop {
550 token: token()?,
551 repo: repo.map(|r| r.to_string()),
552 })? {
553 Response::Ok => Ok(()),
554 other => Err(format!("Agent refused: {other:?}")),
555 }
556}
557
558pub fn status() -> Result<(usize, u64), String> {
559 match request(&Request::Status { token: token()? })? {
560 Response::Status {
561 entries,
562 idle_timeout_secs,
563 } => Ok((entries, idle_timeout_secs)),
564 other => Err(format!("Agent refused: {other:?}")),
565 }
566}
567
568pub fn shutdown() -> Result<(), String> {
569 let result = request(&Request::Shutdown { token: token()? });
570 Endpoint::remove();
573 match result? {
574 Response::Ok => Ok(()),
575 other => Err(format!("Agent refused to stop: {other:?}")),
576 }
577}
578
579#[cfg(test)]
580mod tests {
581 use super::*;
582
583 #[test]
584 fn test_entries_expire_when_idle() {
585 let mut store = Store::new(Duration::from_millis(50));
586 store.put("repo".to_string(), "hunter2".to_string());
587 assert!(store.get("repo").is_some());
588
589 std::thread::sleep(Duration::from_millis(80));
590 assert!(
591 store.get("repo").is_none(),
592 "an entry left alone past the timeout should be gone"
593 );
594 assert!(store.is_empty());
595 }
596
597 #[test]
598 fn test_use_refreshes_the_timeout() {
599 let mut store = Store::new(Duration::from_millis(120));
602 store.put("repo".to_string(), "hunter2".to_string());
603
604 for _ in 0..4 {
605 std::thread::sleep(Duration::from_millis(50));
606 assert!(store.get("repo").is_some(), "use should keep it alive");
607 }
608 }
609
610 #[test]
611 fn test_drop_all_forgets_everything() {
612 let mut store = Store::new(Duration::from_secs(60));
613 store.put("a".to_string(), "one".to_string());
614 store.put("b".to_string(), "two".to_string());
615 assert_eq!(store.len(), 2);
616
617 store.drop_all();
618 assert!(store.is_empty());
619 }
620
621 #[test]
622 fn test_token_comparison_rejects_wrong_and_short_tokens() {
623 let real = generate_token();
624 assert!(token_matches(&real, &real));
625 assert!(!token_matches("", &real));
626 assert!(!token_matches(&real[..real.len() - 1], &real));
627
628 let mut wrong = real.clone();
629 let last = if wrong.ends_with('a') { 'b' } else { 'a' };
631 wrong.pop();
632 wrong.push(last);
633 assert!(!token_matches(&wrong, &real));
634 }
635
636 #[test]
637 fn test_generated_tokens_differ() {
638 assert_ne!(generate_token(), generate_token());
639 }
640
641 #[test]
644 fn test_wrong_token_is_denied_and_changes_nothing() {
645 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
646 let real = generate_token();
647
648 let (resp, stop) = respond_to(
649 Request::Put {
650 token: "not-the-token".to_string(),
651 repo: "repo".to_string(),
652 passphrase: "hunter2".to_string(),
653 },
654 &real,
655 &store,
656 );
657
658 assert!(matches!(resp, Response::Denied));
659 assert!(!stop);
660 assert!(
661 store.lock().unwrap().is_empty(),
662 "an unauthenticated Put must store nothing"
663 );
664 }
665
666 #[test]
669 fn test_handshake_proves_the_peer_holds_the_token() {
670 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
671 let real = generate_token();
672 let nonce = generate_token();
673
674 let (resp, _) = respond_to(
675 Request::Hello {
676 nonce: nonce.clone(),
677 },
678 &real,
679 &store,
680 );
681
682 let proof = match resp {
683 Response::Hello { proof } => proof,
684 other => panic!("expected a proof, got {other:?}"),
685 };
686 assert!(token_matches(&proof, &proof_for(&real, &nonce)));
687
688 let impostor = generate_token();
690 assert!(!token_matches(&proof, &proof_for(&impostor, &nonce)));
691
692 let other_nonce = generate_token();
694 assert!(!token_matches(&proof, &proof_for(&real, &other_nonce)));
695 }
696
697 #[test]
700 fn test_hello_carries_no_token_but_grants_nothing() {
701 assert!(token_of(&Request::Hello {
702 nonce: "n".to_string()
703 })
704 .is_none());
705
706 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
707 let (resp, stop) = apply(
708 Request::Hello {
709 nonce: "n".to_string(),
710 },
711 &store,
712 );
713 assert!(matches!(resp, Response::Denied));
714 assert!(!stop);
715 }
716
717 #[test]
718 fn test_put_then_get_round_trips_through_apply() {
719 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
720
721 let (resp, stop) = apply(
722 Request::Put {
723 token: String::new(),
724 repo: "repo".to_string(),
725 passphrase: "hunter2".to_string(),
726 },
727 &store,
728 );
729 assert!(matches!(resp, Response::Ok));
730 assert!(!stop);
731
732 let (resp, _) = apply(
733 Request::Get {
734 token: String::new(),
735 repo: "repo".to_string(),
736 },
737 &store,
738 );
739 match resp {
740 Response::Passphrase { passphrase } => assert_eq!(passphrase, "hunter2"),
741 other => panic!("expected the passphrase back, got {other:?}"),
742 }
743
744 let (_, stop) = apply(
745 Request::Shutdown {
746 token: String::new(),
747 },
748 &store,
749 );
750 assert!(stop, "shutdown should stop the agent");
751 assert!(
752 store.lock().unwrap().is_empty(),
753 "shutdown should clear what it held"
754 );
755 }
756
757 #[test]
758 fn test_get_for_unknown_repo_is_missing_not_an_error() {
759 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
760 let (resp, _) = apply(
761 Request::Get {
762 token: String::new(),
763 repo: "never-stored".to_string(),
764 },
765 &store,
766 );
767 assert!(matches!(resp, Response::Missing));
768 }
769}