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