fix(gateway): reject schemeless worker URLs at the API boundary (#1977)

Author: ighutake-debug

model_gateway/src/config/mod.rs

@@ -4,7 +4,7 @@ pub(crate) mod validation;
 
 pub use builder::*;
 pub use types::*;
-pub use validation::validate_mesh_server_name;
+pub use validation::{validate_mesh_server_name, validate_worker_url};
 
 #[derive(Debug, thiserror::Error)]
 pub enum ConfigError {

model_gateway/src/config/validation.rs

@@ -18,9 +18,63 @@ pub fn validate_mesh_server_name(name: &str) -> ConfigResult<()> {
     Ok(())
 }
 
+/// Validate a single worker URL: non-empty, an allowed scheme, and a
+/// parseable host. Shared by [`ConfigValidator::validate_urls`] (startup
+/// config) and the worker-management API so both reject schemeless or
+/// unparsable URLs identically. Rejecting at the API boundary prevents
+/// the orphaned `url_to_id` reservation from #1533: the AddWorker
+/// workflow rewrites schemeless input via `normalize_url`, so a
+/// reservation keyed on the raw URL would never match the registered
+/// worker.
+pub fn validate_worker_url(url: &str) -> ConfigResult<()> {
+    if url.is_empty() {
+        return Err(ConfigError::InvalidValue {
+            field: "worker_url".to_string(),
+            value: url.to_string(),
+            reason: "URL cannot be empty".to_string(),
+        });
+    }
+
+    // Exact (lowercase) scheme allow-list. Case-insensitive matching is
+    // tempting but wrong here: the AddWorker workflow's normalize_url
+    // matches schemes case-sensitively, so a mixed-case scheme would be
+    // rewritten downstream and diverge from the reservation key — the
+    // same orphan failure as a schemeless URL.
+    const ALLOWED_SCHEMES: &[&str] = &["http", "https", "grpc", "grpcs"];
+    let scheme = url.split_once("://").map_or("", |(s, _)| s);
+    if !ALLOWED_SCHEMES.contains(&scheme) {
+        return Err(ConfigError::InvalidValue {
+            field: "worker_url".to_string(),
+            value: url.to_string(),
+            reason:
+                "URL must start with a lowercase http://, https://, grpc://, or grpcs:// scheme"
+                    .to_string(),
+        });
+    }
+
+    match ::url::Url::parse(url) {
+        Ok(parsed) => {
+            if parsed.host_str().is_none() {
+                return Err(ConfigError::InvalidValue {
+                    field: "worker_url".to_string(),
+                    value: url.to_string(),
+                    reason: "URL must have a valid host".to_string(),
+                });
+            }
+        }
+        Err(e) => {
+            return Err(ConfigError::InvalidValue {
+                field: "worker_url".to_string(),
+                value: url.to_string(),
+                reason: format!("Invalid URL format: {e}"),
+            });
+        }
+    }
+    Ok(())
+}
+
 /// Configuration validator
 pub(crate) struct ConfigValidator;
-
 impl ConfigValidator {
     pub(crate) fn validate(config: &RouterConfig) -> ConfigResult<()> {
         Self::validate_mode(&config.mode)?;
@@ -958,48 +1012,7 @@ impl ConfigValidator {
 
     fn validate_urls(urls: &[String]) -> ConfigResult<()> {
         for url in urls {
-            if url.is_empty() {
-                return Err(ConfigError::InvalidValue {
-                    field: "worker_url".to_string(),
-                    value: url.clone(),
-                    reason: "URL cannot be empty".to_string(),
-                });
-            }
-
-            // Case-insensitive scheme allow-list. Compare just the scheme
-            // segment so we don't allocate a lowercased copy of the full URL.
-            const ALLOWED_SCHEMES: &[&str] = &["http", "https", "grpc", "grpcs"];
-            let scheme = url.split_once("://").map_or("", |(s, _)| s);
-            if !ALLOWED_SCHEMES
-                .iter()
-                .any(|allowed| scheme.eq_ignore_ascii_case(allowed))
-            {
-                return Err(ConfigError::InvalidValue {
-                    field: "worker_url".to_string(),
-                    value: url.clone(),
-                    reason: "URL must start with http://, https://, grpc://, or grpcs://"
-                        .to_string(),
-                });
-            }
-
-            match ::url::Url::parse(url) {
-                Ok(parsed) => {
-                    if parsed.host_str().is_none() {
-                        return Err(ConfigError::InvalidValue {
-                            field: "worker_url".to_string(),
-                            value: url.clone(),
-                            reason: "URL must have a valid host".to_string(),
-                        });
-                    }
-                }
-                Err(e) => {
-                    return Err(ConfigError::InvalidValue {
-                        field: "worker_url".to_string(),
-                        value: url.clone(),
-                        reason: format!("Invalid URL format: {e}"),
-                    });
-                }
-            }
+            validate_worker_url(url)?;
         }
         Ok(())
     }
@@ -1031,6 +1044,56 @@ mod tests {
         assert!(validate_mesh_server_name("node-a").is_ok());
     }
 
+    #[test]
+    fn worker_url_accepts_allowed_schemes() {
+        for url in [
+            "http://10.0.0.5:8000",
+            "https://worker.example.com",
+            "grpc://10.0.0.5:50051",
+            "grpcs://worker.example.com:443",
+        ] {
+            assert!(validate_worker_url(url).is_ok(), "expected {url} to pass");
+        }
+    }
+
+    #[test]
+    fn worker_url_rejects_non_lowercase_scheme() {
+        // normalize_url in the AddWorker workflow matches schemes
+        // case-sensitively, so `HTTP://…` would be mangled into
+        // `http://HTTP://…` and orphan the reservation just like a
+        // schemeless URL would.
+        assert!(validate_worker_url("HTTP://10.0.0.5:8000").is_err());
+        assert!(validate_worker_url("Grpc://10.0.0.5:50051").is_err());
+    }
+
+    #[test]
+    fn worker_url_rejects_empty() {
+        assert!(matches!(
+            validate_worker_url(""),
+            Err(ConfigError::InvalidValue { ref field, .. }) if field == "worker_url"
+        ));
+    }
+
+    #[test]
+    fn worker_url_rejects_schemeless_host_port() {
+        // The #1533 case: bare host:port input would be rewritten by the
+        // AddWorker workflow, orphaning the API-layer ID reservation.
+        assert!(matches!(
+            validate_worker_url("10.0.0.5:8000"),
+            Err(ConfigError::InvalidValue { ref field, .. }) if field == "worker_url"
+        ));
+    }
+
+    #[test]
+    fn worker_url_rejects_disallowed_scheme() {
+        assert!(validate_worker_url("ftp://example.com:21").is_err());
+    }
+
+    #[test]
+    fn worker_url_rejects_unparsable() {
+        assert!(validate_worker_url("http://").is_err());
+    }
+
     #[test]
     fn test_validate_regular_mode() {
         let config = RouterConfig::new(

model_gateway/src/worker/service.rs

@@ -16,11 +16,21 @@ use serde_json::json;
 use tracing::warn;
 
 use crate::{
-    config::RouterConfig,
+    config::{validate_worker_url, RouterConfig},
     worker::{registry::WorkerId, worker::worker_to_info, WorkerRegistry, WorkerType},
     workflow::{Job, JobQueue, WorkerRegistrationMode},
 };
 
+/// Validate the URL of an incoming worker-management request, translating a
+/// config-layer rejection into `400 Bad Request`. Must run before
+/// `WorkerRegistry::reserve_id_for_url` so rejected input can never leave an
+/// orphaned reservation (#1533).
+fn validate_worker_url_request(url: &str) -> Result<(), WorkerServiceError> {
+    validate_worker_url(url).map_err(|e| WorkerServiceError::BadRequest {
+        message: e.to_string(),
+    })
+}
+
 /// Error types for worker service operations
 #[derive(Debug)]
 pub enum WorkerServiceError {
@@ -256,6 +266,8 @@ impl WorkerService {
         &self,
         config: WorkerSpec,
     ) -> Result<CreateWorkerResult, WorkerServiceError> {
+        validate_worker_url_request(&config.url)?;
+
         if self.router_config.api_key.is_some() && config.api_key.is_none() {
             warn!(
                 "Adding worker {} without API key while router has API key configured. \
@@ -564,4 +576,52 @@ mod tests {
         assert_eq!(result.workers.len(), 1);
         assert_eq!(result.total, 1);
     }
+
+    fn worker_spec(url: &str) -> WorkerSpec {
+        serde_json::from_value(json!({ "url": url })).expect("worker spec")
+    }
+
+    #[tokio::test]
+    async fn create_worker_rejects_schemeless_url_before_reserving() {
+        let registry = Arc::new(WorkerRegistry::new());
+        let service = make_service(registry);
+
+        let err = service
+            .create_worker(worker_spec("10.0.0.5:8000"))
+            .await
+            .expect_err("schemeless URL must be rejected at the boundary");
+
+        assert!(matches!(err, WorkerServiceError::BadRequest { .. }));
+        assert_eq!(err.error_code(), "BAD_REQUEST");
+        assert_eq!(err.status_code(), StatusCode::BAD_REQUEST);
+    }
+
+    #[tokio::test]
+    async fn create_worker_rejects_mixed_case_scheme() {
+        let registry = Arc::new(WorkerRegistry::new());
+        let service = make_service(registry);
+
+        let err = service
+            .create_worker(worker_spec("HTTP://10.0.0.5:8000"))
+            .await
+            .expect_err("mixed-case scheme must be rejected at the boundary");
+
+        assert!(matches!(err, WorkerServiceError::BadRequest { .. }));
+    }
+
+    #[tokio::test]
+    async fn create_worker_accepts_schemed_url() {
+        let registry = Arc::new(WorkerRegistry::new());
+        let service = make_service(registry);
+
+        // The harness leaves the job queue uninitialized, so a valid URL
+        // passes boundary validation and stops at queue submission —
+        // proving validation did not reject it.
+        let err = service
+            .create_worker(worker_spec("grpc://10.0.0.5:8000"))
+            .await
+            .expect_err("queue is uninitialized in the test harness");
+
+        assert!(matches!(err, WorkerServiceError::QueueNotInitialized));
+    }
 }