Session helper class (#328)

* Session helpers

* Add tests

* Rubocop

* Put ruby back where it was

* And the gemfile

* move gems to gemspec

* rubocop again

Gemfile.lock

@@ -2,19 +2,25 @@ PATH
   remote: .
   specs:
     workos (5.6.0)
+      encryptor (~> 3.0)
+      jwt (~> 2.8)
 
 GEM
   remote: https://rubygems.org/
   specs:
     addressable (2.8.6)
       public_suffix (>= 2.0.2, < 6.0)
     ast (2.4.2)
+    base64 (0.2.0)
     bigdecimal (3.1.7)
     crack (1.0.0)
       bigdecimal
       rexml
     diff-lcs (1.5.1)
+    encryptor (3.0.0)
     hashdiff (1.1.0)
+    jwt (2.8.2)
+      base64
     parallel (1.24.0)
     parser (3.3.0.5)
       ast (~> 2.4.1)

lib/workos.rb

@@ -71,6 +71,7 @@ def self.key
   autoload :Profile, 'workos/profile'
   autoload :ProfileAndToken, 'workos/profile_and_token'
   autoload :RefreshAuthenticationResponse, 'workos/refresh_authentication_response'
+  autoload :Session, 'workos/session'
   autoload :SSO, 'workos/sso'
   autoload :Types, 'workos/types'
   autoload :User, 'workos/user'

lib/workos/authentication_response.rb

@@ -6,20 +6,39 @@ module WorkOS
   class AuthenticationResponse
     include HashProvider
 
-    attr_accessor :user, :organization_id, :impersonator, :access_token, :refresh_token
+    attr_accessor :user,
+                  :organization_id,
+                  :impersonator,
+                  :access_token,
+                  :refresh_token,
+                  :authentication_method,
+                  :sealed_session
 
-    def initialize(authentication_response_json)
+    # rubocop:disable Metrics/AbcSize
+    def initialize(authentication_response_json, session = nil)
       json = JSON.parse(authentication_response_json, symbolize_names: true)
+      @access_token = json[:access_token]
+      @refresh_token = json[:refresh_token]
       @user = WorkOS::User.new(json[:user].to_json)
       @organization_id = json[:organization_id]
       @impersonator =
         if (impersonator_json = json[:impersonator])
           Impersonator.new(email: impersonator_json[:email],
                            reason: impersonator_json[:reason],)
         end
-      @access_token = json[:access_token]
-      @refresh_token = json[:refresh_token]
+      @authentication_method = json[:authentication_method]
+      @sealed_session =
+        if session && session[:seal_session]
+          WorkOS::Session.seal_data({
+                                      access_token: access_token,
+                                      refresh_token: refresh_token,
+                                      user: user.to_json,
+                                      organization_id: organization_id,
+                                      impersonator: impersonator.to_json,
+                                    }, session[:cookie_password],)
+        end
     end
+    # rubocop:enable Metrics/AbcSize
 
     def to_json(*)
       {
@@ -28,6 +47,8 @@ def to_json(*)
         impersonator: impersonator.to_json,
         access_token: access_token,
         refresh_token: refresh_token,
+        authentication_method: authentication_method,
+        sealed_session: sealed_session,
       }
     end
   end

lib/workos/refresh_authentication_response.rb

@@ -6,18 +6,41 @@ module WorkOS
   class RefreshAuthenticationResponse
     include HashProvider
 
-    attr_accessor :access_token, :refresh_token
+    attr_accessor :user, :organization_id, :impersonator, :access_token, :refresh_token, :sealed_session
 
-    def initialize(authentication_response_json)
+    # rubocop:disable Metrics/AbcSize
+    def initialize(authentication_response_json, session = nil)
       json = JSON.parse(authentication_response_json, symbolize_names: true)
       @access_token = json[:access_token]
       @refresh_token = json[:refresh_token]
+      @user = WorkOS::User.new(json[:user].to_json)
+      @organization_id = json[:organization_id]
+      @impersonator =
+        if (impersonator_json = json[:impersonator])
+          Impersonator.new(email: impersonator_json[:email],
+                           reason: impersonator_json[:reason],)
+        end
+      @sealed_session =
+        if session && session[:seal_session]
+          WorkOS::Session.seal_data({
+                                      access_token: access_token,
+                                      refresh_token: refresh_token,
+                                      user: user.to_json,
+                                      organization_id: organization_id,
+                                      impersonator: impersonator.to_json,
+                                    }, session[:cookie_password],)
+        end
     end
+    # rubocop:enable Metrics/AbcSize
 
     def to_json(*)
       {
+        user: user.to_json,
+        organization_id: organization_id,
+        impersonator: impersonator.to_json,
         access_token: access_token,
         refresh_token: refresh_token,
+        sealed_session: sealed_session,
       }
     end
   end

lib/workos/session.rb

@@ -0,0 +1,183 @@
+# frozen_string_literal: true
+
+require 'jwt'
+require 'uri'
+require 'net/http'
+require 'encryptor'
+require 'securerandom'
+require 'json'
+require 'uri'
+
+module WorkOS
+  # The Session class provides helper methods for working with WorkOS sessions
+  # This class is not meant to be instantiated in a user space, and is instantiated internally but exposed.
+  # rubocop:disable Metrics/ClassLength
+  class Session
+    attr_accessor :jwks, :jwks_algorithms, :user_management, :cookie_password, :session_data, :client_id
+
+    def initialize(user_management:, client_id:, session_data:, cookie_password:)
+      raise ArgumentError, 'cookiePassword is required' if cookie_password.nil? || cookie_password.empty?
+
+      @user_management = user_management
+      @cookie_password = cookie_password
+      @session_data = session_data
+      @client_id = client_id
+
+      @jwks = create_remote_jwk_set(URI(@user_management.get_jwks_url(client_id)))
+      @jwks_algorithms = @jwks.map { |key| key[:alg] }.compact.uniq
+    end
+
+    # Authenticates the user based on the session data
+    # @return [Hash] A hash containing the authentication response and a reason if the authentication failed
+    def authenticate
+      return { authenticated: false, reason: 'NO_SESSION_COOKIE_PROVIDED' } if @session_data.nil?
+
+      begin
+        session = Session.unseal_data(@session_data, @cookie_password)
+      rescue StandardError
+        return { authenticated: false, reason: 'INVALID_SESSION_COOKIE' }
+      end
+
+      return { authenticated: false, reason: 'INVALID_SESSION_COOKIE' } unless session[:access_token]
+      return { authenticated: false, reason: 'INVALID_JWT' } unless is_valid_jwt(session[:access_token])
+
+      decoded = JWT.decode(session[:access_token], nil, true, algorithms: @jwks_algorithms, jwks: @jwks).first
+
+      {
+        authenticated: true,
+        session_id: decoded['sid'],
+        organization_id: decoded['org_id'],
+        role: decoded['role'],
+        permissions: decoded['permissions'],
+        user: session[:user],
+        impersonator: session[:impersonator],
+        reason: nil,
+      }
+    end
+
+    # Refreshes the session data using the refresh token stored in the session data
+    # @param options [Hash] Options for refreshing the session
+    # @option options [String] :cookie_password The password to use for unsealing the session data
+    # @option options [String] :organization_id The organization ID to use for refreshing the session
+    # @return [Hash] A hash containing a new sealed session, the authentication response,
+    # and a reason if the refresh failed
+    # rubocop:disable Metrics/AbcSize
+    # rubocop:disable Metrics/CyclomaticComplexity
+    # rubocop:disable Metrics/PerceivedComplexity
+    def refresh(options = nil)
+      cookie_password = options.nil? || options[:cookie_password].nil? ? @cookie_password : options[:cookie_password]
+
+      begin
+        session = Session.unseal_data(@session_data, cookie_password)
+      rescue StandardError
+        return { authenticated: false, reason: 'INVALID_SESSION_COOKIE' }
+      end
+
+      return { authenticated: false, reason: 'INVALID_SESSION_COOKIE' } unless session[:refresh_token] && session[:user]
+
+      begin
+        auth_response = @user_management.authenticate_with_refresh_token(
+          client_id: @client_id,
+          refresh_token: session[:refresh_token],
+          organization_id: options.nil? || options[:organization_id].nil? ? nil : options[:organization_id],
+          session: { seal_session: true, cookie_password: cookie_password },
+        )
+
+        @session_data = auth_response.sealed_session
+        @cookie_password = cookie_password
+
+        {
+          authenticated: true,
+          sealed_session: auth_response.sealed_session,
+          session: auth_response,
+          reason: nil,
+        }
+      rescue StandardError => e
+        { authenticated: false, reason: e.message }
+      end
+    end
+    # rubocop:enable Metrics/AbcSize
+    # rubocop:enable Metrics/CyclomaticComplexity
+    # rubocop:enable Metrics/PerceivedComplexity
+
+    # Returns a URL to redirect the user to for logging out
+    # @return [String] The URL to redirect the user to for logging out
+    # rubocop:disable Naming/AccessorMethodName
+    def get_logout_url
+      auth_response = authenticate
+
+      unless auth_response[:authenticated]
+        raise "Failed to extract session ID for logout URL: #{auth_response[:reason]}"
+      end
+
+      @user_management.get_logout_url(session_id: auth_response[:session_id])
+    end
+    # rubocop:enable Naming/AccessorMethodName
+
+    # Encrypts and seals data using AES-256-GCM
+    # @param data [Hash] The data to seal
+    # @param key [String] The key to use for encryption
+    # @return [String] The sealed data
+    def self.seal_data(data, key)
+      iv = SecureRandom.random_bytes(12)
+
+      encrypted_data = Encryptor.encrypt(
+        value: JSON.generate(data),
+        key: key,
+        iv: iv,
+        algorithm: 'aes-256-gcm',
+      )
+      Base64.encode64(iv + encrypted_data) # Combine IV with encrypted data and encode as base64
+    end
+
+    # Decrypts and unseals data using AES-256-GCM
+    # @param sealed_data [String] The sealed data to unseal
+    # @param key [String] The key to use for decryption
+    # @return [Hash] The unsealed data
+    def self.unseal_data(sealed_data, key)
+      decoded_data = Base64.decode64(sealed_data)
+      iv = decoded_data[0..11] # Extract the IV (first 12 bytes)
+      encrypted_data = decoded_data[12..-1] # Extract the encrypted data
+
+      decrypted_data = Encryptor.decrypt(
+        value: encrypted_data,
+        key: key,
+        iv: iv,
+        algorithm: 'aes-256-gcm',
+      )
+
+      JSON.parse(decrypted_data, symbolize_names: true) # Parse the decrypted JSON string back to original data
+    end
+
+    private
+
+    # Creates a JWKS set from a remote JWKS URL
+    # @param uri [URI] The URI of the JWKS
+    # @return [JWT::JWK::Set] The JWKS set
+    def create_remote_jwk_set(uri)
+      # Fetch the JWKS from the remote URL
+      response = Net::HTTP.get(uri)
+
+      jwks_hash = JSON.parse(response)
+      jwks = JWT::JWK::Set.new(jwks_hash)
+
+      # filter jwks so it only returns the keys where 'use' is equal to 'sig'
+      jwks.keys.select! { |key| key[:use] == 'sig' }
+
+      jwks
+    end
+
+    # Validates a JWT token using the JWKS set
+    # @param token [String] The JWT token to validate
+    # @return [Boolean] True if the token is valid, false otherwise
+    # rubocop:disable Naming/PredicateName
+    def is_valid_jwt(token)
+      JWT.decode(token, nil, true, algorithms: @jwks_algorithms, jwks: @jwks)
+      true
+    rescue StandardError
+      false
+    end
+    # rubocop:enable Naming/PredicateName
+  end
+  # rubocop:enable Metrics/ClassLength
+end