feat: fine grained access control policy

This allows a set of rules to be defined for determining access to resources.

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

llama_stack/distribution/access_control/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/distribution/access_control/access_control.py

@@ -0,0 +1,175 @@
+# 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.
+
+from typing import Any, Protocol
+
+from llama_stack.distribution.request_headers import User
+
+from .datatypes import (
+    AccessAttributes,
+    AccessRule,
+    Action,
+    AttributeReference,
+    Condition,
+    Scope,
+)
+
+
+def matches_resource(resource_scope: str, actual_resource: str) -> bool:
+    if resource_scope == actual_resource:
+        return True
+    return resource_scope.endswith("::*") and actual_resource.startswith(resource_scope[:-1])
+
+
+def matches_scope(
+    scope: Scope,
+    action: Action,
+    resource: str,
+    user: str | None,
+) -> bool:
+    if scope.resource and not matches_resource(scope.resource, resource):
+        return False
+    if scope.principal and scope.principal != user:
+        return False
+    return action in scope.actions
+
+
+def user_in_literal(
+    literal: str,
+    user_attributes: dict[str, list[str]] | None,
+) -> bool:
+    for qualifier in ["role::", "team::", "project::", "namespace::"]:
+        if literal.startswith(qualifier):
+            if not user_attributes:
+                return False
+            ref = qualifier.replace("::", "s")
+            if ref in user_attributes:
+                value = literal.removeprefix(qualifier)
+                return value in user_attributes[ref]
+            else:
+                return False
+    return False
+
+
+def user_in(
+    ref: AttributeReference | str,
+    resource_attributes: AccessAttributes | None,
+    user_attributes: dict[str, list[str]] | None,
+) -> bool:
+    if not ref.startswith("resource."):
+        return user_in_literal(ref, user_attributes)
+    name = ref.removeprefix("resource.")
+    required = resource_attributes and getattr(resource_attributes, name)
+    if not required:
+        return True
+    if not user_attributes or name not in user_attributes:
+        return False
+    actual = user_attributes[name]
+    for value in required:
+        if value in actual:
+            return True
+    return False
+
+
+def as_list(obj: Any) -> list[Any]:
+    if isinstance(obj, list):
+        return obj
+    return [obj]
+
+
+def matches_conditions(
+    conditions: list[Condition],
+    resource_attributes: AccessAttributes | None,
+    user_attributes: dict[str, list[str]] | None,
+) -> bool:
+    for condition in conditions:
+        # must match all conditions
+        if not matches_condition(condition, resource_attributes, user_attributes):
+            return False
+    return True
+
+
+def matches_condition(
+    condition: Condition | list[Condition],
+    resource_attributes: AccessAttributes | None,
+    user_attributes: dict[str, list[str]] | None,
+) -> bool:
+    if isinstance(condition, list):
+        return matches_conditions(as_list(condition), resource_attributes, user_attributes)
+    if condition.user_in:
+        for ref in as_list(condition.user_in):
+            # if multiple references are specified, all must match
+            if not user_in(ref, resource_attributes, user_attributes):
+                return False
+        return True
+    if condition.user_not_in:
+        for ref in as_list(condition.user_not_in):
+            # if multiple references are specified, none must match
+            if user_in(ref, resource_attributes, user_attributes):
+                return False
+        return True
+    return True
+
+
+def default_policy() -> list[AccessRule]:
+    # for backwards compatibility, if no rules are provided , assume
+    # full access to all subject to attribute matching rules
+    return [
+        AccessRule(
+            permit=Scope(actions=list(Action)),
+            when=Condition(user_in=list(AttributeReference)),
+        )
+    ]
+
+
+class ProtectedResource(Protocol):
+    type: str
+    identifier: str
+    access_attributes: AccessAttributes
+
+
+def is_action_allowed(
+    policy: list[AccessRule],
+    action: Action,
+    resource: ProtectedResource,
+    user: User | None,
+) -> bool:
+    # If user is not set, assume authentication is not enabled
+    if not user:
+        return True
+
+    if not len(policy):
+        policy = default_policy()
+
+    resource_attributes = AccessAttributes()
+    if hasattr(resource, "access_attributes"):
+        resource_attributes = resource.access_attributes
+    qualified_resource_id = resource.type + "::" + resource.identifier
+    for rule in policy:
+        if rule.forbid and matches_scope(rule.forbid, action, qualified_resource_id, user.principal):
+            if rule.when:
+                if matches_condition(rule.when, resource_attributes, user.attributes):
+                    return False
+            elif rule.unless:
+                if not matches_condition(rule.unless, resource_attributes, user.attributes):
+                    return False
+            else:
+                return False
+        elif rule.permit and matches_scope(rule.permit, action, qualified_resource_id, user.principal):
+            if rule.when:
+                if matches_condition(rule.when, resource_attributes, user.attributes):
+                    return True
+            elif rule.unless:
+                if not matches_condition(rule.unless, resource_attributes, user.attributes):
+                    return True
+            else:
+                return True
+    # assume access is denied unless we find a rule that permits access
+    return False
+
+
+class AccessDeniedError(RuntimeError):
+    pass

