Coverage for polar/auth/middlewares.py: 89%

1import logfire 

2import structlog 

3from fastapi import Request 

4from fastapi.security.utils import get_authorization_scheme_param 

5from starlette.types import ASGIApp, Receive, Send 

6from starlette.types import Scope as ASGIScope 

7 

8from polar.customer_session.service import customer_session as customer_session_service 

9from polar.kit.utils import utc_now 

10from polar.logging import Logger 

11from polar.models import ( 

12 CustomerSession, 

13 OAuth2Token, 

14 OrganizationAccessToken, 

15 PersonalAccessToken, 

16 UserSession, 

17) 

18from polar.oauth2.constants import is_registration_token_prefix 

19from polar.oauth2.exception_handlers import OAuth2Error, oauth2_error_exception_handler 

20from polar.oauth2.exceptions import InvalidTokenError 

21from polar.oauth2.service.oauth2_token import oauth2_token as oauth2_token_service 

22from polar.organization_access_token.service import ( 

23 organization_access_token as organization_access_token_service, 

24) 

25from polar.personal_access_token.service import ( 

26 personal_access_token as personal_access_token_service, 

27) 

28from polar.postgres import AsyncSession 

29from polar.sentry import set_sentry_user 

30from polar.worker._enqueue import enqueue_job 

31 

32from .models import Anonymous, AuthSubject, Subject 

33from .scope import Scope 

34from .service import auth as auth_service 

35 

36log: Logger = structlog.get_logger(__name__) 

37 

38 

39async def get_user_session( 

40 request: Request, session: AsyncSession 

41) -> UserSession | None: 

42 return await auth_service.authenticate(session, request) 

43 

44 

45def get_bearer_token(request: Request) -> str | None: 

46 authorization = request.headers.get("Authorization") 

47 scheme, value = get_authorization_scheme_param(authorization) 

48 if not scheme or not value or scheme.lower() != "bearer": 

49 return None 

50 if not value.isascii(): 50 ↛ 51line 50 didn't jump to line 51 because the condition on line 50 was never true

51 return None 

52 return value 

53 

54 

55async def get_oauth2_token(session: AsyncSession, value: str) -> OAuth2Token | None: 

56 return await oauth2_token_service.get_by_access_token(session, value) 

57 

58 

59async def get_personal_access_token( 

60 session: AsyncSession, value: str 

61) -> PersonalAccessToken | None: 

62 token = await personal_access_token_service.get_by_token(session, value) 

63 

64 if token is not None: 

65 enqueue_job( 

66 "personal_access_token.record_usage", 

67 personal_access_token_id=token.id, 

68 last_used_at=utc_now().timestamp(), 

69 ) 

70 

71 return token 

72 

73 

74async def get_organization_access_token( 

75 session: AsyncSession, value: str 

76) -> OrganizationAccessToken | None: 

77 token = await organization_access_token_service.get_by_token(session, value) 

78 

79 if token is not None: 79 ↛ 80line 79 didn't jump to line 80 because the condition on line 79 was never true

80 enqueue_job( 

81 "organization_access_token.record_usage", 

82 organization_access_token_id=token.id, 

83 last_used_at=utc_now().timestamp(), 

84 ) 

85 

86 return token 

87 

88 

89async def get_customer_session( 

90 session: AsyncSession, value: str 

91) -> CustomerSession | None: 

92 return await customer_session_service.get_by_token(session, value) 

93 

94 

95async def get_auth_subject( 

96 request: Request, session: AsyncSession 

97) -> AuthSubject[Subject]: 

98 token = get_bearer_token(request) 

99 if token is not None: 

100 if is_registration_token_prefix(token): 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true

101 return AuthSubject(Anonymous(), set(), None) 

102 

103 customer_session = await get_customer_session(session, token) 

104 if customer_session: 

105 return AuthSubject( 

106 customer_session.customer, 

107 {Scope.customer_portal_write}, 

108 customer_session, 

109 ) 

110 

111 organization_access_token = await get_organization_access_token(session, token) 

112 if organization_access_token: 112 ↛ 113line 112 didn't jump to line 113 because the condition on line 112 was never true

113 return AuthSubject( 

114 organization_access_token.organization, 

115 organization_access_token.scopes, 

116 organization_access_token, 

117 ) 

118 

119 oauth2_token = await get_oauth2_token(session, token) 

120 if oauth2_token: 120 ↛ 121line 120 didn't jump to line 121 because the condition on line 120 was never true

121 return AuthSubject(oauth2_token.sub, oauth2_token.scopes, oauth2_token) 

122 

123 personal_access_token = await get_personal_access_token(session, token) 

124 if personal_access_token: 

125 return AuthSubject( 

126 personal_access_token.user, 

127 personal_access_token.scopes, 

128 personal_access_token, 

129 ) 

130 

131 raise InvalidTokenError() 

132 

133 user_session = await get_user_session(request, session) 

134 if user_session is not None: 134 ↛ 135line 134 didn't jump to line 135 because the condition on line 134 was never true

135 return AuthSubject(user_session.user, set(user_session.scopes), user_session) 

136 

137 return AuthSubject(Anonymous(), set(), None) 

138 

139 

140class AuthSubjectMiddleware: 

141 def __init__(self, app: ASGIApp) -> None: 

142 self.app = app 

143 

144 async def __call__(self, scope: ASGIScope, receive: Receive, send: Send) -> None: 

145 if scope["type"] != "http": 

146 await self.app(scope, receive, send) 

147 return 

148 

149 session: AsyncSession = scope["state"]["async_session"] 

150 request = Request(scope) 

151 

152 try: 

153 auth_subject = await get_auth_subject(request, session) 

154 except OAuth2Error as e: 

155 response = await oauth2_error_exception_handler(request, e) 

156 return await response(scope, receive, send) 

157 

158 scope["state"]["auth_subject"] = auth_subject 

159 

160 with logfire.set_baggage(**auth_subject.log_context): 

161 log.info("Authenticated subject", **auth_subject.log_context) 

162 set_sentry_user(auth_subject) 

163 await self.app(scope, receive, send)