Coverage for src/lexigram/web/transport/websocket_guards.py: 45%

1""" 

2WebSocket guard support for securing WebSocket connections. 

3 

4Guards run during the handshake phase before connection is accepted. 

5""" 

6 

7from __future__ import annotations 

8 

9from typing import TYPE_CHECKING, Any 

10 

11from lexigram.result import Err, Ok, Result 

12from lexigram.web.security.guards import GuardRejection 

13 

14if TYPE_CHECKING: 

15 from starlette.websockets import WebSocket 

16 

17 from lexigram.web.security.guards import GuardProtocol 

18 

19 

20async def execute_websocket_guards( 

21 guards: list[GuardProtocol], websocket: WebSocket 

22) -> Result[None, GuardRejection]: 

23 """ 

24 Execute guards during WebSocket handshake. 

25 

26 Args: 

27 guards: List of guard instances 

28 websocket: The WebSocket connection 

29 

30 Returns: 

31 Ok(None) if all guards pass, Err(GuardRejection) if any guard fails 

32 """ 

33 for guard in guards: 

34 can_proceed = await guard.can_activate(websocket) # type: ignore[arg-type] 

35 if not can_proceed: 

36 await websocket.close(code=4003, reason="Forbidden") 

37 return Err(GuardRejection()) 

38 

39 return Ok(None) 

40 

41 

42class GuardedWebSocket: 

43 """ 

44 WebSocket wrapper that enforces guards during connection. 

45 

46 Usage: 

47 @websocket("/chat") 

48 @use_guards(AuthGuard) 

49 async def chat(self, ws: WebSocket): 

50 guarded_ws = GuardedWebSocket(ws, guards) 

51 await guarded_ws.accept() 

52 ... 

53 """ 

54 

55 def __init__(self, websocket: WebSocket, guards: list[GuardProtocol] | None = None): 

56 self.websocket = websocket 

57 self.guards = guards or [] 

58 self._accepted = False 

59 

60 async def accept(self, subprotocol: str | None = None) -> Any: 

61 """ 

62 Accept WebSocket connection after running guards. 

63 

64 Args: 

65 subprotocol: Optional WebSocket subprotocol 

66 

67 Raises: 

68 WebSocketDisconnect: If guards fail 

69 """ 

70 if self.guards: 

71 result = await execute_websocket_guards(self.guards, self.websocket) 

72 if result.is_err(): 

73 return # Connection already closed by guard 

74 

75 await self.websocket.accept(subprotocol=subprotocol) 

76 self._accepted = True 

77 

78 async def send_text(self, data: str) -> Any: 

79 """Send text message.""" 

80 await self.websocket.send_text(data) 

81 

82 async def send_bytes(self, data: bytes) -> Any: 

83 """Send binary message.""" 

84 await self.websocket.send_bytes(data) 

85 

86 async def send_json(self, data: Any) -> Any: 

87 """Send JSON message.""" 

88 await self.websocket.send_json(data) 

89 

90 async def receive_text(self) -> str: 

91 """Receive text message.""" 

92 return await self.websocket.receive_text() 

93 

94 async def receive_bytes(self) -> bytes: 

95 """Receive binary message.""" 

96 return await self.websocket.receive_bytes() 

97 

98 async def receive_json(self) -> Any: 

99 """Receive JSON message.""" 

100 return await self.websocket.receive_json() 

101 

102 async def close(self, code: int = 1000, reason: str | None = None) -> Any: 

103 """Close WebSocket connection.""" 

104 await self.websocket.close(code=code, reason=reason) 

105 

106 @property 

107 def client_state(self) -> Any: 

108 """Get WebSocket client state.""" 

109 return self.websocket.client_state 

110 

111 @property 

112 def application_state(self) -> Any: 

113 """Get WebSocket application state.""" 

114 return self.websocket.application_state 

115 

116 

117__all__ = [ 

118 "GuardedWebSocket", 

119 "execute_websocket_guards", 

120]