Skip to main content

arete_server/websocket/
subscription.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use crate::websocket::auth::AuthDeny;
7
8pub const PROTOCOL_VERSION: u8 = 2;
9pub const MAX_SUBSCRIPTION_ID_BYTES: usize = 128;
10
11/// Client message types for subscription management.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13#[serde(tag = "type", rename_all = "snake_case")]
14pub enum ClientMessage {
15    Subscribe(Subscription),
16    Unsubscribe(Unsubscription),
17    Ping,
18    RefreshAuth(RefreshAuthRequest),
19}
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
22#[serde(deny_unknown_fields)]
23pub struct RefreshAuthRequest {
24    pub token: String,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
28#[serde(rename_all = "camelCase")]
29pub struct RefreshAuthResponse {
30    pub success: bool,
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub error: Option<String>,
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub expires_at: Option<u64>,
35}
36
37/// Server-sent protocol, authorization, or quota error.
38#[derive(Debug, Clone, Serialize, Deserialize)]
39#[serde(rename_all = "camelCase")]
40pub struct SocketIssueMessage {
41    #[serde(rename = "type")]
42    pub kind: String,
43    pub protocol_version: u8,
44    /// Null only when a request did not contain a usable subscription ID.
45    pub subscription_id: Option<String>,
46    pub error: String,
47    pub message: String,
48    pub code: String,
49    pub retryable: bool,
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub retry_after: Option<u64>,
52    #[serde(skip_serializing_if = "Option::is_none")]
53    pub suggested_action: Option<String>,
54    #[serde(skip_serializing_if = "Option::is_none")]
55    pub docs_url: Option<String>,
56    pub fatal: bool,
57}
58
59impl SocketIssueMessage {
60    pub fn from_auth_deny(deny: &AuthDeny, fatal: bool, subscription_id: Option<String>) -> Self {
61        let response = deny.to_error_response();
62        Self {
63            kind: "error".to_string(),
64            protocol_version: PROTOCOL_VERSION,
65            subscription_id,
66            error: response.error,
67            message: response.message,
68            code: response.code,
69            retryable: response.retryable,
70            retry_after: response.retry_after,
71            suggested_action: response.suggested_action,
72            docs_url: response.docs_url,
73            fatal,
74        }
75    }
76
77    pub fn protocol(
78        subscription_id: Option<String>,
79        code: impl Into<String>,
80        message: impl Into<String>,
81    ) -> Self {
82        let code = code.into();
83        Self {
84            kind: "error".to_string(),
85            protocol_version: PROTOCOL_VERSION,
86            subscription_id,
87            error: code.clone(),
88            message: message.into(),
89            code,
90            retryable: false,
91            retry_after: None,
92            suggested_action: None,
93            docs_url: None,
94            fatal: false,
95        }
96    }
97}
98
99/// Canonical query identity and behavior for protocol v2.
100#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
101#[serde(rename_all = "camelCase", deny_unknown_fields)]
102pub struct SubscriptionQuery {
103    pub view: String,
104    #[serde(skip_serializing_if = "Option::is_none")]
105    pub key: Option<String>,
106    #[serde(skip_serializing_if = "Option::is_none")]
107    pub partition: Option<String>,
108    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
109    pub filters: BTreeMap<String, Value>,
110    #[serde(skip_serializing_if = "Option::is_none")]
111    pub take: Option<usize>,
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub skip: Option<usize>,
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub after: Option<String>,
116    #[serde(skip_serializing_if = "Option::is_none")]
117    pub snapshot_limit: Option<usize>,
118}
119
120impl SubscriptionQuery {
121    pub fn matches_key(&self, key: &str) -> bool {
122        self.key.as_ref().is_none_or(|expected| expected == key)
123    }
124
125    pub fn validate(&self) -> Result<(), &'static str> {
126        if self.view.trim().is_empty() {
127            return Err("query.view must not be empty");
128        }
129        if self.take == Some(0) {
130            return Err("query.take must be greater than zero");
131        }
132        if self.snapshot_limit == Some(0) {
133            return Err("query.snapshotLimit must be greater than zero");
134        }
135        if self
136            .filters
137            .keys()
138            .any(|path| path.is_empty() || path.split('.').any(|segment| segment.is_empty()))
139        {
140            return Err("query filter paths must contain non-empty dot-path segments");
141        }
142        Ok(())
143    }
144}
145
146#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
147#[serde(rename_all = "camelCase", deny_unknown_fields)]
148pub struct SnapshotOptions {
149    #[serde(default = "default_snapshot_enabled")]
150    pub enabled: bool,
151}
152
153impl Default for SnapshotOptions {
154    fn default() -> Self {
155        Self { enabled: true }
156    }
157}
158
159fn default_snapshot_enabled() -> bool {
160    true
161}
162
163#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
164#[serde(rename_all = "camelCase", deny_unknown_fields)]
165pub struct Subscription {
166    pub protocol_version: u8,
167    pub subscription_id: String,
168    pub query: SubscriptionQuery,
169    #[serde(default)]
170    pub snapshot: SnapshotOptions,
171}
172
173impl Subscription {
174    pub fn validate(&self) -> Result<(), &'static str> {
175        validate_protocol(self.protocol_version)?;
176        validate_subscription_id(&self.subscription_id)?;
177        self.query.validate()
178    }
179}
180
181#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
182#[serde(rename_all = "camelCase", deny_unknown_fields)]
183pub struct Unsubscription {
184    pub protocol_version: u8,
185    pub subscription_id: String,
186}
187
188impl Unsubscription {
189    pub fn validate(&self) -> Result<(), &'static str> {
190        validate_protocol(self.protocol_version)?;
191        validate_subscription_id(&self.subscription_id)
192    }
193}
194
195pub fn validate_protocol(version: u8) -> Result<(), &'static str> {
196    if version == PROTOCOL_VERSION {
197        Ok(())
198    } else {
199        Err("protocolVersion must be 2")
200    }
201}
202
203pub fn validate_subscription_id(subscription_id: &str) -> Result<(), &'static str> {
204    if subscription_id.is_empty() {
205        return Err("subscriptionId must not be empty");
206    }
207    if subscription_id.len() > MAX_SUBSCRIPTION_ID_BYTES {
208        return Err("subscriptionId exceeds 128 bytes");
209    }
210    if subscription_id.trim() != subscription_id || subscription_id.chars().any(char::is_control) {
211        return Err("subscriptionId must be opaque non-whitespace text without control characters");
212    }
213    Ok(())
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use crate::websocket::auth::{AuthDeny, AuthErrorCode};
220    use serde_json::json;
221
222    fn valid_subscription() -> Value {
223        json!({
224            "type": "subscribe",
225            "protocolVersion": 2,
226            "subscriptionId": "rounds:page-1",
227            "query": {
228                "view": "OreRound/latest",
229                "filters": {"state.status": "open"},
230                "take": 10,
231                "skip": 0,
232                "after": "100:000000000001",
233                "snapshotLimit": 10
234            },
235            "snapshot": {"enabled": true}
236        })
237    }
238
239    #[test]
240    fn parses_canonical_v2_subscription() {
241        let message: ClientMessage = serde_json::from_value(valid_subscription()).unwrap();
242        let ClientMessage::Subscribe(subscription) = message else {
243            panic!("expected subscribe");
244        };
245
246        assert!(subscription.validate().is_ok());
247        assert_eq!(subscription.subscription_id, "rounds:page-1");
248        assert_eq!(subscription.query.filters["state.status"], "open");
249        assert!(subscription.snapshot.enabled);
250    }
251
252    #[test]
253    fn rejects_legacy_and_unknown_subscription_fields() {
254        assert!(serde_json::from_value::<ClientMessage>(json!({
255            "type": "subscribe",
256            "view": "OreRound/latest"
257        }))
258        .is_err());
259
260        let mut value = valid_subscription();
261        value["withSnapshot"] = json!(true);
262        assert!(serde_json::from_value::<ClientMessage>(value).is_err());
263    }
264
265    #[test]
266    fn validates_protocol_and_opaque_id() {
267        let mut value = valid_subscription();
268        value["protocolVersion"] = json!(1);
269        let ClientMessage::Subscribe(subscription) =
270            serde_json::from_value::<ClientMessage>(value).unwrap()
271        else {
272            panic!("expected subscribe");
273        };
274        assert_eq!(subscription.validate(), Err("protocolVersion must be 2"));
275
276        assert!(validate_subscription_id("client-selected.id:1").is_ok());
277        assert!(validate_subscription_id("").is_err());
278        assert!(validate_subscription_id(" leading").is_err());
279        assert!(validate_subscription_id(&"x".repeat(129)).is_err());
280    }
281
282    #[test]
283    fn parses_id_only_unsubscribe() {
284        let message: ClientMessage = serde_json::from_value(json!({
285            "type": "unsubscribe",
286            "protocolVersion": 2,
287            "subscriptionId": "rounds:page-1"
288        }))
289        .unwrap();
290        let ClientMessage::Unsubscribe(unsubscription) = message else {
291            panic!("expected unsubscribe");
292        };
293        assert!(unsubscription.validate().is_ok());
294    }
295
296    #[test]
297    fn socket_issue_carries_protocol_and_subscription_identity() {
298        let deny = AuthDeny::new(
299            AuthErrorCode::SubscriptionLimitExceeded,
300            "subscription limit exceeded",
301        );
302        let issue =
303            SocketIssueMessage::from_auth_deny(&deny, false, Some("rounds:page-1".to_string()));
304        assert_eq!(issue.protocol_version, PROTOCOL_VERSION);
305        assert_eq!(issue.subscription_id.as_deref(), Some("rounds:page-1"));
306        assert_eq!(issue.code, "subscription-limit-exceeded");
307    }
308}