From d2ed4524f17ae025fb2e55956d7c82d377872a3e Mon Sep 17 00:00:00 2001
From: John Gallagher <john@oxidecomputer.com>
Date: Mon, 13 May 2024 13:33:40 -0400
Subject: [PATCH] Validate `ByteCount` ranges during deserialization (#5743)

Fixes #5732.

I tried adding a custom `JsonSchema` too but hit a typify issue; left a
link in the TODO comment.
---
 common/src/api/external/mod.rs          | 17 +++++--
 nexus/tests/integration_tests/quotas.rs | 67 +++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/common/src/api/external/mod.rs b/common/src/api/external/mod.rs
index 9e0966f949..9d55522af8 100644
--- a/common/src/api/external/mod.rs
+++ b/common/src/api/external/mod.rs
@@ -527,12 +527,21 @@ impl JsonSchema for RoleName {
 // in the database as an i64.  Constraining it here ensures that we can't fail
 // to serialize the value.
 //
-// TODO: custom JsonSchema and Deserialize impls to enforce i64::MAX limit
-#[derive(
-    Copy, Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq,
-)]
+// TODO: custom JsonSchema impl to describe i64::MAX limit; this is blocked by
+// https://github.com/oxidecomputer/typify/issues/589
+#[derive(Copy, Clone, Debug, Serialize, JsonSchema, PartialEq, Eq)]
 pub struct ByteCount(u64);
 
+impl<'de> Deserialize<'de> for ByteCount {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let bytes = u64::deserialize(deserializer)?;
+        ByteCount::try_from(bytes).map_err(serde::de::Error::custom)
+    }
+}
+
 #[allow(non_upper_case_globals)]
 const KiB: u64 = 1024;
 #[allow(non_upper_case_globals)]
diff --git a/nexus/tests/integration_tests/quotas.rs b/nexus/tests/integration_tests/quotas.rs
index 4e3335d04c..d8f72ba666 100644
--- a/nexus/tests/integration_tests/quotas.rs
+++ b/nexus/tests/integration_tests/quotas.rs
@@ -20,6 +20,7 @@ use nexus_types::external_api::views::{Silo, SiloQuotas};
 use omicron_common::api::external::ByteCount;
 use omicron_common::api::external::IdentityMetadataCreateParams;
 use omicron_common::api::external::InstanceCpuCount;
+use serde_json::json;
 
 type ControlPlaneTestContext =
     nexus_test_utils::ControlPlaneTestContext<omicron_nexus::Server>;
@@ -315,3 +316,69 @@ async fn test_quotas(cptestctx: &ControlPlaneTestContext) {
         .await
         .expect("Disk should be provisioned");
 }
+
+#[nexus_test]
+async fn test_quota_limits(cptestctx: &ControlPlaneTestContext) {
+    let client = &cptestctx.external_client;
+
+    let system = setup_silo_with_quota(
+        &client,
+        "quota-test-silo",
+        params::SiloQuotasCreate::empty(),
+    )
+    .await;
+
+    // Maximal legal limits should be allowed.
+    let quota_limit = params::SiloQuotasUpdate {
+        cpus: Some(i64::MAX),
+        memory: Some(i64::MAX.try_into().unwrap()),
+        storage: Some(i64::MAX.try_into().unwrap()),
+    };
+    system
+        .set_quotas(client, quota_limit.clone())
+        .await
+        .expect("set max quotas");
+    let quotas = system.get_quotas(client).await;
+    assert_eq!(quotas.limits.cpus, quota_limit.cpus.unwrap());
+    assert_eq!(quotas.limits.memory, quota_limit.memory.unwrap());
+    assert_eq!(quotas.limits.storage, quota_limit.storage.unwrap());
+
+    // Construct a value that fits in a u64 but not an i64.
+    let out_of_bounds = u64::try_from(i64::MAX).unwrap() + 1;
+
+    for key in ["cpus", "memory", "storage"] {
+        // We can't construct a `SiloQuotasUpdate` with higher-than-maximal
+        // values, but we can construct the equivalent JSON blob of such a
+        // request.
+        let request = json!({ key: out_of_bounds });
+
+        let err = NexusRequest::expect_failure_with_body(
+            client,
+            http::StatusCode::BAD_REQUEST,
+            http::Method::PUT,
+            "/v1/system/silos/quota-test-silo/quotas",
+            &request,
+        )
+        .authn_as(system.auth.clone())
+        .execute()
+        .await
+        .expect("sent quota update")
+        .parsed_body::<HttpErrorResponseBody>()
+        .expect("parsed error body");
+        assert!(
+            err.message.contains(key)
+                && (err.message.contains("invalid value")
+                    || err
+                        .message
+                        .contains("value is too large for a byte count")),
+            "Unexpected error: {0}",
+            err.message
+        );
+
+        // The quota limits we set above should be unchanged.
+        let quotas = system.get_quotas(client).await;
+        assert_eq!(quotas.limits.cpus, quota_limit.cpus.unwrap());
+        assert_eq!(quotas.limits.memory, quota_limit.memory.unwrap());
+        assert_eq!(quotas.limits.storage, quota_limit.storage.unwrap());
+    }
+}