From fab4be696f1a8a968418400a27a60c522d7b79fc Mon Sep 17 00:00:00 2001 From: Ian Patrick Hines Date: Mon, 22 Jun 2026 07:52:41 -0400 Subject: [PATCH 1/2] feat(security): validate nation_slug + idempotent Stripe webhook (fixes #13) - Add shared/validation.py: nation_slug must match ^[a-z0-9-]{1,63}$, rejected with 400 at every entry point (used in DDB keys, Secrets Manager names, NB API URLs) - Make Stripe webhook idempotent: persist processed event IDs in DynamoDB with TTL; no-op on duplicate delivery (Stripe retries) - Tests cover invalid slugs and replayed webhook events (+23) Co-Authored-By: Claude Opus 4.8 (1M context) --- infrastructure/template.yaml | 41 +++ src/lambdas/nat_agent/handler.py | 16 ++ src/lambdas/nat_agent_streaming/handler.py | 16 ++ src/lambdas/nb_oauth_callback/handler.py | 22 ++ src/lambdas/shared/validation.py | 43 +++ src/lambdas/stripe_checkout/handler.py | 21 ++ src/lambdas/stripe_webhook/handler.py | 143 +++++++++- .../test_stripe_webhook_integration.py | 51 +++- tests/test_nat_agent.py | 17 ++ tests/test_nb_oauth_callback.py | 30 ++ tests/test_stripe_checkout.py | 30 ++ tests/test_stripe_webhook.py | 262 +++++++++++++++++- tests/test_validation.py | 69 +++++ 13 files changed, 739 insertions(+), 22 deletions(-) create mode 100644 src/lambdas/shared/validation.py create mode 100644 tests/test_validation.py diff --git a/infrastructure/template.yaml b/infrastructure/template.yaml index 6ba9042..16268b5 100644 --- a/infrastructure/template.yaml +++ b/infrastructure/template.yaml @@ -126,6 +126,32 @@ Resources: - Key: Project Value: nat + # StripeEventsTable - records processed Stripe webhook event IDs so duplicate + # deliveries (Stripe retries) are no-ops. TTL purges old markers automatically. + StripeEventsTable: + Type: AWS::DynamoDB::Table + DeletionPolicy: Retain + UpdateReplacePolicy: Retain + Properties: + TableName: !Sub 'nat-stripe-events-${Environment}' + BillingMode: PAY_PER_REQUEST + AttributeDefinitions: + - AttributeName: event_id + AttributeType: S + KeySchema: + - AttributeName: event_id + KeyType: HASH + TimeToLiveSpecification: + AttributeName: expires_at + Enabled: true + PointInTimeRecoverySpecification: + PointInTimeRecoveryEnabled: true + Tags: + - Key: Environment + Value: !Ref Environment + - Key: Project + Value: nat + # ============================================================================= # IAM Roles # ============================================================================= @@ -163,6 +189,8 @@ Resources: - !Sub '${TenantsTable.Arn}/index/*' - !GetAtt UsersTable.Arn - !Sub '${UsersTable.Arn}/index/*' + - !GetAtt StripeEventsTable.Arn + - !Sub '${StripeEventsTable.Arn}/index/*' - PolicyName: SecretsManagerAccess PolicyDocument: Version: '2012-10-17' @@ -282,6 +310,7 @@ Resources: NATIONS_TABLE: !Ref NationsTable TENANTS_TABLE: !Ref TenantsTable USERS_TABLE: !Ref UsersTable + STRIPE_EVENTS_TABLE: !Ref StripeEventsTable STRIPE_WEBHOOK_SECRET_NAME: !Sub 'nat/stripe-webhook-secret-${Environment}' Tags: - Key: Environment @@ -872,6 +901,18 @@ Outputs: Export: Name: !Sub '${AWS::StackName}-UsersTableArn' + StripeEventsTableName: + Description: Name of the Stripe events idempotency DynamoDB table + Value: !Ref StripeEventsTable + Export: + Name: !Sub '${AWS::StackName}-StripeEventsTableName' + + StripeEventsTableArn: + Description: ARN of the Stripe events idempotency DynamoDB table + Value: !GetAtt StripeEventsTable.Arn + Export: + Name: !Sub '${AWS::StackName}-StripeEventsTableArn' + LambdaExecutionRoleArn: Description: ARN of the Lambda execution role Value: !GetAtt LambdaExecutionRole.Arn diff --git a/src/lambdas/nat_agent/handler.py b/src/lambdas/nat_agent/handler.py index a40b190..371609f 100644 --- a/src/lambdas/nat_agent/handler.py +++ b/src/lambdas/nat_agent/handler.py @@ -38,6 +38,11 @@ authenticate_request, ) +try: # Resolve in both pytest (repo root) and flattened Lambda packages. + from src.lambdas.shared.validation import is_valid_nation_slug +except ModuleNotFoundError: # pragma: no cover - exercised only in Lambda + from shared.validation import is_valid_nation_slug # type: ignore[no-redef] + logger = logging.getLogger() logger.setLevel(logging.INFO) @@ -373,6 +378,17 @@ def handler(event: dict[str, Any], context: Any) -> LambdaResponse: user_id = session.user_id nation_slug = session.nation_slug + # Defense in depth: the slug is token-attested (minted post-validation at + # OAuth connect), but it is interpolated into Secrets Manager / DynamoDB + # lookups, so reject a malformed value rather than build a bad key. + if not is_valid_nation_slug(nation_slug): + logger.error(f"Invalid nation_slug in session token: {nation_slug!r}") + return { + "statusCode": 400, + "body": json.dumps({"error": "Invalid nation_slug format"}), + "headers": headers, + } + # Extract request payload (identity fields are ignored if present). query = body.get("query") page_context = body.get("context", {}) diff --git a/src/lambdas/nat_agent_streaming/handler.py b/src/lambdas/nat_agent_streaming/handler.py index 1cb226a..50ffc1d 100644 --- a/src/lambdas/nat_agent_streaming/handler.py +++ b/src/lambdas/nat_agent_streaming/handler.py @@ -40,6 +40,11 @@ authenticate_request, ) +try: # Resolve in both pytest (repo root) and flattened Lambda packages. + from src.lambdas.shared.validation import is_valid_nation_slug +except ModuleNotFoundError: # pragma: no cover - exercised only in Lambda + from shared.validation import is_valid_nation_slug # type: ignore[no-redef] + logger = logging.getLogger() logger.setLevel(logging.INFO) @@ -516,6 +521,17 @@ async def process_streaming_request(body: dict[str, Any]) -> AsyncGenerator[str, }) return + # Defense in depth: the slug is token-attested (minted post-validation at + # OAuth connect), but it is interpolated into Secrets Manager / DynamoDB + # lookups, so reject a malformed value rather than build a bad key. + if not is_valid_nation_slug(nation_slug): + logger.error(f"Invalid nation_slug in session token: {nation_slug!r}") + yield format_sse_event(SSE_EVENT_ERROR, { + "error": "Invalid nation_slug format", + "error_code": "BAD_REQUEST", + }) + return + logger.info(f"Processing streaming query for nation {nation_slug}, user {user_id}: {query[:100]}...") # Check if billing cycle has reset for this nation (must run before the diff --git a/src/lambdas/nb_oauth_callback/handler.py b/src/lambdas/nb_oauth_callback/handler.py index 73ed887..8005f52 100644 --- a/src/lambdas/nb_oauth_callback/handler.py +++ b/src/lambdas/nb_oauth_callback/handler.py @@ -13,6 +13,7 @@ import json import logging import os +import re from datetime import datetime, timezone from typing import Any, TypedDict from urllib.parse import parse_qs, urlencode @@ -51,6 +52,19 @@ "ERROR_REDIRECT_URL", "https://natassistant.com/connection-error" ) +# The nation slug arrives in the (client-supplied) OAuth ``state`` and is the +# most security-sensitive entry point: it is interpolated into the Secrets +# Manager path ``nat/nation/{slug}/nb-tokens`` and the NationBuilder token URL +# host ``https://{slug}.nationbuilder.com``. Reject malformed values before +# either is constructed. Mirrors shared.validation.NATION_SLUG_PATTERN (this +# Lambda is packaged without shared/, so the pattern is inlined). +NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}$") + + +def is_valid_nation_slug(slug: Any) -> bool: + """Return True if *slug* is a non-empty, well-formed nation slug.""" + return isinstance(slug, str) and NATION_SLUG_PATTERN.match(slug) is not None + class LambdaResponse(TypedDict): """Lambda response type.""" @@ -434,6 +448,14 @@ def handler(event: dict[str, Any], context: Any) -> LambdaResponse: error_url = f"{ERROR_REDIRECT_URL}?error=invalid_state" return create_redirect_response(error_url) + # Reject a malformed slug before it reaches Secrets Manager / the NB token + # URL. This endpoint is a browser OAuth redirect, so a rejected slug is + # surfaced as an error redirect rather than a 400 JSON body. + if not is_valid_nation_slug(nb_slug): + logger.error(f"Invalid nation_slug in OAuth state: {nb_slug!r}") + error_url = f"{ERROR_REDIRECT_URL}?error=invalid_nation" + return create_redirect_response(error_url) + # Extract optional email from state for admin user creation user_email = state_data.get("email") diff --git a/src/lambdas/shared/validation.py b/src/lambdas/shared/validation.py new file mode 100644 index 0000000..d7d8dbb --- /dev/null +++ b/src/lambdas/shared/validation.py @@ -0,0 +1,43 @@ +""" +Input validation helpers shared across Lambda handlers. + +The canonical home for ``nation_slug`` format validation. A nation slug flows +directly into DynamoDB keys, Secrets Manager secret names +(``nat/nation/{slug}/nb-tokens``) and NationBuilder API URLs +(``https://{slug}.nationbuilder.com``), so every entry point that accepts one +must reject malformed values before they are used. + +Only the Lambdas that bundle ``shared/`` at deploy time (the agent handlers) +import this module. The standalone handlers (stripe_checkout, stripe_webhook, +nb_oauth_callback) are packaged without ``shared/`` and inline an identical +``NATION_SLUG_PATTERN`` instead. +""" + +from __future__ import annotations + +import re +from typing import Any + +# NationBuilder slugs are lowercase alphanumeric plus hyphen. The 63-char ceiling +# matches a DNS label (slugs appear as the subdomain of *.nationbuilder.com). +NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}$") + + +class InvalidNationSlugError(ValueError): + """Raised when a nation_slug does not match the required format.""" + + +def is_valid_nation_slug(slug: Any) -> bool: + """Return True if *slug* is a non-empty, well-formed nation slug.""" + return isinstance(slug, str) and NATION_SLUG_PATTERN.match(slug) is not None + + +def validate_nation_slug(slug: Any) -> str: + """Return *slug* unchanged if valid, else raise :class:`InvalidNationSlugError`.""" + if not is_valid_nation_slug(slug): + raise InvalidNationSlugError( + f"Invalid nation_slug format: {slug!r} (must match ^[a-z0-9-]{{1,63}}$)" + ) + # is_valid_nation_slug guarantees a str; assert narrows it for the type checker. + assert isinstance(slug, str) + return slug diff --git a/src/lambdas/stripe_checkout/handler.py b/src/lambdas/stripe_checkout/handler.py index 9bec170..4a9d597 100644 --- a/src/lambdas/stripe_checkout/handler.py +++ b/src/lambdas/stripe_checkout/handler.py @@ -10,6 +10,7 @@ import json import logging import os +import re from typing import Any, TypedDict import boto3 @@ -18,6 +19,17 @@ logger = logging.getLogger() logger.setLevel(logging.INFO) +# nation_slug is validated here because it is forwarded to Stripe as checkout +# metadata and later used by the webhook as a DynamoDB key. Mirrors +# shared.validation.NATION_SLUG_PATTERN (this Lambda is packaged without +# shared/, so the pattern is inlined). +NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}$") + + +def is_valid_nation_slug(slug: Any) -> bool: + """Return True if *slug* is a non-empty, well-formed nation slug.""" + return isinstance(slug, str) and NATION_SLUG_PATTERN.match(slug) is not None + # Environment variables STRIPE_SECRET_KEY_NAME = os.environ.get( "STRIPE_SECRET_KEY_NAME", "nat/stripe-secret-key" @@ -174,6 +186,15 @@ def handler(event: dict[str, Any], context: Any) -> LambdaResponse: "headers": headers, } + if not is_valid_nation_slug(nation_slug): + return { + "statusCode": 400, + "body": json.dumps({ + "error": "Invalid nation_slug format (must match ^[a-z0-9-]{1,63}$)" + }), + "headers": headers, + } + if plan not in STRIPE_PRICE_IDS: return { "statusCode": 400, diff --git a/src/lambdas/stripe_webhook/handler.py b/src/lambdas/stripe_webhook/handler.py index 32a0275..05f5a61 100644 --- a/src/lambdas/stripe_webhook/handler.py +++ b/src/lambdas/stripe_webhook/handler.py @@ -14,6 +14,7 @@ import json import logging import os +import re import time import uuid from datetime import datetime, timezone @@ -29,10 +30,28 @@ NATIONS_TABLE = os.environ.get("NATIONS_TABLE", "nat-nations-dev") TENANTS_TABLE = os.environ.get("TENANTS_TABLE", "nat-tenants-dev") USERS_TABLE = os.environ.get("USERS_TABLE", "nat-users-dev") +# Table of processed Stripe event IDs (with TTL) used for webhook idempotency. +STRIPE_EVENTS_TABLE = os.environ.get("STRIPE_EVENTS_TABLE", "nat-stripe-events-dev") STRIPE_WEBHOOK_SECRET_NAME = os.environ.get( "STRIPE_WEBHOOK_SECRET_NAME", "nat/stripe-webhook-secret" ) +# Stripe retries a delivery for up to ~3 days until it gets a 2xx. We keep each +# processed event ID a bit beyond that window so replays are reliably deduped, +# then let DynamoDB TTL purge the marker. +EVENT_TTL_SECONDS = int(os.environ.get("STRIPE_EVENT_TTL_SECONDS", str(7 * 24 * 3600))) + +# nation_slug is validated wherever it is accepted because it flows into +# DynamoDB keys, Secrets Manager names, and NationBuilder API URLs. Mirrors +# shared.validation.NATION_SLUG_PATTERN (this Lambda is packaged without +# shared/, so the pattern is inlined). +NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}$") + + +def is_valid_nation_slug(slug: Any) -> bool: + """Return True if *slug* is a non-empty, well-formed nation slug.""" + return isinstance(slug, str) and NATION_SLUG_PATTERN.match(slug) is not None + # NEW PRICING MODEL: Nation-level plans # Nat: $29/mo, 500 queries/month # Nat Pro: $79/mo, unlimited queries (represented as 0 for unlimited) @@ -143,6 +162,66 @@ def get_dynamodb_resource() -> Any: return boto3.resource("dynamodb") +def record_event_id(event_id: str) -> bool: + """ + Atomically record a processed Stripe event ID for idempotency. + + Uses a conditional put so concurrent or retried deliveries of the same event + race on a single write: the first caller wins and should process the event; + every later caller sees the existing marker and skips. This is what prevents + ``handle_checkout_completed`` (and the other handlers) from double-writing + when Stripe legitimately redelivers an event. + + Returns: + True if this call recorded the ID (caller should process the event), + False if the ID was already present (duplicate -- caller should no-op). + """ + if not event_id: + # Genuine Stripe events always carry an ``id``. If one is missing we + # cannot dedupe, so process rather than silently drop the event. + logger.warning("Stripe event has no id; cannot dedupe, processing anyway") + return True + + dynamodb = get_dynamodb_resource() + events_table = dynamodb.Table(STRIPE_EVENTS_TABLE) + expires_at = int(time.time()) + EVENT_TTL_SECONDS + + try: + events_table.put_item( + Item={ + "event_id": event_id, + "processed_at": datetime.now(timezone.utc).isoformat(), + "expires_at": expires_at, + }, + ConditionExpression="attribute_not_exists(event_id)", + ) + return True + except ClientError as e: + if e.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + logger.info(f"Duplicate Stripe event {event_id}; skipping reprocessing") + return False + raise + + +def forget_event_id(event_id: str) -> None: + """ + Best-effort removal of an event marker after processing failed. + + If a handler raises *after* the marker was written, a subsequent Stripe retry + would otherwise be treated as a duplicate and dropped. Removing the marker + lets the retry reprocess the event. Never raises: a cleanup failure must not + mask the original processing error. + """ + if not event_id: + return + try: + dynamodb = get_dynamodb_resource() + events_table = dynamodb.Table(STRIPE_EVENTS_TABLE) + events_table.delete_item(Key={"event_id": event_id}) + except Exception as e: # noqa: BLE001 - best effort cleanup + logger.warning(f"Failed to remove event marker {event_id}: {e}") + + def get_plan_from_price(price_id: str) -> str: """Map Stripe price ID to plan name.""" # Default to extracting plan name from price ID if not in mapping @@ -189,6 +268,14 @@ def handle_checkout_completed(session: dict[str, Any]) -> None: # In production, consider alerting on this error raise ValueError(f"Missing nation_slug in checkout metadata for customer {customer_id}") + if not is_valid_nation_slug(nation_slug): + # A malformed slug would be written straight into the nation primary key; + # reject it rather than persist a corrupt record. + logger.error(f"Invalid nation_slug in checkout metadata: {nation_slug!r}") + raise ValueError( + f"Invalid nation_slug format in checkout metadata for customer {customer_id}" + ) + # If subscription exists, we'll get more details from subscription.updated # For now, create the nation with basic info dynamodb = get_dynamodb_resource() @@ -384,6 +471,16 @@ def handle_subscription_deleted(subscription: dict[str, Any]) -> None: raise +# Maps Stripe event types to their handlers. Used both to dispatch and to decide +# which event types are worth recording for idempotency (others are acked and +# ignored without a DynamoDB write). +EVENT_HANDLERS = { + "checkout.session.completed": handle_checkout_completed, + "customer.subscription.updated": handle_subscription_updated, + "customer.subscription.deleted": handle_subscription_deleted, +} + + def handler(event: dict[str, Any], context: Any) -> LambdaResponse: """ Lambda handler for Stripe webhooks. @@ -422,18 +519,38 @@ def handler(event: dict[str, Any], context: Any) -> LambdaResponse: webhook_event: WebhookEvent = json.loads(body) event_type = webhook_event.get("type", "") event_data = webhook_event.get("data", {}).get("object", {}) + event_id = webhook_event.get("id", "") - logger.info(f"Processing webhook event: {event_type}") + logger.info(f"Processing webhook event: {event_type} ({event_id})") - # Route to appropriate handler - if event_type == "checkout.session.completed": - handle_checkout_completed(event_data) - elif event_type == "customer.subscription.updated": - handle_subscription_updated(event_data) - elif event_type == "customer.subscription.deleted": - handle_subscription_deleted(event_data) - else: + # Short-circuit event types we don't act on. Acknowledging them with 200 + # (so Stripe stops retrying) without recording an idempotency marker + # keeps the events table free of the many types we ignore. + if event_type not in EVENT_HANDLERS: logger.info(f"Ignoring unhandled event type: {event_type}") + return { + "statusCode": 200, + "body": json.dumps({"received": True}), + "headers": headers, + } + + # Idempotency gate: record the event ID before doing any work. If we have + # already processed this delivery (Stripe retry, or near-simultaneous + # duplicate deliveries), skip so handlers do not double-write. + if not record_event_id(event_id): + return { + "statusCode": 200, + "body": json.dumps({"received": True, "duplicate": True}), + "headers": headers, + } + + try: + EVENT_HANDLERS[event_type](event_data) + except Exception: + # Processing failed after the marker was written; remove it so a + # Stripe retry can reprocess instead of being deduped away. + forget_event_id(event_id) + raise return { "statusCode": 200, @@ -448,6 +565,14 @@ def handler(event: dict[str, Any], context: Any) -> LambdaResponse: "body": json.dumps({"error": "Invalid JSON"}), "headers": headers, } + except ValueError as e: + # Raised by handlers for malformed payloads (e.g. invalid nation_slug). + logger.error(f"Webhook validation error: {e}") + return { + "statusCode": 400, + "body": json.dumps({"error": str(e)}), + "headers": headers, + } except ClientError as e: logger.error(f"AWS service error: {e}") return { diff --git a/tests/integration/test_stripe_webhook_integration.py b/tests/integration/test_stripe_webhook_integration.py index 20312b8..dc251cb 100644 --- a/tests/integration/test_stripe_webhook_integration.py +++ b/tests/integration/test_stripe_webhook_integration.py @@ -19,6 +19,7 @@ from unittest.mock import patch import pytest +from botocore.exceptions import ClientError from src.lambdas.stripe_webhook.handler import ( PLAN_QUERY_LIMITS, @@ -75,19 +76,47 @@ def __init__(self, items: list[dict[str, Any]] | None = None) -> None: self._items[key] = item.copy() self.put_calls: list[dict[str, Any]] = [] self.update_calls: list[dict[str, Any]] = [] + self.delete_calls: list[dict[str, Any]] = [] self.query_calls: list[dict[str, Any]] = [] self.query_results: list[dict[str, Any]] = [] @staticmethod def _key(data: dict[str, Any]) -> Any: - return data.get("nation_slug") or data.get("tenant_id") or data.get("user_id") - - def put_item(self, Item: dict[str, Any]) -> None: + return ( + data.get("nation_slug") + or data.get("tenant_id") + or data.get("user_id") + or data.get("event_id") + ) + + def put_item( + self, + Item: dict[str, Any], + ConditionExpression: str | None = None, + ) -> None: key = self._key(Item) + # Simulate a conditional put (idempotency marker): fail if it exists. + if ConditionExpression and "attribute_not_exists" in ConditionExpression: + if key is not None and key in self._items: + raise ClientError( + { + "Error": { + "Code": "ConditionalCheckFailedException", + "Message": "The conditional request failed", + } + }, + "PutItem", + ) if key: self._items[key] = Item.copy() self.put_calls.append(Item) + def delete_item(self, Key: dict[str, Any]) -> None: + key = self._key(Key) + if key is not None and key in self._items: + del self._items[key] + self.delete_calls.append(Key) + def update_item( self, Key: dict[str, Any], @@ -129,14 +158,22 @@ def get_item(self, Key: dict[str, Any]) -> dict[str, Any]: class MockDynamoDBResource: """Mock DynamoDB resource. - The webhook handler operates exclusively on the NationsTable, so a single - backing table is returned regardless of the requested table name. + Routes the Stripe events idempotency table to a dedicated backing table so + event markers don't pollute nation-record assertions; everything else + resolves to the NationsTable. """ - def __init__(self, nations_table: MockDynamoDBTable) -> None: + def __init__( + self, + nations_table: MockDynamoDBTable, + events_table: MockDynamoDBTable | None = None, + ) -> None: self.nations_table = nations_table + self.events_table = events_table or MockDynamoDBTable() def Table(self, name: str) -> MockDynamoDBTable: + if "stripe-events" in name: + return self.events_table return self.nations_table @@ -236,7 +273,7 @@ def test_checkout_with_different_plans(self) -> None: """Test checkout applies the correct query limit per plan for an already-existing nation (the plan/limit are set on the update path).""" for plan in ["starter", "team", "org"]: - nation_slug = f"nation_{plan}" + nation_slug = f"nation-{plan}" nations_table = MockDynamoDBTable([ {"nation_slug": nation_slug, "stripe_customer_id": f"{TEST_CUSTOMER_ID}_{plan}"} ]) diff --git a/tests/test_nat_agent.py b/tests/test_nat_agent.py index 42fccef..2e6deca 100644 --- a/tests/test_nat_agent.py +++ b/tests/test_nat_agent.py @@ -335,6 +335,23 @@ def test_missing_token_returns_401(self) -> None: body = json.loads(response["body"]) assert body["error_code"] == "MISSING_TOKEN" + def test_malformed_slug_in_token_returns_400(self) -> None: + """Defense in depth: a malformed slug in a validly-signed token is rejected. + + The slug is interpolated into Secrets Manager / DynamoDB lookups, so even + a token-attested value is format-checked before use. + """ + token = mint_session_token(TEST_USER_ID, "bad slug!", TEST_JWT_SECRET) + event = create_api_event( + body={"query": "test query"}, + headers={"Authorization": f"Bearer {token}"}, + ) + response = handler(event, None) + + assert response["statusCode"] == 400 + body = json.loads(response["body"]) + assert "Invalid nation_slug" in body["error"] + def test_forged_token_returns_401(self) -> None: """A token signed with the wrong secret is rejected with 401.""" event = create_api_event( diff --git a/tests/test_nb_oauth_callback.py b/tests/test_nb_oauth_callback.py index dc4e31d..11432f3 100644 --- a/tests/test_nb_oauth_callback.py +++ b/tests/test_nb_oauth_callback.py @@ -345,6 +345,36 @@ def test_state_missing_fields_redirects_to_error(self) -> None: assert response["statusCode"] == 302 assert "error=invalid_state" in response["headers"]["Location"] + def test_invalid_nation_slug_redirects_to_error(self) -> None: + """A malformed slug in state is rejected before any token exchange. + + The slug flows into the Secrets Manager path and the NationBuilder token + URL host, so a bad value (here containing a path-traversal segment) must + be rejected up front. This endpoint is a browser redirect, so rejection + is surfaced as an error redirect rather than a 400 body. + """ + malicious_state = base64.urlsafe_b64encode( + json.dumps( + { + "user_id": TEST_USER_ID, + "nb_slug": "../../etc/passwd", + "redirect_uri": TEST_REDIRECT_URI, + } + ).encode() + ).decode() + + event = { + "queryStringParameters": { + "code": TEST_CODE, + "state": malicious_state, + }, + } + + response = handler(event, None) + + assert response["statusCode"] == 302 + assert "error=invalid_nation" in response["headers"]["Location"] + def test_successful_oauth_flow(self) -> None: """Test successful OAuth flow end-to-end.""" users_table = MockDynamoDBTable() diff --git a/tests/test_stripe_checkout.py b/tests/test_stripe_checkout.py index 6cc365f..3b26bde 100644 --- a/tests/test_stripe_checkout.py +++ b/tests/test_stripe_checkout.py @@ -289,6 +289,36 @@ def test_invalid_plan(self) -> None: body = json.loads(response["body"]) assert "Invalid plan" in body["error"] + def test_missing_nation_slug(self) -> None: + """Test error when nation_slug is missing.""" + event: dict[str, Any] = { + "httpMethod": "POST", + "body": json.dumps({"plan": "nat"}), + } + + response = handler(event, None) + + assert response["statusCode"] == 400 + body = json.loads(response["body"]) + assert "Missing required field: nation_slug" in body["error"] + + @pytest.mark.parametrize( + "bad_slug", + ["Bad Slug!", "UPPER", "under_score", "dots.here", "a/b", "a" * 64], + ) + def test_invalid_nation_slug_returns_400(self, bad_slug: str) -> None: + """A malformed nation_slug is rejected with 400 before reaching Stripe.""" + event: dict[str, Any] = { + "httpMethod": "POST", + "body": json.dumps({"plan": "nat", "nation_slug": bad_slug}), + } + + response = handler(event, None) + + assert response["statusCode"] == 400 + body = json.loads(response["body"]) + assert "Invalid nation_slug" in body["error"] + def test_invalid_json_body(self) -> None: """Test error handling for invalid JSON body.""" event: dict[str, Any] = { diff --git a/tests/test_stripe_webhook.py b/tests/test_stripe_webhook.py index 9c9e363..9507ee8 100644 --- a/tests/test_stripe_webhook.py +++ b/tests/test_stripe_webhook.py @@ -12,14 +12,18 @@ from unittest.mock import MagicMock, patch import pytest +from botocore.exceptions import ClientError from src.lambdas.stripe_webhook.handler import ( PLAN_QUERY_LIMITS, + forget_event_id, get_plan_from_price, handle_checkout_completed, handle_subscription_deleted, handle_subscription_updated, handler, + is_valid_nation_slug, + record_event_id, verify_stripe_signature, ) @@ -111,18 +115,47 @@ def __init__(self) -> None: self.items: dict[str, dict[str, Any]] = {} self.put_calls: list[dict[str, Any]] = [] self.update_calls: list[dict[str, Any]] = [] + self.delete_calls: list[dict[str, Any]] = [] self.query_results: list[dict[str, Any]] = [] @staticmethod def _key(data: dict[str, Any]) -> Any: - return data.get("nation_slug") or data.get("tenant_id") or data.get("user_id") - - def put_item(self, Item: dict[str, Any]) -> None: + return ( + data.get("nation_slug") + or data.get("tenant_id") + or data.get("user_id") + or data.get("event_id") + ) + + def put_item( + self, + Item: dict[str, Any], + ConditionExpression: str | None = None, + ) -> None: key = self._key(Item) + # Simulate a conditional put (idempotency marker): fail if the item + # already exists, mirroring DynamoDB's ConditionalCheckFailedException. + if ConditionExpression and "attribute_not_exists" in ConditionExpression: + if key is not None and key in self.items: + raise ClientError( + { + "Error": { + "Code": "ConditionalCheckFailedException", + "Message": "The conditional request failed", + } + }, + "PutItem", + ) if key: self.items[key] = Item self.put_calls.append(Item) + def delete_item(self, Key: dict[str, Any]) -> None: + key = self._key(Key) + if key is not None and key in self.items: + del self.items[key] + self.delete_calls.append(Key) + def get_item(self, Key: dict[str, Any]) -> dict[str, Any]: key = self._key(Key) if key is not None and key in self.items: @@ -154,14 +187,22 @@ def query( class MockDynamoDBResource: """Mock DynamoDB resource for testing. - The webhook handler now operates exclusively on the NationsTable, so a - single backing table is sufficient regardless of the table name requested. + Routes by table name: the Stripe events idempotency table is backed by a + dedicated mock so event markers don't pollute nation-record assertions; + everything else resolves to the NationsTable. """ - def __init__(self, nations_table: MockDynamoDBTable) -> None: + def __init__( + self, + nations_table: MockDynamoDBTable, + events_table: MockDynamoDBTable | None = None, + ) -> None: self.nations_table = nations_table + self.events_table = events_table or MockDynamoDBTable() def Table(self, name: str) -> MockDynamoDBTable: + if "stripe-events" in name: + return self.events_table return self.nations_table @@ -479,6 +520,7 @@ def test_lowercase_signature_header(self) -> None: mock_resource = MockDynamoDBResource(nations_table) event_body = { + "id": "evt_lowercase", "type": "checkout.session.completed", "data": { "object": { @@ -510,3 +552,211 @@ def test_lowercase_signature_header(self) -> None: response = handler(lambda_event, None) assert response["statusCode"] == 200 + + +class TestNationSlugValidation: + """Tests for nation_slug format validation in the webhook.""" + + def test_is_valid_nation_slug_accepts_good_slugs(self) -> None: + assert is_valid_nation_slug("testnation") + assert is_valid_nation_slug("my-nation-123") + assert is_valid_nation_slug("a") + assert is_valid_nation_slug("a" * 63) + + def test_is_valid_nation_slug_rejects_bad_slugs(self) -> None: + assert not is_valid_nation_slug("") + assert not is_valid_nation_slug("UPPER") + assert not is_valid_nation_slug("has space") + assert not is_valid_nation_slug("under_score") + assert not is_valid_nation_slug("dots.here") + assert not is_valid_nation_slug("slash/here") + assert not is_valid_nation_slug("a" * 64) + assert not is_valid_nation_slug(None) # type: ignore[arg-type] + + def test_handle_checkout_invalid_slug_raises(self) -> None: + """A malformed slug in checkout metadata must not be persisted.""" + nations_table = MockDynamoDBTable() + mock_resource = MockDynamoDBResource(nations_table) + + session = { + "customer": TEST_CUSTOMER_ID, + "subscription": TEST_SUBSCRIPTION_ID, + "metadata": {"plan": "nat", "nation_slug": "Bad Slug!"}, + } + + with patch( + "src.lambdas.stripe_webhook.handler.get_dynamodb_resource", + return_value=mock_resource, + ): + with pytest.raises(ValueError, match="Invalid nation_slug"): + handle_checkout_completed(session) + + # Nothing written to the nations table. + assert len(nations_table.put_calls) == 0 + assert len(nations_table.update_calls) == 0 + + def test_handler_invalid_slug_returns_400(self) -> None: + """End-to-end: an invalid slug surfaces as a 400 from the handler.""" + nations_table = MockDynamoDBTable() + mock_resource = MockDynamoDBResource(nations_table) + + event_body = { + "id": "evt_badslug", + "type": "checkout.session.completed", + "data": { + "object": { + "customer": TEST_CUSTOMER_ID, + "subscription": TEST_SUBSCRIPTION_ID, + "metadata": {"nation_slug": "Bad Slug!"}, + } + }, + } + body = json.dumps(event_body) + signature = create_stripe_signature(body, TEST_WEBHOOK_SECRET) + lambda_event = {"body": body, "headers": {"Stripe-Signature": signature}} + + with ( + patch( + "src.lambdas.stripe_webhook.handler.get_stripe_webhook_secret", + return_value=TEST_WEBHOOK_SECRET, + ), + patch( + "src.lambdas.stripe_webhook.handler.get_dynamodb_resource", + return_value=mock_resource, + ), + ): + response = handler(lambda_event, None) + + assert response["statusCode"] == 400 + assert len(nations_table.put_calls) == 0 + + +class TestWebhookIdempotency: + """Tests for Stripe webhook idempotency (duplicate event handling).""" + + def test_record_event_id_first_then_duplicate(self) -> None: + """First call records and returns True; a replay returns False.""" + events_table = MockDynamoDBTable() + mock_resource = MockDynamoDBResource(MockDynamoDBTable(), events_table) + + with patch( + "src.lambdas.stripe_webhook.handler.get_dynamodb_resource", + return_value=mock_resource, + ): + assert record_event_id("evt_dedupe") is True + assert record_event_id("evt_dedupe") is False + + # Marker stored exactly once, carries a TTL. + assert "evt_dedupe" in events_table.items + assert "expires_at" in events_table.items["evt_dedupe"] + + def test_record_event_id_missing_id_processes(self) -> None: + """An event with no id cannot be deduped; we still process it.""" + events_table = MockDynamoDBTable() + mock_resource = MockDynamoDBResource(MockDynamoDBTable(), events_table) + + with patch( + "src.lambdas.stripe_webhook.handler.get_dynamodb_resource", + return_value=mock_resource, + ): + assert record_event_id("") is True + + assert len(events_table.put_calls) == 0 + + def test_forget_event_id_allows_reprocessing(self) -> None: + """Removing a marker lets a later delivery be processed again.""" + events_table = MockDynamoDBTable() + mock_resource = MockDynamoDBResource(MockDynamoDBTable(), events_table) + + with patch( + "src.lambdas.stripe_webhook.handler.get_dynamodb_resource", + return_value=mock_resource, + ): + assert record_event_id("evt_retry") is True + forget_event_id("evt_retry") + # Marker gone -> a retry records it fresh. + assert record_event_id("evt_retry") is True + + def test_duplicate_delivery_does_not_double_write(self) -> None: + """A replayed checkout webhook must not create the nation twice.""" + nations_table = MockDynamoDBTable() + events_table = MockDynamoDBTable() + mock_resource = MockDynamoDBResource(nations_table, events_table) + + event_body = { + "id": "evt_duplicate", + "type": "checkout.session.completed", + "data": { + "object": { + "customer": TEST_CUSTOMER_ID, + "subscription": TEST_SUBSCRIPTION_ID, + "customer_email": TEST_EMAIL, + "metadata": {"nation_slug": TEST_NATION_SLUG}, + } + }, + } + body = json.dumps(event_body) + signature = create_stripe_signature(body, TEST_WEBHOOK_SECRET) + lambda_event = {"body": body, "headers": {"Stripe-Signature": signature}} + + with ( + patch( + "src.lambdas.stripe_webhook.handler.get_stripe_webhook_secret", + return_value=TEST_WEBHOOK_SECRET, + ), + patch( + "src.lambdas.stripe_webhook.handler.get_dynamodb_resource", + return_value=mock_resource, + ), + ): + first = handler(lambda_event, None) + second = handler(lambda_event, None) + + assert first["statusCode"] == 200 + assert json.loads(first["body"]) == {"received": True} + + # Second delivery is acknowledged but treated as a duplicate. + assert second["statusCode"] == 200 + assert json.loads(second["body"]) == {"received": True, "duplicate": True} + + # The nation record was created exactly once. + assert len(nations_table.put_calls) == 1 + + def test_failed_processing_removes_marker_for_retry(self) -> None: + """If a handler raises, the marker is removed so Stripe can retry.""" + nations_table = MockDynamoDBTable() + events_table = MockDynamoDBTable() + mock_resource = MockDynamoDBResource(nations_table, events_table) + + # Missing nation_slug -> handle_checkout_completed raises ValueError. + event_body = { + "id": "evt_fail", + "type": "checkout.session.completed", + "data": { + "object": { + "customer": TEST_CUSTOMER_ID, + "subscription": TEST_SUBSCRIPTION_ID, + "metadata": {"plan": "nat"}, + } + }, + } + body = json.dumps(event_body) + signature = create_stripe_signature(body, TEST_WEBHOOK_SECRET) + lambda_event = {"body": body, "headers": {"Stripe-Signature": signature}} + + with ( + patch( + "src.lambdas.stripe_webhook.handler.get_stripe_webhook_secret", + return_value=TEST_WEBHOOK_SECRET, + ), + patch( + "src.lambdas.stripe_webhook.handler.get_dynamodb_resource", + return_value=mock_resource, + ), + ): + response = handler(lambda_event, None) + + # Validation failure -> 400, and the marker was rolled back. + assert response["statusCode"] == 400 + assert "evt_fail" not in events_table.items + assert events_table.delete_calls == [{"event_id": "evt_fail"}] diff --git a/tests/test_validation.py b/tests/test_validation.py new file mode 100644 index 0000000..029a5ba --- /dev/null +++ b/tests/test_validation.py @@ -0,0 +1,69 @@ +""" +Unit tests for shared nation_slug validation. +""" + +from __future__ import annotations + +import pytest + +from src.lambdas.shared.validation import ( + InvalidNationSlugError, + is_valid_nation_slug, + validate_nation_slug, +) + + +VALID_SLUGS = [ + "testnation", + "my-nation", + "nation123", + "a", + "a" * 63, + "0-9-abc", +] + +INVALID_SLUGS = [ + "", # empty + "UPPER", # uppercase + "Mixed-Case", + "has space", + "under_score", + "dots.here", + "slash/here", + "../../etc/passwd", # path traversal + "nation!", # punctuation + "a" * 64, # too long (DNS label ceiling is 63) + "naïve", # non-ascii +] + + +@pytest.mark.parametrize("slug", VALID_SLUGS) +def test_is_valid_accepts(slug: str) -> None: + assert is_valid_nation_slug(slug) is True + + +@pytest.mark.parametrize("slug", INVALID_SLUGS) +def test_is_valid_rejects(slug: str) -> None: + assert is_valid_nation_slug(slug) is False + + +def test_is_valid_rejects_non_str() -> None: + assert is_valid_nation_slug(None) is False # type: ignore[arg-type] + assert is_valid_nation_slug(123) is False # type: ignore[arg-type] + assert is_valid_nation_slug(["testnation"]) is False # type: ignore[arg-type] + + +@pytest.mark.parametrize("slug", VALID_SLUGS) +def test_validate_returns_slug(slug: str) -> None: + assert validate_nation_slug(slug) == slug + + +@pytest.mark.parametrize("slug", INVALID_SLUGS) +def test_validate_raises(slug: str) -> None: + with pytest.raises(InvalidNationSlugError): + validate_nation_slug(slug) + + +def test_invalid_error_is_value_error() -> None: + """InvalidNationSlugError is a ValueError so handlers can catch broadly.""" + assert issubclass(InvalidNationSlugError, ValueError) From 70bada0b6da93d33be50a10e8d85f17a9dc4a045 Mon Sep 17 00:00:00 2001 From: Ian Patrick Hines Date: Mon, 22 Jun 2026 08:13:27 -0400 Subject: [PATCH 2/2] fix(security): anchor nation_slug regex with \Z to reject trailing newline Independent review caught a bypass: ^[a-z0-9-]{1,63}$ with .match() accepts a trailing newline (Python $ matches before a final \n), so "legit-slug\n" passed validation and flowed into Secrets Manager paths, NB API URLs, and DynamoDB keys. Switched $ -> \Z in all four pattern copies (shared/validation.py + 3 inlined handler copies). Added newline regression cases to INVALID_SLUGS. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/lambdas/nb_oauth_callback/handler.py | 2 +- src/lambdas/shared/validation.py | 2 +- src/lambdas/stripe_checkout/handler.py | 2 +- src/lambdas/stripe_webhook/handler.py | 2 +- tests/test_validation.py | 3 +++ 5 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/lambdas/nb_oauth_callback/handler.py b/src/lambdas/nb_oauth_callback/handler.py index 8005f52..4ac40b0 100644 --- a/src/lambdas/nb_oauth_callback/handler.py +++ b/src/lambdas/nb_oauth_callback/handler.py @@ -58,7 +58,7 @@ # host ``https://{slug}.nationbuilder.com``. Reject malformed values before # either is constructed. Mirrors shared.validation.NATION_SLUG_PATTERN (this # Lambda is packaged without shared/, so the pattern is inlined). -NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}$") +NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}\Z") def is_valid_nation_slug(slug: Any) -> bool: diff --git a/src/lambdas/shared/validation.py b/src/lambdas/shared/validation.py index d7d8dbb..b21a9e0 100644 --- a/src/lambdas/shared/validation.py +++ b/src/lambdas/shared/validation.py @@ -20,7 +20,7 @@ # NationBuilder slugs are lowercase alphanumeric plus hyphen. The 63-char ceiling # matches a DNS label (slugs appear as the subdomain of *.nationbuilder.com). -NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}$") +NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}\Z") class InvalidNationSlugError(ValueError): diff --git a/src/lambdas/stripe_checkout/handler.py b/src/lambdas/stripe_checkout/handler.py index 4a9d597..9900c9c 100644 --- a/src/lambdas/stripe_checkout/handler.py +++ b/src/lambdas/stripe_checkout/handler.py @@ -23,7 +23,7 @@ # metadata and later used by the webhook as a DynamoDB key. Mirrors # shared.validation.NATION_SLUG_PATTERN (this Lambda is packaged without # shared/, so the pattern is inlined). -NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}$") +NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}\Z") def is_valid_nation_slug(slug: Any) -> bool: diff --git a/src/lambdas/stripe_webhook/handler.py b/src/lambdas/stripe_webhook/handler.py index 05f5a61..dac1301 100644 --- a/src/lambdas/stripe_webhook/handler.py +++ b/src/lambdas/stripe_webhook/handler.py @@ -45,7 +45,7 @@ # DynamoDB keys, Secrets Manager names, and NationBuilder API URLs. Mirrors # shared.validation.NATION_SLUG_PATTERN (this Lambda is packaged without # shared/, so the pattern is inlined). -NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}$") +NATION_SLUG_PATTERN = re.compile(r"^[a-z0-9-]{1,63}\Z") def is_valid_nation_slug(slug: Any) -> bool: diff --git a/tests/test_validation.py b/tests/test_validation.py index 029a5ba..200931f 100644 --- a/tests/test_validation.py +++ b/tests/test_validation.py @@ -34,6 +34,9 @@ "nation!", # punctuation "a" * 64, # too long (DNS label ceiling is 63) "naïve", # non-ascii + "legit-slug\n", # trailing newline ($ anchor bypass — must use \Z) + "legit\nslug", # embedded newline + "legit-slug\r\n", # CRLF ]