pb_mapper_protocol/
command.rs

1use serde::{Deserialize, Serialize};
2use snafu::ResultExt;
3
4use pb_mapper_auth::{
5    AuthStatus, IssuedTemporaryKey, KeyPage, LegacyProtocolPolicy, TemporaryKeyMetadata,
6};
7use pb_mapper_core::checksum::AesKeyType;
8use pb_mapper_core::error::{MsgSerializeSnafu, Result};
9
10pub const CONTROL_PROTOCOL_V2: u16 = 2;
11
12pub trait MessageSerializer {
13    fn encode(&self) -> Result<Vec<u8>>;
14    fn decode(msg: &[u8]) -> Result<Self>
15    where
16        Self: Sized;
17}
18
19#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
20pub enum PbConnStatusReq {
21    RemoteId,
22    Keys,
23    Service {
24        key: String,
25        #[serde(default, skip_serializing_if = "Option::is_none")]
26        direct_query: Option<bool>,
27    },
28    Sessions,
29}
30
31#[derive(Debug, Clone, Deserialize, Serialize)]
32pub enum PbConnStatusResp {
33    Sessions(Vec<crate::direct::SessionObservation>),
34    RemoteId {
35        server_map: String,
36        active: String,
37        idle: String,
38    },
39    Keys(Vec<String>),
40    Service {
41        key: String,
42        connections: Vec<PbServiceConnStatus>,
43        #[serde(default, skip_serializing_if = "Option::is_none")]
44        direct: Option<crate::direct::DirectCapability>,
45    },
46}
47
48#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
49pub struct PbServiceConnStatus {
50    pub conn_id: u32,
51    pub generation: u64,
52    pub protocol_version: u16,
53    pub healthy: bool,
54    pub last_rx_age_ms: u64,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
58pub enum PbConnRequest {
59    /// Authenticated discovery; available to administrator and temporary credentials.
60    QuicEndpoint,
61    DirectControl {
62        namespace: Option<u64>,
63    },
64    DirectRelay {
65        id: crate::direct::SessionId,
66        secret: crate::direct::SessionSecret,
67        namespace: Option<u64>,
68        data_protocol: Option<u16>,
69    },
70    Register {
71        need_codec: bool,
72        is_datagram: bool,
73        key: String,
74        #[serde(default, skip_serializing_if = "Option::is_none")]
75        protocol_version: Option<u16>,
76        #[serde(default, skip_serializing_if = "Option::is_none")]
77        client_instance_id: Option<String>,
78        #[serde(default, skip_serializing_if = "Option::is_none")]
79        heartbeat_interval_ms: Option<u64>,
80        #[serde(default, skip_serializing_if = "Option::is_none")]
81        heartbeat_tolerance_ms: Option<u64>,
82        #[serde(default, skip_serializing_if = "Option::is_none")]
83        direct: Option<u16>,
84    },
85    RegisterScoped {
86        need_codec: bool,
87        is_datagram: bool,
88        key: String,
89        namespace: u64,
90        force_namespace: bool,
91        #[serde(default, skip_serializing_if = "Option::is_none")]
92        protocol_version: Option<u16>,
93        #[serde(default, skip_serializing_if = "Option::is_none")]
94        client_instance_id: Option<String>,
95        #[serde(default, skip_serializing_if = "Option::is_none")]
96        heartbeat_interval_ms: Option<u64>,
97        #[serde(default, skip_serializing_if = "Option::is_none")]
98        heartbeat_tolerance_ms: Option<u64>,
99        #[serde(default, skip_serializing_if = "Option::is_none")]
100        direct: Option<u16>,
101    },
102    Subcribe {
103        key: String,
104        #[serde(default, skip_serializing_if = "Option::is_none")]
105        data_protocol: Option<u16>,
106    },
107    SubcribeScoped {
108        key: String,
109        namespace: u64,
110        #[serde(default, skip_serializing_if = "Option::is_none")]
111        data_protocol: Option<u16>,
112    },
113    Status(PbConnStatusReq),
114    StatusScoped {
115        status: PbConnStatusReq,
116        namespace: u64,
117    },
118    Stream {
119        key: String,
120        #[serde(default, skip_serializing_if = "Option::is_none")]
121        data_protocol: Option<u16>,
122        dst_id: u32,
123        #[serde(default)]
124        server_generation: u64,
125    },
126    StreamScoped {
127        key: String,
128        #[serde(default, skip_serializing_if = "Option::is_none")]
129        data_protocol: Option<u16>,
130        namespace: u64,
131        dst_id: u32,
132        #[serde(default)]
133        server_generation: u64,
134    },
135    Admin(AdminRequest),
136}
137
138#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
139pub enum AdminRequest {
140    KeyIssue {
141        ttl_seconds: u64,
142        #[serde(default, skip_serializing_if = "Option::is_none")]
143        label: Option<String>,
144    },
145    KeyList {
146        #[serde(default)]
147        page: u32,
148        #[serde(default = "default_page_size")]
149        page_size: u16,
150    },
151    KeyShow {
152        key_id: u64,
153    },
154    KeyReveal {
155        key_id: u64,
156    },
157    KeyRenew {
158        key_id: u64,
159        ttl_seconds: u64,
160    },
161    KeyRevoke {
162        key_id: u64,
163    },
164    KeyGc,
165    AuthStatus,
166    AuthStateReset {
167        confirm: bool,
168    },
169    RootKeyRotate {
170        new_admin_key: String,
171    },
172    LegacyProtocolSet {
173        policy: LegacyProtocolPolicy,
174    },
175    ConnectionList {
176        #[serde(default, skip_serializing_if = "Option::is_none")]
177        key_id: Option<u64>,
178        #[serde(default)]
179        page: u32,
180        #[serde(default = "default_page_size")]
181        page_size: u16,
182    },
183    ServiceList {
184        #[serde(default, skip_serializing_if = "Option::is_none")]
185        key_id: Option<u64>,
186        #[serde(default)]
187        page: u32,
188        #[serde(default = "default_page_size")]
189        page_size: u16,
190    },
191    /// Drop registered control connections the relay is still holding.
192    ///
193    /// The manual counterpart to the relay's own lease sweep, for the case an
194    /// operator can see but the relay cannot: a registration that answers its
195    /// heartbeat yet no longer forwards, or a service whose connection quota is
196    /// full of connections that should have gone away.
197    ConnectionRetire {
198        /// Namespace owning the service. Absent means the unscoped namespace,
199        /// which is where an administrator's own registrations live.
200        #[serde(default, skip_serializing_if = "Option::is_none")]
201        key_id: Option<u64>,
202        service_name: String,
203        /// Retire only this connection. Absent retires every connection the
204        /// service has, which is what frees a full quota in one call.
205        #[serde(default, skip_serializing_if = "Option::is_none")]
206        conn_id: Option<u32>,
207    },
208}
209
210impl AdminRequest {
211    pub fn is_mutating(&self) -> bool {
212        matches!(
213            self,
214            Self::KeyIssue { .. }
215                | Self::KeyRenew { .. }
216                | Self::KeyRevoke { .. }
217                | Self::KeyGc
218                | Self::AuthStateReset { .. }
219                | Self::RootKeyRotate { .. }
220                | Self::LegacyProtocolSet { .. }
221                | Self::ConnectionRetire { .. }
222        )
223    }
224}
225
226const fn default_page_size() -> u16 {
227    100
228}
229
230#[derive(Debug, Clone, Deserialize, Serialize)]
231pub struct PbErrorResponse {
232    pub code: String,
233    pub message: String,
234    pub retryable: bool,
235    pub server_time: u64,
236}
237
238#[derive(Debug, Clone, Deserialize, Serialize)]
239pub struct AdminServiceInfo {
240    pub key_id: u64,
241    pub namespace: u64,
242    pub service_name: String,
243    pub transport: String,
244    pub codec_enabled: bool,
245    pub connection_count: u32,
246}
247
248#[derive(Debug, Clone, Deserialize, Serialize)]
249pub struct AdminConnectionInfo {
250    pub key_id: u64,
251    pub namespace: u64,
252    pub service_name: String,
253    pub conn_id: u32,
254    pub generation: u64,
255    pub protocol_version: u16,
256    pub healthy: bool,
257    pub transport: String,
258    pub codec_enabled: bool,
259    pub last_rx_age_ms: u64,
260}
261
262#[derive(Debug, Clone, Deserialize, Serialize)]
263pub struct AdminServicePage {
264    pub schema_version: u16,
265    pub items: Vec<AdminServiceInfo>,
266    pub next_page: Option<u32>,
267}
268
269#[derive(Debug, Clone, Deserialize, Serialize)]
270pub struct AdminConnectionPage {
271    pub schema_version: u16,
272    pub items: Vec<AdminConnectionInfo>,
273    pub next_page: Option<u32>,
274}
275
276#[derive(Debug, Clone, Deserialize, Serialize)]
277pub enum AdminResponse {
278    KeyIssued(IssuedTemporaryKey),
279    KeyList(KeyPage),
280    KeyShown(IssuedTemporaryKey),
281    KeyRenewed(IssuedTemporaryKey),
282    KeyRevoked(TemporaryKeyMetadata),
283    KeyGc {
284        removed: u64,
285    },
286    AuthStatus(AuthStatus),
287    Services(AdminServicePage),
288    Connections(AdminConnectionPage),
289    /// How many registered connections `ConnectionRetire` actually dropped. Zero
290    /// is a normal answer: the target may have unwound on its own first.
291    ConnectionsRetired {
292        retired: u32,
293    },
294    Ok {
295        action: String,
296    },
297}
298
299#[derive(Debug, Clone, Deserialize, Serialize)]
300pub enum PbConnResponse {
301    QuicEndpoint(Option<QuicEndpoint>),
302    Direct(crate::direct::DirectEvent),
303    Register(u32),
304    RegisterV2 {
305        conn_id: u32,
306        generation: u64,
307        lease_ttl_ms: u64,
308        #[serde(default, skip_serializing_if = "Option::is_none")]
309        direct: Option<crate::direct::DirectCapability>,
310    },
311    Subcribe {
312        codec_key: Option<AesKeyType>,
313        #[serde(default, skip_serializing_if = "Option::is_none")]
314        data_protocol: Option<u16>,
315        client_id: u32,
316        server_id: u32,
317    },
318    Stream {
319        codec_key: Option<AesKeyType>,
320        #[serde(default, skip_serializing_if = "Option::is_none")]
321        data_protocol: Option<u16>,
322    },
323    Status(PbConnStatusResp),
324    Admin(AdminResponse),
325    Error(PbErrorResponse),
326}
327
328impl PbConnResponse {
329    pub fn error(code: impl Into<String>, message: impl Into<String>, retryable: bool) -> Self {
330        Self::Error(PbErrorResponse {
331            code: code.into(),
332            message: message.into(),
333            retryable,
334            server_time: std::time::SystemTime::now()
335                .duration_since(std::time::UNIX_EPOCH)
336                .unwrap_or_default()
337                .as_secs(),
338        })
339    }
340}
341
342#[derive(Debug, Clone, Deserialize, Serialize)]
343pub enum PbServerRequest {
344    Direct(crate::direct::DirectProviderMessage),
345    Ping,
346    PingV2 {
347        seq: u64,
348    },
349    StreamAck {
350        client_id: u32,
351        #[serde(default)]
352        server_generation: u64,
353    },
354}
355
356#[derive(Debug, Clone, Deserialize, Serialize)]
357pub enum LocalServer {
358    Direct(Box<crate::direct::DirectEvent>),
359    /// pb server makes a stream request to local server
360    Stream {
361        client_id: u32,
362        #[serde(default)]
363        server_generation: u64,
364    },
365    /// pb server response a pong msg when it receive a ping request
366    Pong,
367    PongV2 {
368        seq: u64,
369    },
370    Retire {
371        reason: String,
372        conn_id: u32,
373        #[serde(default)]
374        server_generation: u64,
375    },
376}
377
378macro_rules! gen_impl_msg_serializer {
379    ($struct_name:ident) => {
380        impl MessageSerializer for $struct_name {
381            fn encode(&self) -> Result<Vec<u8>> {
382                serde_json::to_vec(self).with_context(|_| MsgSerializeSnafu {
383                    action: "encode",
384                    struct_name: stringify!($struct_name),
385                    content: "payload redacted".to_string(),
386                })
387            }
388
389            fn decode(msg: &[u8]) -> Result<Self> {
390                serde_json::from_slice(msg).with_context(|_| MsgSerializeSnafu {
391                    action: "decode",
392                    struct_name: stringify!($struct_name),
393                    content: format!("{}-byte payload redacted", msg.len()),
394                })
395            }
396        }
397    };
398}
399
400gen_impl_msg_serializer!(PbConnRequest);
401gen_impl_msg_serializer!(PbConnResponse);
402gen_impl_msg_serializer!(PbServerRequest);
403gen_impl_msg_serializer!(LocalServer);
404
405#[cfg(test)]
406mod tests {
407    use super::PbConnRequest;
408
409    /// The wire form of `Register` is load-bearing: a running peer on the other
410    /// side of an upgrade has to keep parsing it. The `None` fields must stay
411    /// absent from the JSON rather than serialise as null.
412    #[test]
413    fn test_serde_mapper_header() {
414        let mapper = PbConnRequest::Register {
415            direct: None,
416            key: "test".into(),
417            need_codec: false,
418            is_datagram: false,
419            protocol_version: None,
420            client_instance_id: None,
421            heartbeat_interval_ms: None,
422            heartbeat_tolerance_ms: None,
423        };
424        let json_value = serde_json::to_string(&mapper).unwrap();
425        let raw_json_str =
426            r##"{"Register":{"need_codec":false,"is_datagram":false,"key":"test"}}"##;
427        assert_eq!(raw_json_str, json_value);
428
429        let value: PbConnRequest = serde_json::from_str(raw_json_str).unwrap();
430        assert_eq!(mapper, value)
431    }
432}
433
434/// Only a UDP port is advertised; use the IP of the authenticated TCP peer.
435#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
436pub struct QuicEndpoint {
437    pub port: u16,
438    pub certificate: Vec<u8>,
439}
pb_mapper_protocol/command.rs

pb_mapper_protocol/
command.rs
