jam-cloud/websocket-gateway/lib/jam_websockets/router.rb

require 'set'
require 'amqp'
require 'thread'
require 'json'
require 'eventmachine'

include Jampb

# add new field to client connection
module EventMachine
  module WebSocket
    class Connection < EventMachine::Connection
      attr_accessor :encode_json, :client_id # client_id is uuid we give to each client to track them as we like
    end
  end
end

module JamWebsockets

  class Router

    attr_accessor :user_context_lookup

    def initialize()
      @log = Logging.logger[self]
      @pending_clients = Set.new # clients that have connected to server, but not logged in.
      @clients = {} # clients that have logged in
      @user_context_lookup = {} # lookup a set of client_contexts by user_id
      @client_lookup = {} # lookup a client by client_id
      @amqp_connection_manager = nil
      @users_exchange = nil
      @message_factory = JamRuby::MessageFactory.new
      @semaphore = Mutex.new
      @user_topic = nil
      @client_topic = nil
      @thread_pool = nil
      @heartbeat_interval = nil

    end

    def start(connect_time_stale, options={:host => "localhost", :port => 5672}, &block)

      @log.info "startup"

      @heartbeat_interval = connect_time_stale / 2

      begin
        @amqp_connection_manager = AmqpConnectionManager.new(true, 4, :host => options[:host], :port => options[:port])
        @amqp_connection_manager.connect do |channel|
          register_topics(channel)
          block.call
        end

      rescue => e
        @log.error "unable to initialize #{e.to_s}"
        cleanup
        raise e
      end

      @log.info "started"
    end

    def add_client(client_id, client_context)
      @client_lookup[client_id] = client_context
    end

    def remove_client(client_id, client)
      deleted = @client_lookup.delete(client_id)

      if deleted.nil?
        @log.warn "unable to delete #{client_id} from client_lookup"
      elsif deleted.client != client
        # put it back--this is only possible if add_client hit the 'old connection' path
        # so in other words if this happens:
        # add_client(1, clientX)
        # add_client(1, clientY) # but clientX is essentially defunct - this could happen due to a bug in client, or EM doesn't notify always of connection close in time
        # remove_client(1, clientX) -- this check maintains that clientY stays as the current client in the hash
        @client_lookup[client_id] = deleted
        @log.debug "putting back client into @client_lookup for #{client_id} #{client.inspect}"
      else
        @log.debug "cleaned up @client_lookup for #{client_id}"
      end

    end

    def add_user(context)
      user_contexts = @user_context_lookup[context.user.id]
      if user_contexts.nil?
        user_contexts = Hash.new
        @user_context_lookup[context.user.id] = user_contexts
      end

      user_contexts[context.client] = context
    end

    def remove_user(client_context)
      user_contexts = @user_context_lookup[client_context.user.id]

      if user_contexts.nil?
        @log.warn "user can not be removed #{client_context}"
      else
        # delete the context from set of user contexts
        user_contexts.delete(client_context.client)

        # if last user context, delete entire set (memory leak concern)
        if user_contexts.length == 0
          @user_context_lookup.delete(client_context.user.id)
        end

        client_context.user = nil
      end
    end

    # register topic for user messages and session messages
    def register_topics(channel)

      ######################## USER MESSAGING ###########################

      # create user exchange
      @users_exchange = channel.topic('users')
      # create user messaging topic
      @user_topic = channel.queue("", :auto_delete => true)
      @user_topic.bind(@users_exchange, :routing_key => "user.#")
      @user_topic.purge

      # subscribe for any messages to users
      @user_topic.subscribe(:ack => false) do |headers, msg|
        begin
          routing_key = headers.routing_key
          user_id = routing_key["user.".length..-1]

          @semaphore.synchronize do
            contexts = @user_context_lookup[user_id]

            if !contexts.nil?

              @log.debug "received user-directed message for user: #{user_id}"

              msg = Jampb::ClientMessage.parse(msg)

              contexts.each do |client_id, context|
                EM.schedule do
                  @log.debug "sending user message to #{context}"
                  send_to_client(context.client, msg)
                end
              end
            else
              @log.debug "Can't route message: no user connected with id #{user_id}"
            end
          end

        rescue => e
          @log.error "unhandled error in messaging to client"
          @log.error e
        end
      end

      MQRouter.user_exchange = @users_exchange

      ############## CLIENT MESSAGING ###################

      @clients_exchange = channel.topic('clients')

      @client_topic = channel.queue("", :auto_delete => true)
      @client_topic.bind(@clients_exchange, :routing_key => "client.#")
      @client_topic.purge

      # subscribe for any p2p messages to a client
      @client_topic.subscribe(:ack => false) do |headers, msg|
        begin
          routing_key = headers.routing_key
          client_id = routing_key["client.".length..-1]
          @semaphore.synchronize do
            client_context = @client_lookup[client_id]
            client = client_context.client

            msg = Jampb::ClientMessage.parse(msg)

            @log.debug "client-directed message received from #{msg.from} to client #{client_id}"

            unless client.nil?

              EM.schedule do
                @log.debug "sending client-directed down websocket to #{client_id}"
                send_to_client(client, msg)
              end
            else
              @log.debug "client-directed message unroutable to disconnected client #{client_id}"
            end
          end
        rescue => e
          @log.error "unhandled error in messaging to client"
          @log.error e
        end
      end

      MQRouter.client_exchange = @clients_exchange
    end


    def new_client(client)
      @semaphore.synchronize do
        @pending_clients.add(client)
      end

      # default to using json instead of pb
      client.encode_json = true

      client.onopen { |handshake|
        #binding.pry
        @log.debug "client connected #{client}"

        # check for '?pb' or '?pb=true' in url query parameters
        query_pb = handshake.query["pb"]

        if !query_pb.nil? && (query_pb == "" || query_pb == "true")
          client.encode_json = false
        end

      }

      client.onclose {
        @log.debug "Connection closed"
        stale_client(client)
      }

      client.onerror { |error|
        if error.kind_of?(EM::WebSocket::WebSocketError)
          @log.error "websockets error: #{error}"
        else
          @log.error "generic error: #{error} #{error.backtrace}"
        end

        cleanup_client(client)
        client.close_websocket
      }

      client.onmessage { |msg|
        @log.debug("msg received")

        # TODO: set a max message size before we put it through PB?
        # TODO: rate limit?

        pb_msg = nil

        begin
          if client.encode_json
            #example: {"type":"LOGIN", "target":"server", "login" : {"username":"hi"}}
            parse = JSON.parse(msg)
            pb_msg = Jampb::ClientMessage.json_create(parse)
            self.route(pb_msg, client)
          else
            pb_msg = Jampb::ClientMessage.parse(msg.to_s)
            self.route(pb_msg, client)
          end
        rescue SessionError => e
          @log.info "ending client session deliberately due to malformed client behavior. reason=#{e}"
          begin
            # wrap the message up and send it down
            error_msg = @message_factory.server_rejection_error(e.to_s)
            send_to_client(client, error_msg)
          ensure
            client.close_websocket
            cleanup_client(client)
          end
        rescue PermissionError => e
          @log.info "permission error.  reason=#{e.to_s}"
          @log.info e

          # wrap the message up and send it down
          error_msg = @message_factory.server_permission_error(pb_msg.message_id, e.to_s)
          send_to_client(client, error_msg)
        rescue => e
          @log.error "ending client session due to server programming or runtime error. reason=#{e.to_s}"
          @log.error e

          begin
            # wrap the message up and send it down
            error_msg = @message_factory.server_generic_error(e.to_s)
            send_to_client(client, error_msg)
          ensure
            client.close_websocket
            cleanup_client(client)
          end
        end
      }
    end


    def send_to_client(client, msg)
      @log.debug "SEND TO CLIENT (#{@message_factory.get_message_type(msg)})"
      if client.encode_json
        client.send(msg.to_json.to_s)
      else
        # this is so odd that this is necessary from an API perspective.  but searching through the source code... it's all I could find in em-websocket for allowing a binary message to be sent
        client.instance_variable_get(:@handler).send_frame(:binary, msg.to_s)
      end
    end

    def cleanup()
      # shutdown topic listeners and mq connection

      unless @amqp_connection_manager.nil?
        @amqp_connection_manager.disconnect
      end

      # tear down each individual client
      @clients.each do |client, context|
        cleanup_client(client)
      end
    end

    def stop
      @log.info "shutdown"
      cleanup
    end

    # caused a client connection to be marked stale
    def stale_client(client)
      if cid = client.client_id
        ConnectionManager.active_record_transaction do |connection_manager|
          music_session_id = connection_manager.flag_connection_stale_with_client_id(cid)
          # update the session members, letting them know this client went stale
          context = @client_lookup[client.client_id]
          if music_session = MusicSession.find_by_id(music_session_id) 
            Notification.send_musician_session_stale(music_session, client.client_id, context.user)
          end unless music_session_id.nil?
        end
      end
    end

    def cleanup_clients_with_ids(client_ids)
      # @log.debug("*** cleanup_clients_with_ids:  client_ids = #{client_ids.inspect}")
      client_ids.each do |cid|

        client_context = @client_lookup[cid]
        self.cleanup_client(client_context.client) unless client_context.nil?

        # remove this connection from the database
        ConnectionManager.active_record_transaction do |mgr|
          mgr.delete_connection(cid) { |conn, count, music_session_id, user_id|
            Notification.send_friend_update(user_id, false, conn) if count == 0
            music_session = MusicSession.find_by_id(music_session_id) unless music_session_id.nil?
            user = User.find_by_id(user_id) unless user_id.nil?
            Notification.send_musician_session_depart(music_session, cid, user) unless music_session.nil? || user.nil?
          }
        end
      end
    end

    # removes all resources associated with a client
    def cleanup_client(client)
      @semaphore.synchronize do
        # @log.debug("*** cleanup_clients:  client = #{client}")
        pending = @pending_clients.delete?(client)

        if !pending.nil?
          @log.debug "cleaning up not-logged-in client #{client}"
        else

          @log.debug "cleanup up logged-in client #{client}"

          remove_client(client.client_id, client)

          context = @clients.delete(client)

          if !context.nil?
            remove_user(context)
          else
            @log.debug "skipping duplicate cleanup attempt of logged-in client"
          end

        end
      end
    end

    def route(client_msg, client)
      message_type = @message_factory.get_message_type(client_msg)

      raise SessionError, "unknown message type received: #{client_msg.type}" if message_type.nil?

      @log.debug("msg received #{message_type}")

      raise SessionError, 'client_msg.route_to is null' if client_msg.route_to.nil?

      if @pending_clients.include? client and client_msg.type != ClientMessage::Type::LOGIN
        # this client has not logged in and is trying to send a non-login message
        raise SessionError, "must 'Login' first"
      end

      if @message_factory.server_directed? client_msg

        handle_server_directed(client_msg, client)

      elsif @message_factory.client_directed? client_msg

        to_client_id = client_msg.route_to[MessageFactory::CLIENT_TARGET_PREFIX.length..-1]
        handle_client_directed(to_client_id, client_msg, client)

      elsif @message_factory.session_directed? client_msg

        session_id = client_msg.target[MessageFactory::SESSION_TARGET_PREFIX.length..-1]
        handle_session_directed(session_id, client_msg, client)

      elsif @message_factory.user_directed? client_msg

        user_id = client_msg.target[MessageFactory::USER_PREFIX_TARGET.length..-1]
        handle_user_directed(user_id, client_msg, client)

      else
        raise SessionError, "client_msg.route_to is unknown type: #{client_msg.route_to}"
      end

    end

    def handle_server_directed(client_msg, client)
      # @log.info("*** handle_server_directed(#{client_msg.inspect}, #{client})")

      if client_msg.type == ClientMessage::Type::LOGIN

        handle_login(client_msg.login, client)

      elsif client_msg.type == ClientMessage::Type::HEARTBEAT

        handle_heartbeat(client_msg.heartbeat, client_msg.message_id, client)

      else
        raise SessionError, "unknown message type '#{client_msg.type}' for #{client_msg.route_to}-directed message"
      end
    end

    def handle_login(login, client)
      username = login.username if login.value_for_tag(1)
      password = login.password if login.value_for_tag(2)
      token = login.token if login.value_for_tag(3)
      client_id = login.client_id if login.value_for_tag(4)
      reconnect_music_session_id = login.client_id if login.value_for_tag(5)

      @log.info("*** handle_login: token=#{token}; client_id=#{client_id}")
      connection = nil
      reconnected = false

      # you don't have to supply client_id in login--if you don't, we'll generate one
      if client_id.nil? || client_id.empty?
        # give a unique ID to this client.  This is used to prevent session messages
        # from echoing back to the sender, for instance.
        client_id = UUIDTools::UUID.random_create.to_s
      else
        # check if there's a connection for the client... if it's stale, reconnect it
        if connection = JamRuby::Connection.find_by_client_id(client_id)
          # FIXME:  I think connection table needs to updated within connection_manager
          # otherwise this would be 1 line of code (connection.connect!)

          music_session_upon_reentry = connection.music_session

          ConnectionManager.active_record_transaction do |connection_manager|
            music_session_id, reconnected = connection_manager.reconnect(connection, reconnect_music_session_id)
            context = @client_lookup[client_id]
            if music_session_id.nil?
              # if this is a reclaim of a connection, but music_session_id comes back null, then we need to check if this connection was IN a music session before.
              # if so, then we need to tell the others in the session that this user is now departed
              Notification.send_musician_session_depart(music_session_upon_reentry, client.client_id, context.user) unless context.nil? || music_session_upon_reentry.nil? || music_session_upon_reentry.destroyed?
            else
              music_session = MusicSession.find_by_id(music_session_id)
              Notification.send_musician_session_fresh(music_session, client.client_id, context.user) unless context.nil?
            end

          end if connection.stale?
        end
        # if there's a client_id but no connection object, create new client_id
        client_id = UUIDTools::UUID.random_create.to_s if !connection
      end

      client.client_id = client_id

      user = valid_login(username, password, token, client_id)

      if !user.nil?
        @log.debug "user #{user} logged in"

        # respond with LOGIN_ACK to let client know it was successful
        remote_ip = extract_ip(client)


        @semaphore.synchronize do
          # remove from pending_queue
          @pending_clients.delete(client)

          # add a tracker for this user
          context = ClientContext.new(user, client)
          @clients[client] = context
          add_user(context)
          add_client(client_id, context)

          unless connection
            # log this connection in the database
            ConnectionManager.active_record_transaction do |connection_manager|
              connection_manager.create_connection(user.id, client.client_id, remote_ip) do |conn, count|
                if count == 1
                  Notification.send_friend_update(user.id, true, conn)
                end
              end
            end
          end
          login_ack = @message_factory.login_ack(remote_ip,
                                                 client_id,
                                                 user.remember_token,
                                                 @heartbeat_interval,
                                                 connection.try(:music_session_id),
                                                 reconnected)
          send_to_client(client, login_ack)
        end
      else
        raise SessionError, 'invalid login'
      end
    end

    # TODO: deprecated; jam_ruby has routine inspired by this
    def send_friend_update(user, online, client)
      @log.debug "sending friend update for user #{user} online = #{online}"

      if !user.nil? && user.friends.exists?
        @log.debug "user has friends - sending friend updates"

        # create the friend_update message
        friend_update_msg = @message_factory.friend_update(user.id, online)

        # send the friend_update to each friend that has active connections
        user.friends.each do |friend|
          @log.debug "sending friend update message to #{friend}"

          handle_user_directed(friend.id, friend_update_msg, client)
        end
      end
    end

    def handle_heartbeat(heartbeat, heartbeat_message_id, client)
      unless context = @clients[client]
        @log.warn "*** WARNING:   unable to find context due to heartbeat from client: #{client.client_id}; calling cleanup"
        cleanup_client(client)
        raise SessionError, 'context state is gone. please reconnect.'
      else
        connection = Connection.find_by_user_id_and_client_id(context.user.id, context.client.client_id)
        if connection.nil?
          @log.warn "*** WARNING:  unable to find connection due to heartbeat from client: #{context}; calling cleanup_client"
          cleanup_client(client)
          raise SessionError, 'connection state is gone. please reconnect.'
        else
          connection.touch

          ConnectionManager.active_record_transaction do |connection_manager|
            connection_manager.reconnect(connection, connection.music_session_id)
          end if connection.stale?
        end

        heartbeat_ack = @message_factory.heartbeat_ack()

        send_to_client(client, heartbeat_ack)

        # send errors to clients in response to heartbeats if rabbitmq is down
        if !@amqp_connection_manager.connected?
          error_msg = @message_factory.server_bad_state_error(heartbeat_message_id, "messaging system down")
          context.sent_bad_state_previously = true
          send_to_client(client, error_msg)
          return
        elsif context.sent_bad_state_previously
          context.sent_bad_state_previously = false
          recovery_msg = @message_factory.server_bad_state_recovered(heartbeat_message_id)
          send_to_client(client, recovery_msg)
        end
      end
    end

    def valid_login(username, password, token, client_id)

      if !token.nil? && token != ''
        @log.debug "logging in via token"
        # attempt login with token
        user = JamRuby::User.find_by_remember_token(token)

        if user.nil?
          @log.debug "no user found with token #{token}"
          return nil
        else
          @log.debug "#{user} login via token"
          return user
        end

      elsif !username.nil? and !password.nil?

        @log.debug "logging in via user/pass '#{username}' '#{password}'"
        # attempt login with username and password
        user = User.find_by_email(username)

        if !user.nil? && user.valid_password?(password)
          @log.debug "#{user} login via password"
          return user
        else
          @log.debug "#{username} login failure"
          return nil
        end
      else
        raise SessionError, 'no login data was found in Login message'
      end
    end

    def access_music_session(music_session_id, user)
      music_session = MusicSession.find_by_id(music_session_id)

      if music_session.nil?
        raise SessionError, 'specified session not found'
      end

      if !music_session.access? user
        raise SessionError, 'not allowed to join the specified session'
      end

      return music_session
    end

    # client_id = the id of the client being accessed
    # client = the current client
    def access_p2p(client_id, user, msg)

       return nil
      # ping_request and ping_ack messages are special in that they are simply allowed
      if msg.type == ClientMessage::Type::PING_REQUEST || msg.type == ClientMessage::Type::PING_ACK
        return nil
      end

      client_connection = Connection.find_by_client_id(client_id)

      if client_connection.nil?
        raise PermissionError, 'specified client not found'
      end

      if !client_connection.access_p2p? user
        raise SessionError, 'not allowed to message this client'
      end
    end


    def handle_client_directed(to_client_id, client_msg, client)
      context = @clients[client]

      # by not catching any exception here, a PermissionError will be thrown if this isn't valid
      # if for some reason the client is trying to send to a client that it doesn't
      # belong to
      access_p2p(to_client_id, context.user, client_msg)

      # populate routing data
      client_msg.from = client.client_id

      @log.debug "publishing to client #{to_client_id} from client_id #{client.client_id}"

      # put it on the topic exchange for clients
      @clients_exchange.publish(client_msg.to_s, :routing_key => "client.#{to_client_id}", :properties => {:headers => {"client_id" => client.client_id}})
    end

    def handle_user_directed(user_id, client_msg, client)

      @log.debug "publishing to user #{user_id} from client_id #{client.client_id}"

      # put it on the topic exchange for users
      @users_exchange.publish(client_msg.to_s, :routing_key => "user.#{user_id}")
    end

    def handle_session_directed(session_id, client_msg, client)
      context = @clients[client]

      user_publish_to_session(session_id, context.user, client_msg, :client_id => client.client_id)
    end

    # sends a message to a session on behalf of a user
    # if this is originating in the context of a client, it should be specified as :client_id => "value"
    # client_msg should be a well-structure message (jam-pb message)
    def user_publish_to_session(music_session_id, user, client_msg, sender = {:client_id => ""})
      music_session = access_music_session(music_session_id, user)

      # gather up client_ids in the session
      client_ids = music_session.music_session_clients.map { |client| client.client_id }.reject { |client_id| client_id == sender[:client_id] }

      publish_to_session(music_session.id, client_ids, client_msg.to_s, sender)
    end


    # sends a message to a session with no checking of permissions
    # this method deliberately has no database interactivity/active_record objects
    def publish_to_session(music_session_id, client_ids, client_msg, sender = {:client_id => ""})

      EM.schedule do
        sender_client_id = sender[:client_id]

        # iterate over each person in the session, and send a p2p message
        client_ids.each do |client_id|

          @@log.debug "publishing to session:#{music_session_id} client:#{client_id} from client:#{sender_client_id}"
          # put it on the topic exchange3 for clients
          self.class.client_exchange.publish(client_msg, :routing_key => "client.#{music_session_id}")
        end
      end
    end

    def extract_ip(client)
      return Socket.unpack_sockaddr_in(client.get_peername)[1]
    end
  end
end