Skip to main content

volli_agent/
lib.rs

1#![cfg_attr(test, allow(unused_crate_dependencies))]
2
3use ed25519_dalek::VerifyingKey;
4use eyre::Report;
5use eyre::eyre;
6use hex::decode as hex_decode;
7use mac_address::get_mac_address;
8use quinn::{ClientConfig, Endpoint};
9use rustls::{Certificate, RootCertStore, ServerName};
10use sha2::{Digest, Sha256};
11use std::convert::TryFrom;
12use std::fs;
13use std::net::ToSocketAddrs;
14use std::path::PathBuf;
15use std::sync::Arc;
16use tokio::net::TcpStream;
17use tokio_rustls::TlsConnector;
18use volli_core::{DEFAULT_QUIC_PORT, DEFAULT_TCP_PORT, Message};
19use volli_transport::{QuicTransport, TcpTransport, Transport};
20
21#[derive(Clone, Debug)]
22pub enum Protocol {
23    Quic,
24    Tcp,
25}
26
27#[derive(Debug, Clone)]
28pub enum Role {
29    Agent,
30    Coordinator,
31}
32
33#[derive(Clone, Debug)]
34pub struct AgentConfig {
35    pub host: String,
36    pub quic_port: u16,
37    pub tcp_port: u16,
38    pub protocol: Option<Protocol>,
39    pub token: String,
40    pub fingerprint: String,
41    pub cert: Vec<u8>,
42    pub role: Role,
43}
44
45impl Default for AgentConfig {
46    fn default() -> Self {
47        Self {
48            host: "127.0.0.1".into(),
49            quic_port: DEFAULT_QUIC_PORT,
50            tcp_port: DEFAULT_TCP_PORT,
51            protocol: None,
52            token: String::new(),
53            fingerprint: String::new(),
54            cert: Vec::new(),
55            role: Role::Agent,
56        }
57    }
58}
59
60fn default_agent_dir() -> PathBuf {
61    let mut base = volli_core::config_dir();
62    base.push("profiles");
63    base.push("agent");
64    base
65}
66
67pub fn agent_dir(profile: Option<&str>) -> PathBuf {
68    let mut dir = default_agent_dir();
69    if let Some(p) = profile {
70        dir.push(p);
71    }
72    dir
73}
74
75#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
76pub struct PeerEntry {
77    pub host: String,
78    pub quic_port: u16,
79    pub tcp_port: u16,
80    #[serde(default)]
81    pub last_ok: Option<u64>,
82    #[serde(default)]
83    pub last_fail: Option<u64>,
84}
85
86pub fn load_peers(profile: &str) -> Result<Vec<PeerEntry>, Report> {
87    let path = agent_dir(Some(profile)).join("peers.json");
88    match fs::read_to_string(path) {
89        Ok(s) => Ok(serde_json::from_str(&s)?),
90        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Vec::new()),
91        Err(e) => Err(e.into()),
92    }
93}
94
95pub fn save_peers(profile: &str, peers: &[PeerEntry]) -> Result<(), Report> {
96    let dir = agent_dir(Some(profile));
97    fs::create_dir_all(&dir)?;
98    fs::write(dir.join("peers.json"), serde_json::to_string(peers)?)?;
99    Ok(())
100}
101
102pub fn add_peer(profile: &str, peer: PeerEntry) -> Result<(), Report> {
103    let mut peers = load_peers(profile)?;
104    if !peers.iter().any(|p| {
105        p.host == peer.host && p.tcp_port == peer.tcp_port && p.quic_port == peer.quic_port
106    }) {
107        peers.push(peer);
108        save_peers(profile, &peers)?;
109    }
110    Ok(())
111}
112
113fn now_secs() -> u64 {
114    use std::time::{SystemTime, UNIX_EPOCH};
115    SystemTime::now()
116        .duration_since(UNIX_EPOCH)
117        .unwrap()
118        .as_secs()
119}
120
121pub fn save_state(profile: &str, secret: &str) -> Result<(), Report> {
122    let dir = agent_dir(Some(profile));
123    fs::create_dir_all(&dir)?;
124    fs::write(dir.join("agent_state"), secret)?;
125    Ok(())
126}
127
128pub fn load_state(profile: &str) -> Result<Option<String>, Report> {
129    let path = agent_dir(Some(profile)).join("agent_state");
130    match fs::read_to_string(path) {
131        Ok(s) => Ok(Some(s)),
132        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
133        Err(e) => Err(e.into()),
134    }
135}
136
137pub fn list_profiles() -> Result<Vec<String>, Report> {
138    let base = default_agent_dir();
139    let mut profiles = Vec::new();
140    if base.exists() {
141        for entry in fs::read_dir(base)? {
142            let entry = entry?;
143            if entry.path().is_dir() {
144                if let Some(name) = entry.file_name().to_str() {
145                    profiles.push(name.to_string());
146                }
147            }
148        }
149    }
150    profiles.sort();
151    Ok(profiles)
152}
153
154pub fn delete_profile(profile: &str) -> Result<(), Report> {
155    let dir = agent_dir(Some(profile));
156    if dir.exists() {
157        fs::remove_dir_all(dir)?;
158    }
159    Ok(())
160}
161
162pub fn profile_exists(profile: &str) -> bool {
163    agent_dir(Some(profile)).exists()
164}
165
166pub fn rename_profile(old: &str, new: &str) -> Result<(), Report> {
167    let src = agent_dir(Some(old));
168    let dst = agent_dir(Some(new));
169    if !src.exists() {
170        return Err(eyre!("profile not found"));
171    }
172    if dst.exists() {
173        return Err(eyre!("profile exists"));
174    }
175    fs::create_dir_all(dst.parent().unwrap())?;
176    fs::rename(src, dst)?;
177    Ok(())
178}
179
180#[derive(serde::Serialize, serde::Deserialize)]
181pub struct AgentProfileExport {
182    pub name: String,
183    pub secret: String,
184    #[serde(default)]
185    pub peers: Vec<PeerEntry>,
186}
187
188pub fn export_profile(profile: &str) -> Result<String, Report> {
189    let secret = load_state(profile)?.ok_or_else(|| eyre!("profile not found"))?;
190    let peers = load_peers(profile).unwrap_or_default();
191    let exp = AgentProfileExport {
192        name: profile.to_string(),
193        secret,
194        peers,
195    };
196    Ok(serde_yaml::to_string(&exp)?)
197}
198
199pub fn import_profile(yaml: &str, name: Option<&str>, force: bool) -> Result<String, Report> {
200    let mut exp: AgentProfileExport = serde_yaml::from_str(yaml)?;
201    if let Some(n) = name {
202        exp.name = n.to_string();
203    }
204    if profile_exists(&exp.name) && !force {
205        return Err(eyre!("profile exists"));
206    }
207    save_state(&exp.name, &exp.secret)?;
208    if !exp.peers.is_empty() {
209        save_peers(&exp.name, &exp.peers)?;
210    }
211    Ok(exp.name)
212}
213
214fn configure_client(cert: &[u8], alpn: &str) -> Result<ClientConfig, Report> {
215    let mut roots = rustls::RootCertStore::empty();
216    roots.add(&Certificate(cert.to_vec()))?;
217    let mut crypto = rustls::ClientConfig::builder()
218        .with_safe_defaults()
219        .with_root_certificates(roots)
220        .with_no_client_auth();
221    crypto.alpn_protocols = vec![alpn.as_bytes().to_vec()];
222    Ok(ClientConfig::new(Arc::new(crypto)))
223}
224
225pub fn apply_secret(cfg: &mut AgentConfig, profile: &str, secret: &str) -> Result<(), Report> {
226    let bs = volli_core::BootstrapSecret::decode(secret)?;
227    let fp = hex::encode(Sha256::digest(&bs.cert));
228    cfg.host = bs.host.clone();
229    cfg.quic_port = bs.quic_port;
230    cfg.tcp_port = bs.tcp_port;
231    cfg.token = volli_core::token::encode_token(&bs.token)?;
232    cfg.fingerprint = fp;
233    cfg.cert = bs.cert;
234    let peer = PeerEntry {
235        host: cfg.host.clone(),
236        quic_port: cfg.quic_port,
237        tcp_port: cfg.tcp_port,
238        last_ok: None,
239        last_fail: None,
240    };
241    let _ = add_peer(profile, peer);
242    Ok(())
243}
244
245pub async fn run(
246    mut config: AgentConfig,
247    profile: String,
248    on_first_connect: Option<Box<dyn FnOnce() + Send>>,
249) -> Result<(), Report> {
250    let proto_pref = config.protocol.take();
251    let mut backoff = 1u64;
252    let mut connect_cb = on_first_connect;
253    let mut peers = load_peers(&profile).unwrap_or_default();
254    if peers.is_empty() {
255        peers.push(PeerEntry {
256            host: config.host.clone(),
257            quic_port: config.quic_port,
258            tcp_port: config.tcp_port,
259            last_ok: None,
260            last_fail: None,
261        });
262        save_peers(&profile, &peers).ok();
263    }
264    let mut idx = 0usize;
265    loop {
266        let peer = peers.get(idx).cloned().unwrap();
267        config.host = peer.host.clone();
268        config.quic_port = peer.quic_port;
269        config.tcp_port = peer.tcp_port;
270        let res = match proto_pref.as_ref().unwrap_or(&Protocol::Quic) {
271            Protocol::Quic => match connect_quic(&config).await {
272                Ok((tr, p)) => {
273                    peers[idx].last_ok = Some(now_secs());
274                    save_peers(&profile, &peers).ok();
275                    handle_agent(tr, &config, &p, &mut connect_cb).await
276                }
277                Err(e) => {
278                    tracing::warn!("quic connect error: {}", e);
279                    match connect_tcp(&config).await {
280                        Ok((tr, p)) => {
281                            peers[idx].last_ok = Some(now_secs());
282                            save_peers(&profile, &peers).ok();
283                            handle_agent(tr, &config, &p, &mut connect_cb).await
284                        }
285                        Err(e) => Err(e),
286                    }
287                }
288            },
289            Protocol::Tcp => match connect_tcp(&config).await {
290                Ok((tr, p)) => {
291                    peers[idx].last_ok = Some(now_secs());
292                    save_peers(&profile, &peers).ok();
293                    handle_agent(tr, &config, &p, &mut connect_cb).await
294                }
295                Err(e) => Err(e),
296            },
297        };
298
299        match res {
300            Ok(_) => {
301                backoff = 1;
302                idx = 0;
303            }
304            Err(e) => {
305                tracing::error!("connection error: {}", e);
306                peers[idx].last_fail = Some(now_secs());
307                save_peers(&profile, &peers).ok();
308                backoff = (backoff * 2).min(32);
309                idx = (idx + 1) % peers.len();
310            }
311        }
312
313        tokio::time::sleep(std::time::Duration::from_secs(backoff)).await;
314    }
315}
316
317async fn handle_agent(
318    mut transport: Box<dyn Transport>,
319    cfg: &AgentConfig,
320    peer: &str,
321    on_first_connect: &mut Option<Box<dyn FnOnce() + Send>>,
322) -> Result<(), Report> {
323    transport
324        .send(&Message::Auth {
325            token: cfg.token.clone(),
326        })
327        .await?;
328    let mut authed = false;
329    while let Some(msg) = transport.recv().await? {
330        match msg {
331            Message::AuthOk => {
332                authed = true;
333            }
334            Message::Hello {
335                coord_id,
336                nonce,
337                sig,
338            } => {
339                if !authed {
340                    return Err(eyre!("handshake before auth"));
341                }
342                let pk_bytes = hex_decode(&coord_id)?;
343                let arr: [u8; 32] = pk_bytes
344                    .as_slice()
345                    .try_into()
346                    .map_err(|_| eyre!("bad coord id"))?;
347                let pk = VerifyingKey::from_bytes(&arr)?;
348                volli_core::handshake::verify_nonce(&pk, &nonce, &sig)?;
349                transport
350                    .send(&Message::Welcome {
351                        coord_id,
352                        nonce,
353                        sig,
354                    })
355                    .await?;
356                tracing::info!(target: "connection", %peer, role=?cfg.role, "authenticated");
357                if let Some(cb) = on_first_connect.take() {
358                    cb();
359                }
360            }
361            Message::AuthErr => return Err(eyre!("authentication failed")),
362            Message::Ping => {
363                tracing::info!(target: "connection", %peer, "received ping");
364                let mac = get_mac_address()
365                    .ok()
366                    .flatten()
367                    .map(|m| m.to_string())
368                    .unwrap_or_default();
369                tracing::info!(target: "connection", %peer, "sending pong");
370                transport.send(&Message::Pong { mac }).await?;
371            }
372            _ => {}
373        }
374    }
375    Ok(())
376}
377
378pub async fn connect_tcp(cfg: &AgentConfig) -> Result<(Box<dyn Transport>, String), Report> {
379    let addr = format!("{}:{}", cfg.host, cfg.tcp_port);
380    let mut addrs = addr.to_socket_addrs()?;
381    let addr = addrs
382        .find(|a| a.is_ipv4())
383        .or_else(|| addrs.next())
384        .ok_or_else(|| eyre!("invalid addr"))?;
385    let stream = TcpStream::connect(addr).await?;
386    let alpn = match cfg.role {
387        Role::Agent => "volli/agent",
388        Role::Coordinator => "volli/coord",
389    };
390    let mut roots = RootCertStore::empty();
391    roots.add(&Certificate(cfg.cert.clone()))?;
392    let mut root = rustls::ClientConfig::builder()
393        .with_safe_defaults()
394        .with_root_certificates(roots)
395        .with_no_client_auth();
396    root.alpn_protocols = vec![alpn.as_bytes().to_vec()];
397    let connector = TlsConnector::from(Arc::new(root));
398    let tls = connector
399        .connect(ServerName::try_from("volli")?, stream)
400        .await?;
401    if let Some(certs) = tls.get_ref().1.peer_certificates() {
402        if let Some(cert) = certs.first() {
403            let hash = Sha256::digest(&cert.0);
404            if hex::encode(hash) != cfg.fingerprint {
405                return Err(eyre!("server fingerprint mismatch"));
406            }
407        }
408    }
409    let peer = tls.get_ref().0.peer_addr()?.to_string();
410    Ok((Box::new(TcpTransport::new(tls)), peer))
411}
412
413pub async fn connect_quic(cfg: &AgentConfig) -> Result<(Box<dyn Transport>, String), Report> {
414    let addr = format!("{}:{}", cfg.host, cfg.quic_port);
415    let mut addrs = addr.to_socket_addrs()?;
416    let addr = addrs
417        .find(|a| a.is_ipv4())
418        .or_else(|| addrs.next())
419        .ok_or_else(|| eyre!("invalid addr"))?;
420    let mut endpoint = Endpoint::client("0.0.0.0:0".parse()?)?;
421    let alpn = match cfg.role {
422        Role::Agent => "volli/agent",
423        Role::Coordinator => "volli/coord",
424    };
425    let quinn_cfg = configure_client(&cfg.cert, alpn)?;
426    endpoint.set_default_client_config(quinn_cfg);
427    let connection = endpoint.connect(addr, "volli")?.await?;
428    if let Some(identity) = connection.peer_identity() {
429        if let Ok(certs) = identity.downcast::<Vec<Certificate>>() {
430            if let Some(cert) = certs.first() {
431                let hash = Sha256::digest(&cert.0);
432                if hex::encode(hash) != cfg.fingerprint {
433                    return Err(eyre!("server fingerprint mismatch"));
434                }
435            }
436        }
437    }
438    let peer = connection.remote_address().to_string();
439    let (send, recv) = connection.open_bi().await?;
440    Ok((Box::new(QuicTransport::new(send, recv)), peer))
441}