Add support for non-binary WebMKS websocket

lib/websocket_proxy.rb

@@ -6,11 +6,17 @@ def initialize(env, console, logger)
     @id = SecureRandom.uuid
     @console = console
     @logger = logger
+    @protocol = 'binary'
 
     secure = Rack::Request.new(env).ssl?
     scheme = secure ? 'wss:' : 'ws:'
     @url = scheme + '//' + env['HTTP_HOST'] + env['REQUEST_URI']
-    @driver = WebSocket::Driver.rack(self, :protocols => %w(binary))
+
+    # VMware vCloud WebMKS SDK (console access) grabs 'binary' if offered, but then fails because in fact it's not
+    # able to use it :) We workaround this by only forcing the 'uint8utf8' protocol which actually works.
+    @protocol = 'uint8utf8' if console.protocol.to_s.end_with?('uint8utf8')
+
+    @driver = WebSocket::Driver.rack(self, :protocols => [@protocol])
 
     begin
       # Hijack the socket from the Rack middleware
@@ -24,6 +30,8 @@ def initialize(env, console, logger)
                   @console.ssl ? WebsocketSSLSocket : WebsocketSocket
                 when 'webmks'
                   WebsocketWebmks
+                when 'webmks-uint8utf8'
+                  WebsocketWebmksUint8utf8
                 end
       @right = adapter.new(@sock, @console)
     rescue => ex
@@ -33,7 +41,12 @@ def initialize(env, console, logger)
 
     @driver.on(:open) { @console.update(:opened => true) }
 
-    @driver.on(:message) { |msg| @right.issue(msg.data.pack('C*')) }
+    # TODO: Move binary <-> string interpretation into client class, don't do it here (reusability).
+    if binary?
+      @driver.on(:message) { |msg| @right.issue(msg.data.pack('C*')) }
+    else
+      @driver.on(:message) { |msg| @right.issue(msg.data) }
+    end
 
     @driver.on(:close) { cleanup }
   end
@@ -49,10 +62,14 @@ def transmit(sockets, is_ws)
     if is_ws
       data = @ws.recv_nonblock(64.kilobytes)
       @driver.parse(data)
-    else
+    elsif binary?
       @right.fetch(64.kilobytes) do |data|
         @driver.binary(data)
       end
+    else
+      @right.fetch(64.kilobytes) do |data|
+        @driver.frame(data)
+      end
     end
   end
 
@@ -76,4 +93,8 @@ def write(string)
   def vm_id
     @console ? @console.vm_id : 'unknown'
   end
+
+  def binary?
+    @protocol == 'binary'
+  end
 end

lib/websocket_webmks_uint8utf8.rb

@@ -0,0 +1,31 @@
+class WebsocketWebmksUint8utf8 < WebsocketSSLSocket
+  attr_accessor :url
+
+  def initialize(socket, model)
+    super(socket, model)
+    @url = URI::Generic.build(:scheme => 'wss',
+                              :host   => @model.host_name,
+                              :port   => @model.port,
+                              :path   => @model.url).to_s
+
+    @driver = WebSocket::Driver.client(self, :protocols => ['uint8utf8'])
+    @driver.on(:close) { socket.close unless socket.closed? }
+    @driver.start
+  end
+
+  def fetch(length)
+    # WebSocket::Driver requires an event handler that should be registered only once
+    @driver.on(:message) { |msg| yield(msg.data) } if @driver.listeners(:message).length.zero?
+
+    data = @ssl.sysread(length)
+    @driver.parse(data)
+  end
+
+  def issue(data)
+    @driver.frame(data)
+  end
+
+  def write(data)
+    @ssl.syswrite(data)
+  end
+end

spec/lib/websocket_proxy_spec.rb

@@ -8,9 +8,97 @@
   subject { described_class.new(env, console, logger) }
 
   describe '#initialize' do
+    before do
+      allow(TCPSocket).to receive(:open) # prevent real sockets from opening
+    end
+
     it 'sets the URL' do
       expect(subject.url).to eq("ws://#{host}#{uri}")
     end
