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)]
66#[serde(tag = "op", rename_all = "snake_case")]
67pub enum Request {
68 Hello { nonce: String },
71 Put {
73 token: String,
74 repo: String,
75 passphrase: String,
76 },
77 Get { token: String, repo: String },
79 Drop { token: String, repo: Option<String> },
81 Status { token: String },
83 Shutdown { token: String },
85}
86
87#[derive(Serialize, Deserialize, Debug)]
89#[serde(tag = "result", rename_all = "snake_case")]
90pub enum Response {
91 Hello {
93 proof: String,
94 },
95 Passphrase {
96 passphrase: String,
97 },
98 Missing,
100 Ok,
101 Status {
102 entries: usize,
103 idle_timeout_secs: u64,
104 },
105 Denied,
106 Malformed {
107 message: String,
108 },
109}
110
111struct Entry {
112 passphrase: Zeroizing<String>,
113 last_used: Instant,
114}
115
116pub struct Store {
121 entries: HashMap<String, Entry>,
122 idle_timeout: Duration,
123}
124
125impl Store {
126 pub fn new(idle_timeout: Duration) -> Self {
127 Store {
128 entries: HashMap::new(),
129 idle_timeout,
130 }
131 }
132
133 pub fn put(&mut self, repo: String, passphrase: String) {
134 self.entries.insert(
135 repo,
136 Entry {
137 passphrase: Zeroizing::new(passphrase),
138 last_used: Instant::now(),
139 },
140 );
141 }
142
143 pub fn get(&mut self, repo: &str) -> Option<Zeroizing<String>> {
145 self.expire();
146 let entry = self.entries.get_mut(repo)?;
147 entry.last_used = Instant::now();
148 Some(entry.passphrase.clone())
149 }
150
151 pub fn drop_one(&mut self, repo: &str) {
152 self.entries.remove(repo);
153 }
154
155 pub fn drop_all(&mut self) {
156 self.entries.clear();
157 }
158
159 pub fn len(&mut self) -> usize {
160 self.expire();
161 self.entries.len()
162 }
163
164 pub fn is_empty(&mut self) -> bool {
165 self.len() == 0
166 }
167
168 fn expire(&mut self) {
169 let timeout = self.idle_timeout;
170 self.entries.retain(|_, e| e.last_used.elapsed() < timeout);
171 }
172}
173
174#[derive(Serialize, Deserialize)]
178pub struct Endpoint {
179 pub port: u16,
180 pub token: String,
181 pub idle_timeout_secs: u64,
182}
183
184pub fn endpoint_path() -> Result<PathBuf, String> {
185 let home = dirs::home_dir().ok_or("Could not determine home directory")?;
186 Ok(home.join(".lit").join("agent.json"))
187}
188
189impl Endpoint {
190 pub fn load() -> Result<Endpoint, String> {
191 Self::load_from(&endpoint_path()?)
192 }
193
194 pub(crate) fn load_from(path: &std::path::Path) -> Result<Endpoint, String> {
195 let raw = std::fs::read(path)
196 .map_err(|_| "No agent is running (start one with `lit agent start`)".to_string())?;
197 serde_json::from_slice(&raw).map_err(|e| format!("Agent endpoint file is unreadable: {e}"))
198 }
199
200 fn save(&self) -> Result<(), String> {
201 self.save_to(&endpoint_path()?)
202 }
203
204 pub(crate) fn save_to(&self, path: &std::path::Path) -> Result<(), String> {
218 if let Some(parent) = path.parent() {
219 std::fs::create_dir_all(parent)
220 .map_err(|e| format!("Failed to create agent directory: {e}"))?;
221 let _ = crate::crypto::encryption::restrict_dir_to_owner(parent);
224 }
225
226 let raw =
227 serde_json::to_vec(self).map_err(|e| format!("Failed to encode endpoint: {e}"))?;
228
229 let temp = path.with_extension("tmp");
230 std::fs::write(&temp, raw).map_err(|e| format!("Failed to write endpoint: {e}"))?;
231 restrict_to_owner(&temp)?;
232
233 std::fs::rename(&temp, path).map_err(|e| {
234 let _ = std::fs::remove_file(&temp);
235 format!("Failed to write endpoint: {e}")
236 })?;
237
238 Ok(())
239 }
240
241 fn remove() {
242 if let Ok(path) = endpoint_path() {
243 let _ = std::fs::remove_file(path);
244 }
245 }
246}
247
248fn generate_token() -> String {
250 use aes_gcm::aead::rand_core::RngCore;
251 use aes_gcm::aead::OsRng;
252
253 let mut bytes = [0u8; 32];
254 OsRng.fill_bytes(&mut bytes);
255 hex::encode(bytes)
256}
257
258fn token_matches(presented: &str, expected: &str) -> bool {
261 let a = presented.as_bytes();
262 let b = expected.as_bytes();
263 if a.len() != b.len() {
264 return false;
265 }
266 a.ct_eq(b).into()
267}
268
269fn token_of(req: &Request) -> Option<&str> {
270 match req {
271 Request::Put { token, .. }
272 | Request::Get { token, .. }
273 | Request::Drop { token, .. }
274 | Request::Status { token }
275 | Request::Shutdown { token } => Some(token),
276 Request::Hello { .. } => None,
278 }
279}
280
281fn proof_for(token: &str, nonce: &str) -> String {
292 use hmac::{Hmac, Mac};
293 use sha2::Sha256;
294
295 let mut mac = <Hmac<Sha256> as Mac>::new_from_slice(token.as_bytes())
296 .expect("HMAC accepts keys of any length");
297 mac.update(nonce.as_bytes());
298 hex::encode(mac.finalize().into_bytes())
299}
300
301fn apply(req: Request, store: &Arc<Mutex<Store>>) -> (Response, bool) {
305 let mut store = match store.lock() {
306 Ok(s) => s,
307 Err(_) => {
308 return (
309 Response::Malformed {
310 message: "agent state is poisoned".to_string(),
311 },
312 false,
313 )
314 }
315 };
316
317 match req {
318 Request::Hello { .. } => (Response::Denied, false),
320 Request::Put {
321 repo, passphrase, ..
322 } => {
323 store.put(repo, passphrase);
324 (Response::Ok, false)
325 }
326 Request::Get { repo, .. } => match store.get(&repo) {
327 Some(p) => (
328 Response::Passphrase {
329 passphrase: p.to_string(),
330 },
331 false,
332 ),
333 None => (Response::Missing, false),
334 },
335 Request::Drop { repo, .. } => {
336 match repo {
337 Some(r) => store.drop_one(&r),
338 None => store.drop_all(),
339 }
340 (Response::Ok, false)
341 }
342 Request::Status { .. } => (
343 Response::Status {
344 entries: store.len(),
345 idle_timeout_secs: store.idle_timeout.as_secs(),
346 },
347 false,
348 ),
349 Request::Shutdown { .. } => {
350 store.drop_all();
351 (Response::Ok, true)
352 }
353 }
354}
355
356fn respond_to(req: Request, expected_token: &str, store: &Arc<Mutex<Store>>) -> (Response, bool) {
358 match req {
359 Request::Hello { nonce } => (
360 Response::Hello {
361 proof: proof_for(expected_token, &nonce),
362 },
363 false,
364 ),
365 other => match token_of(&other) {
366 Some(t) if token_matches(t, expected_token) => apply(other, store),
370 _ => (Response::Denied, false),
371 },
372 }
373}
374
375fn handle_connection(
379 stream: &mut TcpStream,
380 expected_token: &str,
381 store: &Arc<Mutex<Store>>,
382) -> bool {
383 let _ = stream.set_read_timeout(Some(Duration::from_secs(5)));
385 let _ = stream.set_write_timeout(Some(Duration::from_secs(5)));
386
387 let Ok(peer) = stream.try_clone() else {
388 return false;
389 };
390
391 let mut reader = BufReader::new(peer.take(MAX_REQUEST_BYTES));
394
395 for _ in 0..2 {
398 let mut line = String::new();
399 match reader.read_line(&mut line) {
400 Ok(0) | Err(_) => return false,
401 Ok(_) => {}
402 }
403
404 let (response, shutdown) = match serde_json::from_str::<Request>(line.trim()) {
405 Ok(req) => respond_to(req, expected_token, store),
406 Err(e) => (
407 Response::Malformed {
408 message: e.to_string(),
409 },
410 false,
411 ),
412 };
413
414 let closing = !matches!(response, Response::Hello { .. });
415
416 if let Ok(mut body) = serde_json::to_vec(&response) {
417 body.push(b'\n');
418 if stream.write_all(&body).is_err() {
419 return false;
420 }
421 let _ = stream.flush();
422 }
423
424 if closing {
426 return shutdown;
427 }
428 }
429
430 false
431}
432
433pub fn serve(idle_timeout: Duration) -> Result<(), String> {
435 if Endpoint::load().is_ok() && ping().is_ok() {
436 return Err("An agent is already running (`lit agent stop` to replace it)".to_string());
437 }
438
439 let listener = TcpListener::bind(SocketAddr::from((Ipv4Addr::LOCALHOST, 0)))
442 .map_err(|e| format!("Failed to bind agent socket: {e}"))?;
443 let port = listener
444 .local_addr()
445 .map_err(|e| format!("Failed to read agent port: {e}"))?
446 .port();
447
448 let token = generate_token();
449 Endpoint {
450 port,
451 token: token.clone(),
452 idle_timeout_secs: idle_timeout.as_secs(),
453 }
454 .save()?;
455
456 let store = Arc::new(Mutex::new(Store::new(idle_timeout)));
457
458 for incoming in listener.incoming() {
459 let mut stream = match incoming {
460 Ok(s) => s,
461 Err(_) => continue,
462 };
463 if handle_connection(&mut stream, &token, &store) {
464 break;
465 }
466 }
467
468 if let Ok(mut s) = store.lock() {
469 s.drop_all();
470 }
471 Endpoint::remove();
472 Ok(())
473}
474
475fn write_line(stream: &mut TcpStream, req: &Request) -> Result<(), String> {
476 let mut body = serde_json::to_vec(req).map_err(|e| format!("Failed to encode request: {e}"))?;
477 body.push(b'\n');
478 stream
479 .write_all(&body)
480 .map_err(|e| format!("Failed to reach agent: {e}"))
481}
482
483fn read_response(reader: &mut impl BufRead) -> Result<Response, String> {
484 let mut line = String::new();
485 reader
486 .read_line(&mut line)
487 .map_err(|e| format!("Failed to read agent reply: {e}"))?;
488 serde_json::from_str(line.trim()).map_err(|e| format!("Agent sent an unreadable reply: {e}"))
489}
490
491fn request(req: &Request) -> Result<Response, String> {
499 let endpoint = Endpoint::load()?;
500 let mut stream = TcpStream::connect(SocketAddr::from((Ipv4Addr::LOCALHOST, endpoint.port)))
501 .map_err(|_| "No agent is running (start one with `lit agent start`)".to_string())?;
502 stream
503 .set_read_timeout(Some(Duration::from_secs(5)))
504 .map_err(|e| format!("Failed to configure agent socket: {e}"))?;
505 stream
506 .set_write_timeout(Some(Duration::from_secs(5)))
507 .map_err(|e| format!("Failed to configure agent socket: {e}"))?;
508
509 let peer = stream
510 .try_clone()
511 .map_err(|e| format!("Failed to read from agent: {e}"))?;
512 let mut reader = BufReader::new(peer.take(MAX_REQUEST_BYTES));
513
514 let nonce = generate_token();
515 write_line(
516 &mut stream,
517 &Request::Hello {
518 nonce: nonce.clone(),
519 },
520 )?;
521
522 let proved = matches!(
527 read_response(&mut reader),
528 Ok(Response::Hello { ref proof }) if token_matches(proof, &proof_for(&endpoint.token, &nonce))
529 );
530
531 if !proved {
532 return Err(
533 "Whatever is listening on the agent's port could not prove it is the agent; \
534 refusing to send anything to it. Run `lit agent stop` and start a new one."
535 .to_string(),
536 );
537 }
538
539 write_line(&mut stream, req)?;
542 read_response(&mut reader)
543}
544
545fn token() -> Result<String, String> {
546 Ok(Endpoint::load()?.token)
547}
548
549pub fn ping() -> Result<(), String> {
551 match request(&Request::Status { token: token()? })? {
552 Response::Status { .. } => Ok(()),
553 _ => Err("Agent did not answer a status request".to_string()),
554 }
555}
556
557pub fn get(repo: &str) -> Option<Zeroizing<String>> {
561 let token = token().ok()?;
562 match request(&Request::Get {
563 token,
564 repo: repo.to_string(),
565 })
566 .ok()?
567 {
568 Response::Passphrase { passphrase } => Some(Zeroizing::new(passphrase)),
569 _ => None,
570 }
571}
572
573pub fn put(repo: &str, passphrase: &str) -> Result<(), String> {
574 match request(&Request::Put {
575 token: token()?,
576 repo: repo.to_string(),
577 passphrase: passphrase.to_string(),
578 })? {
579 Response::Ok => Ok(()),
580 other => Err(format!("Agent refused to store the passphrase: {other:?}")),
581 }
582}
583
584pub fn drop_entry(repo: Option<&str>) -> Result<(), String> {
585 match request(&Request::Drop {
586 token: token()?,
587 repo: repo.map(|r| r.to_string()),
588 })? {
589 Response::Ok => Ok(()),
590 other => Err(format!("Agent refused: {other:?}")),
591 }
592}
593
594pub fn status() -> Result<(usize, u64), String> {
595 match request(&Request::Status { token: token()? })? {
596 Response::Status {
597 entries,
598 idle_timeout_secs,
599 } => Ok((entries, idle_timeout_secs)),
600 other => Err(format!("Agent refused: {other:?}")),
601 }
602}
603
604pub fn shutdown() -> Result<(), String> {
605 match request(&Request::Shutdown { token: token()? }) {
606 Ok(Response::Ok) => {
607 Endpoint::remove();
610 Ok(())
611 }
612 Ok(other) => Err(format!("Agent refused to stop: {other:?}")),
613
614 Err(e) if e.contains("No agent is running") || e.contains("could not prove") => {
618 Endpoint::remove();
619 Err(e)
620 }
621
622 Err(e) => Err(format!(
628 "{e}. The agent may still be running and holding a passphrase; \
629 its endpoint file has been left in place so it can be reached again."
630 )),
631 }
632}
633
634#[cfg(test)]
635mod tests {
636 use super::*;
637
638 #[test]
639 fn test_entries_expire_when_idle() {
640 let mut store = Store::new(Duration::from_millis(50));
641 store.put("repo".to_string(), "hunter2".to_string());
642 assert!(store.get("repo").is_some());
643
644 std::thread::sleep(Duration::from_millis(80));
645 assert!(
646 store.get("repo").is_none(),
647 "an entry left alone past the timeout should be gone"
648 );
649 assert!(store.is_empty());
650 }
651
652 #[test]
653 fn test_use_refreshes_the_timeout() {
654 let mut store = Store::new(Duration::from_millis(120));
657 store.put("repo".to_string(), "hunter2".to_string());
658
659 for _ in 0..4 {
660 std::thread::sleep(Duration::from_millis(50));
661 assert!(store.get("repo").is_some(), "use should keep it alive");
662 }
663 }
664
665 #[test]
666 fn test_drop_all_forgets_everything() {
667 let mut store = Store::new(Duration::from_secs(60));
668 store.put("a".to_string(), "one".to_string());
669 store.put("b".to_string(), "two".to_string());
670 assert_eq!(store.len(), 2);
671
672 store.drop_all();
673 assert!(store.is_empty());
674 }
675
676 #[test]
677 fn test_token_comparison_rejects_wrong_and_short_tokens() {
678 let real = generate_token();
679 assert!(token_matches(&real, &real));
680 assert!(!token_matches("", &real));
681 assert!(!token_matches(&real[..real.len() - 1], &real));
682
683 let mut wrong = real.clone();
684 let last = if wrong.ends_with('a') { 'b' } else { 'a' };
686 wrong.pop();
687 wrong.push(last);
688 assert!(!token_matches(&wrong, &real));
689 }
690
691 #[test]
692 fn test_generated_tokens_differ() {
693 assert_ne!(generate_token(), generate_token());
694 }
695
696 #[test]
699 fn test_wrong_token_is_denied_and_changes_nothing() {
700 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
701 let real = generate_token();
702
703 let (resp, stop) = respond_to(
704 Request::Put {
705 token: "not-the-token".to_string(),
706 repo: "repo".to_string(),
707 passphrase: "hunter2".to_string(),
708 },
709 &real,
710 &store,
711 );
712
713 assert!(matches!(resp, Response::Denied));
714 assert!(!stop);
715 assert!(
716 store.lock().unwrap().is_empty(),
717 "an unauthenticated Put must store nothing"
718 );
719 }
720
721 #[test]
724 fn test_handshake_proves_the_peer_holds_the_token() {
725 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
726 let real = generate_token();
727 let nonce = generate_token();
728
729 let (resp, _) = respond_to(
730 Request::Hello {
731 nonce: nonce.clone(),
732 },
733 &real,
734 &store,
735 );
736
737 let proof = match resp {
738 Response::Hello { proof } => proof,
739 other => panic!("expected a proof, got {other:?}"),
740 };
741 assert!(token_matches(&proof, &proof_for(&real, &nonce)));
742
743 let impostor = generate_token();
745 assert!(!token_matches(&proof, &proof_for(&impostor, &nonce)));
746
747 let other_nonce = generate_token();
749 assert!(!token_matches(&proof, &proof_for(&real, &other_nonce)));
750 }
751
752 #[test]
755 fn test_hello_carries_no_token_but_grants_nothing() {
756 assert!(token_of(&Request::Hello {
757 nonce: "n".to_string()
758 })
759 .is_none());
760
761 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
762 let (resp, stop) = apply(
763 Request::Hello {
764 nonce: "n".to_string(),
765 },
766 &store,
767 );
768 assert!(matches!(resp, Response::Denied));
769 assert!(!stop);
770 }
771
772 #[test]
773 fn test_put_then_get_round_trips_through_apply() {
774 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
775
776 let (resp, stop) = apply(
777 Request::Put {
778 token: String::new(),
779 repo: "repo".to_string(),
780 passphrase: "hunter2".to_string(),
781 },
782 &store,
783 );
784 assert!(matches!(resp, Response::Ok));
785 assert!(!stop);
786
787 let (resp, _) = apply(
788 Request::Get {
789 token: String::new(),
790 repo: "repo".to_string(),
791 },
792 &store,
793 );
794 match resp {
795 Response::Passphrase { passphrase } => assert_eq!(passphrase, "hunter2"),
796 other => panic!("expected the passphrase back, got {other:?}"),
797 }
798
799 let (_, stop) = apply(
800 Request::Shutdown {
801 token: String::new(),
802 },
803 &store,
804 );
805 assert!(stop, "shutdown should stop the agent");
806 assert!(
807 store.lock().unwrap().is_empty(),
808 "shutdown should clear what it held"
809 );
810 }
811
812 #[test]
820 fn test_endpoint_is_written_restricted_and_by_rename() {
821 let dir = tempfile::tempdir().unwrap();
822 let path = dir.path().join("nested").join("agent.json");
823
824 let endpoint = Endpoint {
825 port: 4242,
826 token: generate_token(),
827 idle_timeout_secs: 900,
828 };
829 endpoint.save_to(&path).unwrap();
830
831 let read_back = Endpoint::load_from(&path).unwrap();
832 assert_eq!(read_back.port, endpoint.port);
833 assert_eq!(read_back.token, endpoint.token);
834
835 assert!(
836 !path.with_extension("tmp").exists(),
837 "the temporary endpoint file was left behind"
838 );
839
840 #[cfg(unix)]
841 {
842 use std::os::unix::fs::PermissionsExt;
843 let mode = std::fs::metadata(&path).unwrap().permissions().mode();
844 assert_eq!(mode & 0o777, 0o600, "the token file is readable by others");
845
846 let dir_mode = std::fs::metadata(path.parent().unwrap())
847 .unwrap()
848 .permissions()
849 .mode();
850 assert_eq!(dir_mode & 0o777, 0o700, "the agent directory is listable");
851 }
852 }
853
854 #[test]
857 fn test_endpoint_can_be_replaced() {
858 let dir = tempfile::tempdir().unwrap();
859 let path = dir.path().join("agent.json");
860
861 for port in [1111u16, 2222] {
862 Endpoint {
863 port,
864 token: generate_token(),
865 idle_timeout_secs: 60,
866 }
867 .save_to(&path)
868 .unwrap();
869 assert_eq!(Endpoint::load_from(&path).unwrap().port, port);
870 }
871 }
872
873 #[test]
874 fn test_get_for_unknown_repo_is_missing_not_an_error() {
875 let store = Arc::new(Mutex::new(Store::new(Duration::from_secs(60))));
876 let (resp, _) = apply(
877 Request::Get {
878 token: String::new(),
879 repo: "never-stored".to_string(),
880 },
881 &store,
882 );
883 assert!(matches!(resp, Response::Missing));
884 }
885}