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

@jesster2k10

class Api::V1::AuthController < Api::V1::ApplicationController
  skip_before_action :authenticate_request, only: [:refresh_access_token]

  def refresh_access_token
    user = User.find_by(id: params[:id])
    refresh_token = user.refresh_tokens.last
    access_token = Jwt::Authenticator.authenticate_header(
      request.headers
    )

    begin
      # The `verify: false` option is used to skip the verification of the token's signature.
      # This is not good since if someone has an expired token, they can refresh it without verifying the signature.
      # The thing is that the access token is expired at this point and the code expects it to be expired here. The
      # issue is that JWT throws a JWT::ExpiredSignature error and haults the entire operation if `verify: true`.
      # I still think that I need to revisit this area of the code to see if there is a way to hand this solution.
      decoded_access_token = Jwt::Decoder.decode!(access_token, verify: false)

      new_access_token = Jwt::Refresher.refresh!(
        refresh_token: refresh_token,
        decoded_token: decoded_access_token,
        user: user,
      )

      render json: { access_token: new_access_token }, status: :ok
    rescue JWT::DecodeError => e
      render json: { error: "Invalid token: #{e.message}" }, status: :unauthorized
    rescue StandardError => e
      render json: { error: e.message }, status: :unauthorized
    end
  end
end

Thanks for this gist! I am using it for nextjs/rails app.

One area that stumped me was refreshing the access token using the Refresher module. Above, I kept running into JWT::SignatureExpired error whenever I tried decoding an expired access token. The only way to refresh the token was to set verify to false while decoding it, otherwise the error wouldn't let me continue.

Did you face this issue? If so, did you come up with a solution?