Integrate JWT into Rails Grape API for Multi-Login

jwt_in_rails_1_model_setup.rb

# Prepare the schema for the additional session records.
# XXXXXXXXXXXXXXXX_create_user_session.rb
class CreateUserSessions < ActiveRecord::Migration
  def change
    create_table :user_sessions do |t|
      t.references :user, :foreign_key => true, :index => true, :null => false
      t.text :user_agent
      t.datetime :expire_at # use expire_at to control the expiration after issuing JWT

      t.timestamps null: false
    end
  end
end

# Setup the UserSession model
# app/models/user_session.rb
class UserSession < ActiveRecord::Base
  belongs_to :user
  validates :user, :presence => true

  store :user_agent, :accessors => [:device_idid, :request_platform], :coder => JSON

  before_create :multi_login_count

  private

  # We're not going to allow the user manage their sessions here
  # will just expire the earliest session before creating a new one
  def multi_login_count
    existed_valid_sessions = user.user_sessions.where(expire_at: nil).order('created_at ASC')
    existed_valid_sessions.first.update_columns(expire_at: Time.now) if existed_valid_sessions.count == 5
  end
end

# Setup User model relationship
class User < ActiveRecord::Base
  ...
  has_many :user_sessions, :dependent => :destroy
  ...
end

jwt_in_rails_2_wrapper_services.rb

# In Gemfile.rb
  gem 'jwt'

# Create your own JasonWebToken.rb for better encapsulation
# I would like to make the error handle specific from my own JasonWebToken class
# app/services/jason_web_token.rb
require 'jwt'
require 'json/ext'

class JasonWebToken
  attr_reader :header, :payload

  JWT_SETTING = {
    protocol: 'Bearer',
    algorithm: 'HS256',
    issuer: TOKEN_ISS[Rails.env.to_sym],
    secret: Rails.application.secrets.secret_key_base
  }

  def initialize(user_session, options = {})
    # Will feed the JasonWebToken with valid UserSession object, so that the
    # payload info can use the user_session information for further encoding
    raise JasonWebTokenError, "Must provide valid UserSession object." unless user_session.kind_of? UserSession

    @header = {
      typ: 'JWT',
      alg: options[:algorithm] || JWT_SETTING[:algorithm]
    }
    @payload = {
      user_session_id: user_session.id,
      user_agent: user_session.user_agent,
      iss: 'http://lwstudio.org'
    }
    @payload[:exp] = user_session.expire_at if user_session.expire_at.present?
  end

  def encode
    JWT.encode(payload, JWT_SETTING[:secret], header[:alg])
  end

  def self.decode(request_header_hash) # return user object
    raise JasonWebTokenError, "Missing Authentication Info in Header." unless request_header_hash['Authorization'].present?

    protocol, token = request_header_hash['Authorization'].split(' ')
    raise JasonWebTokenError, "JWT Protocol is Wrong." unless protocol == JWT_SETTING[:protocol]

    begin
      payload, header = JWT.decode(token, JWT_SETTING[:secret], true, { :algorithm => JWT_SETTING[:algorithm] })
    rescue JWT::ExpiredSignature => e
      # rescue the jwt error and replace it with our custom error
      # we'll take the expiration token handled by custom error
      raise JasonWebTokenError, "Token is expired."
    end

    raise JasonWebTokenError, "Token issuer is not valid." unless payload['iss'] == JWT_SETTING[:issuer]

    user_session = UserSession.find_by_id(payload['user_session_id'].to_i)
    raise JasonWebTokenError, "UserSession Not Found." unless user_session.present?

    # payload 'user_agent' decoded from jwt will be a hash, request_header contains json string, need .to_json to compare
    raise JasonWebTokenError, "UserAgent info in token does not match." unless payload['user_agent'].to_json == request_header_hash['Session-User-Agent']
    # payload 'user_agent' can compare with the rails model directly
    raise JasonWebTokenError, "UserAgent info in token does not match." unless payload['user_agent'] == user_session.user_agent

    # Handle if the token was expired by the UserSession Model
    raise JasonWebTokenError, "Token is expired." unless user_session.expire_at.nil? || user_session.expire_at > Time.now

    user_session.user
  end
end

class JasonWebTokenError < StandardError; end

jwt_in_rails_3_grape_api_login_output_jwt.rb

# Given a valid user object from the login API endpoint,
# create an UserSession object for the login.
# And present the jwt with feeding our JasonWebToken the new
# UserSession object in the login API output.
require 'json/ext'

resource :login do
  desc "Return logged in user and jwt."
  params do
    requires :account, desc: 'Username', type: String
    requires :password, desc: 'user password', type: String
  end
  post do
    user = User.find_by_username(params[:account])
    error!('User Not Found', 403) unless user.present?
    error!('Password invalid', 403) unless user.valid_password?(params[:password])
    
    user_agent = headers['Session-User-Agent'] || request.env['Session-User-Agent']
    user_session = UserSession.create(user: user, user_agent: JSON.parse(user_agent, :symbolize_names => true))
    
    present :user, user, with: API::Entities::EntityUser
    present :jwt, JasonWebToken.new(user_session).encode
  end
end

jwt_in_rails_4_grape_api_helpers.rb

# add authentication helper method, then we can use this helper method
# in a before block for those API endpoints requiring logged in user

def authenticate!
  # In RSpec test, the grape headers method cannot get out the request header
  # parameters, have to use request.env to get the information from tests.
  # However, in real request, the request.env does not work properly, only
  # headers method works. To work around this we'll reassemble the auth header first
  request_headers = {}
  request_headers['Authorization'] = headers['Authorization'] || request.env['Authorization']
  request_headers['Session-User-Agent'] = headers['Session-User-Agent'] || request.env['Session-User-Agent']
  begin
    @user ||= JasonWebToken.decode(request_headers)
  rescue JasonWebTokenError => e
    error!(e.message, 403)
  end
end