JWT Auth + Refresh Tokens in Rails

README.md

JWT Auth + Refresh Tokens in Rails

This is just some code I recently used in my development application in order to add token-based authentication for my api-only rails app. The api-client was to be consumed by a mobile application, so I needed an authentication solution that would keep the user logged in indefinetly and the only way to do this was either using refresh tokens or sliding sessions.

I also needed a way to both blacklist and whitelist tokens based on a unique identifier (jti)

Before trying it out DIY, I considered using:

• devise-jwt which unfortunately does not support refresh tokens
• devise_token_auth I ran into issues when it came to the changing headers on request on mobile, disabling this meant users would have to sign in periodically
• doorkeeper This was pretty close to what I needed, however, it was quite complicated and I considered it wasn't worth the extra effort of implmeneting OAuth2 (for now)
• api_guard This was great, almost everything I needed but it didn't play too nicely with GraphQL and I needed to implement token whitelisting also.

So, since I couldn't find any widely-used gem to meet my needs; I decided to just go DIY, and the end result works pretty well. And overview of how things works is so:

• You call on the Jwt::Issuer module to create an access_token and refresh_token pair.
• You call on the Jwt::Authenticator module to authenticate the access_token get the current_user and the decoeded_token
• You call on the Jwt::Revoker module to revoke (blacklist/remove whitelist) a token
• You call on the Jwt::Refresher module to refresh an access_token based on a refresh_token

There are more modules, but you can preview them for yourself.

There are some prequistes you need in order to use this code:

1. You need to create a blacklisted tokens table like so: rails g model BlacklistedToken jti:string:uniq:index user:belongs_to exp:datetime

2. If you want to use whitelisting to, create a tokens table like so: rails g model WhitelistedToken jti:string:uniq:index user:belongs_to exp:datetime

3. Create a refresh tokens table like & model so: rails g model RefreshToken crypted_token:string:uniq user:belongs_to

class RefreshToken < ApplicationRecord
  belongs_to :user
  before_create :set_crypted_token

  attr_accessor :token

  def self.find_by_token(token)
    crypted_token = Digest::SHA256.hexdigest token
    RefreshToken.find_by(crypted_token: crypted_token)
  end

  private

  def set_crypted_token
    self.token = SecureRandom.hex
    self.crypted_token = Digest::SHA256.hexdigest(token)
  end
end

4. Update the user model to include the associations

  has_many :refresh_tokens, dependent: :delete_all
  has_many :whitelisted_tokens, dependent: :delete_all
  has_many :blacklisted_tokens, dependent: :delete_all

Then you are pretty much ready to go!

In the future, I might make this into a gem or add redis support or similar.

I hope this gist helps someone!

application_controller.rb

class ApplicationControler < ActionController::Base
  before_action :authenticate
  
  private
  
  def authenticate
    current_user, decoded_token = Jwt::Authenticator.call(
       headers: request.headers,
       access_token: params[:access_token] # authenticate from header OR params
     )

     @current_user = current_user
     @decoded_token = decoded_token
  end
end

authenticator.rb

module Jwt
  module Authenticator
    module_function

    def call(headers:, access_token:)
      token = access_token || Jwt::Authenticator.authenticate_header(
        headers
      )
      raise Errors::Jwt::MissingToken unless token.present?

      decoded_token = Jwt::Decoder.decode!(token)
      user = Jwt::Authenticator.authenticate_user_from_token(decoded_token)
      raise Errors::Unauthorized unless user.present?

      [user, decoded_token]
    end

    def authenticate_header(headers)
      headers['Authorization']&.split('Bearer ')&.last
    end

    def authenticate_user_from_token(decoded_token)
      raise Errors::Jwt::InvalidToken unless decoded_token[:jti].present? && decoded_token[:user_id].present?

      user = User.find(decoded_token.fetch(:user_id))
      blacklisted = Jwt::Blacklister.blacklisted?(jti: decoded_token[:jti])
      whitelisted = Jwt::Whitelister.whitelisted?(jti: decoded_token[:jti])
      valid_issued_at = Jwt::Authenticator.valid_issued_at?(user, decoded_token)

      return user if !blacklisted && whitelisted && valid_issued_at
    end

    def valid_issued_at?(user, decoded_token)
      !user.token_issued_at || decoded_token[:iat] >= user.token_issued_at.to_i
    end

    module Helpers
      extend ActiveSupport::Concern

      def logout!(user:, decoded_token:)
        Jwt::Revoker.revoke(
          decoded_token: decoded_token,
          user: user
        )
      end
    end
  end
