tests: add test_keep_going.nit
[nit.git] / lib / websocket / websocket.nit
1 # This file is part of NIT ( http://www.nitlanguage.org ).
2 #
3 # Copyright 2014 Lucas Bajolet <r4pass@hotmail.com>
4 #
5 # Licensed under the Apache License, Version 2.0 (the "License");
6 # you may not use this file except in compliance with the License.
7 # You may obtain a copy of the License at
8 #
9 # http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing, software
12 # distributed under the License is distributed on an "AS IS" BASIS,
13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 # See the License for the specific language governing permissions and
15 # limitations under the License.
16
17 # Adds support for a websocket connection in Nit
18 # Uses standard sockets
19 module websocket
20
21 import socket
22 import sha1
23 import base64
24
25 intrude import standard::stream
26
27 # Websocket compatible listener
28 #
29 # Produces Websocket client-server connections
30 class WebSocketListener
31 super Socket
32
33 # Socket listening to connections on a defined port
34 var listener: TCPServer
35
36 # Creates a new Websocket server listening on given port with `max_clients` slots available
37 init(port: Int, max_clients: Int)
38 do
39 listener = new TCPServer(port)
40 listener.listen max_clients
41 end
42
43 # Accepts an incoming connection
44 fun accept: WebsocketConnection
45 do
46 assert not listener.closed
47
48 var client = listener.accept
49 assert client != null
50
51 return new WebsocketConnection(listener.port, "", client)
52 end
53
54 # Stop listening for incoming connections
55 fun close
56 do
57 listener.close
58 end
59 end
60
61 # Connection to a websocket client
62 #
63 # Can be used to communicate with a client
64 class WebsocketConnection
65 super TCPStream
66
67 init do
68 _buffer = new FlatBuffer
69 _buffer_pos = 0
70 var headers = parse_handshake
71 var resp = handshake_response(headers)
72
73 client.write(resp)
74 end
75
76 # Client connection to the server
77 var client: TCPStream
78
79 # Disconnect from a client
80 redef fun close
81 do
82 client.close
83 end
84
85 # Parses the input handshake sent by the client
86 # See RFC 6455 for information
87 private fun parse_handshake: Map[String,String]
88 do
89 var recved = read_http_frame(new FlatBuffer)
90 var headers = recved.split("\r\n")
91 var headmap = new HashMap[String,String]
92 for i in headers do
93 var temp_head = i.split(" ")
94 var head = temp_head.shift
95 if head.is_empty or head.length == 1 then continue
96 if head.chars.last == ':' then
97 head = head.substring(0, head.length - 1)
98 end
99 var body = temp_head.join(" ")
100 headmap[head] = body
101 end
102 return headmap
103 end
104
105 # Generates the handshake
106 private fun handshake_response(heads: Map[String,String]): String
107 do
108 var resp_map = new HashMap[String,String]
109 resp_map["HTTP/1.1"] = "101 Switching Protocols"
110 resp_map["Upgrade:"] = "websocket"
111 resp_map["Connection:"] = "Upgrade"
112 var key = heads["Sec-WebSocket-Key"]
113 key += "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
114 key = key.sha1.encode_base64
115 resp_map["Sec-WebSocket-Accept:"] = key
116 var resp = resp_map.join("\r\n", " ")
117 resp += "\r\n\r\n"
118 return resp
119 end
120
121 # Frames a text message to be sent to a client
122 private fun frame_message(msg: String): String
123 do
124 var ans_buffer = new FlatBuffer
125 # Flag for final frame set to 1
126 # opcode set to 1 (for text)
127 ans_buffer.add(129.ascii)
128 if msg.length < 126 then
129 ans_buffer.add(msg.length.ascii)
130 end
131 if msg.length >= 126 and msg.length <= 65535 then
132 ans_buffer.add(126.ascii)
133 ans_buffer.add(msg.length.rshift(8).ascii)
134 ans_buffer.add(msg.length.ascii)
135 end
136 ans_buffer.append(msg)
137 return ans_buffer.to_s
138 end
139
140 # Reads an HTTP frame
141 protected fun read_http_frame(buf: Buffer): String
142 do
143 buf.append client.read_line
144 buf.append("\r\n")
145 if buf.has_suffix("\r\n\r\n") then return buf.to_s
146 return read_http_frame(buf)
147 end
148
149 # Gets the message from the client, unpads it and reconstitutes the message
150 private fun unpad_message do
151 var fin = false
152 while not fin do
153 var fst_byte = client.read_byte
154 var snd_byte = client.read_byte
155 if fst_byte == null or snd_byte == null then
156 last_error = new IOError("Error: bad frame")
157 client.close
158 return
159 end
160 # First byte in msg is formatted this way :
161 # |(fin - 1bit)|(RSV1 - 1bit)|(RSV2 - 1bit)|(RSV3 - 1bit)|(opcode - 4bits)
162 # fin = Flag indicating if current frame is the last one
163 # RSV1/2/3 = Extension flags, unsupported
164 # Opcode values :
165 # %x0 denotes a continuation frame
166 # %x1 denotes a text frame
167 # %x2 denotes a binary frame
168 # %x3-7 are reserved for further non-control frames
169 # %x8 denotes a connection close
170 # %x9 denotes a ping
171 # %xA denotes a pong
172 # %xB-F are reserved for further control frames
173 var fin_flag = fst_byte.bin_and(128)
174 if fin_flag != 0 then fin = true
175 var opcode = fst_byte.bin_and(15)
176 if opcode == 9 then
177 _buffer.add(138.ascii)
178 _buffer.add('\0')
179 client.write(_buffer.to_s)
180 _buffer_pos += 2
181 return
182 end
183 if opcode == 8 then
184 self.client.close
185 return
186 end
187 # Second byte is formatted this way :
188 # |(mask - 1bit)|(payload length - 7 bits)
189 # As specified, if the payload length is 126 or 127
190 # The next 16 or 64 bits contain an extended payload length
191 var mask_flag = snd_byte.bin_and(128)
192 var len = snd_byte.bin_and(127)
193 var payload_ext_len = 0
194 if len == 126 then
195 var tmp = client.read(2)
196 if tmp.length != 2 then
197 last_error = new IOError("Error: received interrupted frame")
198 client.close
199 return
200 end
201 payload_ext_len = tmp[1].ascii + tmp[0].ascii.lshift(8)
202 else if len == 127 then
203 # 64 bits for length are not supported,
204 # only the last 32 will be interpreted as a Nit Integer
205 var tmp = client.read(8)
206 if tmp.length != 8 then
207 last_error = new IOError("Error: received interrupted frame")
208 client.close
209 return
210 end
211 for pos in [0 .. tmp.length[ do
212 var i = tmp[pos].ascii
213 payload_ext_len += i.lshift(8 * (7 - pos))
214 end
215 end
216 if mask_flag != 0 then
217 if payload_ext_len != 0 then
218 var msg = client.read(payload_ext_len+4)
219 var mask = msg.substring(0,4)
220 _buffer.append(unmask_message(mask, msg.substring(4, payload_ext_len)))
221 else
222 if len == 0 then
223 return
224 end
225 var msg = client.read(len+4)
226 var mask = msg.substring(0,4)
227 _buffer.append(unmask_message(mask, msg.substring(4, len)))
228 end
229 end
230 end
231 end
232
233 # Unmasks a message sent by a client
234 private fun unmask_message(key: String, message: String): String
235 do
236 var return_message = new FlatBuffer.with_capacity(message.length)
237 var msg_iter = message.chars.iterator
238
239 while msg_iter.is_ok do
240 return_message.chars[msg_iter.index] = msg_iter.item.ascii.bin_xor(key.chars[msg_iter.index%4].ascii).ascii
241 msg_iter.next
242 end
243
244 return return_message.to_s
245 end
246
247 # Checks if a connection to a client is available
248 redef fun connected do return client.connected
249
250 redef fun write(msg)
251 do
252 client.write(frame_message(msg.to_s))
253 end
254
255 redef fun is_writable do return client.connected
256
257 redef fun fill_buffer
258 do
259 _buffer.clear
260 _buffer_pos = 0
261 unpad_message
262 end
263
264 redef fun end_reached do return client._buffer_pos >= client._buffer.length and client.end_reached
265
266 # Is there some data available to be read ?
267 fun can_read(timeout: Int): Bool do return client.ready_to_read(timeout)
268
269 redef fun poll_in do return client.poll_in
270 end