x0x/
dm_capability_service.rs

1//! Runtime service that publishes this agent's DM capability advert to the
2//! mesh-wide `x0x/caps/v1` topic and consumes peers' adverts into a
3//! shared [`crate::dm_capability::CapabilityStore`].
4
5use crate::dm::DmCapabilities;
6use crate::dm_capability::{
7    now_unix_ms, CapabilityAdvert, CapabilityStore, ADVERT_PUBLISH_INTERVAL_SECS,
8    DM_CAPABILITY_TOPIC,
9};
10use crate::error::{NetworkError, NetworkResult};
11use crate::gossip::{PubSubManager, SigningContext};
12use crate::identity::{AgentId, MachineId};
13use bytes::Bytes;
14use std::sync::Arc;
15use std::time::Duration;
16use tokio::task::JoinHandle;
17
18pub const ADVERT_PROTOCOL_VERSION: u16 = 1;
19
20const FIRST_PUBLISH_DELAY_MS: u64 = 250;
21
22/// Startup-burst schedule so late-joining peers catch our advert quickly.
23const STARTUP_BURST_INTERVALS_MS: &[u64] = &[5_000, 10_000, 20_000, 45_000];
24
25pub struct CapabilityAdvertService {
26    publisher: JoinHandle<()>,
27    subscriber: JoinHandle<()>,
28}
29
30impl CapabilityAdvertService {
31    #[allow(clippy::too_many_arguments)]
32    pub async fn spawn(
33        pubsub: Arc<PubSubManager>,
34        signing: Arc<SigningContext>,
35        self_agent_id: AgentId,
36        self_machine_id: MachineId,
37        caps_rx: tokio::sync::watch::Receiver<DmCapabilities>,
38        store: Arc<CapabilityStore>,
39        publish_interval: Duration,
40    ) -> NetworkResult<Self> {
41        let mut subscription = pubsub.subscribe(DM_CAPABILITY_TOPIC.to_string()).await;
42        let store_sub = Arc::clone(&store);
43        let self_agent_for_sub = self_agent_id;
44
45        let subscriber = tokio::spawn(async move {
46            while let Some(message) = subscription.recv().await {
47                let (pubsub_sender, sender_pubkey) =
48                    match (message.sender, message.sender_public_key.as_deref()) {
49                        (Some(s), Some(pk)) if message.verified => (s, pk.to_vec()),
50                        _ => continue,
51                    };
52                if pubsub_sender == self_agent_for_sub {
53                    continue;
54                }
55                let advert: CapabilityAdvert = match postcard::from_bytes(&message.payload) {
56                    Ok(a) => a,
57                    Err(_) => continue,
58                };
59                if advert.protocol_version != ADVERT_PROTOCOL_VERSION {
60                    continue;
61                }
62                if advert.agent_id != *pubsub_sender.as_bytes() {
63                    continue;
64                }
65                if !verify_advert_signature(&advert, &sender_pubkey) {
66                    continue;
67                }
68                store_sub.insert(
69                    AgentId(advert.agent_id),
70                    MachineId(advert.machine_id),
71                    advert.capabilities,
72                );
73                tracing::debug!(
74                    "cached capability advert from {}",
75                    hex::encode(advert.agent_id)
76                );
77            }
78            tracing::debug!("capability advert subscriber exited");
79        });
80
81        let publisher_pubsub = Arc::clone(&pubsub);
82        let publisher_signing = Arc::clone(&signing);
83        let mut publisher_caps_rx = caps_rx;
84        let publisher = tokio::spawn(async move {
85            tokio::time::sleep(Duration::from_millis(FIRST_PUBLISH_DELAY_MS)).await;
86            let mut burst_idx: usize = 0;
87            loop {
88                let caps_snapshot = publisher_caps_rx.borrow().clone();
89                match build_signed_advert(
90                    &publisher_signing,
91                    self_agent_id,
92                    self_machine_id,
93                    caps_snapshot,
94                ) {
95                    Ok(bytes) => {
96                        if let Err(e) = publisher_pubsub
97                            .publish(DM_CAPABILITY_TOPIC.to_string(), Bytes::from(bytes))
98                            .await
99                        {
100                            tracing::warn!("capability advert publish failed: {e}");
101                        } else {
102                            tracing::debug!("capability advert published");
103                        }
104                    }
105                    Err(e) => tracing::warn!("capability advert build failed: {e}"),
106                }
107                let next_delay = if burst_idx < STARTUP_BURST_INTERVALS_MS.len() {
108                    let d = Duration::from_millis(STARTUP_BURST_INTERVALS_MS[burst_idx]);
109                    burst_idx += 1;
110                    d
111                } else {
112                    publish_interval
113                };
114                tokio::select! {
115                    _ = tokio::time::sleep(next_delay) => {}
116                    res = publisher_caps_rx.changed() => {
117                        if res.is_ok() {
118                            tracing::debug!("capability advert upgraded; republishing");
119                            burst_idx = 0;
120                        }
121                    }
122                }
123            }
124        });
125
126        Ok(Self {
127            publisher,
128            subscriber,
129        })
130    }
131
132    pub async fn spawn_default(
133        pubsub: Arc<PubSubManager>,
134        signing: Arc<SigningContext>,
135        self_agent_id: AgentId,
136        self_machine_id: MachineId,
137        caps_rx: tokio::sync::watch::Receiver<DmCapabilities>,
138        store: Arc<CapabilityStore>,
139    ) -> NetworkResult<Self> {
140        Self::spawn(
141            pubsub,
142            signing,
143            self_agent_id,
144            self_machine_id,
145            caps_rx,
146            store,
147            Duration::from_secs(ADVERT_PUBLISH_INTERVAL_SECS),
148        )
149        .await
150    }
151
152    pub fn abort(&self) {
153        self.publisher.abort();
154        self.subscriber.abort();
155    }
156}
157
158impl Drop for CapabilityAdvertService {
159    fn drop(&mut self) {
160        self.abort();
161    }
162}
163
164pub fn build_signed_advert(
165    signing: &SigningContext,
166    self_agent_id: AgentId,
167    self_machine_id: MachineId,
168    capabilities: DmCapabilities,
169) -> NetworkResult<Vec<u8>> {
170    let mut advert = CapabilityAdvert {
171        protocol_version: ADVERT_PROTOCOL_VERSION,
172        agent_id: *self_agent_id.as_bytes(),
173        machine_id: *self_machine_id.as_bytes(),
174        created_at_unix_ms: now_unix_ms(),
175        capabilities,
176        signature: Vec::new(),
177    };
178    let signed_bytes = advert
179        .signed_bytes()
180        .map_err(|e| NetworkError::SerializationError(format!("advert sign-bytes: {e}")))?;
181    advert.signature = signing.sign(&signed_bytes)?;
182    postcard::to_stdvec(&advert)
183        .map_err(|e| NetworkError::SerializationError(format!("advert encode: {e}")))
184}
185
186pub fn verify_advert_signature(advert: &CapabilityAdvert, public_key_bytes: &[u8]) -> bool {
187    let signed_bytes = match advert.signed_bytes() {
188        Ok(b) => b,
189        Err(_) => return false,
190    };
191    let public_key = match ant_quic::MlDsaPublicKey::from_bytes(public_key_bytes) {
192        Ok(pk) => pk,
193        Err(_) => return false,
194    };
195    let derived = crate::identity::AgentId::from_public_key(&public_key);
196    if derived.0 != advert.agent_id {
197        return false;
198    }
199    let signature =
200        match ant_quic::crypto::raw_public_keys::pqc::MlDsaSignature::from_bytes(&advert.signature)
201        {
202            Ok(s) => s,
203            Err(_) => return false,
204        };
205    ant_quic::crypto::raw_public_keys::pqc::verify_with_ml_dsa(
206        &public_key,
207        &signed_bytes,
208        &signature,
209    )
210    .is_ok()
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use crate::identity::AgentKeypair;
217
218    #[test]
219    fn build_and_verify_advert_roundtrip() {
220        let kp = AgentKeypair::generate().expect("keygen");
221        let signing = SigningContext::from_keypair(&kp);
222        let agent_id = kp.agent_id();
223        let machine_id = MachineId([9u8; 32]);
224        let encoded = build_signed_advert(
225            &signing,
226            agent_id,
227            machine_id,
228            DmCapabilities::v1_gossip_ready(vec![0u8; 1184]),
229        )
230        .expect("build");
231        let advert: CapabilityAdvert = postcard::from_bytes(&encoded).expect("decode");
232        assert!(verify_advert_signature(&advert, &signing.public_key_bytes));
233    }
234
235    #[test]
236    fn verify_advert_rejects_tampered_signature() {
237        let kp = AgentKeypair::generate().expect("keygen");
238        let signing = SigningContext::from_keypair(&kp);
239        let encoded = build_signed_advert(
240            &signing,
241            kp.agent_id(),
242            MachineId([0u8; 32]),
243            DmCapabilities::v1_gossip_ready(vec![0u8; 1184]),
244        )
245        .expect("build");
246        let mut advert: CapabilityAdvert = postcard::from_bytes(&encoded).expect("decode");
247        advert.signature[0] ^= 0x01;
248        assert!(!verify_advert_signature(&advert, &signing.public_key_bytes));
249    }
250}
x0x/dm_capability_service.rs

x0x/
dm_capability_service.rs
