jwt_session_store.rb

require 'active_support/key_generator'

module ActionDispatch
  module Session
    class JWTStore < AbstractStore
      GENERATOR_KEY = "action_dispatch.key_generator".freeze
      JWT_SECRET_SALT = "jwt secret salt".freeze

      def initialize
        super
        key_generator = env[GENERATOR_KEY]
        @jwt_secret = key_generator.generate_key(JWT_SECRET_SALT)
      end

      private

      def commit_session(env, status, headers, body)
        [status, headers, body] super
      end

      def get_session(env, session_id)
        decoded_jwt_data(env) || {}
      end

      def set_session(env, session_id, session, options) 
        session['id'] = session_id
        session['iat'] = Time.now
        session['exp'] = Time.now + options[:expire_after] if options[:expire_after]
        session['sub'] = 'session'
        jwe = JSON::JWT.new(session).sign(@jwt_secret).encrypt(@jwt_secret)
        jwe.to_s
      end

      def extract_session_id(env)
        data = decoded_jwt_data(env)
        data["id"] if data
      end

      def decoded_jwt_data(env)
        request = ActionDispatch::Request.new(env)
        jwt_string ||= if request.headers.key?('Authorization')
                         request.headers['Authorization'].split(' ').last
                       else
                         request.cookie_jar[key]
                       end
        begin
          jwt = JSON::JWT.decode(jwt_string, @jwt_secret)
          return jwt if jwt['exp'] < Time.now.to_i
        rescue JSON::JWT::InvalidFormat, JSON::JWS::VerificationFailed
          nil
        end
      end
    end
  end
end