+
+    describe 'based on console type' do
+      before do
+        [WebsocketSocket, WebsocketWebmks, WebsocketWebmksUint8utf8].each do |k|
+          allow(k).to receive(:new).and_return(k.allocate)
+        end
+      end
+
+      let(:proxy)      { described_class.new(env, FactoryGirl.create(:system_console, :protocol => console_type), logger) }
+      let(:proto)      { proxy.instance_variable_get(:@driver).instance_variable_get(:@options)[:protocols] }
+      let(:adapter)    { proxy.instance_variable_get(:@right) }
+      let(:on_message) { proxy.instance_variable_get(:@driver).listeners(:message).first }
+      let(:right)      { double('right socket') }
+
+      context 'when vnc' do
+        let(:console_type) { 'vnc' }
+
+        it 'uses binary protocol' do
+          expect(proto).to eq(['binary'])
+        end
+
+        it 'uses WebsocketSocket adapter' do
+          expect(adapter).to be_an_instance_of(WebsocketSocket)
+        end
+
+        it 'decodes message' do
+          assert_message_transformation('BANANA'.unpack('C*'), 'BANANA')
+        end
+      end
+
+      context 'when spice' do
+        let(:console_type) { 'spice' }
+
+        it 'uses binary protocol' do
+          expect(proto).to eq(['binary'])
+        end
+
+        it 'uses WebsocketSocket adapter' do
+          expect(adapter).to be_an_instance_of(WebsocketSocket)
+        end
+
+        it 'decodes message' do
+          assert_message_transformation('BANANA'.unpack('C*'), 'BANANA')
+        end
+      end
+
+      context 'when webmks' do
+        let(:console_type) { 'webmks' }
+
+        it 'uses binary protocol' do
+          expect(proto).to eq(['binary'])
+        end
+
+        it 'uses WebsocketWebmks adapter' do
+          expect(adapter).to be_an_instance_of(WebsocketWebmks)
+        end
+
+        it 'decodes message' do
+          assert_message_transformation('BANANA'.unpack('C*'), 'BANANA')
+        end
+      end
+
+      context 'when webmks-uint8utf8' do
+        let(:console_type) { 'webmks-uint8utf8' }
+
+        it 'uses uint8utf8 protocol' do
+          expect(proto).to eq(['uint8utf8'])
+        end
+
+        it 'uses WebsocketWebmksUint8utf8 adapter' do
+          expect(adapter).to be_an_instance_of(WebsocketWebmksUint8utf8)
+        end
+
+        it 'does not decode message' do
+          assert_message_transformation('BANANA', 'BANANA')
+        end
+      end
+    end
+  end
+
+  def assert_message_transformation(input, output)
+    proxy.instance_variable_set(:@right, right)
+    expect(right).to receive(:issue).with(output)
+    on_message.call(double('message', :data => input))
   end
 
   describe '#cleanup' do
@@ -67,11 +155,23 @@
     context 'socket to websocket' do
       let(:is_ws) { false }
 
-      it 'reads from the socket and sends the result to the driver' do
-        expect(sock).to receive(:recv_nonblock).and_return(123)
-        expect(driver).to receive(:binary).with(123)
+      context 'binary' do
+        it 'reads from the socket and sends the result to the driver' do
+          expect(sock).to receive(:recv_nonblock).and_return(123)
+          expect(driver).to receive(:binary).with(123)
+
+          subject.transmit([sock], is_ws)
+        end
+      end
+
+      context 'non-binary' do
+        it 'reads from the socket and sends the result to the driver' do
+          allow(subject).to receive(:binary?).and_return(false)
+          expect(sock).to receive(:recv_nonblock).and_return(123)
+          expect(driver).to receive(:frame).with(123)
 
-        subject.transmit([sock], is_ws)
+          subject.transmit([sock], is_ws)
+        end
       end
     end
   end