Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
169 changes: 169 additions & 0 deletions tests/sentry/seer/endpoints/test_organization_agent_token.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import re
from collections.abc import Callable
from contextlib import ExitStack, nullcontext
from dataclasses import dataclass
from datetime import timedelta
Expand Down Expand Up @@ -38,9 +39,11 @@
from sentry.organizations.services.organization import organization_service
from sentry.replays.lib.storage import FilestoreBlob, RecordingSegmentStorageMeta
from sentry.replays.testutils import mock_replay, mock_replay_viewed
from sentry.roles import organization_roles
from sentry.seer import agent_token
from sentry.seer.models.agent_write_grant import SeerAgentWriteGrant
from sentry.seer.signed_seer_api import _resolve_viewer_context
from sentry.service_accounts.models import GRANTABLE_SCOPES
from sentry.silo.base import SiloMode
from sentry.snuba.dataset import Dataset
from sentry.snuba.models import SnubaQuery, SnubaQueryEventType
Expand Down Expand Up @@ -329,6 +332,7 @@ class MatrixAuthentication(StrEnum):
AGENT_TOKEN = "agent_token"
SCOPED_DOWN_AGENT_TOKEN = "scoped_down_agent_token"
APPROVED_AGENT_TOKEN = "approved_agent_token"
SERVICE_ACCOUNT_VIEWER_CONTEXT = "service_account_viewer_context"


@dataclass(frozen=True)
Expand Down Expand Up @@ -586,6 +590,60 @@ def _public_mutation_endpoints() -> tuple[PublicMutationEndpoint, ...]:
AGENT_SUDO_BOUNDARY_MUTATIONS = frozenset({("TeamDetailsEndpoint", "DELETE")})


# Oracle: an owner token with the account's scopes. The account is linked to the project
# in an open-membership org, so it should reach everything the owner does. User-global
# routes must deny it, since it has no user.
SERVICE_ACCOUNT_SCOPES = frozenset(organization_roles.get("owner").scopes & GRANTABLE_SCOPES)


class ServiceAccountGap(StrEnum):
CONTROL_SILO = "service account auth is cell-only"
REQUIRES_USER = "rejects, or stores state for, an anonymous request.user"
USER_ACTOR_ERROR = "500s using an anonymous request.user as an actor"
NO_ORG_ROLE = "checks the caller's org role"


# Rows that don't match the oracle yet. Strict, so a fixed row fails until removed here.
SERVICE_ACCOUNT_KNOWN_GAPS: dict[str, ServiceAccountGap] = {
"OrganizationIntegrationChannelValidate": ServiceAccountGap.CONTROL_SILO,
"OrganizationIntegrationChannels": ServiceAccountGap.CONTROL_SILO,
"OrganizationIntegrationDetails": ServiceAccountGap.CONTROL_SILO,
"OrganizationIntegrationDetails_delete": ServiceAccountGap.CONTROL_SILO,
"OrganizationIntegrations": ServiceAccountGap.CONTROL_SILO,
"OrganizationSentryAppComponents": ServiceAccountGap.CONTROL_SILO,
"OrganizationSentryApps": ServiceAccountGap.CONTROL_SILO,
"SentryAppDetails": ServiceAccountGap.CONTROL_SILO,
"SentryAppDetails_delete": ServiceAccountGap.CONTROL_SILO,
"SentryAppDetails_put": ServiceAccountGap.CONTROL_SILO,
"SentryAppInstallationExternalIssueActions_post": ServiceAccountGap.CONTROL_SILO,
"SentryAppInstallationExternalRequests": ServiceAccountGap.CONTROL_SILO,
"GroupIntegrationDetails": ServiceAccountGap.REQUIRES_USER,
"GroupIntegrationDetails_post": ServiceAccountGap.REQUIRES_USER,
"GroupIntegrationDetails_put": ServiceAccountGap.REQUIRES_USER,
"OrganizationDashboards": ServiceAccountGap.REQUIRES_USER,
"OrganizationDashboards_post": ServiceAccountGap.REQUIRES_USER,
"OrganizationDetectorIndex": ServiceAccountGap.REQUIRES_USER,
"OrganizationDetectorIndex_delete": ServiceAccountGap.REQUIRES_USER,
"OrganizationDetectorIndex_put": ServiceAccountGap.REQUIRES_USER,
"OrganizationGroupSearchViews": ServiceAccountGap.REQUIRES_USER,
"OrganizationGroupSearchViews_post": ServiceAccountGap.REQUIRES_USER,
"OrganizationMemberDetails": ServiceAccountGap.REQUIRES_USER,
"OrganizationMemberDetails_delete": ServiceAccountGap.REQUIRES_USER,
"OrganizationMemberDetails_put": ServiceAccountGap.REQUIRES_USER,
"OrganizationMemberTeamDetails_delete": ServiceAccountGap.REQUIRES_USER,
"OrganizationMemberTeamDetails_post": ServiceAccountGap.REQUIRES_USER,
"OrganizationMemberTeamDetails_put": ServiceAccountGap.REQUIRES_USER,
"OrganizationProjects_post": ServiceAccountGap.REQUIRES_USER,
"OrganizationStarredServiceSpans_delete": ServiceAccountGap.REQUIRES_USER,
"ProjectDetails_put": ServiceAccountGap.REQUIRES_USER,
"OrganizationStarredServiceSpans_post": ServiceAccountGap.USER_ACTOR_ERROR,
"ProjectDetails_delete": ServiceAccountGap.USER_ACTOR_ERROR,
"TeamDetails_delete": ServiceAccountGap.USER_ACTOR_ERROR,
"EventAttachmentDetails": ServiceAccountGap.NO_ORG_ROLE,
"OrganizationMemberIndex_post": ServiceAccountGap.NO_ORG_ROLE,
}


