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_char = client.read_char
 154                         var snd_char = client.read_char
 155                         # First byte in msg is formatted this way :
 156                         # |(fin - 1bit)|(RSV1 - 1bit)|(RSV2 - 1bit)|(RSV3 - 1bit)|(opcode - 4bits)
 157                         # fin = Flag indicating if current frame is the last one
 158                         # RSV1/2/3 = Extension flags, unsupported
 159                         # Opcode values :
 160                         #       %x0 denotes a continuation frame
 161                         #       %x1 denotes a text frame
 162                         #       %x2 denotes a binary frame
 163                         #       %x3-7 are reserved for further non-control frames
 164                         #       %x8 denotes a connection close
 165                         #       %x9 denotes a ping
 166                         #       %xA denotes a pong
 167                         #       %xB-F are reserved for further control frames
 168                         var fin_flag = fst_char.bin_and(128)
 169                         if fin_flag != 0 then fin = true
 170                         var opcode = fst_char.bin_and(15)
 171                         if opcode == 9 then
 172                                 _buffer.add(138.ascii)
 173                                 _buffer.add('\0')
 174                                 client.write(_buffer.to_s)
 175                                 _buffer_pos += 2
 176                                 return
 177                         end
 178                         if opcode == 8 then
 179                                 self.client.close
 180                                 return
 181                         end
 182                         # Second byte is formatted this way :
 183                         # |(mask - 1bit)|(payload length - 7 bits)
 184                         # As specified, if the payload length is 126 or 127
 185                         # The next 16 or 64 bits contain an extended payload length
 186                         var mask_flag = snd_char.bin_and(128)
 187                         var len = snd_char.bin_and(127)
 188                         var payload_ext_len = 0
 189                         if len == 126 then
 190                                 payload_ext_len = client.read_char.lshift(8)
 191                                 payload_ext_len += client.read_char
 192                         else if len == 127 then
 193                                 # 64 bits for length are not supported,
 194                                 # only the last 32 will be interpreted as a Nit Integer
 195                                 for i in [0..4[ do client.read_char
 196                                 payload_ext_len = client.read_char.lshift(24)
 197                                 payload_ext_len += client.read_char.lshift(16)
 198                                 payload_ext_len += client.read_char.lshift(8)
 199                                 payload_ext_len += client.read_char
 200                         end
 201                         if mask_flag != 0 then
 202                                 if payload_ext_len != 0 then
 203                                         var msg = client.read(payload_ext_len+4)
 204                                         var mask = msg.substring(0,4)
 205                                         _buffer.append(unmask_message(mask, msg.substring(4, payload_ext_len)))
 206                                 else
 207                                         if len == 0 then
 208                                                 return
 209                                         end
 210                                         var msg = client.read(len+4)
 211                                         var mask = msg.substring(0,4)
 212                                         _buffer.append(unmask_message(mask, msg.substring(4, len)))
 213                                 end
 214                         end
 215                 end
 216         end
 217
 218         # Unmasks a message sent by a client
 219         private fun unmask_message(key: String, message: String): String
 220         do
 221                 var return_message = new FlatBuffer.with_capacity(message.length)
 222                 var msg_iter = message.chars.iterator
 223
 224                 while msg_iter.is_ok do
 225                         return_message.chars[msg_iter.index] = msg_iter.item.ascii.bin_xor(key.chars[msg_iter.index%4].ascii).ascii
 226                         msg_iter.next
 227                 end
 228
 229                 return return_message.to_s
 230         end
 231
 232         # Checks if a connection to a client is available
 233         redef fun connected do return client.connected
 234
 235         redef fun write(msg)
 236         do
 237                 client.write(frame_message(msg.to_s))
 238         end
 239
 240         redef fun is_writable do return client.connected
 241
 242         redef fun fill_buffer
 243         do
 244                 _buffer.clear
 245                 _buffer_pos = 0
 246                 unpad_message
 247         end
 248
 249         redef fun end_reached do return client._buffer_pos >= client._buffer.length and client.end_reached
 250
 251         # Is there some data available to be read ?
 252         fun can_read(timeout: Int): Bool do return client.ready_to_read(timeout)
 253
 254         redef fun poll_in do return client.poll_in
 255 end