Allow pluggable encryptors for session seal/unseal (#432)

Author: gjtorikian

lib/workos.rb

@@ -56,6 +56,7 @@ def self.key
   autoload :DirectorySync, 'workos/directory_sync'
   autoload :DirectoryUser, 'workos/directory_user'
   autoload :EmailVerification, 'workos/email_verification'
+  autoload :Encryptors, 'workos/encryptors'
   autoload :Event, 'workos/event'
   autoload :Events, 'workos/events'
   autoload :Factor, 'workos/factor'

lib/workos/authentication_response.rb

@@ -31,13 +31,17 @@ def initialize(authentication_response_json, session = nil)
       @oauth_tokens = json[:oauth_tokens] ? WorkOS::OAuthTokens.new(json[:oauth_tokens].to_json) : nil
       @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],)
+          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],
+            encryptor: session[:encryptor],
+          )
         end
     end
     # rubocop:enable Metrics/AbcSize

lib/workos/encryptors.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module WorkOS
+  # Encryptors module provides pluggable encryption implementations for session data.
+  # The default encryptor is AesGcm, which uses AES-256-GCM encryption.
+  module Encryptors
+    autoload :AesGcm, 'workos/encryptors/aes_gcm'
+  end
+end

lib/workos/encryptors/aes_gcm.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require 'encryptor'
+require 'securerandom'
+require 'json'
+require 'base64'
+
+module WorkOS
+  module Encryptors
+    # Default encryptor using AES-256-GCM.
+    # Implements the encryptor interface: #seal(data, key) and #unseal(sealed_data, key)
+    class AesGcm
+      # Encrypts and seals data using AES-256-GCM
+      # @param data [Hash] The data to seal
+      # @param key [String] The encryption key
+      # @return [String] Base64-encoded sealed data
+      def seal(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)
+      end
+
+      # Decrypts and unseals data using AES-256-GCM
+      # @param sealed_data [String] The sealed data to unseal
+      # @param key [String] The decryption key
+      # @return [Hash] The unsealed data with symbolized keys
+      def unseal(sealed_data, key)
+        decoded_data = Base64.decode64(sealed_data)
+        iv = decoded_data[0..11]
+        encrypted_data = decoded_data[12..]
+
+        decrypted_data = Encryptor.decrypt(
+          value: encrypted_data,
+          key: key,
+          iv: iv,
+          algorithm: 'aes-256-gcm',
+        )
+
+        JSON.parse(decrypted_data, symbolize_names: true)
+      end
+    end
+  end
+end

lib/workos/refresh_authentication_response.rb

@@ -22,13 +22,17 @@ def initialize(authentication_response_json, session = nil)
         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],)
+          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],
+            encryptor: session[:encryptor],
+          )
         end
     end
     # rubocop:enable Metrics/AbcSize

lib/workos/session.rb

@@ -12,11 +12,14 @@ 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.
   class Session
-    attr_accessor :jwks, :jwks_algorithms, :user_management, :cookie_password, :session_data, :client_id
+    attr_accessor :jwks, :jwks_algorithms, :user_management, :cookie_password, :session_data, :client_id, :encryptor
 
-    def initialize(user_management:, client_id:, session_data:, cookie_password:)
+    def initialize(user_management:, client_id:, session_data:, cookie_password:, encryptor: nil)
       raise ArgumentError, 'cookiePassword is required' if cookie_password.nil? || cookie_password.empty?
 
+      @encryptor = encryptor || WorkOS::Encryptors::AesGcm.new
+      validate_encryptor!(@encryptor)
+
       @user_management = user_management
       @cookie_password = cookie_password
       @session_data = session_data
@@ -37,7 +40,7 @@ def authenticate(include_expired: false, &claim_extractor)
       return { authenticated: false, reason: 'NO_SESSION_COOKIE_PROVIDED' } if @session_data.nil?
 
       begin
-        session = Session.unseal_data(@session_data, @cookie_password)
+        session = Session.unseal_data(@session_data, @cookie_password, encryptor: @encryptor)
       rescue StandardError
         return { authenticated: false, reason: 'INVALID_SESSION_COOKIE' }
       end
@@ -92,7 +95,7 @@ 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)
+        session = Session.unseal_data(@session_data, cookie_password, encryptor: @encryptor)
       rescue StandardError
         return { authenticated: false, reason: 'INVALID_SESSION_COOKIE' }
       end
@@ -104,7 +107,7 @@ def refresh(options = nil)
           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: { seal_session: true, cookie_password: cookie_password, encryptor: @encryptor },
         )
 
         @session_data = auth_response.sealed_session
@@ -137,43 +140,34 @@ def get_logout_url(return_to: nil)
       @user_management.get_logout_url(session_id: auth_response[:session_id], return_to: return_to)
     end
 
-    # Encrypts and seals data using AES-256-GCM
+    # Encrypts and seals data using the provided encryptor (defaults to AES-256-GCM)
     # @param data [Hash] The data to seal
     # @param key [String] The key to use for encryption
+    # @param encryptor [Object] Optional encryptor that responds to #seal(data, key)
     # @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
+    def self.seal_data(data, key, encryptor: nil)
+      enc = encryptor || WorkOS::Encryptors::AesGcm.new
+      enc.seal(data, key)
     end
 
-    # Decrypts and unseals data using AES-256-GCM
+    # Decrypts and unseals data using the provided encryptor (defaults to AES-256-GCM)
     # @param sealed_data [String] The sealed data to unseal
     # @param key [String] The key to use for decryption