def _matrix_cases() -> tuple[tuple[PublicGetEndpoint, MatrixAuthentication], ...]:
cases: list[tuple[PublicGetEndpoint, MatrixAuthentication]] = []
for endpoint in PUBLIC_GET_ENDPOINTS:
Expand All @@ -595,6 +653,7 @@ def _matrix_cases() -> tuple[tuple[PublicGetEndpoint, MatrixAuthentication], ...
MatrixAuthentication.VIEWER_CONTEXT,
MatrixAuthentication.AGENT_TOKEN,
MatrixAuthentication.SCOPED_DOWN_AGENT_TOKEN,
MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT,
):
cases.append((endpoint, authentication))
return tuple(cases)
Expand All @@ -620,6 +679,7 @@ def _private_helper_get_matrix_cases() -> tuple[
MatrixAuthentication.VIEWER_CONTEXT,
MatrixAuthentication.AGENT_TOKEN,
MatrixAuthentication.SCOPED_DOWN_AGENT_TOKEN,
MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT,
)
)

Expand Down Expand Up @@ -2147,6 +2207,51 @@ def _mutation_payload(self, endpoint: PublicMutationEndpoint) -> dict[str, Any]:
}
return payloads.get(key, {})

def _create_matrix_service_account(self) -> None:
# Before any control-silo override, so the cell replica is written.
service_account = self.create_service_account(
organization=self.org,
name="permission-matrix",
scope_list=sorted(SERVICE_ACCOUNT_SCOPES),
project_ids=[self.project.id],
)
self._service_account_context = encode_viewer_context(
ViewerContext(
actor=Actor(type=ActorType.SERVICE_ACCOUNT, actor_id=service_account.id),
organization_id=self.org.id,
project_id=self.project.id,
),
key=SECRET,
)

def _assert_service_account_outcome(
self,
endpoint: PublicGetEndpoint | PublicMutationEndpoint,
authorization_oracle: Response,
response: Response,
) -> None:
matches_oracle = response.status_code == authorization_oracle.status_code and (
authorization_oracle.status_code < 400
or response.content == authorization_oracle.content
)
assert matches_oracle, (
endpoint.path_template,
authorization_oracle.content,
response.content,
)

def _assert_service_account_user_global_get_denied(self, endpoint: PublicGetEndpoint) -> None:
self._create_matrix_service_account()
assert endpoint.silo_mode is not None
with assume_test_silo_mode(endpoint.silo_mode):
session_response = self.client.get(endpoint.path_template)
response = self._request(
endpoint.path_template, MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT
)

assert session_response.status_code == 200, session_response.content
assert response.status_code in {401, 403}, response.content

