litecable 0.4.0

data/lib/lite_cable/server.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+module LiteCable
+  # Rack middleware to hijack sockets.
+  #
+  # Uses thread-per-connection model (thus recommended only for development and test usage).
+  #
+  # Inspired by https://github.com/ngauthier/tubesock/blob/master/lib/tubesock.rb
+  module Server
+    require "websocket"
+    require "lite_cable/server/subscribers_map"
+    require "lite_cable/server/client_socket"
+    require "lite_cable/server/heart_beat"
+    require "lite_cable/server/middleware"
+
+    class << self
+      attr_accessor :subscribers_map
+
+      # Broadcast encoded message to the stream
+      def broadcast(stream, message, coder: nil)
+        coder ||= LiteCable.config.coder
+        subscribers_map.broadcast stream, message, coder
+      end
+    end
+
+    self.subscribers_map = SubscribersMap.new
+  end
+end

data/lib/lite_cable/server/client_socket.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+module LiteCable
+  module Server
+    module ClientSocket # :nodoc:
+      require "lite_cable/server/client_socket/subscriptions"
+      require "lite_cable/server/client_socket/base"
+    end
+  end
+end

data/lib/lite_cable/server/client_socket/base.rb

@@ -0,0 +1,163 @@
+# frozen_string_literal: true
+module LiteCable
+  module Server
+    module ClientSocket
+      # Wrapper over web socket
+      # rubocop:disable Metrics/ClassLength
+      class Base
+        include Logging
+        include Subscriptions
+
+        attr_reader :version, :active
+
+        def initialize(env, socket, version)
+          log(:debug, "WebSocket version #{version}")
+          @env = env
+          @socket = socket
+          @version = version
+          @active = true
+
+          @open_handlers    = []
+          @message_handlers = []
+          @close_handlers   = []
+          @error_handlers   = []
+
+          @close_on_error = true
+        end
+
+        def prevent_close_on_error
+          @close_on_error = false
+        end
+
+        def transmit(data, type: :text)
+          frame = WebSocket::Frame::Outgoing::Server.new(
+            version: version,
+            data: data,
+            type: type
+          )
+          socket.write frame.to_s
+        rescue IOError, Errno::EPIPE, Errno::ETIMEDOUT => e
+          log(:error, "Socket send failed: #{e}")
+          close
+        end
+
+        def request
+          @request ||= Rack::Request.new(@env)
+        end
+
+        def onopen(&block)
+          @open_handlers << block
+        end
+
+        def onmessage(&block)
+          @message_handlers << block
+        end
+
+        def onclose(&block)
+          @close_handlers << block
+        end
+
+        def onerror(&block)
+          @error_handlers << block
+        end
+
+        def listen
+          keepalive
+          Thread.new do
+            Thread.current.abort_on_exception = true
+            begin
+              @open_handlers.each(&:call)
+              each_frame do |data|
+                @message_handlers.each do |h|
+                  begin
+                    h.call(data)
+                  rescue => e
+                    log(:error, "Socket receive failed: #{e}")
+                    @error_handlers.each { |eh| eh.call(e, data) }
+                    close if close_on_error
+                  end
+                end
+              end
+            ensure
+              close
+            end
+          end
+        end
+
+        def close
+          return unless @active
+
+          @close_handlers.each(&:call)
+          close!
+
+          @active = false
+        end
+
+        def closed?
+          @socket.closed?
+        end
+
+        private
+
+        attr_reader :socket, :close_on_error
+
+        def close!
+          if @socket.respond_to?(:closed?)
+            close_socket unless @socket.closed?
+          else
+            close_socket
+          end
+        end
+
+        def close_socket
+          frame = WebSocket::Frame::Outgoing::Server.new(version: version, type: :close, code: 1000)
+          @socket.write(frame.to_s) if frame.supported?
+          @socket.close
+        end
+
+        def keepalive
+          thread = Thread.new do
+            Thread.current.abort_on_exception = true
+            loop do
+              sleep 5
+              transmit nil, type: :ping
+            end
+          end
+
+          onclose do
+            thread.kill
+          end
+        end
+
+        # rubocop:disable Metrics/AbcSize
+        # rubocop:disable Metrics/CyclomaticComplexity
+        # rubocop:disable Metrics/PerceivedComplexity
+        # rubocop:disable Metrics/MethodLength
+        def each_frame
+          framebuffer = WebSocket::Frame::Incoming::Server.new(version: version)
+
+          while IO.select([socket])
+            if socket.respond_to?(:recvfrom)
+              data, _addrinfo = socket.recvfrom(2000)
+            else
+              data, _addrinfo = socket.readpartial(2000), socket.peeraddr
+            end
+            break if data.empty?
+            framebuffer << data
+            while frame = framebuffer.next
+              case frame.type
+              when :close
+                return
+              when :text, :binary
+                yield frame.data
+              end
+            end
+          end
+        rescue Errno::EHOSTUNREACH, Errno::ETIMEDOUT, Errno::ECONNRESET, IOError, Errno::EBADF => e
+          log(:debug, "Socket frame error: #{e}")
+          nil # client disconnected or timed out
+        end
+      end
+    end
+  end
+end