end

blacklister.rb

module Jwt
  module Blacklister
    module_function

    def blacklist!(jti:, exp:, user:)
      user.blacklisted_tokens.create!(
        jti: jti,
        exp: Time.at(exp)
      )
    end

    def blacklisted?(jti:)
      BlacklistedToken.exists?(jti: jti)
    end
  end
end

decoder.rb

module Jwt
  module Decoder
    module_function

    def decode!(access_token, verify: true)
      decoded = JWT.decode(access_token, Jwt::Secret.secret, verify, verify_iat: true)[0]
      raise Errors::Jwt::InvalidToken unless decoded.present?

      decoded.symbolize_keys
    end

    def decode(access_token, verify: true)
      decode!(access_token, verify: verify)
    rescue StandardError
      nil
    end
  end
end

encoder.rb

module Jwt
  module Encoder
    module_function

    def call(user)
      jti = SecureRandom.hex
      exp = Jwt::Encoder.token_expiry
      access_token = JWT.encode(
        {
          user_id: user.id,
          jti: jti,
          iat: Jwt::Encoder.token_issued_at.to_i,
          exp: exp
        },
        Jwt::Secret.secret
      )

      [access_token, jti, exp]
    end

    def token_expiry
      (Jwt::Encoder.token_issued_at + Jwt::Expiry.expiry).to_i
    end

    def token_issued_at
      Time.now
    end
  end
end

expiry.rb

module Jwt
  module Expiry
    module_function

    def expiry
      2.hours
    end
  end
end

issuer.rb

module Jwt
  module Issuer
    module_function

    def call(user)
      access_token, jti, exp = Jwt::Encoder.call(user)
      refresh_token = user.refresh_tokens.create!
      Jwt::Whitelister.whitelist!(
        jti: jti,
        exp: exp,
        user: user
      )

      [access_token, refresh_token]
    end
  end
end

refresher.rb

module Jwt
  module Refresher
    module_function

    def refresh!(refresh_token:, decoded_token:, user:)
      raise Errors::Jwt::MissingToken, token: 'refresh' unless refresh_token.present? || decoded_token.nil?

      existing_refresh_token = user.refresh_tokens.find_by_token(
        refresh_token
      )
      raise Errors::Jwt::InvalidToken, token: 'refresh' unless existing_refresh_token.present?

      jti = decoded_token.fetch(:jti)

      new_access_token, new_refresh_token = Jwt::Issuer.call(user)
      existing_refresh_token.destroy!

      Jwt::Blacklister.blacklist!(jti: jti, exp: decoded_token.fetch(:exp), user: user)
      Jwt::Whitelister.remove_whitelist!(jti: jti)

      [new_access_token, new_refresh_token]
    end
  end
end

revoker.rb

module Jwt
  module Revoker
    module_function

    def revoke(decoded_token:, user:)
      jti = decoded_token.fetch(:jti)
      exp = decoded_token.fetch(:exp)

      Jwt::Whitelister.remove_whitelist!(jti: jti)
      Jwt::Blacklister.blacklist!(
        jti: jti,
        exp: exp,
        user: user
      )
    rescue StandardError
      raise Errors::Jwt::InvalidToken
    end
  end
end

secret.rb

module Jwt
  module Secret
    module_function

    def secret
      Rails.application.secrets.secret_key_base
    end
  end
end

whitelisted.rb

module Jwt
  module Whitelister
    module_function

    def whitelist!(jti:, exp:, user:)
      user.whitelisted_tokens.create!(
        jti: jti,
        exp: Time.at(exp)
      )
    end

    def remove_whitelist!(jti:)
      whitelist = WhitelistedToken.find_by(
        jti: jti
      )
      whitelist.destroy if whitelist.present?
    end

    def whitelisted?(jti:)
      WhitelistedToken.exists?(jti: jti)
    end
  end
end