def _user_token(self, scopes: frozenset[str] | None = None) -> str:
with assume_test_silo_mode(SiloMode.CONTROL):
token = self.create_user_auth_token(
Expand Down Expand Up @@ -2182,6 +2287,8 @@ def _request(
key=SECRET,
)
return client.get(path, HTTP_X_VIEWER_CONTEXT=viewer_context)
if authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
return client.get(path, HTTP_X_VIEWER_CONTEXT=self._service_account_context)
if authentication in {
MatrixAuthentication.AGENT_TOKEN,
MatrixAuthentication.SCOPED_DOWN_AGENT_TOKEN,
Expand Down Expand Up @@ -2213,6 +2320,8 @@ def _mutation_request(
),
key=SECRET,
)
elif authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
headers["HTTP_X_VIEWER_CONTEXT"] = self._service_account_context
elif authentication in {
MatrixAuthentication.AGENT_TOKEN,
MatrixAuthentication.SCOPED_DOWN_AGENT_TOKEN,
Expand Down Expand Up @@ -2300,6 +2409,8 @@ def _assert_public_get_authentication(
with self.feature(FLAG):
agent_bearer = self._mint_agent_token(requested_scopes=requested_scopes)
agent_scopes = frozenset(agent_token.decode_agent_token(agent_bearer)["scopes"])
if authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
self._create_matrix_service_account()
if endpoint.endpoint_name == "SeerModelsEndpoint":
cache.set(SEER_MODELS_CACHE_KEY, {"models": ["matrix-model"]})
self.addCleanup(cache.delete, SEER_MODELS_CACHE_KEY)
Expand Down Expand Up @@ -2369,6 +2480,12 @@ def _assert_public_get_authentication(
MatrixAuthentication.USER_TOKEN,
user_token_scopes=agent_scopes,
)
elif authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
authorization_oracle = self._request(
path,
MatrixAuthentication.USER_TOKEN,
user_token_scopes=SERVICE_ACCOUNT_SCOPES,
)
response = (
baseline
if authentication is MatrixAuthentication.SESSION
Expand All @@ -2378,6 +2495,9 @@ def _assert_public_get_authentication(
assert baseline.status_code < 500, baseline.content
expected_baseline_status = OUTER_BOUNDARY_GET_STATUSES.get(endpoint.endpoint_name, 200)
assert baseline.status_code == expected_baseline_status, baseline.content
if authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
self._assert_service_account_outcome(endpoint, authorization_oracle, response)
return

# Successful representations can legitimately differ by credential (for
# example, an ``access`` field must reflect narrower token scopes), and
Expand Down Expand Up @@ -2439,6 +2559,8 @@ def _assert_public_mutation_authentication(
requested_scopes=requested_scopes,
)
agent_scopes = frozenset(agent_token.decode_agent_token(agent_bearer)["scopes"])
if authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
self._create_matrix_service_account()

silo_scope = (
assume_test_silo_mode(endpoint.silo_mode)
Expand Down Expand Up @@ -2503,6 +2625,14 @@ def _assert_public_mutation_authentication(
MatrixAuthentication.USER_TOKEN,
user_token_scopes=agent_scopes,
)
elif authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
authorization_oracle = self._rolled_back_mutation_request(
endpoint,
path,
payload,
MatrixAuthentication.USER_TOKEN,
user_token_scopes=SERVICE_ACCOUNT_SCOPES,
)

response = (
baseline
Expand All @@ -2528,6 +2658,9 @@ def _assert_public_mutation_authentication(
)
else:
assert baseline.status_code == outer_status, baseline.content
if authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
self._assert_service_account_outcome(endpoint, authorization_oracle, response)
return

if (
authentication is MatrixAuthentication.APPROVED_AGENT_TOKEN
Expand Down Expand Up @@ -2562,6 +2695,16 @@ def _assert_public_mutation_authentication(
)


def _mark_service_account_gap[F: Callable[..., None]](test: F, test_id: str) -> F:
# Until service account auth lands (#128779), service account requests are anonymous.
if not hasattr(access, "ServiceAccountAccess"):
return pytest.mark.xfail(reason="service account auth isn't supported yet")(test)
gap = SERVICE_ACCOUNT_KNOWN_GAPS.get(test_id)
if gap is None:
return test
return pytest.mark.xfail(strict=True, reason=gap.value)(test)


def _install_public_get_matrix_tests() -> None:
"""Generate real unittest methods because pytest cannot parametrize APITestCase."""

Expand All @@ -2577,6 +2720,8 @@ def test_matrix_cell(
test_matrix_cell.__name__ = (
f"test_public_get_{index:03d}_{endpoint.test_id}_{authentication.value}"
)
if authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
test_matrix_cell = _mark_service_account_gap(test_matrix_cell, endpoint.test_id)
setattr(AgentTokenPublicGetMatrixTest, test_matrix_cell.__name__, test_matrix_cell)


Expand All @@ -2596,6 +2741,8 @@ def test_matrix_cell(
test_matrix_cell.__name__ = (
f"test_public_mutation_{index:03d}_{endpoint.test_id}_{authentication.value}"
)
if authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
test_matrix_cell = _mark_service_account_gap(test_matrix_cell, endpoint.test_id)
setattr(AgentTokenPublicGetMatrixTest, test_matrix_cell.__name__, test_matrix_cell)


Expand All @@ -2619,6 +2766,10 @@ def test_get_matrix_cell(
test_get_matrix_cell.__name__ = (
f"test_private_helper_get_{index:03d}_{get_endpoint.test_id}_{get_authentication.value}"
)
if get_authentication is MatrixAuthentication.SERVICE_ACCOUNT_VIEWER_CONTEXT:
test_get_matrix_cell = _mark_service_account_gap(
test_get_matrix_cell, get_endpoint.test_id
)
setattr(AgentTokenPublicGetMatrixTest, test_get_matrix_cell.__name__, test_get_matrix_cell)


Expand Down Expand Up @@ -2663,3 +2814,21 @@ def test_matrix_cell(


_install_user_global_boundary_matrix_tests()


def _install_service_account_user_global_matrix_tests() -> None:
for index, endpoint in enumerate(USER_GLOBAL_GET_ENDPOINTS):

def test_matrix_cell(
self: AgentTokenPublicGetMatrixTest,
endpoint: PublicGetEndpoint = endpoint,
) -> None:
self._assert_service_account_user_global_get_denied(endpoint)

test_matrix_cell.__name__ = (
f"test_user_global_get_{index:03d}_{endpoint.test_id}_service_account_viewer_context"
)
setattr(AgentTokenPublicGetMatrixTest, test_matrix_cell.__name__, test_matrix_cell)


_install_service_account_user_global_matrix_tests()
Loading