data/lib/lite_cable/server/client_socket/subscriptions.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+module LiteCable
+  module Server
+    module ClientSocket
+      # Handle socket subscriptions
+      module Subscriptions
+        def subscribe(channel, broadcasting)
+          LiteCable::Server.subscribers_map
+                           .add_subscriber(broadcasting, self, channel)
+        end
+
+        def unsubscribe(channel, broadcasting)
+          LiteCable::Server.subscribers_map
+                           .remove_subscriber(broadcasting, self, channel)
+        end
+
+        def unsubscribe_from_all(channel)
+          LiteCable::Server.subscribers_map.remove_socket(self, channel)
+        end
+      end
+    end
+  end
+end

data/lib/lite_cable/server/heart_beat.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+module LiteCable
+  module Server
+    # Sends pings to sockets
+    class HeartBeat
+      BEAT_INTERVAL = 3
+
+      def initialize
+        @sockets = []
+        run
+      end
+
+      def add(socket)
+        @sockets << socket
+      end
+
+      def remove(socket)
+        @sockets.delete(socket)
+      end
+
+      def stop
+        @stopped = true
+      end
+
+      def run
+        Thread.new do
+          Thread.current.abort_on_exception = true
+          loop do
+            break if @stopped
+
+            unless @sockets.empty?
+              msg = ping_message Time.now.to_i
+              @sockets.each do |socket|
+                socket.transmit msg
+              end
+            end
+
+            sleep BEAT_INTERVAL
+          end
+        end
+      end
+
+      private
+
+      def ping_message(time)
+        { type: LiteCable::INTERNAL[:message_types][:ping], message: time }.to_json
+      end
+    end
+  end
+end

data/lib/lite_cable/server/middleware.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+module LiteCable
+  module Server
+    require "lite_cable/server/websocket_ext/protocols"
+    # Rack middleware to hijack the socket
+    class Middleware
+      class HijackNotAvailable < RuntimeError; end
+
+      def initialize(_app, connection_class:)
+        @connection_class = connection_class
+        @heart_beat = HeartBeat.new
+      end
+
+      def call(env)
+        return [404, { 'Content-Type' => 'text/plain' }, ['Not Found']] unless
+          env["HTTP_UPGRADE"] == 'websocket'
+
+        raise HijackNotAvailable unless env['rack.hijack']
+
+        env['rack.hijack'].call
+        handshake = send_handshake(env)
+
+        socket = ClientSocket::Base.new env, env['rack.hijack_io'], handshake.version
+        init_connection socket
+        init_heartbeat socket
+        socket.listen
+        [-1, {}, []]
+      end
+
+      private
+
+      def send_handshake(env)
+        handshake = WebSocket::Handshake::Server.new
+        handshake.from_rack env
+        handshake.protocols LiteCable::INTERNAL[:protocols]
+
+        env['rack.hijack_io'].write handshake.to_s
+        handshake
+      end
+
+      def init_connection(socket)
+        connection = @connection_class.new(socket)
+
+        socket.onopen { connection.handle_open }
+        socket.onclose { connection.handle_close }
+        socket.onmessage { |data| connection.handle_command(data) }
+      end
+
+      def init_heartbeat(socket)
+        @heart_beat.add(socket)
+        socket.onclose { @heart_beat.remove(socket) }
+      end
+    end
+  end
+end

