Coverage for src/taipanstack/bridges/web_bridge.py: 100%

1""" 

2Web Bridge — ASGI middleware for rate limiting and security headers. 

3 

4Provides a framework-agnostic ASGI middleware that integrates 

5TaipanStack's rate limiter and security headers into any ASGI 

6application (FastAPI, Litestar, Starlette, etc.). 

7""" 

8 

9from __future__ import annotations 

10 

11import json 

12import logging 

13from collections.abc import Awaitable, Callable, MutableMapping 

14from dataclasses import dataclass 

15from typing import Literal, TypeAlias, TypeVar 

16 

17from taipanstack.core.result import Ok, Result 

18from taipanstack.utils.rate_limit import RateLimiter 

19 

20logger = logging.getLogger("taipanstack.bridges.web") 

21 

22T = TypeVar("T") 

23 

24# ASGI type aliases 

25Scope: TypeAlias = MutableMapping[str, object] 

26Receive: TypeAlias = Callable[[], Awaitable[MutableMapping[str, object]]] 

27Send: TypeAlias = Callable[[MutableMapping[str, object]], Awaitable[None]] 

28ASGIApp: TypeAlias = Callable[[Scope, Receive, Send], Awaitable[None]] 

29 

30 

31@dataclass(frozen=True) 

32class SecurityHeadersConfig: 

33 """Configuration for security response headers. 

34 

35 Attributes: 

36 x_content_type_options: Value for X-Content-Type-Options. 

37 x_frame_options: Value for X-Frame-Options. 

38 x_xss_protection: Value for X-XSS-Protection. 

39 strict_transport_security: Value for Strict-Transport-Security. 

40 referrer_policy: Value for Referrer-Policy. 

41 content_security_policy: Value for Content-Security-Policy. 

42 

43 """ 

44 

45 x_content_type_options: Literal["nosniff"] = "nosniff" 

46 x_frame_options: Literal["DENY", "SAMEORIGIN"] = "DENY" 

47 x_xss_protection: Literal["1; mode=block", "0"] = "1; mode=block" 

48 strict_transport_security: str = "max-age=31536000; includeSubDomains" 

49 referrer_policy: Literal[ 

50 "no-referrer", 

51 "no-referrer-when-downgrade", 

52 "origin", 

53 "origin-when-cross-origin", 

54 "same-origin", 

55 "strict-origin", 

56 "strict-origin-when-cross-origin", 

57 "unsafe-url", 

58 ] = "strict-origin-when-cross-origin" 

59 content_security_policy: str = "default-src 'self'" 

60 

61 def to_headers(self) -> list[tuple[bytes, bytes]]: 

62 """Convert config to ASGI header pairs. 

63 

64 Returns: 

65 List of (name, value) byte tuples. 

66 

67 """ 

68 return [ 

69 (b"x-content-type-options", self.x_content_type_options.encode()), 

70 (b"x-frame-options", self.x_frame_options.encode()), 

71 (b"x-xss-protection", self.x_xss_protection.encode()), 

72 ( 

73 b"strict-transport-security", 

74 self.strict_transport_security.encode(), 

75 ), 

76 (b"referrer-policy", self.referrer_policy.encode()), 

77 (b"content-security-policy", self.content_security_policy.encode()), 

78 ] 

79 

80 

81def result_to_response( 

82 result: Result[T, Exception], 

83 *, 

84 status_ok: int = 200, 

85 status_err: int = 500, 

86) -> dict[str, object]: 

87 """Convert a ``Result`` to a JSON-friendly response dict. 

88 

89 Args: 

90 result: The Result to convert. 

91 status_ok: HTTP status for ``Ok`` values. 

92 status_err: HTTP status for ``Err`` values. 

93 

94 Returns: 

95 Dict with ``status``, ``data``/``error`` keys. 

96 

97 Example: 

98 >>> result_to_response(Ok({"id": 1})) 

99 {"status": 200, "data": {"id": 1}} 

100 

101 """ 

102 if isinstance(result, Ok): 

103 return {"status": status_ok, "data": result.ok_value} 

104 return {"status": status_err, "error": str(result.err_value)} 

105 

106 

107async def _send_json_response( 

108 send: Send, 

109 *, 

110 status: int, 

111 body: dict[str, object], 

112 extra_headers: list[tuple[bytes, bytes]] | None = None, 

113) -> None: 

114 """Send a JSON response via ASGI send. 

115 

116 Args: 

117 send: ASGI send callable. 

118 status: HTTP status code. 

119 body: JSON-serializable body. 

120 extra_headers: Additional headers to include. 

121 

122 """ 

