feat: allow access attributes for resources to be configured

This allows a set of rules to be defined for determining the access attributes to apply to
a particular resource. It also checks that the attributes determined for a new resource to
be registered are matched by attributes associated with the request context.

Signed-off-by: Gordon Sim <gsim@redhat.com>

tests/unit/distribution/test_resource_attributes.py

@@ -0,0 +1,209 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import pytest
+
+from llama_stack.apis.datasets import URIDataSource
+from llama_stack.distribution.datatypes import (
+    AuthenticationConfig,
+    DatasetWithACL,
+    ModelWithACL,
+    ShieldWithACL,
+)
+from llama_stack.distribution.resource_attributes import ResourceAccessAttributes, match_access_attributes_rule
+
+
+@pytest.fixture
+def rules():
+    config = """{
+    "provider_type": "custom",
+    "config": {},
+    "resource_attribute_rules": [
+        {
+            "resource_type": "model",
+            "resource_id": "my-model",
+            "provider_id": "my-provider",
+            "attributes": {
+                "roles": ["role1"],
+                "teams": ["team1"],
+                "projects": ["project1"],
+                "namespaces": ["namespace1"]
+            }
+        },
+        {
+            "resource_type": "dataset",
+            "provider_id": "my-provider",
+            "attributes": {
+                "roles": ["role2"],
+                "teams": ["team2"],
+                "projects": ["project2"],
+                "namespaces": ["namespace2"]
+            }
+        },
+        {
+            "provider_id": "my-provider",
+            "attributes": {
+                "roles": ["role3"],
+                "teams": ["team3"],
+                "projects": ["project3"],
+                "namespaces": ["namespace3"]
+            }
+        },
+        {
+            "resource_type": "model",
+            "attributes": {
+                "roles": ["role4"],
+                "teams": ["team4"],
+                "projects": ["project4"],
+                "namespaces": ["namespace4"]
+            }
+        },
+        {
+            "attributes": {
+                "roles": ["role5"],
+                "teams": ["team5"],
+                "projects": ["project5"],
+                "namespaces": ["namespace5"]
+            }
+        }
+    ]
+}"""
+    return AuthenticationConfig.model_validate_json(config).resource_attribute_rules
+
+
+def test_match_access_attributes_rule(rules):
+    assert match_access_attributes_rule(rules[0], "model", "my-model", "my-provider")
+    assert not match_access_attributes_rule(rules[0], "model", "another-model", "my-provider")
+    assert not match_access_attributes_rule(rules[0], "model", "my-model", "another-provider")
+    assert not match_access_attributes_rule(rules[0], "dataset", "my-model", "my-provider")
+
+    assert match_access_attributes_rule(rules[1], "dataset", "my-data", "my-provider")
+    assert match_access_attributes_rule(rules[1], "dataset", "different-data", "my-provider")
+    assert match_access_attributes_rule(rules[1], "dataset", "any-data", "my-provider")
+    assert not match_access_attributes_rule(rules[1], "model", "a-model", "my-provider")
+    assert not match_access_attributes_rule(rules[1], "dataset", "any-data", "another-provider")
+    assert not match_access_attributes_rule(rules[1], "model", "my-model", "my-provider")
+
+    assert match_access_attributes_rule(rules[2], "dataset", "foo", "my-provider")
+    assert match_access_attributes_rule(rules[2], "model", "foo", "my-provider")
+    assert match_access_attributes_rule(rules[2], "vector_db", "bar", "my-provider")
+    assert not match_access_attributes_rule(rules[2], "dataset", "foo", "another-provider")
+
+    assert match_access_attributes_rule(rules[3], "model", "foo", "my-provider")
+    assert match_access_attributes_rule(rules[3], "model", "bar", "my-provider")
+    assert match_access_attributes_rule(rules[3], "model", "foo", "another-provider")
+    assert not match_access_attributes_rule(rules[3], "dataset", "bar", "my-provider")
+
+    assert match_access_attributes_rule(rules[4], "model", "foo", "my-provider")
+    assert match_access_attributes_rule(rules[4], "model", "bar", "my-provider")
+    assert match_access_attributes_rule(rules[4], "model", "foo", "another-provider")
+    assert match_access_attributes_rule(rules[4], "dataset", "bar", "my-provider")
+    assert match_access_attributes_rule(rules[4], "vector_db", "baz", "any-provider")
+
+
+@pytest.fixture
+def resource_access_attributes(rules):
+    return ResourceAccessAttributes(rules)
+
+
+def dataset(identifier: str, provider_id: str) -> DatasetWithACL:
+    return DatasetWithACL(
+        identifier=identifier,
+        provider_id=provider_id,
+        purpose="eval/question-answer",
+        source=URIDataSource(uri="https://a.com/a.jsonl"),
+    )
+
+
+@pytest.mark.parametrize(
+    "resource,expected_roles",
+    [
+        (ModelWithACL(identifier="my-model", provider_id="my-provider"), ["role1"]),
+        (ModelWithACL(identifier="another-model", provider_id="my-provider"), ["role3"]),
+        (ModelWithACL(identifier="third-model", provider_id="another-provider"), ["role4"]),
+        (dataset(identifier="my-dataset", provider_id="my-provider"), ["role2"]),
+        (dataset(identifier="another-dataset", provider_id="another-provider"), ["role5"]),
+        (ShieldWithACL(identifier="my-shield", provider_id="my-provider"), ["role3"]),
+        (ShieldWithACL(identifier="another-shield", provider_id="another-provider"), ["role5"]),
+    ],
+)
+def test_apply(resource_access_attributes, resource, expected_roles):
+    assert resource_access_attributes.apply(resource, None)
+    assert resource.access_attributes.roles == expected_roles
+
+
+@pytest.fixture
+def alternate_access_attributes_rules():
+    config = """{
+    "provider_type": "custom",
+    "config": {},
+    "resource_attribute_rules": [
+        {
+            "resource_type": "model",
+            "resource_id": "my-model",
+            "provider_id": "my-provider",
+            "attributes": {
+                "roles": ["roleA"]
+            }
+        },
+        {
+            "resource_type": "model",
+            "attributes": {
+                "roles": ["roleB"]
+            }
+        }
+    ]
+}"""
+    return AuthenticationConfig.model_validate_json(config).resource_attribute_rules
+
+
+@pytest.fixture
+def alternate_access_attributes(alternate_access_attributes_rules):
+    return ResourceAccessAttributes(alternate_access_attributes_rules)
+
+
+@pytest.mark.parametrize(
+    "resource,expected_roles",
+    [
+        (ModelWithACL(identifier="my-model", provider_id="my-provider"), ["roleA"]),
+        (ModelWithACL(identifier="another-model", provider_id="another-provider"), ["roleB"]),
+        (dataset(identifier="my-dataset", provider_id="my-provider"), None),
+        (dataset(identifier="another-dataset", provider_id="another-provider"), None),
+    ],
+)
+def test_apply_alternate(alternate_access_attributes, resource, expected_roles):
+    if expected_roles:
+        assert alternate_access_attributes.apply(resource, None)
+        assert resource.access_attributes.roles == expected_roles
+    else:
+        assert not alternate_access_attributes.apply(resource, None)
+        assert not resource.access_attributes
+
+
+@pytest.fixture
+def checked_attributes(alternate_access_attributes_rules):
+    attributes = ResourceAccessAttributes(alternate_access_attributes_rules)
+    attributes.enable_access_checks()
+    return attributes
+
+
+@pytest.mark.parametrize(
+    "resource,user_attributes,fails,result",
+    [
+        (ModelWithACL(identifier="my-model", provider_id="my-provider"), {"roles": ["roleA"]}, False, True),
+        (ModelWithACL(identifier="my-model", provider_id="my-provider"), {"roles": ["somethingelse"]}, True, False),
+        (dataset(identifier="my-dataset", provider_id="my-provider"), {"roles": ["somethingelse"]}, False, False),
+        (dataset(identifier="my-dataset", provider_id="my-provider"), None, False, False),
+    ],
+)
+def test_access_check_on_apply(checked_attributes, resource, user_attributes, fails, result):
+    if fails:
+        with pytest.raises(ValueError) as e:
+            checked_attributes.apply(resource, user_attributes)
+        assert "Access denied" in str(e.value)
+        assert not resource.access_attributes
+    else:
+        assert checked_attributes.apply(resource, user_attributes) == result