Coverage for app/backend/src/couchers/middleware/sanitize.py: 92%

1""" 

2Stripping fields marked sensitive (see proto/annotations.proto) out of the requests and responses we log. 

3""" 

4 

5from copy import deepcopy 

6from dataclasses import dataclass 

7from functools import cache 

8from typing import overload 

9 

10from google.protobuf.descriptor import Descriptor 

11from google.protobuf.message import Message 

12 

13from couchers.proto import annotations_pb2 

14 

15 

16@cache 

17def _descriptor_has_sensitive(descriptor: Descriptor) -> bool: 

18 """Whether this message type transitively contains any field marked sensitive.""" 

19 seen: set[Descriptor] = set() 

20 stack = [descriptor] 

21 while stack: 

22 d = stack.pop() 

23 if d in seen: 

24 continue 

25 seen.add(d) 

26 for f in d.fields: 

27 if f.GetOptions().Extensions[annotations_pb2.sensitive]: 

28 return True 

29 if f.message_type is not None: 

30 stack.append(f.message_type) 

31 return False 

32 

33 

34@dataclass(frozen=True, slots=True) 

35class _SanitizePlan: 

36 fields_to_clear: tuple[str, ...] 

37 fields_to_recurse: tuple[tuple[str, bool], ...] # (field name, is_repeated) 

38 

39 

40@cache 

41def _sanitize_plan(descriptor: Descriptor) -> _SanitizePlan: 

42 """For a message type, the fields to clear and the subfields worth recursing into.""" 

43 clear = [] 

44 recurse = [] 

45 for f in descriptor.fields: 

46 if f.GetOptions().Extensions[annotations_pb2.sensitive]: 

47 clear.append(f.name) 

48 elif f.message_type is not None and _descriptor_has_sensitive(f.message_type): 

49 recurse.append((f.name, f.is_repeated)) 

50 return _SanitizePlan(fields_to_clear=tuple(clear), fields_to_recurse=tuple(recurse)) 

51 

52 

53def _sanitize_message(message: Message) -> None: 

54 plan = _sanitize_plan(message.DESCRIPTOR) 

55 for name in plan.fields_to_clear: 

56 message.ClearField(name) 

57 for name, is_repeated in plan.fields_to_recurse: 

58 submessage = getattr(message, name) 

59 if not submessage: 59 ↛ 60line 59 didn't jump to line 60 because the condition on line 59 was never true

60 continue 

61 if is_repeated: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true

62 for msg in submessage: 

63 _sanitize_message(msg) 

64 else: 

65 _sanitize_message(submessage) 

66 

67 

68@overload 

69def sanitized_bytes(proto: Message) -> bytes: ... 

70@overload 

71def sanitized_bytes(proto: None) -> None: ... 

72def sanitized_bytes(proto: Message | None) -> bytes | None: 

73 """ 

74 Remove fields marked sensitive and return serialized bytes. 

75 

76 Sensitivity is static per message type, so the descriptor analysis is cached: messages whose type has no 

77 sensitive field anywhere serialize directly without a copy or walk. 

78 """ 

79 if not proto: 

80 return None 

81 

82 if not _descriptor_has_sensitive(proto.DESCRIPTOR): 

83 return proto.SerializeToString() 

84 

85 new_proto = deepcopy(proto) 

86 _sanitize_message(new_proto) 

87 return new_proto.SerializeToString()