[Fix] User States

src/omega/capabilities.rs

@@ -0,0 +1,145 @@
+use std::collections::BTreeSet;
+
+const OMIKRON_PREFIX: &str = "omikron;caps=";
+const OMEGA_PREFIX: &str = "omega;caps=";
+const SET_USER_STATE: &str = "set_user_state_v1";
+const STATE_SUBSCRIBE: &str = "state_subscribe_v1";
+const SESSION_SNAPSHOT: &str = "session_snapshot_v1";
+const CLIENT_STATE_PUSH: &str = "client_state_push_v1";
+
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct PeerCapabilities {
+    pub set_user_state_v1: bool,
+    pub state_subscribe_v1: bool,
+    pub session_snapshot_v1: bool,
+    pub client_state_push_v1: bool,
+}
+
+impl PeerCapabilities {
+    pub fn current() -> Self {
+        Self {
+            set_user_state_v1: true,
+            state_subscribe_v1: true,
+            session_snapshot_v1: true,
+            client_state_push_v1: true,
+        }
+    }
+
+    pub fn identification_description(&self) -> String {
+        let mut names = Vec::new();
+        if self.set_user_state_v1 {
+            names.push(SET_USER_STATE);
+        }
+        if self.state_subscribe_v1 {
+            names.push(STATE_SUBSCRIBE);
+        }
+        if self.session_snapshot_v1 {
+            names.push(SESSION_SNAPSHOT);
+        }
+        if self.client_state_push_v1 {
+            names.push(CLIENT_STATE_PUSH);
+        }
+        format!("{OMIKRON_PREFIX}{}", names.join(","))
+    }
+
+    pub fn from_identification_description(description: Option<&str>) -> Result<Self, ()> {
+        parse_capabilities(description, OMIKRON_PREFIX)
+    }
+
+    pub fn from_omega_identification_description(description: Option<&str>) -> Result<Self, ()> {
+        // A missing Omega descriptor selects the documented legacy wire
+        // contracts: tuple route snapshots, GetStates-only subscriptions,
+        // and ClientChanged pushes.
+        parse_capabilities(description, OMEGA_PREFIX)
+    }
+}
+
+fn parse_capabilities(description: Option<&str>, prefix: &str) -> Result<PeerCapabilities, ()> {
+    let Some(description) = description else {
+        return Ok(PeerCapabilities::default());
+    };
+    if description == "omega" {
+        return Ok(PeerCapabilities::default());
+    }
+    let Some(capabilities) = description.strip_prefix(prefix) else {
+        return Err(());
+    };
+    let mut seen = BTreeSet::new();
+    for capability in capabilities.split(',') {
+        if capability.is_empty() || !seen.insert(capability) {
+            return Err(());
+        }
+    }
+    if seen.iter().any(|capability| {
+        !matches!(
+            *capability,
+            SET_USER_STATE | STATE_SUBSCRIBE | SESSION_SNAPSHOT | CLIENT_STATE_PUSH
+        )
+    }) {
+        return Err(());
+    }
+    Ok(PeerCapabilities {
+        set_user_state_v1: seen.contains(SET_USER_STATE),
+        state_subscribe_v1: seen.contains(STATE_SUBSCRIBE),
+        session_snapshot_v1: seen.contains(SESSION_SNAPSHOT),
+        client_state_push_v1: seen.contains(CLIENT_STATE_PUSH),
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn capability_description_round_trips() {
+        let capabilities = PeerCapabilities::current();
+        assert_eq!(
+            PeerCapabilities::from_identification_description(Some(
+                &capabilities.identification_description()
+            )),
+            Ok(capabilities)
+        );
+    }
+
+    #[test]
+    fn malformed_capabilities_are_rejected() {
+        assert!(
+            PeerCapabilities::from_identification_description(Some("omikron;caps=unknown_v9"))
+                .is_err()
+        );
+        assert!(
+            PeerCapabilities::from_identification_description(Some(
+                "omikron;caps=state_subscribe_v1,state_subscribe_v1"
+            ))
+            .is_err()
+        );
+        assert!(PeerCapabilities::from_identification_description(Some("omikron;caps=")).is_err());
+    }
+
+    #[test]
+    fn missing_capabilities_mean_legacy_peer() {
+        assert_eq!(
+            PeerCapabilities::from_identification_description(None),
+            Ok(PeerCapabilities::default())
+        );
+        assert_eq!(
+            PeerCapabilities::from_omega_identification_description(Some("omega")),
+            Ok(PeerCapabilities::default())
+        );
+    }
+
+    #[test]
+    fn omega_capabilities_use_the_omega_namespace() {
+        assert!(
+            PeerCapabilities::from_omega_identification_description(Some(
+                "omega;caps=set_user_state_v1"
+            ))
+            .unwrap()
+            .set_user_state_v1
+        );
+        assert!(
+            PeerCapabilities::from_identification_description(Some("omega;caps=set_user_state_v1"))
+                .is_err()
+        );
+    }
+}