+    # @param encryptor [Object] Optional encryptor that responds to #unseal(sealed_data, key)
     # @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
+    def self.unseal_data(sealed_data, key, encryptor: nil)
+      enc = encryptor || WorkOS::Encryptors::AesGcm.new
+      enc.unseal(sealed_data, key)
     end
 
     private
 
+    def validate_encryptor!(enc)
+      return if enc.respond_to?(:seal) && enc.respond_to?(:unseal)
+
+      raise ArgumentError, 'encryptor must respond to #seal(data, key) and #unseal(sealed_data, key)'
+    end
+
     # Creates a JWKS set from a remote JWKS URL
     # @param uri [URI] The URI of the JWKS
     # @return [JWT::JWK::Set] The JWKS set

lib/workos/user_management.rb

@@ -42,14 +42,16 @@ class << self
       # @param [String] client_id The WorkOS client ID for the environment
       # @param [String] session_data The sealed session data
       # @param [String] cookie_password The password used to seal the session
+      # @param [Object] encryptor Optional custom encryptor that responds to #seal and #unseal
       #
       # @return WorkOS::Session
-      def load_sealed_session(client_id:, session_data:, cookie_password:)
+      def load_sealed_session(client_id:, session_data:, cookie_password:, encryptor: nil)
         WorkOS::Session.new(
           user_management: self,
           client_id: client_id,
           session_data: session_data,
           cookie_password: cookie_password,
+          encryptor: encryptor,
         )
       end
 

spec/lib/workos/encryptors/aes_gcm_spec.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+RSpec.describe WorkOS::Encryptors::AesGcm do
+  subject(:encryptor) { described_class.new }
+
+  let(:key) { 'a' * 32 }
+  let(:data) { { access_token: 'tok_123', user: { id: 'user_01' } } }
+
+  describe '#seal' do
+    it 'returns a base64-encoded string' do
+      sealed = encryptor.seal(data, key)
+      expect(sealed).to be_a(String)
+      expect { Base64.decode64(sealed) }.not_to raise_error
+    end
+
+    it 'produces different output each time (random IV)' do
+      sealed1 = encryptor.seal(data, key)
+      sealed2 = encryptor.seal(data, key)
+      expect(sealed1).not_to eq(sealed2)
+    end
+  end
+
+  describe '#unseal' do
+    it 'round-trips data correctly' do
+      sealed = encryptor.seal(data, key)
+      unsealed = encryptor.unseal(sealed, key)
+      expect(unsealed).to eq(data)
+    end
+
+    it 'returns hash with symbolized keys' do
+      sealed = encryptor.seal({ 'string_key' => 'value' }, key)
+      unsealed = encryptor.unseal(sealed, key)
+      expect(unsealed.keys.first).to be_a(Symbol)
+    end
+
+    it 'raises error with wrong key' do
+      sealed = encryptor.seal(data, key)
+      expect { encryptor.unseal(sealed, 'b' * 32) }.to raise_error(OpenSSL::Cipher::CipherError)
+    end
+  end
+end

spec/lib/workos/session_spec.rb

@@ -408,4 +408,68 @@
       end
     end
   end
+
+  describe 'custom encryptor' do
+    let(:user_management) { instance_double('UserManagement') }
+    let(:custom_encryptor) do
+      Class.new do
+        def seal(data, _key)
+          "CUSTOM:#{JSON.generate(data)}"
+        end
+
+        def unseal(sealed_data, _key)
+          json = sealed_data.sub('CUSTOM:', '')
+          JSON.parse(json, symbolize_names: true)
+        end
+      end.new
+    end
+
+    before do
+      allow(user_management).to receive(:get_jwks_url).with(client_id).and_return(jwks_url)
+    end
+
+    it 'uses custom encryptor for seal_data' do
+      sealed = WorkOS::Session.seal_data({ foo: 'bar' }, 'key', encryptor: custom_encryptor)
+      expect(sealed).to start_with('CUSTOM:')
+    end
+
+    it 'uses custom encryptor for unseal_data' do
+      sealed = 'CUSTOM:{"foo":"bar"}'
+      unsealed = WorkOS::Session.unseal_data(sealed, 'key', encryptor: custom_encryptor)
+      expect(unsealed).to eq({ foo: 'bar' })
+    end
+
+    it 'accepts custom encryptor in initialize' do
+      session = WorkOS::Session.new(
+        user_management: user_management,
+        client_id: client_id,
+        session_data: session_data,
+        cookie_password: cookie_password,
+        encryptor: custom_encryptor,
+      )
+      expect(session.encryptor).to eq(custom_encryptor)
+    end
+
+    it 'defaults to AesGcm encryptor when none provided' do
+      session = WorkOS::Session.new(
+        user_management: user_management,
+        client_id: client_id,
+        session_data: session_data,
+        cookie_password: cookie_password,
+      )
+      expect(session.encryptor).to be_a(WorkOS::Encryptors::AesGcm)
+    end
+
+    it 'raises ArgumentError for invalid encryptor' do
+      expect do
+        WorkOS::Session.new(
+          user_management: user_management,
+          client_id: client_id,
+          session_data: session_data,
+          cookie_password: cookie_password,
+          encryptor: Object.new,
+        )
+      end.to raise_error(ArgumentError, /must respond to/)
+    end
+  end
 end