lib/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 # Websocket compatible server, works as an extra layer to the original Sockets
  26 class WebSocket
  27         super BufferedIStream
  28         super OStream
  29         super PollableIStream
  30
  31         # Client connection to the server
  32         var client: Socket
  33
  34         # Socket listening to connections on a defined port
  35         var listener: Socket
  36
  37         # Creates a new Websocket server listening on given port with `max_clients` slots available
  38         init(port: Int, max_clients: Int)
  39         do
  40                 _buffer = new FlatBuffer
  41                 _buffer_pos = 0
  42                 listener = new Socket.server(port, max_clients)
  43         end
  44
  45         # Accept an incoming connection and initializes the handshake
  46         fun accept
  47         do
  48                 assert not listener.eof
  49
  50                 client = listener.accept
  51
  52                 var headers = parse_handshake
  53                 var resp = handshake_response(headers)
  54
  55                 client.write(resp)
  56         end
  57
  58         # Disconnect from a client
  59         fun disconnect_client
  60         do
  61                 client.close
  62         end
  63
  64         # Disconnects the client if one is connected
  65         # And stops the server
  66         redef fun close
  67         do
  68                 client.close
  69                 listener.close
  70         end
  71
  72         # Parses the input handshake sent by the client
  73         # See RFC 6455 for information
  74         private fun parse_handshake: Map[String,String]
  75         do
  76                 var recved = read_http_frame(new FlatBuffer)
  77                 var headers = recved.split("\r\n")
  78                 var headmap = new HashMap[String,String]
  79                 for i in headers do
  80                         var temp_head = i.split(" ")
  81                         var head = temp_head.shift
  82                         if head.is_empty or head.length == 1 then continue
  83                         if head.chars.last == ':' then
  84                                 head = head.substring(0, head.length - 1)
  85                         end
  86                         var body = temp_head.join(" ")
  87                         headmap[head] = body
  88                 end
  89                 return headmap
  90         end
  91
  92         # Generates the handshake
  93         private fun handshake_response(heads: Map[String,String]): String
  94         do
  95                 var resp_map = new HashMap[String,String]
  96                 resp_map["HTTP/1.1"] = "101 Switching Protocols"
  97                 resp_map["Upgrade:"] = "websocket"
  98                 resp_map["Connection:"] = "Upgrade"
  99                 var key = heads["Sec-WebSocket-Key"]
 100                 key += "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
 101                 key = key.sha1.encode_base64
 102                 resp_map["Sec-WebSocket-Accept:"] = key
 103                 var resp = resp_map.join("\r\n", " ")
 104                 resp += "\r\n\r\n"
 105                 return resp
 106         end
 107
 108         # Frames a text message to be sent to a client
 109         private fun frame_message(msg: String): String
 110         do
 111                 var ans_buffer = new FlatBuffer
 112                 # Flag for final frame set to 1
 113                 # opcode set to 1 (for text)
 114                 ans_buffer.add(129.ascii)
 115                 if msg.length < 126 then
 116                         ans_buffer.add(msg.length.ascii)
 117                 end
 118                 if msg.length >= 126 and msg.length <= 65535 then
 119                         ans_buffer.add(126.ascii)
 120                         ans_buffer.add(msg.length.rshift(8).ascii)
 121                         ans_buffer.add(msg.length.ascii)
 122                 end
 123                 ans_buffer.append(msg)
 124                 return ans_buffer.to_s
 125         end
 126
 127         # Reads an HTTP frame
 128         protected fun read_http_frame(buf: Buffer): String
 129         do
 130                 client.append_line_to(buf)
 131                 buf.chars.add('\n')
 132                 if buf.has_substring("\r\n\r\n", buf.length - 4) then return buf.to_s
 133                 return read_http_frame(buf)
 134         end
 135
 136         # Gets the message from the client, unpads it and reconstitutes the message
 137         private fun unpad_message do
 138                 var fin = false
 139                 while not fin do
 140                         var fst_char = client.read_char
 141                         var snd_char = client.read_char
 142                         # First byte in msg is formatted this way :
 143                         # |(fin - 1bit)|(RSV1 - 1bit)|(RSV2 - 1bit)|(RSV3 - 1bit)|(opcode - 4bits)
 144                         # fin = Flag indicating if current frame is the last one
 145                         # RSV1/2/3 = Extension flags, unsupported
 146                         # Opcode values :
 147                         #       %x0 denotes a continuation frame
 148                         #       %x1 denotes a text frame
 149                         #       %x2 denotes a binary frame
 150                         #       %x3-7 are reserved for further non-control frames
 151                         #       %x8 denotes a connection close
 152                         #       %x9 denotes a ping
 153                         #       %xA denotes a pong
 154                         #       %xB-F are reserved for further control frames
 155                         var fin_flag = fst_char.bin_and(128)
 156                         if fin_flag != 0 then fin = true
 157                         var opcode = fst_char.bin_and(15)
 158                         if opcode == 9 then
 159                                 _buffer.add(138.ascii)
 160                                 _buffer.add('\0')
 161                                 client.write(_buffer.to_s)
 162                                 _buffer_pos += 2
 163                                 return
 164                         end
 165                         if opcode == 8 then
 166                                 self.client.close
 167                                 return
 168                         end
 169                         # Second byte is formatted this way :
 170                         # |(mask - 1bit)|(payload length - 7 bits)
 171                         # As specified, if the payload length is 126 or 127
 172                         # The next 16 or 64 bits contain an extended payload length
 173                         var mask_flag = snd_char.bin_and(128)
 174                         var len = snd_char.bin_and(127)
 175                         var payload_ext_len = 0
 176                         if len == 126 then
 177                                 payload_ext_len = client.read_char.lshift(8)
 178                                 payload_ext_len += client.read_char
 179                         else if len == 127 then
 180                                 # 64 bits for length are not supported,
 181                                 # only the last 32 will be interpreted as a Nit Integer
 182                                 for i in [0..4[ do client.read_char
 183                                 payload_ext_len = client.read_char.lshift(24)
 184                                 payload_ext_len += client.read_char.lshift(16)
 185                                 payload_ext_len += client.read_char.lshift(8)
 186                                 payload_ext_len += client.read_char
 187                         end
 188                         if mask_flag != 0 then
 189                                 if payload_ext_len != 0 then
 190                                         var msg = client.read(payload_ext_len+4)
 191                                         var mask = msg.substring(0,4)
 192                                         _buffer.append(unmask_message(mask, msg.substring(4, payload_ext_len)))
 193                                 else
 194                                         if len == 0 then
 195                                                 return
 196                                         end
 197                                         var msg = client.read(len+4)
 198                                         var mask = msg.substring(0,4)
 199                                         _buffer.append(unmask_message(mask, msg.substring(4, len)))
 200                                 end
 201                         end
 202                 end
 203         end
 204
 205         # Unmasks a message sent by a client
 206         private fun unmask_message(key: String, message: String): String
 207         do
 208                 var return_message = new FlatBuffer.with_capacity(message.length)
 209                 var msg_iter = message.chars.iterator
 210
 211                 while msg_iter.is_ok do
 212                         return_message.chars[msg_iter.index] = msg_iter.item.ascii.bin_xor(key.chars[msg_iter.index%4].ascii).ascii
 213                         msg_iter.next
 214                 end
 215
 216                 return return_message.to_s
 217         end
 218
 219         # Checks if a connection to a client is available
 220         fun connected: Bool do return client.connected
 221
 222         redef fun write(msg: Text)
 223         do
 224                 client.write(frame_message(msg.to_s))
 225         end
 226
 227         redef fun is_writable do return client.connected
 228
 229         redef fun fill_buffer
 230         do
 231                 _buffer.clear
 232                 _buffer_pos = 0
 233                 unpad_message
 234         end
 235
 236         redef fun end_reached do return _buffer_pos >= _buffer.length and client.eof
 237
 238         # Is there some data available to be read ?
 239         fun can_read(timeout: Int): Bool do return client.ready_to_read(timeout)
 240
 241         redef fun poll_in do return client.poll_in
 242
 243 end
 244