llama_stack/distribution/access_control/datatypes.py

@@ -0,0 +1,136 @@
+# 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.
+
+from enum import Enum
+
+from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
+
+
+class AccessAttributes(BaseModel):
+    """Structured representation of user attributes for access control.
+
+    This model defines a structured approach to representing user attributes
+    with common standard categories for access control.
+
+    Standard attribute categories include:
+    - roles: Role-based attributes (e.g., admin, data-scientist)
+    - teams: Team-based attributes (e.g., ml-team, infra-team)
+    - projects: Project access attributes (e.g., llama-3, customer-insights)
+    - namespaces: Namespace-based access control for resource isolation
+    """
+
+    # Standard attribute categories - the minimal set we need now
+    roles: list[str] | None = Field(
+        default=None, description="Role-based attributes (e.g., 'admin', 'data-scientist', 'user')"
+    )
+
+    teams: list[str] | None = Field(default=None, description="Team-based attributes (e.g., 'ml-team', 'nlp-team')")
+
+    projects: list[str] | None = Field(
+        default=None, description="Project-based access attributes (e.g., 'llama-3', 'customer-insights')"
+    )
+
+    namespaces: list[str] | None = Field(
+        default=None, description="Namespace-based access control for resource isolation"
+    )
+
+
+class Action(str, Enum):
+    CREATE = "create"
+    READ = "read"
+    UPDATE = "update"
+    DELETE = "delete"
+
+
+class Scope(BaseModel):
+    principal: str | None = None
+    actions: Action | list[Action]
+    resource: str | None = None
+
+
+def _mutually_exclusive(obj, a: str, b: str):
+    if getattr(obj, a) and getattr(obj, b):
+        raise ValueError(f"{a} and {b} are mutually exclusive")
+
+
+def _require_one_of(obj, a: str, b: str):
+    if not getattr(obj, a) and not getattr(obj, b):
+        raise ValueError(f"on of {a} or {b} is required")
+
+
+class AttributeReference(str, Enum):
+    RESOURCE_ROLES = "resource.roles"
+    RESOURCE_TEAMS = "resource.teams"
+    RESOURCE_PROJECTS = "resource.projects"
+    RESOURCE_NAMESPACES = "resource.namespaces"
+
+
+class Condition(BaseModel):
+    user_in: AttributeReference | list[AttributeReference] | str | None = None
+    user_not_in: AttributeReference | list[AttributeReference] | str | None = None
+
+
+class AccessRule(BaseModel):
+    """Access rule based loosely on cedar policy language
+
+    A rule defines a list of action either to permit or to forbid. It may specify a
+    principal or a resource that must match for the rule to take effect. The resource
+    to match should be specified in the form of a type qualified identifier, e.g.
+    model::my-model or vector_db::some-db, or a wildcard for all resources of a type,
+    e.g. model::*. If the principal or resource are not specified, they will match all
+    requests.
+
+    A rule may also specify a condition, either a 'when' or an 'unless', with additional
+    constraints as to where the rule applies. The constraints at present are whether the
+    user requesting access is in or not in some set. This set can either be a particular
+    set of attributes on the resource e.g. resource.roles or a literal value of some
+    notion of group, e.g. role::admin or namespace::foo.
+
+    Rules are tested in order to find a match. If a match is found, the request is
+    permitted or forbidden depending on the type of rule. If no match is found, the
+    request is denied. If no rules are specified, a rule that allows any action as
+    long as the resource attributes match the user attributes is added
+    (i.e. the previous behaviour is the default).
+
+    Some examples in yaml:
+
+    - permit:
+      principal: user-1
+      actions: [create, read, delete]
+      resource: model::*
+      description: user-1 has full access to all models
+    - permit:
+      principal: user-2
+      actions: [read]
+      resource: model::model-1
+      description: user-2 has read access to model-1 only
+    - permit:
+      actions: [read]
+      when:
+        user_in: resource.namespaces
+      description: any user has read access to any resource with matching attributes
+    - forbid:
+      actions: [create, read, delete]
+      resource: vector_db::*
+      unless:
+        user_in: role::admin
+      description: only user with admin role can use vector_db resources
+
+    """
+
+    permit: Scope | None = None
+    forbid: Scope | None = None
+    when: Condition | list[Condition] | None = None
+    unless: Condition | list[Condition] | None = None
+    description: str | None = None
+
+    @model_validator(mode="after")
+    def validate_rule_format(self) -> Self:
+        _require_one_of(self, "permit", "forbid")
+        _mutually_exclusive(self, "permit", "forbid")
+        _mutually_exclusive(self, "when", "unless")
+        return self