123 payload = json.dumps(body).encode("utf-8") 

124 headers: list[tuple[bytes, bytes]] = [ 

125 (b"content-type", b"application/json"), 

126 (b"content-length", str(len(payload)).encode()), 

127 ] 

128 if extra_headers: 

129 headers.extend(extra_headers) 

130 

131 await send( 

132 { 

133 "type": "http.response.start", 

134 "status": status, 

135 "headers": headers, 

136 }, 

137 ) 

138 await send( 

139 { 

140 "type": "http.response.body", 

141 "body": payload, 

142 }, 

143 ) 

144 

145 

146class TaipanMiddleware: 

147 """ASGI middleware providing rate limiting and security headers. 

148 

149 Args: 

150 app: The wrapped ASGI application. 

151 rate_limiter: Optional rate limiter instance. 

152 security_headers: Whether to inject security headers. 

153 headers_config: Custom security headers configuration. 

154 

155 Example: 

156 >>> from taipanstack.utils.rate_limit import RateLimiter 

157 >>> app = TaipanMiddleware( 

158 ... my_asgi_app, 

159 ... rate_limiter=RateLimiter(max_calls=100, time_window=60), 

160 ... security_headers=True, 

161 ... ) 

162 

163 """ 

164 

165 def __init__( 

166 self, 

167 app: ASGIApp, 

168 *, 

169 rate_limiter: RateLimiter | None = None, 

170 security_headers: bool = True, 

171 headers_config: SecurityHeadersConfig | None = None, 

172 ) -> None: 

173 """Initialize the middleware. 

174 

175 Args: 

176 app: ASGI application to wrap. 

177 rate_limiter: Optional rate limiter. 

178 security_headers: Inject security headers. 

179 headers_config: Custom headers config. 

180 

181 """ 

182 self._app = app 

183 self._rate_limiter = rate_limiter 

184 self._security_headers = security_headers 

185 self._headers_config = headers_config or SecurityHeadersConfig() 

186 

187 def _wrap_send_with_security_headers(self, send: Send) -> Send: 

188 """Wrap the send callable to inject security headers if enabled. 

189 

190 Args: 

191 send: The original ASGI send callable. 

192 

193 Returns: 

194 The wrapped ASGI send callable. 

195 

196 """ 

197 if not self._security_headers: 

198 return send 

199 

200 extra_headers = self._headers_config.to_headers() 

201 

202 async def send_with_headers(message: MutableMapping[str, object]) -> None: 

203 if message.get("type") == "http.response.start": 

204 headers = message.get("headers") 

205 existing = list(headers) if isinstance(headers, (list, tuple)) else [] 

206 existing.extend(extra_headers) 

207 message["headers"] = existing 

208 await send(message) 

209 

210 return send_with_headers 

211 

212 async def _handle_rate_limit(self, send: Send) -> bool: 

213 """Apply rate limiting and send response if exceeded. 

214 

215 Args: 

216 send: ASGI send callable. 

217 

218 Returns: 

219 True if rate limit was exceeded, False otherwise. 

220 

221 """ 

222 if self._rate_limiter is None or self._rate_limiter.consume(): 

223 return False 

224 

225 logger.warning("Rate limit exceeded for request") 

226 security_hdrs = ( 

227 self._headers_config.to_headers() if self._security_headers else None 

228 ) 

229 await _send_json_response( 

230 send, 

231 status=429, 

232 body={"error": "Rate limit exceeded", "retry_after": 1}, 

233 extra_headers=security_hdrs, 

234 ) 

235 return True 

236 

237 async def __call__( 

238 self, 

239 scope: Scope, 

240 receive: Receive, 

241 send: Send, 

242 ) -> None: 

243 """Process an ASGI request. 

244 

245 Args: 

246 scope: ASGI scope dict. 

247 receive: ASGI receive callable. 

248 send: ASGI send callable. 

249 

250 """ 

251 # Only handle HTTP requests 

252 if scope.get("type") != "http": 

253 await self._app(scope, receive, send) 

254 return 

255 

256 # Rate limiting 

257 if await self._handle_rate_limit(send): 

258 return 

259 

260 # Wrap send to inject security headers 

261 send = self._wrap_send_with_security_headers(send) 

262 

263 # Call the actual application 

264 try: 

265 await self._app(scope, receive, send) 

266 except Exception as exc: 

267 logger.exception("Unhandled exception in ASGI app", exc_info=exc) 

268 await _send_json_response( 

269 send, 

270 status=500, 

271 body={"error": "Internal server error"}, 

272 )