1# This file is part of Xpra. 

2# Copyright (C) 2019-2020 Antoine Martin <antoine@xpra.org> 

3# Xpra is released under the terms of the GNU GPL v2, or, at your option, any 

4# later version. See the file COPYING for details. 

5 

6import os 

7import struct 

8 

9from xpra.net.websockets.header import encode_hybi_header, decode_hybi 

10from xpra.net.protocol import Protocol 

11from xpra.util import first_time, envbool 

12from xpra.os_util import memoryview_to_bytes 

13from xpra.log import Logger 

14 

15log = Logger("websocket") 

16 

17OPCODE_CONTINUE = 0 

18OPCODE_TEXT = 1 

19OPCODE_BINARY = 2 

20OPCODE_CLOSE = 8 

21OPCODE_PING = 9 

22OPCODE_PONG = 10 

23OPCODES = { 

24 OPCODE_CONTINUE : "continuation frame", 

25 OPCODE_TEXT : "text frame", 

26 OPCODE_BINARY : "binary frame", 

27 OPCODE_CLOSE : "connection close", 

28 OPCODE_PING : "ping", 

29 OPCODE_PONG : "pong", 

30 } 

31 

32MASK = envbool("XPRA_WEBSOCKET_MASK", False) 

33 

34 

35class WebSocketProtocol(Protocol): 

36 

37 STATE_FIELDS = tuple(list(Protocol.STATE_FIELDS)+["legacy_frame_per_chunk"]) 

38 

39 TYPE = "websocket" 

40 

41 def __init__(self, *args, **kwargs): 

42 super().__init__(*args, **kwargs) 

43 self.ws_data = b"" 

44 self.ws_payload = [] 

45 self.ws_payload_opcode = 0 

46 self.ws_mask = MASK 

47 self._process_read = self.parse_ws_frame 

48 self.make_chunk_header = self.make_xpra_header 

49 self.make_frame_header = self.make_wsframe_header 

50 

51 def __repr__(self): 

52 return "WebSocket(%s)" % self._conn 

53 

54 def close(self): 

55 Protocol.close(self) 

56 self.ws_data = b"" 

57 self.ws_payload = [] 

58 

59 

60 def make_wsframe_header(self, packet_type, items): 

61 payload_len = sum(len(item) for item in items) 

62 log("make_wsframe_header(%s, %i items) %i bytes, ms_mask=%s", 

63 packet_type, len(items), payload_len, self.ws_mask) 

64 header = encode_hybi_header(OPCODE_BINARY, payload_len, self.ws_mask) 

65 if self.ws_mask: 

66 from xpra.buffers.cyxor import hybi_mask #@UnresolvedImport 

67 mask = os.urandom(4) 

68 #now mask all the items: 

69 for i, item in enumerate(items): 

70 items[i] = hybi_mask(mask, item) 

71 return header+mask 

72 return header 

73 

74 def parse_ws_frame(self, buf): 

75 if not buf: 

76 self._read_queue_put(buf) 

77 return 

78 if self.ws_data: 

79 ws_data = self.ws_data+buf 

80 self.ws_data = b"" 

81 else: 

82 ws_data = buf 

83 log("parse_ws_frame(%i bytes) total buffer is %i bytes", len(buf), len(ws_data)) 

84 while ws_data and not self._closed: 

85 parsed = decode_hybi(ws_data) 

86 if parsed is None: 

87 log("parse_ws_frame(%i bytes) not enough data", len(ws_data)) 

88 #not enough data to get a full websocket frame, 

89 #save it for later: 

90 self.ws_data = ws_data 

91 return 

92 opcode, payload, processed, fin = parsed 

93 ws_data = ws_data[processed:] 

94 log("parse_ws_frame(%i bytes) payload=%i bytes, processed=%i, remaining=%i, opcode=%s, fin=%s", 

95 len(buf), len(payload), processed, len(ws_data), OPCODES.get(opcode, opcode), fin) 

96 if opcode==OPCODE_CONTINUE: 

97 assert self.ws_payload_opcode and self.ws_payload, "continuation frame does not follow a partial frame" 

98 self.ws_payload.append(payload) 

99 if not fin: 

100 #wait for more 

101 continue 

102 #join all the frames and process the payload: 

103 full_payload = b"".join(memoryview_to_bytes(v) for v in self.ws_payload) 

104 self.ws_payload = [] 

105 opcode = self.ws_payload_opcode 

106 self.ws_payload_opcode = 0 

107 else: 

108 if self.ws_payload and self.ws_payload_opcode: 

109 raise Exception("expected a continuation frame not %s" % OPCODES.get(opcode, opcode)) 

110 full_payload = payload 

111 if not fin: 

112 if opcode not in (OPCODE_BINARY, OPCODE_TEXT): 

113 raise Exception("cannot handle fragmented '%s' frames" % OPCODES.get(opcode, opcode)) 

114 #fragmented, keep this payload for later 

115 self.ws_payload_opcode = opcode 

116 self.ws_payload.append(payload) 

117 continue 

118 if opcode==OPCODE_BINARY: 

119 self._read_queue_put(full_payload) 

120 elif opcode==OPCODE_TEXT: 

121 if first_time("ws-text-frame-from-%s" % self._conn): 

122 log.warn("Warning: handling text websocket frame as binary") 

123 self._read_queue_put(full_payload) 

124 elif opcode==OPCODE_CLOSE: 

125 self._process_ws_close(full_payload) 

126 elif opcode==OPCODE_PING: 

127 self._process_ws_ping(full_payload) 

128 elif opcode==OPCODE_PONG: 

129 self._process_ws_pong(full_payload) 

130 else: 

131 log.warn("Warning unhandled websocket opcode '%s'", OPCODES.get(opcode, "%#x" % opcode)) 

132 log("payload=%r", payload) 

133 

134 def _process_ws_ping(self, payload): 

135 log("_process_ws_ping(%r)", payload) 

136 item = encode_hybi_header(OPCODE_PONG, len(payload)) + memoryview_to_bytes(payload) 

137 items = (item, ) 

138 with self._write_lock: 

139 self.raw_write("ws-ping", items) 

140 

141 def _process_ws_pong(self, payload): 

142 log("_process_ws_pong(%r)", payload) 

143 

144 def _process_ws_close(self, payload): 

145 log("_process_ws_close(%r)", payload) 

146 if len(payload)<2: 

147 self._connection_lost("unknown reason") 

148 return 

149 code = struct.unpack(">H", payload[:2])[0] 

150 try: 

151 reason = memoryview_to_bytes(payload[2:]).decode("utf8") 

152 except UnicodeDecodeError: 

153 reason = str(reason) 

154 self._connection_lost("code %i: %r" % (code, reason))