mesh_llm_protocol/protocol/
mod.rs

1pub mod convert;
2pub mod v0;
3use anyhow::Result;
4pub use convert::*;
5use iroh::endpoint::{ConnectOptions, Connection};
6use iroh::{Endpoint, EndpointAddr};
7use prost::Message;
8pub use v0::*;
9pub const ALPN_CONTROL_V1: &[u8] = b"mesh-llm-control/1";
10pub const ALPN_V1: &[u8] = b"mesh-llm/1";
11pub const NODE_PROTOCOL_GENERATION: u32 = 1;
12pub const MAX_CONTROL_FRAME_BYTES: usize = 8 * 1024 * 1024;
13
14pub const STREAM_GOSSIP: u8 = 0x01;
15pub const STREAM_TUNNEL: u8 = 0x02;
16pub const STREAM_TUNNEL_MAP: u8 = 0x03;
17pub const STREAM_TUNNEL_HTTP: u8 = 0x04;
18pub const STREAM_ROUTE_REQUEST: u8 = 0x05;
19pub const STREAM_PEER_DOWN: u8 = 0x06;
20pub const STREAM_PEER_LEAVING: u8 = 0x07;
21pub const STREAM_PLUGIN_CHANNEL: u8 = 0x08;
22pub const STREAM_PLUGIN_BULK_TRANSFER: u8 = 0x09;
23/// Reserved legacy mesh-plane config subscription stream ID.
24///
25/// Config and inventory control now live exclusively on `mesh-llm-control/1`;
26/// keep 0x0b reserved so old wire values are not accidentally reused.
27pub const STREAM_CONFIG_SUBSCRIBE: u8 = 0x0b;
28/// Reserved legacy mesh-plane config push stream ID.
29///
30/// Config and inventory control now live exclusively on `mesh-llm-control/1`;
31/// keep 0x0c reserved so old wire values are not accidentally reused.
32pub const STREAM_CONFIG_PUSH: u8 = 0x0c;
33pub const STREAM_SUBPROTOCOL: u8 = 0x0d;
34pub const STREAM_DIRECT_PATH_REQUEST: u8 = 0x0e;
35const _: () = {
36    let _ = STREAM_CONFIG_SUBSCRIBE;
37    let _ = STREAM_CONFIG_PUSH;
38    let _ = STREAM_SUBPROTOCOL;
39};
40
41#[derive(Clone, Copy, Debug, Eq, PartialEq)]
42pub enum ControlProtocol {
43    ProtoV1,
44    JsonV0,
45}
46
47#[derive(Debug, PartialEq)]
48pub enum ControlFrameError {
49    OversizeFrame { size: usize },
50    BadGeneration { got: u32 },
51    InvalidEndpointId { got: usize },
52    InvalidSenderId { got: usize },
53    MissingDirectPathAddress,
54    MissingHttpPort,
55    MissingOwnerId,
56    MissingControlOwnerId,
57    InvalidConfigHashLength { got: usize },
58    InvalidSubprotocol,
59    InvalidPublicKeyLength { got: usize },
60    MissingSignature,
61    InvalidSignatureLength { got: usize },
62    MissingConfig,
63    MissingControlEnvelope,
64    MissingControlCommand,
65    MissingControlResult,
66    MissingControlOwnership,
67    MissingRequestId,
68    InvalidOwnerControlErrorCode { got: i32 },
69    InvalidInventoryDisposition { got: i32 },
70    MissingInventoryModelRef,
71    MissingModelRef,
72    InvalidModelRefCombination,
73    InvalidInventoryOrder,
74    DecodeError(String),
75    WrongStreamType { expected: u8, got: u8 },
76    ForgedSender,
77}
78
79impl std::fmt::Display for ControlFrameError {
80    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
81        match self {
82            ControlFrameError::OversizeFrame { size } => write!(
83                f,
84                "control frame too large: {} bytes (max {})",
85                size, MAX_CONTROL_FRAME_BYTES
86            ),
87            ControlFrameError::BadGeneration { got } => write!(
88                f,
89                "bad protocol generation: expected {}, got {}",
90                NODE_PROTOCOL_GENERATION, got
91            ),
92            ControlFrameError::InvalidEndpointId { got } => {
93                write!(f, "invalid endpoint_id length: expected 32, got {}", got)
94            }
95            ControlFrameError::InvalidSenderId { got } => {
96                write!(f, "invalid sender_id length: expected 32, got {}", got)
97            }
98            ControlFrameError::MissingDirectPathAddress => {
99                write!(f, "direct path request missing endpoint address")
100            }
101            ControlFrameError::MissingHttpPort => {
102                write!(f, "HOST-role peer annotation missing http_port")
103            }
104            ControlFrameError::MissingOwnerId => write!(f, "config frame missing owner_id"),
105            ControlFrameError::MissingControlOwnerId => {
106                write!(f, "owner control handshake missing owner_id")
107            }
108            ControlFrameError::InvalidConfigHashLength { got } => {
109                write!(f, "invalid config_hash length: expected 32, got {}", got)
110            }
111            ControlFrameError::InvalidSubprotocol => {
112                write!(f, "subprotocol entries require a non-empty name and major")
113            }
114            ControlFrameError::InvalidPublicKeyLength { got } => {
115                write!(f, "invalid public key length: expected 32, got {}", got)
116            }
117            ControlFrameError::MissingSignature => write!(f, "config push missing signature"),
118            ControlFrameError::InvalidSignatureLength { got } => {
119                write!(f, "invalid signature length: expected 64, got {got}")
120            }
121            ControlFrameError::MissingConfig => {
122                write!(f, "config field is required but missing")
123            }
124            ControlFrameError::MissingControlEnvelope => {
125                write!(f, "owner control envelope requires exactly one payload")
126            }
127            ControlFrameError::MissingControlCommand => {
128                write!(
129                    f,
130                    "owner control request requires exactly one command variant"
131                )
132            }
133            ControlFrameError::MissingControlResult => {
134                write!(
135                    f,
136                    "owner control response requires exactly one result variant"
137                )
138            }
139            ControlFrameError::MissingControlOwnership => {
140                write!(f, "owner control handshake missing ownership attestation")
141            }
142            ControlFrameError::MissingRequestId => {
143                write!(f, "owner control request_id must be non-zero")
144            }
145            ControlFrameError::InvalidOwnerControlErrorCode { got } => {
146                write!(f, "invalid owner control error code: {got}")
147            }
148            ControlFrameError::InvalidInventoryDisposition { got } => {
149                write!(f, "invalid inventory scan disposition: {got}")
150            }
151            ControlFrameError::MissingInventoryModelRef => {
152                write!(f, "inventory entry requires a canonical model ref")
153            }
154            ControlFrameError::MissingModelRef => {
155                write!(
156                    f,
157                    "model lifecycle command requires exactly one model reference"
158                )
159            }
160            ControlFrameError::InvalidModelRefCombination => {
161                write!(
162                    f,
163                    "model lifecycle command has an invalid model identifier combination"
164                )
165            }
166            ControlFrameError::InvalidInventoryOrder => {
167                write!(
168                    f,
169                    "inventory entries must be strictly sorted by canonical model ref"
170                )
171            }
172            ControlFrameError::DecodeError(msg) => write!(f, "protobuf decode error: {}", msg),
173            ControlFrameError::WrongStreamType { expected, got } => write!(
174                f,
175                "wrong stream type: expected {:#04x}, got {:#04x}",
176                expected, got
177            ),
178            ControlFrameError::ForgedSender => {
179                write!(f, "frame peer_id does not match QUIC connection identity")
180            }
181        }
182    }
183}
184
185impl std::error::Error for ControlFrameError {}
186
187pub trait ValidateControlFrame: prost::Message + Default + Sized {
188    fn validate_frame(&self) -> Result<(), ControlFrameError> {
189        Ok(())
190    }
191}
192
193impl ValidateControlFrame for crate::proto::node::GossipFrame {
194    fn validate_frame(&self) -> Result<(), ControlFrameError> {
195        if self.r#gen != NODE_PROTOCOL_GENERATION {
196            return Err(ControlFrameError::BadGeneration { got: self.r#gen });
197        }
198        if self.sender_id.len() != 32 {
199            return Err(ControlFrameError::InvalidSenderId {
200                got: self.sender_id.len(),
201            });
202        }
203        for pa in &self.peers {
204            validate_peer_announcement(pa)?;
205        }
206        Ok(())
207    }
208}
209
210impl ValidateControlFrame for crate::proto::node::TunnelMap {
211    fn validate_frame(&self) -> Result<(), ControlFrameError> {
212        if self.owner_peer_id.len() != 32 {
213            return Err(ControlFrameError::InvalidEndpointId {
214                got: self.owner_peer_id.len(),
215            });
216        }
217        for entry in &self.entries {
218            if entry.target_peer_id.len() != 32 {
219                return Err(ControlFrameError::InvalidEndpointId {
220                    got: entry.target_peer_id.len(),
221                });
222            }
223        }
224        Ok(())
225    }
226}
227impl ValidateControlFrame for crate::proto::node::RouteTableRequest {
228    fn validate_frame(&self) -> Result<(), ControlFrameError> {
229        if self.r#gen != NODE_PROTOCOL_GENERATION {
230            return Err(ControlFrameError::BadGeneration { got: self.r#gen });
231        }
232        if !self.requester_id.is_empty() && self.requester_id.len() != 32 {
233            return Err(ControlFrameError::InvalidEndpointId {
234                got: self.requester_id.len(),
235            });
236        }
237        Ok(())
238    }
239}
240impl ValidateControlFrame for crate::proto::node::RouteTable {
241    fn validate_frame(&self) -> Result<(), ControlFrameError> {
242        if self.r#gen != NODE_PROTOCOL_GENERATION {
243            return Err(ControlFrameError::BadGeneration { got: self.r#gen });
244        }
245        for entry in &self.entries {
246            if entry.endpoint_id.len() != 32 {
247                return Err(ControlFrameError::InvalidEndpointId {
248                    got: entry.endpoint_id.len(),
249                });
250            }
251        }
252        Ok(())
253    }
254}
255impl ValidateControlFrame for crate::proto::node::PeerDown {
256    fn validate_frame(&self) -> Result<(), ControlFrameError> {
257        if self.r#gen != NODE_PROTOCOL_GENERATION {
258            return Err(ControlFrameError::BadGeneration { got: self.r#gen });
259        }
260        if self.peer_id.len() != 32 {
261            return Err(ControlFrameError::InvalidEndpointId {
262                got: self.peer_id.len(),
263            });
264        }
265        Ok(())
266    }
267}
268impl ValidateControlFrame for crate::proto::node::PeerLeaving {
269    fn validate_frame(&self) -> Result<(), ControlFrameError> {
270        if self.r#gen != NODE_PROTOCOL_GENERATION {
271            return Err(ControlFrameError::BadGeneration { got: self.r#gen });
272        }
273        if self.peer_id.len() != 32 {
274            return Err(ControlFrameError::InvalidEndpointId {
275                got: self.peer_id.len(),
276            });
277        }
278        Ok(())
279    }
280}
281
282impl ValidateControlFrame for crate::proto::node::DirectPathRequest {
283    fn validate_frame(&self) -> Result<(), ControlFrameError> {
284        if self.r#gen != NODE_PROTOCOL_GENERATION {
285            return Err(ControlFrameError::BadGeneration { got: self.r#gen });
286        }
287        if self.requester_id.len() != 32 {
288            return Err(ControlFrameError::InvalidEndpointId {
289                got: self.requester_id.len(),
290            });
291        }
292        if self.serialized_addr.is_empty() {
293            return Err(ControlFrameError::MissingDirectPathAddress);
294        }
295        Ok(())
296    }
297}
298
299impl ValidateControlFrame for crate::proto::node::OwnerControlEnvelope {
300    fn validate_frame(&self) -> Result<(), ControlFrameError> {
301        if self.r#gen != NODE_PROTOCOL_GENERATION {
302            return Err(ControlFrameError::BadGeneration { got: self.r#gen });
303        }
304        let payloads = [
305            self.handshake.is_some(),
306            self.request.is_some(),
307            self.response.is_some(),
308            self.error.is_some(),
309        ];
310        if payloads.into_iter().filter(|present| *present).count() != 1 {
311            return Err(ControlFrameError::MissingControlEnvelope);
312        }
313        if let Some(handshake) = &self.handshake {
314            handshake.validate_frame()?;
315        }
316        if let Some(request) = &self.request {
317            request.validate_frame()?;
318        }
319        if let Some(response) = &self.response {
320            response.validate_frame()?;
321        }
322        if let Some(error) = &self.error {
323            error.validate_frame()?;
324        }
325        Ok(())
326    }
327}
328
329impl ValidateControlFrame for crate::proto::node::OwnerControlHandshake {
330    fn validate_frame(&self) -> Result<(), ControlFrameError> {
331        let ownership = self
332            .ownership
333            .as_ref()
334            .ok_or(ControlFrameError::MissingControlOwnership)?;
335        if ownership.owner_id.trim().is_empty() {
336            return Err(ControlFrameError::MissingControlOwnerId);
337        }
338        validate_public_key_length(ownership.owner_sign_public_key.len())?;
339        validate_endpoint_id_length(ownership.node_endpoint_id.len())?;
340        if ownership.signature.is_empty() {
341            return Err(ControlFrameError::MissingSignature);
342        }
343        if ownership.signature.len() != 64 {
344            return Err(ControlFrameError::InvalidSignatureLength {
345                got: ownership.signature.len(),
346            });
347        }
348        Ok(())
349    }
350}
351
352impl ValidateControlFrame for crate::proto::node::OwnerControlRequest {
353    fn validate_frame(&self) -> Result<(), ControlFrameError> {
354        if self.request_id == 0 {
355            return Err(ControlFrameError::MissingRequestId);
356        }
357        let commands = [
358            self.get_config.is_some(),
359            self.watch_config.is_some(),
360            self.apply_config.is_some(),
361            self.refresh_inventory.is_some(),
362            self.load_model.is_some(),
363            self.unload_model.is_some(),
364            self.ensure_model.is_some(),
365            self.drain_model.is_some(),
366        ];
367        if commands.into_iter().filter(|present| *present).count() != 1 {
368            return Err(ControlFrameError::MissingControlCommand);
369        }
370        if let Some(request) = &self.get_config {
371            request.validate_frame()?;
372        }
373        if let Some(request) = &self.watch_config {
374            request.validate_frame()?;
375        }
376        if let Some(request) = &self.apply_config {
377            request.validate_frame()?;
378        }
379        if let Some(request) = &self.refresh_inventory {
380            request.validate_frame()?;
381        }
382        if let Some(request) = &self.load_model {
383            request.validate_frame()?;
384        }
385        if let Some(request) = &self.unload_model {
386            request.validate_frame()?;
387        }
388        if let Some(request) = &self.ensure_model {
389            request.validate_frame()?;
390        }
391        if let Some(request) = &self.drain_model {
392            request.validate_frame()?;
393        }
394        Ok(())
395    }
396}
397
398impl ValidateControlFrame for crate::proto::node::OwnerControlResponse {
399    fn validate_frame(&self) -> Result<(), ControlFrameError> {
400        if self.request_id == 0 {
401            return Err(ControlFrameError::MissingRequestId);
402        }
403        let results = [
404            self.get_config.is_some(),
405            self.watch_config.is_some(),
406            self.apply_config.is_some(),
407            self.refresh_inventory.is_some(),
408            self.load_model.is_some(),
409            self.unload_model.is_some(),
410            self.ensure_model.is_some(),
411            self.drain_model.is_some(),
412        ];
413        if results.into_iter().filter(|present| *present).count() != 1 {
414            return Err(ControlFrameError::MissingControlResult);
415        }
416        if let Some(response) = &self.get_config {
417            response.validate_frame()?;
418        }
419        if let Some(response) = &self.watch_config {
420            response.validate_frame()?;
421        }
422        if let Some(response) = &self.apply_config {
423            response.validate_frame()?;
424        }
425        if let Some(response) = &self.refresh_inventory {
426            response.validate_frame()?;
427        }
428        if let Some(response) = &self.load_model {
429            response.validate_frame()?;
430        }
431        if let Some(response) = &self.unload_model {
432            response.validate_frame()?;
433        }
434        if let Some(response) = &self.ensure_model {
435            response.validate_frame()?;
436        }
437        if let Some(response) = &self.drain_model {
438            response.validate_frame()?;
439        }
440        Ok(())
441    }
442}
443
444impl ValidateControlFrame for crate::proto::node::OwnerControlError {
445    fn validate_frame(&self) -> Result<(), ControlFrameError> {
446        if matches!(
447            crate::proto::node::OwnerControlErrorCode::try_from(self.code),
448            Err(_) | Ok(crate::proto::node::OwnerControlErrorCode::Unspecified)
449        ) {
450            return Err(ControlFrameError::InvalidOwnerControlErrorCode { got: self.code });
451        }
452        Ok(())
453    }
454}
455
456impl ValidateControlFrame for crate::proto::node::OwnerControlGetConfigRequest {
457    fn validate_frame(&self) -> Result<(), ControlFrameError> {
458        validate_endpoint_id_length(self.requester_node_id.len())?;
459        validate_endpoint_id_length(self.target_node_id.len())?;
460        Ok(())
461    }
462}
463
464impl ValidateControlFrame for crate::proto::node::OwnerControlGetConfigResponse {
465    fn validate_frame(&self) -> Result<(), ControlFrameError> {
466        self.snapshot
467            .as_ref()
468            .ok_or(ControlFrameError::MissingConfig)?
469            .validate_frame()
470    }
471}
472
473impl ValidateControlFrame for crate::proto::node::OwnerControlWatchConfigRequest {
474    fn validate_frame(&self) -> Result<(), ControlFrameError> {
475        validate_endpoint_id_length(self.requester_node_id.len())?;
476        validate_endpoint_id_length(self.target_node_id.len())?;
477        Ok(())
478    }
479}
480
481impl ValidateControlFrame for crate::proto::node::OwnerControlWatchConfigResponse {
482    fn validate_frame(&self) -> Result<(), ControlFrameError> {
483        let results = [
484            self.accepted.is_some(),
485            self.snapshot.is_some(),
486            self.update.is_some(),
487        ];
488        if results.into_iter().filter(|present| *present).count() != 1 {
489            return Err(ControlFrameError::MissingControlResult);
490        }
491        if let Some(accepted) = &self.accepted {
492            accepted.validate_frame()?;
493        }
494        if let Some(snapshot) = &self.snapshot {
495            snapshot.validate_frame()?;
496        }
497        if let Some(update) = &self.update {
498            update.validate_frame()?;
499        }
500        Ok(())
501    }
502}
503
504impl ValidateControlFrame for crate::proto::node::OwnerControlWatchAccepted {
505    fn validate_frame(&self) -> Result<(), ControlFrameError> {
506        validate_endpoint_id_length(self.target_node_id.len())?;
507        Ok(())
508    }
509}
510
511impl ValidateControlFrame for crate::proto::node::OwnerControlApplyConfigRequest {
512    fn validate_frame(&self) -> Result<(), ControlFrameError> {
513        validate_endpoint_id_length(self.requester_node_id.len())?;
514        validate_endpoint_id_length(self.target_node_id.len())?;
515        if self.config.is_none() {
516            return Err(ControlFrameError::MissingConfig);
517        }
518        Ok(())
519    }
520}
521
522impl ValidateControlFrame for crate::proto::node::OwnerControlApplyConfigResponse {
523    fn validate_frame(&self) -> Result<(), ControlFrameError> {
524        if self.success || !self.config_hash.is_empty() {
525            validate_config_hash_length(self.config_hash.len())?;
526        }
527        Ok(())
528    }
529}
530
531impl ValidateControlFrame for crate::proto::node::OwnerControlRefreshInventoryRequest {
532    fn validate_frame(&self) -> Result<(), ControlFrameError> {
533        validate_endpoint_id_length(self.requester_node_id.len())?;
534        validate_endpoint_id_length(self.target_node_id.len())?;
535        Ok(())
536    }
537}
538
539impl ValidateControlFrame for crate::proto::node::OwnerControlRefreshInventoryResponse {
540    fn validate_frame(&self) -> Result<(), ControlFrameError> {
541        self.snapshot
542            .as_ref()
543            .ok_or(ControlFrameError::MissingConfig)?
544            .validate_frame()?;
545        if let Some(inventory) = &self.inventory {
546            inventory.validate_frame()?;
547        }
548        Ok(())
549    }
550}
551
552impl ValidateControlFrame for crate::proto::node::OwnerControlLoadModelRequest {
553    fn validate_frame(&self) -> Result<(), ControlFrameError> {
554        validate_endpoint_id_length(self.requester_node_id.len())?;
555        validate_endpoint_id_length(self.target_node_id.len())?;
556        let model = self
557            .model
558            .as_ref()
559            .ok_or(ControlFrameError::MissingModelRef)?;
560        validate_owner_control_model_for_load_or_ensure(model)
561    }
562}
563
564impl ValidateControlFrame for crate::proto::node::OwnerControlUnloadModelRequest {
565    fn validate_frame(&self) -> Result<(), ControlFrameError> {
566        validate_endpoint_id_length(self.requester_node_id.len())?;
567        validate_endpoint_id_length(self.target_node_id.len())?;
568        let model = self
569            .model
570            .as_ref()
571            .ok_or(ControlFrameError::MissingModelRef)?;
572        validate_owner_control_model_for_unload_or_drain(model)
573    }
574}
575
576impl ValidateControlFrame for crate::proto::node::OwnerControlEnsureModelRequest {
577    fn validate_frame(&self) -> Result<(), ControlFrameError> {
578        validate_endpoint_id_length(self.requester_node_id.len())?;
579        validate_endpoint_id_length(self.target_node_id.len())?;
580        let model = self
581            .model
582            .as_ref()
583            .ok_or(ControlFrameError::MissingModelRef)?;
584        validate_owner_control_model_for_load_or_ensure(model)
585    }
586}
587
588impl ValidateControlFrame for crate::proto::node::OwnerControlDrainModelRequest {
589    fn validate_frame(&self) -> Result<(), ControlFrameError> {
590        validate_endpoint_id_length(self.requester_node_id.len())?;
591        validate_endpoint_id_length(self.target_node_id.len())?;
592        let model = self
593            .model
594            .as_ref()
595            .ok_or(ControlFrameError::MissingModelRef)?;
596        validate_owner_control_model_for_unload_or_drain(model)
597    }
598}
599
600impl ValidateControlFrame for crate::proto::node::OwnerControlLoadModelResponse {
601    fn validate_frame(&self) -> Result<(), ControlFrameError> {
602        validate_owner_control_model_for_load_or_ensure(
603            self.target
604                .as_ref()
605                .ok_or(ControlFrameError::MissingModelRef)?,
606        )
607    }
608}
609
610impl ValidateControlFrame for crate::proto::node::OwnerControlUnloadModelResponse {
611    fn validate_frame(&self) -> Result<(), ControlFrameError> {
612        validate_owner_control_model_for_unload_or_drain(
613            self.target
614                .as_ref()
615                .ok_or(ControlFrameError::MissingModelRef)?,
616        )
617    }
618}
619
620impl ValidateControlFrame for crate::proto::node::OwnerControlEnsureModelResponse {
621    fn validate_frame(&self) -> Result<(), ControlFrameError> {
622        validate_owner_control_model_for_load_or_ensure(
623            self.target
624                .as_ref()
625                .ok_or(ControlFrameError::MissingModelRef)?,
626        )
627    }
628}
629
630impl ValidateControlFrame for crate::proto::node::OwnerControlDrainModelResponse {
631    fn validate_frame(&self) -> Result<(), ControlFrameError> {
632        validate_owner_control_model_for_unload_or_drain(
633            self.target
634                .as_ref()
635                .ok_or(ControlFrameError::MissingModelRef)?,
636        )
637    }
638}
639
640impl ValidateControlFrame for crate::proto::node::OwnerControlRefreshInventory {
641    fn validate_frame(&self) -> Result<(), ControlFrameError> {
642        use crate::proto::node::OwnerControlRefreshInventoryDisposition;
643
644        if !matches!(
645            OwnerControlRefreshInventoryDisposition::try_from(self.disposition),
646            Ok(OwnerControlRefreshInventoryDisposition::Executed)
647                | Ok(OwnerControlRefreshInventoryDisposition::Coalesced)
648        ) {
649            return Err(ControlFrameError::InvalidInventoryDisposition {
650                got: self.disposition,
651            });
652        }
653        let mut previous = None;
654        for entry in &self.entries {
655            let canonical = entry.canonical_model_ref.trim();
656            if canonical.is_empty() {
657                return Err(ControlFrameError::MissingInventoryModelRef);
658            }
659            if previous.is_some_and(|value| value >= canonical) {
660                return Err(ControlFrameError::InvalidInventoryOrder);
661            }
662            previous = Some(canonical);
663        }
664        Ok(())
665    }
666}
667
668impl ValidateControlFrame for crate::proto::node::OwnerControlConfigSnapshot {
669    fn validate_frame(&self) -> Result<(), ControlFrameError> {
670        validate_endpoint_id_length(self.node_id.len())?;
671        validate_config_hash_length(self.config_hash.len())?;
672        if self.config.is_none() {
673            return Err(ControlFrameError::MissingConfig);
674        }
675        Ok(())
676    }
677}
678
679impl ValidateControlFrame for crate::proto::node::OwnerControlConfigUpdate {
680    fn validate_frame(&self) -> Result<(), ControlFrameError> {
681        validate_endpoint_id_length(self.node_id.len())?;
682        validate_config_hash_length(self.config_hash.len())?;
683        if self.config.is_none() {
684            return Err(ControlFrameError::MissingConfig);
685        }
686        Ok(())
687    }
688}
689
690impl ValidateControlFrame for crate::proto::node::MeshSubprotocolOpen {
691    fn validate_frame(&self) -> Result<(), ControlFrameError> {
692        if self.r#gen != NODE_PROTOCOL_GENERATION {
693            return Err(ControlFrameError::BadGeneration { got: self.r#gen });
694        }
695        if self.name.trim().is_empty() || self.major == 0 {
696            return Err(ControlFrameError::InvalidSubprotocol);
697        }
698        Ok(())
699    }
700}
701
702pub fn validate_peer_announcement(
703    pa: &crate::proto::node::PeerAnnouncement,
704) -> Result<(), ControlFrameError> {
705    if pa.endpoint_id.len() != 32 {
706        return Err(ControlFrameError::InvalidEndpointId {
707            got: pa.endpoint_id.len(),
708        });
709    }
710    if pa.role == crate::proto::node::NodeRole::Host as i32 && pa.http_port.is_none() {
711        return Err(ControlFrameError::MissingHttpPort);
712    }
713    for subprotocol in &pa.subprotocols {
714        if subprotocol.name.trim().is_empty() || subprotocol.major == 0 {
715            return Err(ControlFrameError::InvalidSubprotocol);
716        }
717    }
718    Ok(())
719}
720
721fn validate_endpoint_id_length(len: usize) -> Result<(), ControlFrameError> {
722    if len != 32 {
723        return Err(ControlFrameError::InvalidEndpointId { got: len });
724    }
725    Ok(())
726}
727
728fn validate_config_hash_length(len: usize) -> Result<(), ControlFrameError> {
729    if len != 32 {
730        return Err(ControlFrameError::InvalidConfigHashLength { got: len });
731    }
732    Ok(())
733}
734
735pub fn validate_owner_control_model_for_load_or_ensure(
736    model: &crate::proto::node::OwnerControlModelRef,
737) -> Result<(), ControlFrameError> {
738    let canonical = !model.canonical_model_ref.trim().is_empty();
739    let instance = model
740        .instance_id
741        .as_deref()
742        .is_some_and(|id| !id.trim().is_empty());
743    match (canonical, instance) {
744        (true, false) => Ok(()),
745        (false, false) => Err(ControlFrameError::MissingModelRef),
746        _ => Err(ControlFrameError::InvalidModelRefCombination),
747    }
748}
749
750pub fn validate_owner_control_model_for_unload_or_drain(
751    model: &crate::proto::node::OwnerControlModelRef,
752) -> Result<(), ControlFrameError> {
753    let canonical = !model.canonical_model_ref.trim().is_empty();
754    let instance = model
755        .instance_id
756        .as_deref()
757        .is_some_and(|id| !id.trim().is_empty());
758    if canonical ^ instance {
759        Ok(())
760    } else if canonical || instance {
761        Err(ControlFrameError::InvalidModelRefCombination)
762    } else {
763        Err(ControlFrameError::MissingModelRef)
764    }
765}
766
767fn validate_public_key_length(len: usize) -> Result<(), ControlFrameError> {
768    if len != 32 {
769        return Err(ControlFrameError::InvalidPublicKeyLength { got: len });
770    }
771    Ok(())
772}
773
774pub fn protocol_from_alpn(alpn: &[u8]) -> ControlProtocol {
775    if alpn == ALPN_V0 {
776        ControlProtocol::JsonV0
777    } else {
778        ControlProtocol::ProtoV1
779    }
780}
781
782pub fn connection_protocol(conn: &Connection) -> ControlProtocol {
783    protocol_from_alpn(conn.alpn())
784}
785
786pub async fn connect_mesh(endpoint: &Endpoint, addr: EndpointAddr) -> Result<Connection> {
787    let opts = ConnectOptions::new().with_additional_alpns(vec![ALPN_V0.to_vec()]);
788    let connecting = endpoint.connect_with_opts(addr, ALPN_V1, opts).await?;
789    Ok(connecting.await?)
790}
791
792pub async fn write_len_prefixed(send: &mut iroh::endpoint::SendStream, body: &[u8]) -> Result<()> {
793    ensure_control_frame_size(body)?;
794    send.write_all(&(body.len() as u32).to_le_bytes()).await?;
795    send.write_all(body).await?;
796    Ok(())
797}
798
799pub fn ensure_control_frame_size(body: &[u8]) -> Result<(), ControlFrameError> {
800    if body.len() > MAX_CONTROL_FRAME_BYTES {
801        return Err(ControlFrameError::OversizeFrame { size: body.len() });
802    }
803    Ok(())
804}
805
806pub async fn read_len_prefixed(recv: &mut iroh::endpoint::RecvStream) -> Result<Vec<u8>> {
807    let mut len_buf = [0u8; 4];
808    recv.read_exact(&mut len_buf).await?;
809    let len = u32::from_le_bytes(len_buf) as usize;
810    if len > MAX_CONTROL_FRAME_BYTES {
811        anyhow::bail!("control frame too large: {} bytes", len);
812    }
813    let mut buf = vec![0u8; len];
814    recv.read_exact(&mut buf).await?;
815    Ok(buf)
816}
817
818pub fn encode_control_frame(stream_type: u8, msg: &impl prost::Message) -> Vec<u8> {
819    let proto_bytes = msg.encode_to_vec();
820    let len = proto_bytes.len() as u32;
821    let mut buf = Vec::with_capacity(1 + 4 + proto_bytes.len());
822    buf.push(stream_type);
823    buf.extend_from_slice(&len.to_le_bytes());
824    buf.extend_from_slice(&proto_bytes);
825    buf
826}
827
828pub fn decode_control_frame<T: ValidateControlFrame>(
829    expected_stream_type: u8,
830    data: &[u8],
831) -> Result<T, ControlFrameError> {
832    const HEADER_LEN: usize = 5;
833    if data.len() < HEADER_LEN {
834        return Err(ControlFrameError::DecodeError(format!(
835            "frame too short: {} bytes (minimum {})",
836            data.len(),
837            HEADER_LEN
838        )));
839    }
840    let actual_type = data[0];
841    if actual_type != expected_stream_type {
842        return Err(ControlFrameError::WrongStreamType {
843            expected: expected_stream_type,
844            got: actual_type,
845        });
846    }
847    let len = u32::from_le_bytes(data[1..5].try_into().unwrap()) as usize;
848    if len > MAX_CONTROL_FRAME_BYTES {
849        return Err(ControlFrameError::OversizeFrame { size: len });
850    }
851    let proto_bytes = data.get(5..5 + len).ok_or_else(|| {
852        ControlFrameError::DecodeError(format!(
853            "frame truncated: header says {} bytes but only {} available",
854            len,
855            data.len().saturating_sub(5)
856        ))
857    })?;
858    let msg = T::decode(proto_bytes).map_err(|e| ControlFrameError::DecodeError(e.to_string()))?;
859    msg.validate_frame()?;
860    Ok(msg)
861}
862
863pub fn encode_owner_control_envelope(msg: &crate::proto::node::OwnerControlEnvelope) -> Vec<u8> {
864    msg.encode_to_vec()
865}
866
867pub fn decode_owner_control_envelope(
868    data: &[u8],
869) -> Result<crate::proto::node::OwnerControlEnvelope, ControlFrameError> {
870    let msg = crate::proto::node::OwnerControlEnvelope::decode(data)
871        .map_err(|e| ControlFrameError::DecodeError(e.to_string()))?;
872    msg.validate_frame()?;
873    Ok(msg)
874}
875
876pub fn owner_control_error_envelope(
877    code: crate::proto::node::OwnerControlErrorCode,
878    request_id: Option<u64>,
879    message: impl Into<String>,
880) -> crate::proto::node::OwnerControlEnvelope {
881    crate::proto::node::OwnerControlEnvelope {
882        r#gen: NODE_PROTOCOL_GENERATION,
883        handshake: None,
884        request: None,
885        response: None,
886        error: Some(crate::proto::node::OwnerControlError {
887            code: code as i32,
888            message: message.into(),
889            request_id,
890            current_revision: None,
891        }),
892    }
893}
894
895pub fn owner_control_rejection_envelope(
896    data: &[u8],
897    request_id: Option<u64>,
898    err: &ControlFrameError,
899) -> crate::proto::node::OwnerControlEnvelope {
900    let code = if matches!(err, ControlFrameError::MissingControlCommand) {
901        crate::proto::node::OwnerControlErrorCode::UnknownCommand
902    } else if serde_json::from_slice::<serde_json::Value>(data).is_ok() {
903        crate::proto::node::OwnerControlErrorCode::LegacyJsonUnsupported
904    } else {
905        crate::proto::node::OwnerControlErrorCode::BadRequest
906    };
907    owner_control_error_envelope(code, request_id, err.to_string())
908}
909
910#[cfg(test)]
911mod tests {
912    use super::*;
913    use crate::proto::node::{
914        CompactModelMetadata, ConfigApplyMode, GossipFrame, InferenceAdmissionState,
915        NodeConfigSnapshot, NodeGpuConfig, NodeModelEntry, NodeRole,
916        OwnerControlApplyConfigRequest, OwnerControlApplyConfigResponse,
917        OwnerControlConfigSnapshot, OwnerControlConfigUpdate, OwnerControlEnvelope,
918        OwnerControlError, OwnerControlErrorCode, OwnerControlGetConfigRequest,
919        OwnerControlGetConfigResponse, OwnerControlHandshake, OwnerControlInventoryEntry,
920        OwnerControlRefreshInventory, OwnerControlRefreshInventoryDisposition,
921        OwnerControlRefreshInventoryRequest, OwnerControlRefreshInventoryResponse,
922        OwnerControlRequest, OwnerControlResponse, OwnerControlWatchAccepted,
923        OwnerControlWatchConfigResponse, PeerAnnouncement, SignedNodeOwnership,
924    };
925    use prost::Message;
926
927    fn control_plane_test_config() -> NodeConfigSnapshot {
928        NodeConfigSnapshot {
929            version: 1,
930            gpu: Some(NodeGpuConfig {
931                assignment: crate::proto::node::GpuAssignment::Auto as i32,
932            }),
933            models: vec![NodeModelEntry {
934                model: "Qwen3-8B".to_string(),
935                mmproj: None,
936                ctx_size: Some(8192),
937                gpu_id: None,
938                model_ref: None,
939                mmproj_ref: None,
940            }],
941            plugins: vec![],
942            config_toml: None,
943            mesh_requirements: None,
944        }
945    }
946
947    fn control_plane_test_snapshot() -> OwnerControlConfigSnapshot {
948        OwnerControlConfigSnapshot {
949            node_id: vec![0x55; 32],
950            revision: 7,
951            config_hash: vec![0xA5; 32],
952            config: Some(control_plane_test_config()),
953            hostname: Some("node-01".to_string()),
954        }
955    }
956
957    fn control_plane_test_handshake() -> OwnerControlEnvelope {
958        OwnerControlEnvelope {
959            r#gen: NODE_PROTOCOL_GENERATION,
960            handshake: Some(OwnerControlHandshake {
961                ownership: Some(SignedNodeOwnership {
962                    version: 1,
963                    cert_id: "cert-1".to_string(),
964                    owner_id: "owner-1".to_string(),
965                    owner_sign_public_key: vec![0x11; 32],
966                    node_endpoint_id: vec![0x22; 32],
967                    issued_at_unix_ms: 1,
968                    expires_at_unix_ms: 2,
969                    node_label: Some("node-01".to_string()),
970                    hostname_hint: Some("node-01".to_string()),
971                    signature: vec![0x33; 64],
972                }),
973            }),
974            request: None,
975            response: None,
976            error: None,
977        }
978    }
979
980    #[test]
981    fn control_plane_messages_constants_are_stable() {
982        assert_eq!(ALPN_CONTROL_V1, b"mesh-llm-control/1");
983        assert_eq!(ALPN_V1, b"mesh-llm/1");
984        assert_eq!(ALPN_V0, b"mesh-llm/0");
985        assert_eq!(STREAM_CONFIG_SUBSCRIBE, 0x0b);
986        assert_eq!(STREAM_CONFIG_PUSH, 0x0c);
987        assert_eq!(STREAM_SUBPROTOCOL, 0x0d);
988        assert_eq!(STREAM_DIRECT_PATH_REQUEST, 0x0e);
989    }
990
991    #[test]
992    fn control_plane_messages_roundtrip_commands_and_responses() {
993        let handshake = control_plane_test_handshake();
994        let decoded = decode_owner_control_envelope(&encode_owner_control_envelope(&handshake))
995            .expect("handshake must decode");
996        assert!(decoded.handshake.is_some());
997
998        let get_request = OwnerControlEnvelope {
999            r#gen: NODE_PROTOCOL_GENERATION,
1000            handshake: None,
1001            request: Some(OwnerControlRequest {
1002                request_id: 10,
1003                get_config: Some(OwnerControlGetConfigRequest {
1004                    requester_node_id: vec![0x10; 32],
1005                    target_node_id: vec![0x20; 32],
1006                }),
1007                watch_config: None,
1008                apply_config: None,
1009                refresh_inventory: None,
1010                load_model: None,
1011                unload_model: None,
1012                ensure_model: None,
1013                drain_model: None,
1014            }),
1015            response: None,
1016            error: None,
1017        };
1018        let decoded = decode_owner_control_envelope(&encode_owner_control_envelope(&get_request))
1019            .expect("get-config request must decode");
1020        assert_eq!(decoded.request.unwrap().request_id, 10);
1021
1022        let watch_response = OwnerControlEnvelope {
1023            r#gen: NODE_PROTOCOL_GENERATION,
1024            handshake: None,
1025            request: None,
1026            response: Some(OwnerControlResponse {
1027                request_id: 11,
1028                get_config: None,
1029                watch_config: Some(OwnerControlWatchConfigResponse {
1030                    accepted: Some(OwnerControlWatchAccepted {
1031                        target_node_id: vec![0x21; 32],
1032                    }),
1033                    snapshot: None,
1034                    update: None,
1035                }),
1036                apply_config: None,
1037                refresh_inventory: None,
1038                load_model: None,
1039                unload_model: None,
1040                ensure_model: None,
1041                drain_model: None,
1042            }),
1043            error: None,
1044        };
1045        decode_owner_control_envelope(&encode_owner_control_envelope(&watch_response))
1046            .expect("watch-config response must decode");
1047
1048        let apply_request = OwnerControlEnvelope {
1049            r#gen: NODE_PROTOCOL_GENERATION,
1050            handshake: None,
1051            request: Some(OwnerControlRequest {
1052                request_id: 12,
1053                get_config: None,
1054                watch_config: None,
1055                apply_config: Some(OwnerControlApplyConfigRequest {
1056                    requester_node_id: vec![0x30; 32],
1057                    target_node_id: vec![0x40; 32],
1058                    expected_revision: 7,
1059                    config: Some(control_plane_test_config()),
1060                }),
1061                refresh_inventory: None,
1062                load_model: None,
1063                unload_model: None,
1064                ensure_model: None,
1065                drain_model: None,
1066            }),
1067            response: None,
1068            error: None,
1069        };
1070        decode_owner_control_envelope(&encode_owner_control_envelope(&apply_request))
1071            .expect("apply-config request must decode");
1072
1073        let apply_response = OwnerControlEnvelope {
1074            r#gen: NODE_PROTOCOL_GENERATION,
1075            handshake: None,
1076            request: None,
1077            response: Some(OwnerControlResponse {
1078                request_id: 12,
1079                get_config: None,
1080                watch_config: None,
1081                apply_config: Some(OwnerControlApplyConfigResponse {
1082                    success: true,
1083                    current_revision: 8,
1084                    config_hash: vec![0x99; 32],
1085                    error: None,
1086                    apply_mode: ConfigApplyMode::Live as i32,
1087                    diagnostics: Vec::new(),
1088                }),
1089                refresh_inventory: None,
1090                load_model: None,
1091                unload_model: None,
1092                ensure_model: None,
1093                drain_model: None,
1094            }),
1095            error: None,
1096        };
1097        decode_owner_control_envelope(&encode_owner_control_envelope(&apply_response))
1098            .expect("apply-config response must decode");
1099
1100        let refresh_request = OwnerControlEnvelope {
1101            r#gen: NODE_PROTOCOL_GENERATION,
1102            handshake: None,
1103            request: Some(OwnerControlRequest {
1104                request_id: 13,
1105                get_config: None,
1106                watch_config: None,
1107                apply_config: None,
1108                refresh_inventory: Some(OwnerControlRefreshInventoryRequest {
1109                    requester_node_id: vec![0x50; 32],
1110                    target_node_id: vec![0x60; 32],
1111                }),
1112                load_model: None,
1113                unload_model: None,
1114                ensure_model: None,
1115                drain_model: None,
1116            }),
1117            response: None,
1118            error: None,
1119        };
1120        decode_owner_control_envelope(&encode_owner_control_envelope(&refresh_request))
1121            .expect("refresh-inventory request must decode");
1122
1123        let refresh_response = OwnerControlEnvelope {
1124            r#gen: NODE_PROTOCOL_GENERATION,
1125            handshake: None,
1126            request: None,
1127            response: Some(OwnerControlResponse {
1128                request_id: 13,
1129                get_config: None,
1130                watch_config: None,
1131                apply_config: None,
1132                refresh_inventory: Some(OwnerControlRefreshInventoryResponse {
1133                    snapshot: Some(control_plane_test_snapshot()),
1134                    inventory: None,
1135                }),
1136                load_model: None,
1137                unload_model: None,
1138                ensure_model: None,
1139                drain_model: None,
1140            }),
1141            error: None,
1142        };
1143        decode_owner_control_envelope(&encode_owner_control_envelope(&refresh_response))
1144            .expect("refresh-inventory response must decode");
1145
1146        let get_response = OwnerControlEnvelope {
1147            r#gen: NODE_PROTOCOL_GENERATION,
1148            handshake: None,
1149            request: None,
1150            response: Some(OwnerControlResponse {
1151                request_id: 14,
1152                get_config: Some(OwnerControlGetConfigResponse {
1153                    snapshot: Some(control_plane_test_snapshot()),
1154                }),
1155                watch_config: None,
1156                apply_config: None,
1157                refresh_inventory: None,
1158                load_model: None,
1159                unload_model: None,
1160                ensure_model: None,
1161                drain_model: None,
1162            }),
1163            error: None,
1164        };
1165        decode_owner_control_envelope(&encode_owner_control_envelope(&get_response))
1166            .expect("get-config response must decode");
1167
1168        let update_response = OwnerControlEnvelope {
1169            r#gen: NODE_PROTOCOL_GENERATION,
1170            handshake: None,
1171            request: None,
1172            response: Some(OwnerControlResponse {
1173                request_id: 15,
1174                get_config: None,
1175                watch_config: Some(OwnerControlWatchConfigResponse {
1176                    accepted: None,
1177                    snapshot: None,
1178                    update: Some(OwnerControlConfigUpdate {
1179                        node_id: vec![0x55; 32],
1180                        revision: 8,
1181                        config_hash: vec![0x77; 32],
1182                        config: Some(control_plane_test_config()),
1183                    }),
1184                }),
1185                apply_config: None,
1186                refresh_inventory: None,
1187                load_model: None,
1188                unload_model: None,
1189                ensure_model: None,
1190                drain_model: None,
1191            }),
1192            error: None,
1193        };
1194        decode_owner_control_envelope(&encode_owner_control_envelope(&update_response))
1195            .expect("watch update response must decode");
1196    }
1197
1198    #[test]
1199    fn control_plane_messages_unknown_command_rejects_with_structured_error() {
1200        let envelope = OwnerControlEnvelope {
1201            r#gen: NODE_PROTOCOL_GENERATION,
1202            handshake: None,
1203            request: Some(OwnerControlRequest {
1204                request_id: 42,
1205                get_config: None,
1206                watch_config: None,
1207                apply_config: None,
1208                refresh_inventory: None,
1209                load_model: None,
1210                unload_model: None,
1211                ensure_model: None,
1212                drain_model: None,
1213            }),
1214            response: None,
1215            error: None,
1216        };
1217        let bytes = encode_owner_control_envelope(&envelope);
1218        let err = decode_owner_control_envelope(&bytes)
1219            .expect_err("missing command variant must be rejected");
1220        assert!(matches!(err, ControlFrameError::MissingControlCommand));
1221
1222        let rejection = owner_control_rejection_envelope(&bytes, Some(42), &err);
1223        let error = rejection
1224            .error
1225            .expect("structured rejection must carry an error");
1226        assert_eq!(
1227            crate::proto::node::OwnerControlErrorCode::try_from(error.code).unwrap(),
1228            OwnerControlErrorCode::UnknownCommand
1229        );
1230        assert_eq!(error.request_id, Some(42));
1231    }
1232
1233    #[test]
1234    fn owner_control_handshake_empty_owner_id_uses_handshake_error() {
1235        let mut envelope = control_plane_test_handshake();
1236        envelope
1237            .handshake
1238            .as_mut()
1239            .and_then(|handshake| handshake.ownership.as_mut())
1240            .expect("test handshake must include ownership")
1241            .owner_id = "   ".to_string();
1242
1243        let err = decode_owner_control_envelope(&encode_owner_control_envelope(&envelope))
1244            .expect_err("handshake with blank owner_id must be rejected");
1245        assert!(matches!(err, ControlFrameError::MissingControlOwnerId));
1246        assert_eq!(err.to_string(), "owner control handshake missing owner_id");
1247    }
1248
1249    #[test]
1250    fn owner_control_error_rejects_invalid_error_code() {
1251        for code in [OwnerControlErrorCode::Unspecified as i32, 9999] {
1252            let err = OwnerControlError {
1253                code,
1254                message: "invalid".to_string(),
1255                request_id: Some(1),
1256                current_revision: None,
1257            }
1258            .validate_frame()
1259            .expect_err("invalid owner-control error code must be rejected");
1260            assert!(matches!(
1261                err,
1262                ControlFrameError::InvalidOwnerControlErrorCode { got } if got == code
1263            ));
1264            assert_eq!(
1265                err.to_string(),
1266                format!("invalid owner control error code: {code}")
1267            );
1268        }
1269    }
1270
1271    #[test]
1272    fn control_plane_messages_legacy_json_rejects_with_structured_error() {
1273        let legacy_json = br#"{"owner_id":"legacy","command":"GetConfig"}"#;
1274        let err = decode_owner_control_envelope(legacy_json)
1275            .expect_err("legacy json must not decode on protobuf-only control plane");
1276        let rejection = owner_control_rejection_envelope(legacy_json, Some(99), &err);
1277        let error = rejection
1278            .error
1279            .expect("structured rejection must carry an error");
1280        assert_eq!(
1281            crate::proto::node::OwnerControlErrorCode::try_from(error.code).unwrap(),
1282            OwnerControlErrorCode::LegacyJsonUnsupported
1283        );
1284        assert_eq!(error.request_id, Some(99));
1285    }
1286
1287    #[test]
1288    fn outbound_control_frame_size_rejects_before_write() {
1289        let oversized = vec![0u8; MAX_CONTROL_FRAME_BYTES + 1];
1290
1291        let err = ensure_control_frame_size(&oversized)
1292            .expect_err("oversize outbound frame must fail before length/body write");
1293
1294        assert!(matches!(
1295            err,
1296            ControlFrameError::OversizeFrame { size } if size == MAX_CONTROL_FRAME_BYTES + 1
1297        ));
1298    }
1299
1300    #[test]
1301    fn refresh_inventory_snapshot_only_frozen_bytes_remain_compatible() {
1302        let mut frozen = vec![0x0a, 0x48, 0x0a, 0x20];
1303        frozen.extend_from_slice(&[0x55; 32]);
1304        frozen.extend_from_slice(&[0x10, 0x07, 0x1a, 0x20]);
1305        frozen.extend_from_slice(&[0xa5; 32]);
1306        frozen.extend_from_slice(&[0x22, 0x00]);
1307        frozen.extend_from_slice(&[0x98, 0x06, 0x01]);
1308
1309        let decoded = OwnerControlRefreshInventoryResponse::decode(frozen.as_slice())
1310            .expect("frozen snapshot-only response must decode");
1311
1312        decoded
1313            .validate_frame()
1314            .expect("snapshot-only response from an old server must remain valid");
1315        assert!(decoded.inventory.is_none());
1316        assert_eq!(decoded.snapshot.expect("snapshot").revision, 7);
1317    }
1318
1319    #[test]
1320    fn refresh_inventory_rich_response_roundtrips_in_sorted_order() {
1321        let response = OwnerControlRefreshInventoryResponse {
1322            snapshot: Some(control_plane_test_snapshot()),
1323            inventory: Some(OwnerControlRefreshInventory {
1324                entries: vec![
1325                    inventory_entry("hf://mesh/alpha-GGUF:Q4_K_M", "alpha", 1024),
1326                    inventory_entry("hf://mesh/beta-GGUF:Q8_0", "beta", 2048),
1327                ],
1328                disposition: OwnerControlRefreshInventoryDisposition::Executed as i32,
1329            }),
1330        };
1331
1332        response
1333            .validate_frame()
1334            .expect("rich response must validate");
1335        let encoded = response.encode_to_vec();
1336        let decoded = OwnerControlRefreshInventoryResponse::decode(encoded.as_slice())
1337            .expect("rich response must decode");
1338
1339        assert_eq!(decoded, response);
1340        assert_eq!(decoded.inventory.expect("inventory").entries.len(), 2);
1341    }
1342
1343    #[test]
1344    fn peer_announcement_admission_round_trip() {
1345        let mut peer = PeerAnnouncement {
1346            endpoint_id: vec![0x42; 32],
1347            role: NodeRole::Worker as i32,
1348            inference_admission_state: Some(InferenceAdmissionState::Accepting as i32),
1349            ..Default::default()
1350        };
1351
1352        let frame_with_state = GossipFrame {
1353            r#gen: NODE_PROTOCOL_GENERATION,
1354            sender_id: vec![0x11; 32],
1355            peers: vec![peer.clone()],
1356        };
1357
1358        let encoded = frame_with_state.encode_to_vec();
1359        let decoded = GossipFrame::decode(encoded.as_slice())
1360            .expect("peer-announcement frame should decode after encode");
1361        let decoded_peer = &decoded.peers[0];
1362
1363        assert_eq!(
1364            decoded_peer.inference_admission_state,
1365            peer.inference_admission_state
1366        );
1367        assert!(decoded.validate_frame().is_ok());
1368
1369        peer.inference_admission_state = None;
1370        let frame_without_state = GossipFrame {
1371            peers: vec![peer],
1372            ..frame_with_state
1373        };
1374        let reencoded = frame_without_state.encode_to_vec();
1375        let redecoded = GossipFrame::decode(reencoded.as_slice())
1376            .expect("peer-announcement frame should decode without admission state");
1377
1378        assert!(redecoded.peers[0].inference_admission_state.is_none());
1379    }
1380
1381    #[test]
1382    fn refresh_inventory_rejects_unsorted_or_unspecified_details() {
1383        let mut inventory = OwnerControlRefreshInventory {
1384            entries: vec![
1385                inventory_entry("hf://mesh/beta", "beta", 2),
1386                inventory_entry("hf://mesh/alpha", "alpha", 1),
1387            ],
1388            disposition: OwnerControlRefreshInventoryDisposition::Executed as i32,
1389        };
1390        assert!(matches!(
1391            inventory.validate_frame(),
1392            Err(ControlFrameError::InvalidInventoryOrder)
1393        ));
1394
1395        inventory
1396            .entries
1397            .sort_by(|left, right| left.canonical_model_ref.cmp(&right.canonical_model_ref));
1398        inventory.disposition = OwnerControlRefreshInventoryDisposition::Unspecified as i32;
1399        assert!(matches!(
1400            inventory.validate_frame(),
1401            Err(ControlFrameError::InvalidInventoryDisposition { got: 0 })
1402        ));
1403    }
1404
1405    #[test]
1406    fn metadata_rich_inventory_has_deterministic_bounded_size() {
1407        let entries = (0..128)
1408            .map(|index| {
1409                let mut entry = inventory_entry(
1410                    &format!("hf://mesh/model-{index:03}"),
1411                    &format!("model-{index:03}"),
1412                    index,
1413                );
1414                entry.metadata.as_mut().expect("metadata").model_key = "x".repeat(70_000);
1415                entry
1416            })
1417            .collect();
1418        let response = OwnerControlRefreshInventoryResponse {
1419            snapshot: Some(control_plane_test_snapshot()),
1420            inventory: Some(OwnerControlRefreshInventory {
1421                entries,
1422                disposition: OwnerControlRefreshInventoryDisposition::Coalesced as i32,
1423            }),
1424        };
1425
1426        response
1427            .validate_frame()
1428            .expect("synthetic response is valid");
1429        let encoded = response.encode_to_vec();
1430        assert_eq!(encoded.len(), response.encoded_len());
1431        assert_eq!(encoded.len(), response.clone().encoded_len());
1432        assert!(encoded.len() > MAX_CONTROL_FRAME_BYTES);
1433        assert!(matches!(
1434            ensure_control_frame_size(&encoded),
1435            Err(ControlFrameError::OversizeFrame { size }) if size == encoded.len()
1436        ));
1437    }
1438
1439    fn inventory_entry(
1440        canonical_model_ref: &str,
1441        display_name: &str,
1442        total_size_bytes: u64,
1443    ) -> OwnerControlInventoryEntry {
1444        OwnerControlInventoryEntry {
1445            canonical_model_ref: canonical_model_ref.to_string(),
1446            display_name: Some(display_name.to_string()),
1447            total_size_bytes,
1448            metadata: Some(CompactModelMetadata {
1449                model_key: canonical_model_ref.to_string(),
1450                architecture: "llama".to_string(),
1451                quantization_type: "Q4_K_M".to_string(),
1452                ..Default::default()
1453            }),
1454        }
1455    }
1456}
mesh_llm_protocol/protocol/mod.rs

mesh_llm_protocol/protocol/
mod.rs
