diff --git a/src/backend/chatbot_backend.py b/src/backend/chatbot_backend.py index 9f879fc..d93709f 100644 --- a/src/backend/chatbot_backend.py +++ b/src/backend/chatbot_backend.py @@ -49,6 +49,25 @@ def get_conversation_history(session_id: str) -> List[Dict[str, str]]: return list(reversed(messages[-max_turns:])) +def get_session_owner(session_id: str) -> Optional[str]: + """Return the owner_sub of an existing session. + + Returns None if the session has no stored items yet (a brand-new session) or + predates ownership tracking (legacy rows written before this fix). The oldest + item is authoritative — it was written by the session's creator. + """ + table = dynamodb.Table(os.getenv("CONVERSATION_TABLE")) + response = table.query( + KeyConditionExpression=Key("session_id").eq(session_id), + Limit=1, + ProjectionExpression="owner_sub", + ) + items = response.get("Items", []) + if not items: + return None + return items[0].get("owner_sub") + + def _prepare_for_dynamodb(value: Any) -> Any: """Recursively clean a value tree for DynamoDB put_item: - drop None values and empty strings (rejected by the high-level Table API) @@ -76,6 +95,7 @@ def save_message( document_ids: List[str] = None, conversation_turn: str = None, sources: Optional[List[Dict[str, Any]]] = None, + owner_sub: Optional[str] = None, ) -> int: """Save a message to conversation history and return timestamp.""" table = dynamodb.Table(os.getenv("CONVERSATION_TABLE")) @@ -88,6 +108,9 @@ def save_message( "content": content, } + if owner_sub: + item["owner_sub"] = owner_sub + if document_ids: item["document_ids"] = document_ids @@ -444,6 +467,16 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]: user_query: str = body_data["query"] session_id: str = body_data.get("session_id", str(uuid.uuid4())) + # Identity is injected by the authenticated proxy from the validated JWT. + caller_sub: Optional[str] = body_data.get("owner_sub") + if not caller_sub: + return {"statusCode": 401, "body": json.dumps("Unauthorized")} + + # IDOR guard: refuse to read or write a session owned by another user. + existing_owner = get_session_owner(session_id) + if existing_owner is not None and existing_owner != caller_sub: + return {"statusCode": 403, "body": json.dumps("Forbidden")} + # Get conversation history history = get_conversation_history(session_id) @@ -538,7 +571,14 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]: conversation_turn = str(uuid.uuid4()) # Save conversation to history - save_message(session_id, "user", user_query, None, conversation_turn) + save_message( + session_id, + "user", + user_query, + None, + conversation_turn, + owner_sub=caller_sub, + ) assistant_timestamp = save_message( session_id, "assistant", @@ -546,6 +586,7 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]: document_ids, conversation_turn, sources=sources, + owner_sub=caller_sub, ) return { diff --git a/src/backend/feedback.py b/src/backend/feedback.py index 896c962..75ce542 100644 --- a/src/backend/feedback.py +++ b/src/backend/feedback.py @@ -4,17 +4,26 @@ from typing import Any, Dict import boto3 +from botocore.exceptions import ClientError logger = logging.getLogger() logger.setLevel(logging.INFO) dynamodb = boto3.resource("dynamodb") +# Guards every feedback write: the target message row must already exist +# (attribute_exists stops update_item from upserting a brand-new row) AND must +# be owned by the caller (stops cross-user tampering). Fails closed on legacy +# rows written before ownership tracking, which have no owner_sub. +_OWNERSHIP_CONDITION = "attribute_exists(session_id) AND owner_sub = :owner" -def save_feedback(session_id: str, timestamp: int, rating: str, feedback_text: str = "") -> None: - """Save feedback for a message.""" + +def save_feedback( + session_id: str, timestamp: int, rating: str, owner_sub: str, feedback_text: str = "" +) -> None: + """Save feedback for a message the caller owns.""" table = dynamodb.Table(os.getenv("CONVERSATION_TABLE")) - + if rating == "thumbs_up": # Save thumb feedback and clear any prior thumbs-down reason table.update_item( @@ -23,8 +32,10 @@ def save_feedback(session_id: str, timestamp: int, rating: str, feedback_text: s "timestamp": timestamp }, UpdateExpression="SET thumb_rating = :rating REMOVE feedback_text", + ConditionExpression=_OWNERSHIP_CONDITION, ExpressionAttributeValues={ - ":rating": rating + ":rating": rating, + ":owner": owner_sub } ) elif rating == "thumbs_down": @@ -36,9 +47,11 @@ def save_feedback(session_id: str, timestamp: int, rating: str, feedback_text: s "timestamp": timestamp }, UpdateExpression="SET thumb_rating = :rating, feedback_text = :text", + ConditionExpression=_OWNERSHIP_CONDITION, ExpressionAttributeValues={ ":rating": rating, - ":text": feedback_text + ":text": feedback_text, + ":owner": owner_sub } ) else: @@ -49,8 +62,10 @@ def save_feedback(session_id: str, timestamp: int, rating: str, feedback_text: s "timestamp": timestamp }, UpdateExpression="SET thumb_rating = :rating REMOVE feedback_text", + ConditionExpression=_OWNERSHIP_CONDITION, ExpressionAttributeValues={ - ":rating": rating + ":rating": rating, + ":owner": owner_sub } ) @@ -64,13 +79,34 @@ def feedback_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]: rating: str = body_data["rating"] # "thumbs_up" or "thumbs_down" feedback_text: str = body_data.get("feedback_text", "") - save_feedback(session_id, timestamp, rating, feedback_text) + # Identity is injected by the authenticated proxy from the validated JWT. + owner_sub: str = body_data.get("owner_sub") + if not owner_sub: + return { + "statusCode": 401, + "body": json.dumps("Unauthorized") + } + + save_feedback(session_id, timestamp, rating, owner_sub, feedback_text) return { "statusCode": 200, "body": json.dumps({"message": "Feedback saved"}) } + except ClientError as e: + # ConditionExpression failed: the row doesn't exist or isn't the caller's. + if e.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + return { + "statusCode": 403, + "body": json.dumps("Forbidden") + } + logger.error(f"Error in feedback_handler: {e}") + return { + "statusCode": 500, + "body": json.dumps("Error saving feedback") + } + except Exception as e: logger.error(f"Error in feedback_handler: {e}") return { diff --git a/src/proxy/proxy_handler.py b/src/proxy/proxy_handler.py index 56f78f6..54e96a2 100644 --- a/src/proxy/proxy_handler.py +++ b/src/proxy/proxy_handler.py @@ -102,6 +102,38 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]: }) } + if not isinstance(request_data, dict): + return { + 'statusCode': 400, + 'headers': { + 'Content-Type': 'application/json', + 'Access-Control-Allow-Origin': '*' + }, + 'body': json.dumps({ + 'error': 'Request body must be a JSON object' + }) + } + + # Inject the authenticated caller's identity from the validated JWT claims. + # The Cognito authorizer on the proxy API guarantees these claims are present + # for any request that reaches here. Overwrite any client-supplied value so a + # caller cannot spoof another user's identity — the backend trusts owner_sub + # only because it arrives via this keyed, authenticated hop. + claims = (event.get('requestContext', {}).get('authorizer') or {}).get('claims') or {} + caller_sub = claims.get('sub') + if not caller_sub: + return { + 'statusCode': 401, + 'headers': { + 'Content-Type': 'application/json', + 'Access-Control-Allow-Origin': '*' + }, + 'body': json.dumps({ + 'error': 'Unauthorized' + }) + } + request_data['owner_sub'] = caller_sub + # Forward the request to the backend API with the API key try: print(f"Forwarding request to: {full_url}")