lib/websocket/websocket.nit

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