Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1from base64 import encode as base64_encode 

2from fiber import Fiber 

3from hash.sha1 import sha1 

4from net.tcp.server import Client as TcpClient 

5from net.tcp.server import Server as TcpServer 

6from . import WebsocketError 

7from .common import HEADER_FIN 

8from .common import OpCode 

9 

10func _send_frame(tcp_client: TcpClient, data: bytes, op_code: OpCode): 

11 header = b"" 

12 data_size = data.length() 

13 header += (HEADER_FIN | u8(op_code)) 

14 

15 if data_size < 126: 

16 header += u8(data_size) 

17 elif data_size < 65536: 

18 header += 126 

19 header += u8((data_size >> 8) & 0xff) 

20 header += u8((data_size >> 0) & 0xff) 

21 else: 

22 header += 127 

23 header += 0 

24 header += 0 

25 header += 0 

26 header += 0 

27 header += u8((data_size >> 24) & 0xff) 

28 header += u8((data_size >> 16) & 0xff) 

29 header += u8((data_size >> 8) & 0xff) 

30 header += u8((data_size >> 0) & 0xff) 

31 

32 tcp_client.write(header) 

33 tcp_client.write(data) 

34 

35class Client: 

36 """An accepted websocket client. 

37 

38 `path` is the path received in the HTTP request from the client. 

39 

40 """ 

41 

42 path: string? 

43 _tcp_client: TcpClient 

44 _connected: bool 

45 

46 func __init__(self, tcp_client: TcpClient): 

47 self.path = None 

48 self._tcp_client = tcp_client 

49 self._connected = True 

50 sec_websocket_key: string? = None 

51 

52 while True: 

53 line = self._read_line() 

54 

55 if line == "": 

56 break 

57 

58 mo = line.match(re"^GET (.*) HTTP/1.1$") 

59 

60 if mo is not None: 

61 self.path = mo.group(1) 

62 continue 

63 

64 if line.starts_with("Sec-WebSocket-Key: "): 

65 sec_websocket_key = line[19:] 

66 continue 

67 

68 if self.path is None or sec_websocket_key is None: 

69 raise WebsocketError("Invalid HTTP request.") 

70 

71 sec_websocket_key += "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" 

72 sec_websocket_accept = base64_encode(sha1(sec_websocket_key.to_utf8())) 

73 

74 tcp_client.write("HTTP/1.1 101 Switching Protocols\r\n" 

75 "Upgrade: websocket\r\n" 

76 "Connection: Upgrade\r\n" 

77 f"Sec-WebSocket-Accept: {sec_websocket_accept}\r\n" 

78 "\r\n".to_utf8()) 

79 

80 func _read_line(self) -> string: 

81 line = b"" 

82 

83 while True: 

84 byte = self._tcp_client.read(1) 

85 

86 if byte.length() == 0: 

87 raise WebsocketError("Handshake failed.") 

88 

89 line += byte 

90 

91 if line.length() < 2: 

92 continue 

93 

94 if line[-2] == u8('\r') and line[-1] == u8('\n'): 

95 break 

96 

97 return string(line)[:-2] 

98 

99 func is_connected(self) -> bool: 

100 return self._connected and self._tcp_client.is_connected() 

101 

102 func send_binary(self, data: bytes): 

103 """Send given data to the client as a binary message. 

104 

105 """ 

106 

107 raise NotImplementedError() 

108 

109 func send_text(self, data: string): 

110 """Send given data to the client as a text message. 

111 

112 """ 

113 

114 _send_frame(self._tcp_client, data.to_utf8(), OpCode.Text) 

115 

116 func receive_binary(self) -> bytes?: 

117 """Receive a binary message from the client. Returns None if 

118 disconnected. 

119 

120 """ 

121 

122 raise NotImplementedError() 

123 

124 func receive_text(self) -> string?: 

125 """Receive a text message from the client. Returns None if 

126 disconnected. 

127 

128 """ 

129 

130 raise NotImplementedError() 

131 

132class Server: 

133 """A websocket server, used to communicate with websocket clients. 

134 

135 """ 

136 

137 _tcp_server: TcpServer 

138 

139 func __init__(self): 

140 self._tcp_server = TcpServer() 

141 

142 func listen(self, port: i64): 

143 """Start listening for clients to connect to given `port` on any 

144 interface. 

145 

146 """ 

147 

148 self._tcp_server.listen(port) 

149 

150 func accept(self) -> Client: 

151 """Wait for a client to connect and accept it. 

152 

153 """ 

154 

155 return Client(self._tcp_server.accept())