Coverage for src/lilbee/server/auth.py: 100%

1"""Session token auth middleware with decorator-based read-only marking.""" 

2 

3from __future__ import annotations 

4 

5import hmac 

6import json 

7import logging 

8import secrets 

9from collections.abc import Callable 

10from pathlib import Path 

11from typing import Any, TypeVar 

12 

13from litestar.exceptions import NotAuthorizedException 

14from litestar.types import ASGIApp, Receive, Scope, Send 

15 

16from lilbee.core.config import cfg 

17from lilbee.core.security import file_lock_or_warn, harden_private_file, write_private_text 

18 

19log = logging.getLogger(__name__) 

20 

21_TOKEN_BYTES = 32 

22 

23_BOOT_LOCK_TIMEOUT_S = 30 

24 

25# Character floor for a persisted token. secrets.token_urlsafe(n) base64url-encodes 

26# n bytes, so the entropy count is not a length: reusing _TOKEN_BYTES here accepted 

27# tokens roughly a quarter weaker than the ones this module mints. 

28_MIN_TOKEN_CHARS = len(secrets.token_urlsafe(_TOKEN_BYTES)) 

29 

30F = TypeVar("F", bound=Callable[..., Any]) 

31 

32 

33# Route handlers AuthMiddleware skips, registered at import by the decorator 

34# below. A module-level set rather than an attribute on the function object, 

35# which mypy cannot see and would put a # type: ignore on every check. 

36_SELF_AUTHENTICATING_HANDLERS: set[Callable[..., Any]] = set() 

37 

38 

39def auth_checked_in_handler(fn: F) -> F: 

40 """Mark a route whose token check runs inside the handler, not in middleware. 

41 

42 Not an exemption: the route must still reject an unauthenticated caller 

43 itself. Only ``/v1/*`` uses this, to answer a bad token with the OpenAI 

44 error envelope instead of Litestar's 401 shape. 

45 ``test_every_route_is_authenticated`` holds the line. 

46 

47 Must sit *below* the route decorator so it receives the raw function, which 

48 is what ``AuthMiddleware`` looks up via ``handler.fn``; stacked the other 

49 way the lookup misses. Enforced below. 

50 """ 

51 if hasattr(fn, "fn"): 

52 raise TypeError( 

53 "@auth_checked_in_handler must be applied below the route decorator, " 

54 "so it sees the function rather than the route handler." 

55 ) 

56 _SELF_AUTHENTICATING_HANDLERS.add(fn) 

57 return fn 

58 

59 

60def authenticates_itself(fn: Callable[..., Any]) -> bool: 

61 """True iff *fn* was decorated with :func:`auth_checked_in_handler`.""" 

62 return fn in _SELF_AUTHENTICATING_HANDLERS 

63 

64 

65def server_json_path() -> Path: 

66 """Return the path to the server session file.""" 

67 return cfg.data_dir / "server.json" 

68 

69 

70class SessionManager: 

71 """Manages the server session token lifecycle. 

72 Replaces the old module-level ``_session_token`` global so that auth 

73 state is explicit and injectable rather than hidden mutable state. 

74 """ 

75 

76 def __init__(self) -> None: 

77 self.token: str | None = None 

78 # False until load_or_generate() or disable() runs. validate() fails 

79 # closed while unset so an app served without its lifespan (or after 

80 # cleanup) never silently accepts unauthenticated mutating requests. 

81 self._initialized: bool = False 

82 

83 def load_or_generate(self) -> str: 

84 """Return the persisted token if shape-valid; generate a new one otherwise. 

85 

86 Read and write happen under a file lock so concurrent boots converge on 

87 one token: without it both see no file, both mint, and the last write 

88 rejects every client that read the other's file. 

89 """ 

90 path = server_json_path() 

91 with file_lock_or_warn(path, _BOOT_LOCK_TIMEOUT_S): 

92 existing = self._read_persisted_token(path) 

93 if existing is not None: 

94 # The token is reused indefinitely and the file can arrive 

95 # world-readable (backup, older release), so narrow on every load. 

96 harden_private_file(path) 

97 self.token = existing 

98 self._initialized = True 

99 return existing 

100 self.token = secrets.token_urlsafe(_TOKEN_BYTES) 

101 write_private_text(path, json.dumps({"token": self.token})) 

102 self._initialized = True 

103 return self.token 

104 

105 def disable(self) -> None: 

106 """Explicitly turn auth off (test harness / embedded read-only use). 

107 

108 Distinct from the uninitialized state: validate() accepts any request 

109 once disabled, but denies until either this or load_or_generate() runs. 

110 """ 

111 self.token = None 

112 self._initialized = True 

113 

114 @staticmethod 

115 def _read_persisted_token(path: Path) -> str | None: 

116 """Return a previously-persisted token if shape-valid, else None. 

117 

118 Total by design: every way the file can be unusable returns None so the 

119 caller mints a fresh token. A corrupt server.json must never be the 

120 reason the server refuses to boot, since nothing would point the user at 

121 the file to delete. 

122 """ 

123 try: 

124 raw = path.read_text(encoding="utf-8") 

125 except OSError: 

126 return None 

127 except UnicodeDecodeError: 

128 # Not an OSError: a truncated write or a file clobbered by another 

129 # tool leaves bytes that are not valid UTF-8. 

130 return None 

131 try: 

132 data = json.loads(raw) 

133 except json.JSONDecodeError: 

134 return None 

135 if not isinstance(data, dict): 

136 return None 

137 token = data.get("token") 

138 if not isinstance(token, str) or len(token) < _MIN_TOKEN_CHARS: 

139 return None 

140 return token 

141 

142 def cleanup(self) -> None: 

143 """Remove server.json on shutdown and reset to the uninitialized state.""" 

144 self.token = None 

145 self._initialized = False 

146 path = server_json_path() 

147 try: 

148 path.unlink(missing_ok=True) 

149 except OSError: 

150 # A still-open handle on Windows makes unlink raise; the token is 

151 # already invalidated above and the file is rewritten on next boot. 

152 log.debug("Could not remove %s at shutdown.", path, exc_info=True) 

153 

154 def validate(self, auth_header: str) -> bool: 

155 """Check whether *auth_header* carries a valid bearer token. 

156 

157 Fails closed until initialized: a request reaching auth before the 

158 lifespan ran (or after cleanup) is denied rather than allowed. 

159 """ 

160 if not self._initialized: 

161 raise NotAuthorizedException("Server authentication is not initialized") 

162 if self.token is None: 

163 return True # auth explicitly disabled via disable() 

164 return hmac.compare_digest(auth_header, f"Bearer {self.token}") 

165 

166 

167# Singleton instance: used by AuthMiddleware and the app lifespan. 

168session_manager = SessionManager() 

169 

170 

171class AuthMiddleware: 

172 """Bearer token auth middleware for mutating endpoints.""" 

173 

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

175 self.app = app 

176 

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

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

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

180 return 

181 

182 method = scope.get("method", "") 

183 if method == "OPTIONS": 

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

185 return 

186 

187 handler = scope.get("route_handler") 

188 if handler and authenticates_itself(handler.fn): 

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

190 return 

191 

192 headers = dict(scope.get("headers", [])) 

193 auth_header = headers.get(b"authorization", b"").decode() 

194 if session_manager.validate(auth_header): 

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

196 return 

197 raise NotAuthorizedException("Missing or invalid bearer token")