From 4be29808a49c899e7d5286dd9277f49c6456e00b Mon Sep 17 00:00:00 2001
From: Francesco Guardiani <me@slinkydeveloper.com>
Date: Mon, 3 Jun 2024 17:08:50 +0200
Subject: [PATCH] Validate that incoming requests to mutate state are sent to
 keyed services only (virtual objects or workflows) (#1581)

---
 crates/admin/src/rest_api/error.rs    |  7 ++++++-
 crates/admin/src/rest_api/services.rs | 10 ++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/crates/admin/src/rest_api/error.rs b/crates/admin/src/rest_api/error.rs
index 93de96b67..38897a70d 100644
--- a/crates/admin/src/rest_api/error.rs
+++ b/crates/admin/src/rest_api/error.rs
@@ -21,6 +21,7 @@ use okapi_operation::okapi::openapi3::Responses;
 use okapi_operation::{okapi, Components, ToMediaTypes, ToResponses};
 use restate_core::ShutdownError;
 use restate_types::identifiers::{DeploymentId, SubscriptionId};
+use restate_types::invocation::ServiceType;
 use schemars::JsonSchema;
 use serde::Serialize;
 
@@ -41,6 +42,8 @@ pub enum MetaApiError {
     },
     #[error("The requested subscription '{0}' does not exist")]
     SubscriptionNotFound(SubscriptionId),
+    #[error("Cannot {0} for service type {1}")]
+    UnsupportedOperation(&'static str, ServiceType),
     #[error(transparent)]
     Schema(#[from] SchemaError),
     #[error(transparent)]
@@ -68,7 +71,9 @@ impl IntoResponse for MetaApiError {
             | MetaApiError::HandlerNotFound { .. }
             | MetaApiError::DeploymentNotFound(_)
             | MetaApiError::SubscriptionNotFound(_) => StatusCode::NOT_FOUND,
-            MetaApiError::InvalidField(_, _) => StatusCode::BAD_REQUEST,
+            MetaApiError::InvalidField(_, _) | MetaApiError::UnsupportedOperation(_, _) => {
+                StatusCode::BAD_REQUEST
+            }
             MetaApiError::Schema(schema_error) => match schema_error {
                 SchemaError::NotFound(_) => StatusCode::NOT_FOUND,
                 SchemaError::Override(_)
diff --git a/crates/admin/src/rest_api/services.rs b/crates/admin/src/rest_api/services.rs
index f7da97854..76a5651d7 100644
--- a/crates/admin/src/rest_api/services.rs
+++ b/crates/admin/src/rest_api/services.rs
@@ -155,6 +155,16 @@ pub async fn modify_service_state<V>(
         new_state,
     }): Json<ModifyServiceStateRequest>,
 ) -> Result<StatusCode, MetaApiError> {
+    let svc = state
+        .task_center
+        .run_in_scope_sync("get-service", None, || {
+            state.schema_registry.get_service(&service_name)
+        })
+        .ok_or_else(|| MetaApiError::ServiceNotFound(service_name.clone()))?;
+    if !svc.ty.has_state() {
+        return Err(MetaApiError::UnsupportedOperation("modify state", svc.ty));
+    }
+
     let service_id = ServiceId::new(service_name, object_key);
 
     let new_state = new_state