fix(network): align REST and WebSocket endpoint handling (#5336)

Signed-off-by: John Myers <9696606+johntmyers@users.noreply.github.com>

Author: johntmyers

crates/openshell-policy-schema/src/lib.rs

@@ -285,7 +285,7 @@ pub struct NetworkEndpoint {
     /// Defaults to false (strict).
     #[serde(default, skip_serializing_if = "std::ops::Not::not")]
     pub allow_encoded_slash: bool,
-    /// When true, client-to-server WebSocket text messages on this REST
+    /// When true, client-to-server WebSocket text messages on a `protocol: websocket`
     /// endpoint rewrite credential placeholders after an allowed 101 upgrade.
     /// Defaults to false.
     #[serde(default, skip_serializing_if = "std::ops::Not::not")]

crates/openshell-policy/src/l7_validate.rs

@@ -223,6 +223,9 @@ pub struct L7EndpointFields<'a> {
     /// `true` when the endpoint has a non-empty rules list.
     pub has_rules: bool,
 
+    /// Whether WebSocket credential rewriting was enabled.
+    pub websocket_credential_rewrite: bool,
+
     /// `true` when the endpoint has a non-empty `deny_rules` list.
     pub has_deny_rules: bool,
 
@@ -244,6 +247,9 @@ pub fn validate_l7_endpoint_semantics(ep: &L7EndpointFields<'_>) -> Vec<String>
     let mut errors = Vec::new();
     let protocol = ep.protocol;
     let l7_protocol = L7Protocol::parse(protocol);
+    if matches!(l7_protocol, Some(L7Protocol::Rest)) && ep.websocket_credential_rewrite {
+        errors.push("websocket_credential_rewrite is not supported for protocol rest; use protocol: websocket with explicit message rules".to_string());
+    }
     let explicit_tcp = is_explicit_tcp_protocol(protocol);
     let jsonrpc_family = l7_protocol.is_some_and(L7Protocol::is_jsonrpc_family);
     let is_mcp = matches!(l7_protocol, Some(L7Protocol::Mcp));
@@ -335,6 +341,7 @@ mod tests {
 
     fn valid_rest_endpoint() -> L7EndpointFields<'static> {
         L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "rest",
             access: "read-only",
             has_rules: false,
@@ -344,6 +351,27 @@ mod tests {
         }
     }
 
+    #[test]
+    fn rest_rejects_websocket_credential_rewrite_with_migration_hint() {
+        let mut endpoint = valid_rest_endpoint();
+        endpoint.websocket_credential_rewrite = true;
+        for protocol in ["rest", "REST", "Rest"] {
+            endpoint.protocol = protocol;
+            let errors = validate_l7_endpoint_semantics(&endpoint);
+            assert!(
+                errors.iter().any(|e| e.contains("use protocol: websocket")),
+                "{protocol}: {errors:?}"
+            );
+        }
+        for protocol in ["websocket", "WEBSOCKET", "WebSocket"] {
+            endpoint.protocol = protocol;
+            assert!(
+                validate_l7_endpoint_semantics(&endpoint).is_empty(),
+                "{protocol}"
+            );
+        }
+    }
+
     #[test]
     fn valid_endpoint_produces_no_errors() {
         let errors = validate_l7_endpoint_semantics(&valid_rest_endpoint());
@@ -383,6 +411,7 @@ mod tests {
     #[test]
     fn rejects_unknown_protocol() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "ftp",
             access: "",
             has_rules: false,
@@ -397,6 +426,7 @@ mod tests {
     #[test]
     fn rejects_rules_and_access_together() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "rest",
             access: "full",
             has_rules: true,
@@ -411,6 +441,7 @@ mod tests {
     #[test]
     fn rejects_jsonrpc_with_access_presets() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "json-rpc",
             access: "full",
             has_rules: false,
@@ -429,6 +460,7 @@ mod tests {
     #[test]
     fn rejects_mcp_with_access_presets() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "mcp",
             access: "full",
             has_rules: false,
@@ -447,6 +479,7 @@ mod tests {
     #[test]
     fn rejects_jsonrpc_without_rules() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "json-rpc",
             access: "",
             has_rules: false,
@@ -461,6 +494,7 @@ mod tests {
     #[test]
     fn rejects_protocol_without_rules_or_access() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "rest",
             access: "",
             has_rules: false,
@@ -479,6 +513,7 @@ mod tests {
     #[test]
     fn rejects_mcp_without_rules_when_allow_all_false() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "mcp",
             access: "",
             has_rules: false,
@@ -493,6 +528,7 @@ mod tests {
     #[test]
     fn accepts_mcp_with_allow_all_true() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "mcp",
             access: "",
             has_rules: false,
@@ -507,6 +543,7 @@ mod tests {
     #[test]
     fn rejects_rules_that_deny_all() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "rest",
             access: "",
             has_rules: true,
@@ -521,6 +558,7 @@ mod tests {
     #[test]
     fn mcp_allow_all_with_empty_allow_not_rejected() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "mcp",
             access: "",
             has_rules: true,
@@ -538,6 +576,7 @@ mod tests {
     #[test]
     fn rejects_deny_rules_without_protocol() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "",
             access: "",
             has_rules: false,
@@ -556,6 +595,7 @@ mod tests {
     #[test]
     fn rejects_deny_rules_without_allow_base() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "rest",
             access: "",
             has_rules: false,
@@ -574,6 +614,7 @@ mod tests {
     #[test]
     fn no_protocol_no_errors() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "",
             access: "",
             has_rules: false,
@@ -588,6 +629,7 @@ mod tests {
     #[test]
     fn explicit_tcp_is_valid_without_l7_fields() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "tcp",
             access: "",
             has_rules: false,
@@ -604,6 +646,7 @@ mod tests {
     #[test]
     fn explicit_tcp_rejects_l7_fields() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "tcp",
             access: "full",
             has_rules: false,
@@ -648,6 +691,7 @@ mod tests {
     #[test]
     fn jsonrpc_with_access_emits_single_diagnostic() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "json-rpc",
             access: "full",
             has_rules: false,
@@ -669,6 +713,7 @@ mod tests {
     #[test]
     fn jsonrpc_without_rules_or_access_emits_single_diagnostic() {
         let ep = L7EndpointFields {
+            websocket_credential_rewrite: false,
             protocol: "json-rpc",
             access: "",
             has_rules: false,

crates/openshell-policy/src/lib.rs

@@ -1562,6 +1562,7 @@ fn validate_sandbox_policy_with_mcp_presence(
                     })
                 });
             let fields = L7EndpointFields {
+                websocket_credential_rewrite: ep.websocket_credential_rewrite,
                 protocol: &ep.protocol,
                 access,
                 has_rules: !ep.rules.is_empty(),
@@ -4786,6 +4787,18 @@ network_policies:
         );
     }
 
+    #[test]
+    fn validate_rest_rejects_websocket_credential_rewrite() {
+        for protocol in ["rest", "REST", "Rest"] {
+            let yaml = "version: 1\nnetwork_policies:\n  api:\n    endpoints:\n      - { host: api.example.com, port: 443, protocol: rest, access: read-only, websocket_credential_rewrite: true }\n";
+            let yaml = yaml.replace("protocol: rest", &format!("protocol: {protocol}"));
+            let policy = parse_sandbox_policy(&yaml).unwrap();
+            let errors = validate_sandbox_policy(&policy)
+                .expect_err("REST cannot rewrite WebSocket messages");
+            assert!(format!("{errors:?}").contains("use protocol: websocket"));
+        }
+    }
+
     #[test]
     fn round_trip_preserves_websocket_credential_rewrite() {
         let yaml = r"
@@ -4796,7 +4809,7 @@ network_policies:
     endpoints:
       - host: gateway.example.com
         port: 443
-        protocol: rest
+        protocol: websocket
         enforcement: enforce
         access: full
         websocket_credential_rewrite: true
@@ -4808,7 +4821,7 @@ network_policies:
         let proto2 = parse_sandbox_policy(&yaml_out).expect("re-parse failed");
 
         let ep = &proto2.network_policies["discord_gateway"].endpoints[0];
-        assert_eq!(ep.protocol, "rest");
+        assert_eq!(ep.protocol, "websocket");
         assert!(ep.websocket_credential_rewrite);
         assert!(yaml_out.contains("websocket_credential_rewrite: true"));
     }

crates/openshell-policy/src/merge.rs

@@ -3818,8 +3818,8 @@ mod tests {
     /// agent is told its approved rule never landed.
     #[test]
     fn merged_policy_covers_the_rule_it_merged() {
-        // Every field `merge_endpoint` widens rather than replaces, set on the
-        // existing endpoint and absent from the proposal.
+        // REST-compatible fields that `merge_endpoint` widens rather than
+        // replaces, set on the existing endpoint and absent from the proposal.
         let existing_endpoint = NetworkEndpoint {
             protocol: "rest".to_string(),
             rules: vec![rest_rule("GET", "/health")],
@@ -3830,7 +3830,6 @@ mod tests {
             }],
             allowed_ips: vec!["10.0.0.1".to_string()],
             allow_encoded_slash: true,
-            websocket_credential_rewrite: true,
             request_body_credential_rewrite: true,
             advisor_proposed: true,
             ..endpoint("api.example.com", 443)

crates/openshell-providers/src/profiles.rs

@@ -2759,6 +2759,7 @@ pub fn validate_profile_set(
             }
 
             let l7_fields = L7EndpointFields {
+                websocket_credential_rewrite: endpoint.websocket_credential_rewrite,
                 protocol: &endpoint.protocol,
                 access: &endpoint.access,
                 has_rules: endpoint.rules.as_ref().is_some_and(|r| !r.is_empty()),
@@ -6330,6 +6331,42 @@ binaries: ["", /usr/bin/broken]
         assert!(messages.contains(&"binary path must not be empty"));
     }
 
+    #[test]
+    fn validate_profile_rejects_rest_websocket_credential_rewrite() {
+        for protocol in ["rest", "REST", "Rest"] {
+            let profile = parse_profile_yaml(
+                r"
+id: rest-rewrite
+display_name: REST rewrite
+credentials:
+  - name: api_key
+    env_vars: [API_KEY]
+    auth_style: bearer
+    header_name: authorization
+discovery:
+  credentials: [api_key]
+endpoints:
+  - host: api.example.com
+    port: 443
+    protocol: rest
+    access: read-only
+    websocket_credential_rewrite: true
+binaries:
+  - /usr/bin/curl
+"
+                .replace("protocol: rest", &format!("protocol: {protocol}"))
+                .as_str(),
+            )
+            .unwrap();
+            let diagnostics = validate_profile_set(&[("profile.yaml".to_string(), profile)]);
+            assert!(
+                diagnostics
+                    .iter()
+                    .any(|d| d.message.contains("use protocol: websocket"))
+            );
+        }
+    }
+
     #[test]
     fn validate_profile_set_rejects_noncanonical_profile_ids() {
         let profiles = [

crates/openshell-supervisor-network/src/l7/mod.rs

@@ -379,6 +379,9 @@ pub fn parse_l7_config(val: &regorus::Value) -> Option<L7EndpointConfig> {
     let allow_encoded_slash = get_object_bool(val, "allow_encoded_slash").unwrap_or(false);
     let websocket_credential_rewrite =
         get_object_bool(val, "websocket_credential_rewrite").unwrap_or(false);
+    if protocol == L7Protocol::Rest && websocket_credential_rewrite {
+        return None;
+    }
     let request_body_credential_rewrite =
         get_object_bool(val, "request_body_credential_rewrite").unwrap_or(false);
     let allow_uninspected_credentials =
@@ -1393,6 +1396,10 @@ pub fn validate_l7_policies(data_json: &serde_json::Value) -> (Vec<String>, Vec<
                 .and_then(serde_json::Value::as_bool)
                 .unwrap_or(false);
             let l7_fields = L7EndpointFields {
+                websocket_credential_rewrite: ep
+                    .get("websocket_credential_rewrite")
+                    .and_then(serde_json::Value::as_bool)
+                    .unwrap_or(false),
                 protocol,
                 access,
                 has_rules,
@@ -1504,7 +1511,7 @@ pub fn validate_l7_policies(data_json: &serde_json::Value) -> (Vec<String>, Vec<
                 && protocol != "websocket"
             {
                 warnings.push(format!(
-                    "{loc}: websocket_credential_rewrite is ignored unless protocol is rest or websocket"
+                    "{loc}: websocket_credential_rewrite is ignored unless protocol is websocket"
                 ));
             }
 
@@ -2470,10 +2477,24 @@ mod tests {
         assert!(!config.websocket_credential_rewrite);
     }
 
+    #[test]
+    fn rest_websocket_credential_rewrite_fails_validation_and_materialization() {
+        let endpoint = serde_json::json!({
+            "protocol": "rest", "host": "gateway.example.com", "port": 443,
+            "access": "read-only", "websocket_credential_rewrite": true
+        });
+        let val = regorus::Value::from_json_str(&endpoint.to_string()).unwrap();
+        assert!(parse_l7_config(&val).is_none());
+        let (errors, _) = validate_l7_policies(&serde_json::json!({
+            "network_policies": {"test": {"endpoints": [endpoint], "binaries": []}}
+        }));
+        assert!(errors.iter().any(|e| e.contains("use protocol: websocket")));
+    }
+
     #[test]
     fn parse_l7_config_websocket_credential_rewrite_opt_in() {
         let val = regorus::Value::from_json_str(
-            r#"{"protocol": "rest", "host": "gateway.example.com", "port": 443, "websocket_credential_rewrite": true}"#,
+            r#"{"protocol": "websocket", "host": "gateway.example.com", "port": 443, "websocket_credential_rewrite": true}"#,
         )
         .unwrap();
         let config = parse_l7_config(&val).unwrap();
@@ -2521,7 +2542,7 @@ mod tests {
     }
 
     #[test]
-    fn validate_websocket_credential_rewrite_warns_unless_rest_or_websocket() {
+    fn validate_websocket_credential_rewrite_warns_without_inspected_protocol() {
         let data = serde_json::json!({
             "network_policies": {
                 "test": {