heimdall/server.rb

require_relative "proto"
require_relative "hyde/hyde"
require 'webrick/websocket'
require "json"

Users = Heimdall::UserCache.new
SocketsMap = {}
NotifyList = {}

def _require_keys(dict,key_dict)
  raise KeyError, "not a dict" unless dict.kind_of? Hash
  key_dict.each_pair { |k,v|
    unless (dict.has_key? k.to_s) and (dict[k.to_s].kind_of? v) then
      raise KeyError, "key #{k} of type #{v} required"
    end
  }
end

def _send_json(res,data,code: 200)
  res.body = JSON::fast_generate(data)
  res['Content-Type'] = "application/json"
  res.status = code
end

def _throw_error(res,error)
  _send_json(res,{
    error: "#{error}"
  },code: 400)
end

def _parse_json(body,key_dict)
  data = JSON::Parser.new(body).parse
  _require_keys(data,key_dict)
  return data
end


server = Hyde::Server.new Port: 8000 do
  path "user" do
    path "find" do
      index ["by-username"]
      
      get "by-username" do |ctx|
        req,res = ctx.request,ctx.response
        begin
          _require_keys(req.query, {
            username: String
          })
          _send_json(res, {
            "results": Users.search_by_username(req.query['username'])
          })
        rescue KeyError => keyerror
          _throw_error(res,keyerror)
        end
      end
    
      get "by-protoid" do |ctx|
        req,res = ctx.request,ctx.response
        begin
          _require_keys(req.query, {
            protocol_id: String
          })
          _send_json(res, {
            "results": Users.search_by_protoid(req.query['protocol_id'])
          })
        rescue KeyError => keyerror
          _throw_error(res,keyerror)
        end
      end
    end

    get "exists" do |ctx|
      req,res = ctx.request,ctx.response
      begin 
        _require_keys(req.query, {
          protocol_id: String
        })
        _send_json(res,{
          exists: (Users.get(req.query["protocol_id"]) != nil)
        })
      rescue KeyError => keyerror
        _throw_error(res,keyerror)
      rescue Heimdall::ProtocolError => protoerr
        _send_json(res, {
          exists: false
        })
      end
    end

    post "new" do |ctx|
      req,res = ctx.request,ctx.response
      begin
        data = _parse_json(req.body,{
          username: String,
          protocol_id: String
        })
        new_user = Heimdall::User.new(data)
        Users.add(new_user)
        _send_json(res,{
          status: true
        })
        NotifyList[new_user.UID] = []
      rescue JSON::ParserError => jsonerror
        _throw_error(res,jsonerror)
      rescue KeyError => keyerror
        _throw_error(res,keyerror)
      end
    end

    post "send" do |ctx|
      req,res = ctx.request,ctx.response
      begin 
        data = _parse_json(req.body, {
          content: String,
          from: String,
          to: String
        })
        new_message = Heimdall::Message.new(data)
        user = Users.get(new_message.to)
        if NotifyList[user.UID].length != 0 then
          NotifyList[user.UID].each { |sockid|
            sock = SocketsMap[sockid]
            msg = new_message.to_struct
            msg["user"] = Users.get(msg["from"]).to_card
            sock.puts(JSON::fast_generate(msg))
          }
          user.direct.send_silent(new_message)
        else
          user.direct.send(new_message)
        end
        _send_json(res,{
          status: true
        })
      rescue JSON::ParserError => jsonerror
        _throw_error(res,jsonerror)
      rescue KeyError => keyerror
        _throw_error(res,keyerror)
      rescue Heimdall::ProtocolError => protoerr
        _throw_error(res,protoerr)
      end
    end

    get "get" do |ctx|
      req,res = ctx.request,ctx.response
      begin 
        _require_keys(req.query,{
          n: Integer,
          protocol_id: String
        })
        number = req.query[:n]
        id = req.query["protocol_id"]
        user = Users.get(id)
        messages = user.direct.get(number)
        _send_json(res, {
          messages: messages.map { |x| 
            x = x.to_struct
            x["user"] = Users.get(x["from"]).to_card
            x
          }
        })
      rescue Heimdall::ProtocolError => protoerr
        _throw_error(res,protoerr)
      end
    end

    get "read" do |ctx|
      req,res = ctx.request,ctx.response
      begin 
        _require_keys(req.query,{
          protocol_id: String
        })
        id = req.query["protocol_id"]
        user = Users.get(id)
        messages = user.direct.read
        _send_json(res, {
          messages: messages.map { |x| 
            x = x.to_struct 
            x["user"] = Users.get(x["from"]).to_card
            x
          }
        })
      rescue Heimdall::ProtocolError => protoerr
        _throw_error(res,protoerr)
      end
    end

    post "listen" do |ctx|
      req,res = ctx.request, ctx.response
      begin
        data = _parse_json(req.body, {
          websocket: String,
          protocol_id: String
        })
        uid = Users.get(data["protocol_id"]).UID
        raise KeyError, "websocket does not exist" unless SocketsMap.has_key? data["websocket"] 
        NotifyList[uid].append data["websocket"]
        _send_json(res, {
          status: true
        })
      rescue KeyError => keyerror
        _throw_error(res,keyerror)
      end
    end

    post "delete" do |ctx|
      req,res = ctx.request, ctx.response
      begin
        data = _parse_json(req.body,{
          protocol_id: String
        })
        id = data["protocol_id"]
        _send_json(res, {
          status: true
        })
        user = Users.get(id) 
        NotifyList.delete(user.UID)
        Users.delete(id)
      rescue Heimdall::ProtocolError => protoerr
        _throw_error(res,protoerr)
      end
    end
  end
  
  get "version" do |ctx|
    ctx.response.body = "{\"version\":\"#{Heimdall::VERSION}\"}"
    ctx.response['Content-Type'] = "application/json"
  end
    
end

at_exit do
  server.shutdown
end

class WebsocketUID < Heimdall::UID
  def initialize
    @UID_prefix = "websocket"
    super
  end
end

class EventServlet < WEBrick::Websocket::Servlet
  def socket_open(sock)
    @UID = WebsocketUID.new
    @connected_listeners = []
    SocketsMap[@UID.UID] = sock
    sock.puts(JSON::fast_generate({
      websocket: @UID.UID.to_s
    }))
  end

  def socket_close(sock)
    SocketsMap.delete @UID.UID
    NotifyList.each do |k,v| 
      if v.include? @UID.UID
        v.delete @UID.UID
      end
    end
  end

  def socket_text(sock,text)
    # do nothing
  end
end

Thread.new do
  websocket_server = WEBrick::Websocket::HTTPServer.new Port:8001
  websocket_server.mount "/", EventServlet
  websocket_server.start
  at_exit do
    websocket_server.shutdown
  end
end

server.start