1
2# Python WebSocket library with support for "wss://" encryption.
3# Copyright 2011 Joel Martin
4# Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3)
5#
6# Supports following protocol versions:
7#     - http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-75
8#     - http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76
9#     - http://tools.ietf.org/html/draft-ietf-hybi-thewebsocketprotocol-10
10
11require 'gserver'
12require 'openssl'
13require 'stringio'
14require 'digest/md5'
15require 'digest/sha1'
16require 'base64'
17
18unless OpenSSL::SSL::SSLSocket.instance_methods.index("read_nonblock")
19  module OpenSSL
20    module SSL
21      class SSLSocket
22        alias :read_nonblock :readpartial
23      end
24    end
25  end
26end
27
28class EClose < Exception
29end
30
31class WebSocketServer < GServer
32  @@Buffer_size = 65536
33
34  #
35  # WebSocket constants
36  #
37  @@Server_handshake_hixie = "HTTP/1.1 101 Web Socket Protocol Handshake\r
38Upgrade: WebSocket\r
39Connection: Upgrade\r
40%sWebSocket-Origin: %s\r
41%sWebSocket-Location: %s://%s%s\r
42"
43
44  @@Server_handshake_hybi = "HTTP/1.1 101 Switching Protocols\r
45Upgrade: websocket\r
46Connection: Upgrade\r
47Sec-WebSocket-Accept: %s\r
48"
49  @@GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
50
51
52  def initialize(opts)
53    vmsg "in WebSocketServer.initialize"
54    port = opts['listen_port']
55    host = opts['listen_host'] || GServer::DEFAULT_HOST
56
57    super(port, host)
58    msg opts.inspect
59    if opts['server_cert']
60      msg "creating ssl context"
61      @sslContext = OpenSSL::SSL::SSLContext.new
62      @sslContext.cert = OpenSSL::X509::Certificate.new(File.open(opts['server_cert']))
63      @sslContext.key = OpenSSL::PKey::RSA.new(File.open(opts['server_key']))
64      @sslContext.ca_file = opts['server_cert']
65      @sslContext.verify_mode = OpenSSL::SSL::VERIFY_NONE
66      @sslContext.verify_depth = 0
67    end
68
69    @@client_id = 0  # Track client number total on class
70
71    @verbose = opts['verbose']
72    @opts = opts
73  end
74
75  def serve(io)
76    @@client_id += 1
77    msg self.inspect
78    if @sslContext
79      msg "Enabling SSL context"
80      ssl = OpenSSL::SSL::SSLSocket.new(io, @sslContext)
81      #ssl.sync_close = true
82      #ssl.sync = true
83      msg "SSL accepting"
84      ssl.accept
85      io = ssl # replace the unencrypted handle with the encrypted one
86    end
87
88    msg "initializing thread"
89
90    # Initialize per thread state
91    t = Thread.current
92    t[:my_client_id] = @@client_id
93    t[:send_parts] = []
94    t[:recv_part] = nil
95    t[:base64] = nil
96
97    puts "in serve, client: #{t[:my_client_id].inspect}"
98
99    begin
100      t[:client] = do_handshake(io)
101      new_websocket_client(t[:client])
102    rescue EClose => e
103      msg "Client closed: #{e.message}"
104      return
105    rescue Exception => e
106      msg "Uncaught exception: #{e.message}"
107      msg "Trace: #{e.backtrace}"
108      return
109    end
110
111    msg "Client disconnected"
112  end
113
114  #
115  # WebSocketServer logging/output functions
116  #
117  def traffic(token)
118    if @verbose then print token; STDOUT.flush; end
119  end
120
121  def msg(m)
122    printf("% 3d: %s\n", Thread.current[:my_client_id] || 0, m)
123  end
124
125  def vmsg(m)
126    if @verbose then msg(m) end
127  end
128
129  #
130  # WebSocketServer general support routines
131  #
132  def gen_md5(h)
133    key1 = h['sec-websocket-key1']
134    key2 = h['sec-websocket-key2']
135    key3 = h['key3']
136    spaces1 = key1.count(" ")
137    spaces2 = key2.count(" ")
138    num1 = key1.scan(/[0-9]/).join('').to_i / spaces1
139    num2 = key2.scan(/[0-9]/).join('').to_i / spaces2
140
141    return Digest::MD5.digest([num1, num2, key3].pack('NNa8'))
142  end
143
144  def unmask(buf, hlen, length)
145    pstart = hlen + 4
146    mask = buf[hlen...hlen+4].each_byte.map{|b|b}
147    data = buf[pstart...pstart+length]
148    #data = data.bytes.zip(mask.bytes.cycle(length)).map { |d,m| d^m }
149    i=-1
150    data = data.each_byte.map{|b| i+=1; (b ^ mask[i % 4]).chr}.join("")
151    return data
152  end
153
154  def encode_hybi(buf, opcode, base64=false)
155    if base64
156      buf = Base64.encode64(buf).gsub(/\n/, '')
157    end
158
159    b1 = 0x80 | (opcode & 0x0f) # FIN + opcode
160    payload_len = buf.length
161    if payload_len <= 125
162      header = [b1, payload_len].pack('CC')
163    elsif payload_len > 125 && payload_len < 65536
164      header = [b1, 126, payload_len].pack('CCn')
165    elsif payload_len >= 65536
166      header = [b1, 127, payload_len >> 32,
167                         payload_len & 0xffffffff].pack('CCNN')
168    end
169
170    return [header + buf, header.length, 0]
171  end
172
173  def decode_hybi(buf, base64=false)
174    f = {'fin'          => 0,
175         'opcode'       => 0,
176         'hlen'         => 2,
177         'length'       => 0,
178         'payload'      => nil,
179         'left'         => 0,
180         'close_code'   => nil,
181         'close_reason' => nil}
182
183    blen = buf.length
184    f['left'] = blen
185
186    if blen < f['hlen'] then return f end # incomplete frame
187
188    b1, b2 = buf.unpack('CC')
189    f['opcode'] = b1 & 0x0f
190    f['fin']    = (b1 & 0x80) >> 7
191    has_mask    = (b2 & 0x80) >> 7
192
193    f['length'] = b2 & 0x7f
194
195    if f['length'] == 126
196      f['hlen'] = 4
197      if blen < f['hlen'] then return f end # incomplete frame
198      f['length'] = buf.unpack('xxn')[0]
199    elsif f['length'] == 127
200      f['hlen'] = 10
201      if blen < f['hlen'] then return f end # incomplete frame
202      top, bottom = buf.unpack('xxNN')
203      f['length'] = (top << 32) & bottom
204    end
205
206    full_len = f['hlen'] + has_mask * 4 + f['length']
207
208    if blen < full_len then return f end # incomplete frame
209
210    # number of bytes that are part of the next frame(s)
211    f['left'] = blen - full_len
212
213    if has_mask > 0
214      f['payload'] = unmask(buf, f['hlen'], f['length'])
215    else
216      f['payload'] = buf[f['hlen']...full_len]
217    end
218
219    if base64 and [1, 2].include?(f['opcode'])
220      f['payload'] = Base64.decode64(f['payload'])
221    end
222
223    # close frame
224    if f['opcode'] == 0x08
225      if f['length'] >= 2
226        f['close_code'] = f['payload'].unpack('n')
227      end
228      if f['length'] > 3
229        f['close_reason'] = f['payload'][2...f['payload'].length]
230      end
231    end
232
233    return f
234  end
235
236  def encode_hixie(buf)
237    return ["\x00" + Base64.encode64(buf).gsub(/\n/, '') + "\xff", 1, 1]
238  end
239
240  def decode_hixie(buf)
241    last = buf.index("\377")
242    return {'payload' => Base64.decode64(buf[1...last]),
243            'hlen' => 1,
244            'length' => last - 1,
245            'left' => buf.length - (last + 1)}
246  end
247
248  def send_frames(bufs)
249    t = Thread.current
250    if bufs.length > 0
251      encbuf = ""
252      bufs.each do |buf|
253        if t[:version].start_with?("hybi")
254          if t[:base64]
255            encbuf, lenhead, lentail = encode_hybi(
256              buf, opcode=1, base64=true)
257          else
258            encbuf, lenhead, lentail = encode_hybi(
259              buf, opcode=2, base64=false)
260          end
261        else
262          encbuf, lenhead, lentail = encode_hixie(buf)
263        end
264
265        t[:send_parts] << encbuf
266      end
267
268    end
269
270    while t[:send_parts].length > 0
271      buf = t[:send_parts].shift
272      sent = t[:client].write(buf)
273
274      if sent == buf.length
275        traffic "<"
276      else
277        traffic "<."
278        t[:send_parts].unshift(buf[sent...buf.length])
279      end
280    end
281
282    return t[:send_parts].length
283  end
284
285  # Receive and decode Websocket frames
286  # Returns: [bufs_list, closed_string]
287  def recv_frames()
288    t = Thread.current
289    closed = false
290    bufs = []
291
292    buf = t[:client].read_nonblock(@@Buffer_size)
293
294    if buf.length == 0
295      return bufs, "Client closed abrubtly"
296    end
297
298    if t[:recv_part]
299      buf = t[:recv_part] + buf
300      t[:recv_part] = nil
301    end
302
303    while buf.length > 0
304      if t[:version].start_with?("hybi")
305        frame = decode_hybi(buf, base64=t[:base64])
306
307        if frame['payload'] == nil
308          traffic "}."
309          if frame['left'] > 0
310            t[:recv_part] = buf[-frame['left']...buf.length]
311          end
312          break
313        else
314          if frame['opcode'] == 0x8
315            closed = "Client closed, reason: %s - %s" % [
316              frame['close_code'], frame['close_reason']]
317            break
318          end
319        end
320      else
321        if buf[0...2] == "\xff\x00"
322          closed = "Client sent orderly close frame"
323          break
324        elsif buf[0...2] == "\x00\xff"
325          buf = buf[2...buf.length]
326          continue # No-op frame
327        elsif buf.count("\xff") == 0
328          # Partial frame
329          traffic "}."
330          t[:recv_part] = buf
331          break
332        end
333
334        frame = decode_hixie(buf)
335      end
336
337      #msg "Receive frame: #{frame.inspect}"
338
339      traffic "}"
340
341      bufs << frame['payload']
342
343      if frame['left'] > 0
344        buf = buf[-frame['left']...buf.length]
345      else
346        buf = ''
347      end
348    end
349
350    return bufs, closed
351  end
352
353
354  def send_close(code=nil, reason='')
355    t = Thread.current
356    if t[:version].start_with?("hybi")
357      msg = ''
358      if code
359        msg = [reason.length, code].pack("na8")
360      end
361
362      buf, lenh, lent = encode_hybi(msg, opcode=0x08, base64=false)
363      t[:client].write(buf)
364    elsif t[:version] == "hixie-76"
365      buf = "\xff\x00"
366      t[:client].write(buf)
367    end
368  end
369
370  def do_handshake(sock)
371
372    t = Thread.current
373    stype = ""
374
375    if !IO.select([sock], nil, nil, 3)
376      raise EClose, "ignoring socket not ready"
377    end
378
379    handshake = ""
380    msg "About to read from sock [#{sock.inspect}]"
381    handshake = sock.read_nonblock(1024)
382    msg "Handshake [#{handshake.inspect}]"
383
384    if handshake == nil or handshake == ""
385      raise(EClose, "ignoring empty handshake")
386    else
387      stype = "Plain non-SSL (ws://)"
388      scheme = "ws"
389      if sock.class == OpenSSL::SSL::SSLSocket
390        stype = "SSL (wss://)"
391        scheme = "wss"
392      end
393      retsock = sock
394    end
395
396    h = t[:headers] = {}
397    hlines = handshake.split("\r\n")
398    req_split = hlines.shift.match(/^(\w+) (\/[^\s]*) HTTP\/1\.1$/)
399    t[:path] = req_split[2].strip
400    hlines.each do |hline|
401      break if hline == ""
402      hsplit = hline.match(/^([^:]+):\s*(.+)$/)
403      h[hsplit[1].strip.downcase] = hsplit[2]
404    end
405    puts "Headers: #{h.inspect}"
406
407    unless h.has_key?('upgrade') &&
408       h['upgrade'].downcase == 'websocket'
409      raise EClose, "Non-WebSocket connection"
410    end
411
412    protocols = h.fetch("sec-websocket-protocol", h["websocket-protocol"])
413    ver = h.fetch('sec-websocket-version', nil)
414
415    if ver
416      # HyBi/IETF vesrion of the protocol
417
418      # HyBi 07 reports version 7
419      # HyBi 08 - 12 report version 8
420      # HyBi 13 and up report version 13
421      if ['7', '8', '13'].include?(ver)
422        t[:version] = "hybi-%02d" % [ver.to_i]
423      else
424        raise EClose, "Unsupported protocol version %s" % [ver]
425      end
426
427      # choose binary if client supports it
428      if protocols.include?('binary')
429        t[:base64] = false
430      elsif protocols.include?('base64')
431        t[:base64] = true
432      else
433        raise EClose, "Client must support 'binary' or 'base64' sub-protocol"
434      end
435
436      key = h['sec-websocket-key']
437
438      # Generate the hash value for the accpet header
439      accept = Base64.encode64(
440        Digest::SHA1.digest(key + @@GUID)).gsub(/\n/, '')
441
442      response = @@Server_handshake_hybi % [accept]
443
444      if t[:base64]
445        response += "Sec-WebSocket-Protocol: base64\r\n"
446      else
447        response += "Sec-WebSocket-Protocol: binary\r\n"
448      end
449      response += "\r\n"
450
451    else
452      # Hixie vesrion of the protocol (75 or 76)
453      body = handshake.match(/\r\n\r\n(........)/)
454      if body
455        h['key3'] = body[1]
456        trailer = gen_md5(h)
457        pre = "Sec-"
458        t[:version] = "hixie-76"
459      else
460        trailer = ""
461        pre = ""
462        t[:version] = "hixie-75"
463      end
464
465      # base64 required for Hixie since payload is only UTF-8
466      t[:base64] = true
467
468      response = @@Server_handshake_hixie % [pre, h['origin'], pre,
469        "ws", h['host'], t[:path]]
470
471      if protocols && protocols.include?('base64')
472        response += "%sWebSocket-Protocol: base64\r\n" % [pre]
473      else
474        msg "Warning: client does not report 'base64' protocol support"
475      end
476
477      response += "\r\n" + trailer
478    end
479
480    msg "%s WebSocket connection" % [stype]
481    msg "Version %s, base64: '%s'" % [t[:version], t[:base64]]
482    if t[:path] then msg "Path: '%s'" % [t[:path]] end
483
484    #puts "sending reponse #{response.inspect}"
485    retsock.write(response)
486
487    # Return the WebSocket socket which may be SSL wrapped
488    return retsock
489  end
490
491end
492
493# vim: sw=2
494