From af2b0576411e85dc5b234d851d565198159f4678 Mon Sep 17 00:00:00 2001
From: Lann Martin <lann.martin@fermyon.com>
Date: Thu, 13 Jul 2023 16:36:21 -0400
Subject: [PATCH] Add 'package_keys' to AuthorizedKeyPolicy

This allows signing keys to be granted write access to a log based on
the full package name.
---
 .../server/src/policy/record/authorization.rs | 85 ++++++++++++++++---
 1 file changed, 74 insertions(+), 11 deletions(-)

diff --git a/crates/server/src/policy/record/authorization.rs b/crates/server/src/policy/record/authorization.rs
index e1d0f0c0..c3392352 100644
--- a/crates/server/src/policy/record/authorization.rs
+++ b/crates/server/src/policy/record/authorization.rs
@@ -11,8 +11,11 @@ use wasmparser::names::KebabStr;
 #[derive(Default, Deserialize)]
 pub struct AuthorizedKeyPolicy {
     #[serde(skip)]
-    keys: HashSet<KeyID>,
+    superuser_keys: HashSet<KeyID>,
+    #[serde(default)]
     namespace_keys: HashMap<String, HashSet<KeyID>>,
+    #[serde(default)]
+    package_keys: HashMap<PackageId, HashSet<KeyID>>,
 }
 
 impl AuthorizedKeyPolicy {
@@ -24,8 +27,8 @@ impl AuthorizedKeyPolicy {
     }
 
     /// Sets an authorized key for publishing to any namespace.
-    pub fn with_key(mut self, key: KeyID) -> Self {
-        self.keys.insert(key);
+    pub fn with_superuser_key(mut self, key: KeyID) -> Self {
+        self.superuser_keys.insert(key);
         self
     }
 
@@ -42,6 +45,33 @@ impl AuthorizedKeyPolicy {
             .insert(key);
         Ok(self)
     }
+
+    pub fn with_package_key(mut self, package_id: impl Into<String>, key: KeyID) -> Result<Self> {
+        let package_id = PackageId::new(package_id)?;
+        self.package_keys.entry(package_id).or_default().insert(key);
+        Ok(self)
+    }
+
+    pub fn key_authorized_for_package(&self, key: &KeyID, package: &PackageId) -> bool {
+        if self.superuser_keys.contains(key) {
+            return true;
+        }
+
+        let namespace_keys = self.namespace_keys.get(package.namespace());
+        if namespace_keys
+            .map(|keys| keys.contains(key))
+            .unwrap_or(false)
+        {
+            return true;
+        }
+
+        let package_keys = self.package_keys.get(package);
+        if package_keys.map(|keys| keys.contains(key)).unwrap_or(false) {
+            return true;
+        }
+
+        false
+    }
 }
 
 impl RecordPolicy for AuthorizedKeyPolicy {
@@ -50,18 +80,51 @@ impl RecordPolicy for AuthorizedKeyPolicy {
         id: &PackageId,
         record: &ProtoEnvelope<PackageRecord>,
     ) -> RecordPolicyResult<()> {
-        if !self.keys.contains(record.key_id())
-            && !self
-                .namespace_keys
-                .get(id.namespace())
-                .map(|keys| keys.contains(record.key_id()))
-                .unwrap_or(false)
-        {
+        let key = record.key_id();
+        if !self.key_authorized_for_package(key, id) {
             return Err(RecordPolicyError::Unauthorized(format!(
                 "key id `{key}` is not authorized to publish to package `{id}`",
-                key = record.key_id()
             )));
         }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_key_authorized_for_package() -> Result<()> {
+        let super_key = KeyID::from("super-key".to_string());
+        let namespace_key = KeyID::from("namespace-key".to_string());
+        let package_key = KeyID::from("package-key".to_string());
+        let other_key = KeyID::from("other-key".to_string());
+
+        let policy = AuthorizedKeyPolicy::new()
+            .with_superuser_key(super_key.clone())
+            .with_namespace_key("my-namespace", namespace_key.clone())?
+            .with_package_key("my-namespace:my-package", package_key.clone())?;
+
+        let my_package: PackageId = "my-namespace:my-package".parse()?;
+        let my_namespace_other_package: PackageId = "my-namespace:other-package".parse()?;
+        let other_namespace: PackageId = "other-namespace:any-package".parse()?;
+
+        assert!(policy.key_authorized_for_package(&super_key, &my_package));
+        assert!(policy.key_authorized_for_package(&super_key, &my_namespace_other_package));
+        assert!(policy.key_authorized_for_package(&super_key, &other_namespace));
+
+        assert!(policy.key_authorized_for_package(&namespace_key, &my_package));
+        assert!(policy.key_authorized_for_package(&namespace_key, &my_namespace_other_package));
+        assert!(!policy.key_authorized_for_package(&namespace_key, &other_namespace));
+
+        assert!(policy.key_authorized_for_package(&package_key, &my_package));
+        assert!(!policy.key_authorized_for_package(&package_key, &my_namespace_other_package));
+        assert!(!policy.key_authorized_for_package(&package_key, &other_namespace));
+
+        assert!(!policy.key_authorized_for_package(&other_key, &my_package));
+        assert!(!policy.key_authorized_for_package(&other_key, &my_namespace_other_package));
+        assert!(!policy.key_authorized_for_package(&other_key, &other_namespace));
 
         Ok(())
     }