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

pb_mapper_protocol/
command.rs
