jam-cloud/lib/jam_websockets/router.rb

require 'set'
require 'hot_bunnies'
require 'thread'

import java.util.concurrent.Executors

include Jampb

module JamWebsockets
  class Router

		attr_accessor :user_context_lookup, :session_context_lookup
		
    def initialize(options={})
      @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
			@session_context_lookup = {} # lookup a set of client_contexts by session_id
      @sessions_exchange = nil
      @connection = nil
      @channel = nil
      @users_exchange = nil
      @message_factory = JamRuby::MessageFactory.new
			@semaphore = Mutex.new
			@user_topic = nil
			@user_subscription = nil
			@session_topic = nil
			@session_subscription = nil
			@thread_pool = nil
    end

    def start(options = {})

      @log.info "startup"

      begin
				@thread_pool = Executors.new_fixed_thread_pool(8)
        @connection = HotBunnies.connect(:host => options[:host], :port => options[:port])
        @channel = @connection.create_channel
        @channel.prefetch = 10

        register_topics
      rescue => e
        cleanup
        raise e
      end

    end

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

			user_contexts.add(context)
		end

		def remove_user(context)
			user_contexts = @user_context_lookup[context.user.id]
			if user_contexts.nil?
				@log.warn "user can not be removed #{context}"
			else
				# delete the context from set of user contexts
				user_contexts.delete(context)

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

		def add_session(context)
			session_contexts = @session_context_lookup[context.session.id]
			if session_contexts.nil?
				session_contexts = Set.new
				@session_context_lookup[context.session.id] = session_contexts
			end

			session_contexts.add(context)
		end

		def remove_session(context)
			session_contexts = @session_context_lookup[context.session.id]
			if session_contexts.nil?
				@log.warn "session can not be removed #{context}"
			else
				# delete the context from set of session contexts
				session_contexts.delete(context)

				# if last session context, delete entire set (memory leak concern)
				if session_contexts.length == 0
					@session_context_lookup.delete(context.session.id)
				end

				context.session = nil
			end
		end


		# register topic for user messages and session messages
		def register_topics
			@users_exchange = @channel.exchange('users', :type => :topic)
      @sessions_exchange = @channel.exchange('sessions', :type => :topic)
				
			# create user messaging topic
      @user_topic = @channel.queue("", :auto_delete => true)
      @user_topic.bind(@users_exchange, :routing_key => "user.#")
      @user_topic.purge

      # TODO: alert friends

			# subscribe for any messages to users
			
			#@user_subscription = @user_topic.subscribe(:ack => false, :blocking => false, :executor => @threadpool) do |headers, msg|
			@user_subscription = @user_topic.subscribe(:ack => false)
			@user_subscription.each(:blocking => false, :executor => @threadpool) do |headers, msg|
				begin 
					routing_key = headers.envelope.routing_key
					user_id = routing_key["user.".length..-1]
					@sempahore.synchronize do
						contexts = @user_context_lookup[user_id]

						unless contexts.nil?
						
							@log.debug "received user-directed message for session: #{user_id}"
					
							contexts.each do |context|
								@log.debug "sending user message to #{context}"
								EM.schedule do
									context.client.instance_variable_get(:@handler).send_frame(:binary, msg)
								end
							end
						end
					end

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

			@session_topic = @channel.queue("", :auto_delete => true)
      @session_topic.bind(@sessions_exchange, :routing_key => "session.#")
      @session_topic.purge

      # subscribe for any messages to session
      #@session_subscription = @session_topic.subscribe(:ack => false, :blocking => false) do |headers, msg|
      @session_subscription = @session_topic.subscribe(:ack => false)
			@session_subscription.each(:blocking => false, :executor => @threadpool) do |headers, msg|
				begin 
					routing_key = headers.envelope.routing_key
					session_id = routing_key["session.".length..-1]
					@semaphore.synchronize do
						contexts = @session_context_lookup[session_id]

						unless contexts.nil?
						
							@log.debug "received session-directed message for session: #{session_id}"
					
							contexts.each do |context|
								@log.debug "sending session message to #{context}"
								EM.schedule do
									@log.debug "ONTUHNOTEHU"
									context.client.instance_variable_get(:@handler).send_frame(:binary, msg)
									@log.debug "gross"
								end
							end
						end
					end

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

    def cleanup()
			# shutdown topic listeners and mq connection
			begin
				if !@user_subscription.nil? && @user_subscription.active?
					@log.debug "cleaning up user subscription"
					@user_subscription.cancel
					@user_subscription.shutdown!
				end

				if !@session_subscription.nil? && @session_subscription.active?
					@log.debug "cleaning up session subscription"
					@session_subscription.cancel
					@session_subscription.shutdown!
				end

			rescue => e
				@log.debug "unable to cancel subscription on cleanup: #{e}"
			end

			@thread_pool.shutdown

			if !@channel.nil?
				@channel.close
			end

			if !@connection.nil?
				@connection.close
			end

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

		def stop
			@log.info "shutdown"
			cleanup
		end

		def new_client(client)

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

			client.onopen {
				@log.debug "client connected #{client}"
			}

			client.onclose {
				@log.debug "Connection closed"

				cleanup_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?

        begin
					pb_msg = Jampb::ClientMessage.parse(msg.to_s)
          self.route(pb_msg, client)
				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
            client.send(@message_factory.server_rejection_error(e.to_s).to_s)
          ensure
						client.close_websocket
            cleanup_client(client)
          end
        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
            client.send(@message_factory.server_generic_error(e.to_s).to_s)
          ensure
						client.close_websocket
            cleanup_client(client)
          end
        end

      }

    end

		# removes all resources associated with a client
    def cleanup_client(client)

			@semaphore.synchronize do
				pending = @pending_clients.delete?(client)

				if !pending.nil?
					@log.debug "cleaning up pending client #{client}"
				else
					context = @clients.delete(client)

					if !context.nil?
	
							remove_user(context)

							if !context.session.nil?
								remove_session(context)
							end
					else
						@log.debug "skipping duplicate cleanup attempt of authorized 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.target is null' if client_msg.target.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.session_directed? client_msg

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

      elsif @message_factory.user_directed? client_msg

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

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

    end

    def handle_server_directed(client_msg, 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)

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

        handle_join_jam_session(client_msg.login_jam_session, client)

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

        handle_leave_jam_session(client_msg.leave_jam_session, client)

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

    def handle_login(login, client)
      username = login.username
      token = login.token
      password = login.password

      user = valid_login(username, password, token)

      if !user.nil?

        @log.debug "user #{user.email} logged in"

        # respond with LOGIN_ACK to let client know it was successful
				client.send(@message_factory.login_ack(client.request["origin"]).to_s)

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

					# add a tracker for this user
					context = ClientContext.new(user, client)
					@clients[client] = context
					add_user(context)
				end
      else
        raise SessionError, 'invalid login'
      end
    end

		def handle_heartbeat(heartbeat, client)
			# todo: manage staleness
		end

    def handle_join_jam_session(join_jam_session, client)
      # verify that the current user has the rights to actually join the jam session
      context = @clients[client]

      session_id = join_jam_session.jam_session

      begin
        session = access_jam_session?(session_id, context.user)
				@log.debug "user #{context} joining new session #{session}"
				@semaphore.synchronize do
					old_session = context.session
					if !old_session.nil?
						@log.debug "#{context} is already in session.  auto-logging out to join new session."	
						remove_session(context)	
					end
					context.session = session
					add_session(context)
				end
      rescue => e
        # send back a failure ack and bail
				@log.debug "client requested non-existent session. client:#{client.request['origin']} user:#{context.user.email}" 
        client.send(@message_factory.login_jam_session_ack(true, e.to_s).to_s)
        return
      end

      # respond with LOGIN_JAM_SESSION_ACK to let client know it was successful
      client.send(@message_factory.login_jam_session_ack(false, nil).to_s)

      # send 'new client' message to other members in the session
      handle_session_directed(session_id, 
															@message_factory.user_joined_jam_session(context.user.id, context.user.name),
														  client)	
    end

		def handle_leave_jam_session(leave_jam_session, client)

			context = @clients[client]

			raise SessionError, "unsupported"
		end

		def valid_login(username, password, token)

			if !username.nil? and !password.nil?
				# attempt login with username and password
				user = User.find_by_email(username)

				if !user.nil? && user.authenticate(password)
					@log.debug "#{username} login via password"
					return user
				else
					@log.debug "#{username} login failure"
					return nil
				end
			elsif !token.nil?
				# attempt login with token
				user = User.find_by_remember_token(token)

				if user.nil?
					@log.debug "no user found with token"
					return false
				else
					@log.debug "#{username} login via token"
					return nil
				end
			else
				raise SessionError, 'no login data was found in Login message'
			end

		end

		def access_jam_session?(jam_session_id, user)
			jam_session = JamSession.find_by_id(jam_session_id)

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

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

			return jam_session
		end

		def handle_session_directed(session_id, client_msg, client)

			context = @clients[client]

			# by not catching any exception here, this will kill the connection
			# if for some reason the client is trying to send to a session that it doesn't
			# belong to
			session = access_jam_session?(session_id, context.user)

			@log.debug "publishing to session #{session}"
			# put it on the topic exchange for sessions
			@sessions_exchange.publish(client_msg.to_s, :routing_key => "session.#{session_id}")
		end

		def handle_user_directed(user, client_msg, client)

			raise SessionError, 'not implemented'
		end
	end
end