fix(server): reject reserved annotation keys (#4626)

Signed-off-by: Eric Curtin <eric.curtin@docker.com>

Author: ericcurtin

crates/openshell-server/src/grpc/policy.rs

@@ -112,7 +112,7 @@ use tonic::{Code, Request, Response, Status};
 use tracing::{debug, info, warn};
 
 use super::validation::{
-    level_matches, source_matches, validate_and_canonicalize_policy, validate_annotations,
+    level_matches, source_matches, validate_and_canonicalize_policy, validate_caller_annotations,
     validate_no_reserved_provider_policy_keys, validate_policy_safety,
     validate_static_fields_unchanged,
 };
@@ -4167,7 +4167,7 @@ async fn handle_update_config_inner(
 ) -> Result<Response<UpdateConfigResponse>, Status> {
     let replay_facts = super::mutation_replay::ordinary::Facts::from_request(&request);
     let req = request.into_inner();
-    validate_annotations(&req.annotations, "annotations")?;
+    validate_caller_annotations(&req.annotations, "annotations")?;
     let sandbox = if req.global {
         if !req.sandbox.is_empty() || req.workspace_scope.is_some() {
             return Err(Status::invalid_argument(
@@ -24751,6 +24751,49 @@ mod tests {
         assert!(err.message().contains("label key"));
     }
 
+    #[tokio::test]
+    async fn update_config_rejects_reserved_annotations() {
+        let state = test_server_state().await;
+        state
+            .store
+            .put_message(&test_sandbox(
+                "sb-reserved-annotation",
+                "reserved-annotation",
+                ProtoSandboxPolicy::default(),
+                Vec::new(),
+            ))
+            .await
+            .unwrap();
+
+        let err = handle_update_config(
+            &state,
+            with_user(Request::new(UpdateConfigRequest {
+                sandbox: "reserved-annotation".to_string(),
+                workspace_scope: Some(openshell_core::proto::workspace_selector(
+                    "default".to_string(),
+                )),
+                policy: Some(ProtoSandboxPolicy::default()),
+                annotations: HashMap::from([(
+                    "internal.openshell.ai/auth-epoch".to_string(),
+                    "1".to_string(),
+                )]),
+                ..Default::default()
+            })),
+        )
+        .await
+        .unwrap_err();
+
+        assert_eq!(err.code(), Code::InvalidArgument);
+        assert!(err.message().contains("internal.openshell.ai/"));
+        let stored = state
+            .store
+            .get_message_by_name::<Sandbox>("default", "reserved-annotation")
+            .await
+            .unwrap()
+            .unwrap();
+        assert!(stored.metadata.unwrap().annotations.is_empty());
+    }
+
     #[tokio::test]
     async fn update_config_user_policy_rejects_reserved_provider_key() {
         let state = test_server_state().await;

crates/openshell-server/src/grpc/sandbox.rs

@@ -462,13 +462,6 @@ async fn handle_create_sandbox_inner(
 
     validate_create_sandbox_request_pre_io(&request, &workload_template_name)?;
 
-    // Validate labels (keys and values must meet Kubernetes requirements).
-    for (key, value) in &request.labels {
-        crate::grpc::validation::validate_label_key(key)?;
-        crate::grpc::validation::validate_label_value(value)?;
-    }
-    crate::grpc::validation::validate_annotations(&request.annotations, "annotations")?;
-
     let authz = authorize_workspace(
         &state.store,
         &state.admin_role,
@@ -791,7 +784,7 @@ fn validate_create_sandbox_request_pre_io(
         crate::grpc::validation::validate_label_key(key)?;
         crate::grpc::validation::validate_label_value(value)?;
     }
-    crate::grpc::validation::validate_annotations(&request.annotations, "annotations")?;
+    crate::grpc::validation::validate_caller_annotations(&request.annotations, "annotations")?;
 
     if request.service_exposures.len() > MAX_CREATE_SERVICE_EXPOSURES {
         return Err(Status::invalid_argument(format!(
@@ -7236,6 +7229,44 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn create_sandbox_rejects_reserved_annotations() {
+        let state = test_server_state().await;
+
+        let err = handle_create_sandbox(
+            &state,
+            authed_request(CreateSandboxRequest {
+                request_id: String::new(),
+                name: "reserved".to_string(),
+                spec: Some(SandboxSpec::default()),
+                labels: HashMap::new(),
+                annotations: HashMap::from([(
+                    "internal.openshell.ai/compute-driver".to_string(),
+                    "x".to_string(),
+                )]),
+                workspace_scope: Some(openshell_core::proto::workspace_selector(
+                    "default".to_string(),
+                )),
+                await_main_process_attachment: false,
+                workload_template: String::new(),
+                service_exposures: Vec::new(),
+            }),
+        )
+        .await
+        .unwrap_err();
+
+        assert_eq!(err.code(), tonic::Code::InvalidArgument);
+        assert!(err.message().contains("internal.openshell.ai/"));
+        assert!(
+            state
+                .store
+                .get_message_by_name::<Sandbox>("default", "reserved")
+                .await
+                .unwrap()
+                .is_none()
+        );
+    }
+
     #[tokio::test]
     async fn single_replica_get_sandbox_sends_no_owner_hint() {
         let state = test_server_state().await;

crates/openshell-server/src/grpc/validation.rs

@@ -481,6 +481,33 @@ pub(super) fn validate_annotations(
     Ok(())
 }
 
+/// Annotation key prefix owned by the gateway.
+const RESERVED_ANNOTATION_PREFIX: &str = "internal.openshell.ai/";
+
+/// Validate annotations supplied by an API caller.
+///
+/// Same as [`validate_annotations`], but also rejects keys under the
+/// gateway-owned `internal.openshell.ai/` prefix. The gateway keeps session
+/// and runtime state there and trusts it on read.
+pub(super) fn validate_caller_annotations(
+    annotations: &std::collections::HashMap<String, String>,
+    field_name: &str,
+) -> Result<(), Status> {
+    validate_annotations(annotations, field_name)?;
+    if let Some(key) = annotations
+        .keys()
+        .find(|key| key.starts_with(RESERVED_ANNOTATION_PREFIX))
+    {
+        return Err(invalid_argument(
+            field_name,
+            format!(
+                "{field_name} key '{key}' uses the reserved '{RESERVED_ANNOTATION_PREFIX}' prefix"
+            ),
+        ));
+    }
+    Ok(())
+}
+
 /// OPENSHELL_* keys that are allowed in exec environment. The Python SDK's
 /// `exec_python()` sends a serialized callable via this key.
 const EXEC_ALLOWED_OPENSHELL_KEYS: &[&str] = &["OPENSHELL_PYFUNC_B64"];

proto/openshell.proto

@@ -1287,7 +1287,8 @@ message CreateSandboxRequest {
   string name = 2;
   // Optional labels for the sandbox (key-value metadata).
   map<string, string> labels = 3;
-  // Optional annotations for the sandbox (non-selector metadata).
+  // Optional annotations for the sandbox (non-selector metadata). Keys under
+  // the `internal.openshell.ai/` prefix are reserved for the gateway.
   map<string, string> annotations = 4;
   // One-shot launch hint indicating that the creating client will attach to
   // the canonical main process. The supervisor keeps the terminal transport
@@ -2844,7 +2845,8 @@ message UpdateConfigRequest {
   // not interpret or verify their semantics. For policy updates, the gateway
   // stores the annotations immutably with the revision and merges them into
   // sandbox metadata as a convenience projection. For setting-only updates, it
-  // only merges them into sandbox metadata.
+  // only merges them into sandbox metadata. Keys under the
+  // `internal.openshell.ai/` prefix are reserved for the gateway.
   map<string, string> annotations = 9;
   // Required for sandbox-scoped updates and empty for global updates.
   string sandbox = 1;