data/lib/lite_cable/server/subscribers_map.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+module LiteCable
+  module Server
+    # From https://github.com/rails/rails/blob/v5.0.1/actioncable/lib/action_cable/subscription_adapter/subscriber_map.rb
+    class SubscribersMap
+      attr_reader :streams, :sockets
+
+      def initialize
+        @streams = Hash.new do |streams, stream_id|
+          streams[stream_id] = Hash.new { |channels, channel_id| channels[channel_id] = [] }
+        end
+        @sockets = Hash.new { |h, k| h[k] = [] }
+        @sync = Mutex.new
+      end
+
+      def add_subscriber(stream, socket, channel)
+        @sync.synchronize do
+          @streams[stream][channel] << socket
+          @sockets[socket] << [channel, stream]
+        end
+      end
+
+      def remove_subscriber(stream, socket, channel)
+        @sync.synchronize do
+          @streams[stream][channel].delete(socket)
+          @sockets[socket].delete([channel, stream])
+          cleanup stream, socket, channel
+        end
+      end
+
+      def remove_socket(socket, channel)
+        list = @sync.synchronize do
+          return unless @sockets.key?(socket)
+          @sockets[socket].dup
+        end
+
+        list.each do |(channel_id, stream)|
+          remove_subscriber(stream, socket, channel) if channel == channel_id
+        end
+      end
+
+      def broadcast(stream, message, coder)
+        list = @sync.synchronize do
+          return unless @streams.key?(stream)
+          @streams[stream].to_a
+        end
+
+        list.each do |(channel_id, sockets)|
+          cmessage = channel_message(channel_id, message, coder)
+          sockets.each { |s| s.transmit cmessage }
+        end
+      end
+
+      private
+
+      def cleanup(stream, socket, channel)
+        @streams[stream].delete(channel) if @streams[stream][channel].empty?
+        @streams.delete(stream) if @streams[stream].empty?
+        @sockets.delete(socket) if @sockets[socket].empty?
+      end
+
+      def channel_message(channel_id, message, coder)
+        coder.encode(identifier: channel_id, message: message)
+      end
+    end
+  end
+end

data/lib/lite_cable/server/websocket_ext/protocols.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+# Add missing protocols support to websocket-ruby
+module WebSocketExt
+  module Protocols # :nodoc:
+    module Handshake # :nodoc:
+      # Specify server protocols
+      def protocols(values)
+        @protocols = values
+      end
+
+      # Return matching protocol
+      def protocol
+        return @protocol if instance_variable_defined?(:@protocol)
+        protos = @headers['sec-websocket-protocol']
+
+        return @protocol = nil unless protos
+        @protocol = begin
+          protos = protos.split(/ *, */) if protos.is_a?(String)
+          protos.find { |p| @protocols.include?(p) }
+        end
+      end
+    end
+
+    module Handler # :nodoc:
+      def handshake_keys
+        return super unless @handshake.protocol
+        super + [
+          [
+            'Sec-WebSocket-Protocol',
+            @handshake.protocol
+          ]
+        ]
+      end
+    end
+  end
+end
+
+WebSocket::Handshake::Server.include WebSocketExt::Protocols::Handshake
+[
+  WebSocket::Handshake::Handler::Server04,
+  WebSocket::Handshake::Handler::Server75,
+  WebSocket::Handshake::Handler::Server76
+].each do |handler|
+  handler.prepend WebSocketExt::Protocols::Handler
+end