Handle token in cookie, fix refresh session bug. (#76)

* Handle token in cookie, fix refresh session bug.

* fix order

* fix tests

* fix lint

* fix tests

* move managemnt key check to management tests

lib/descope/api/v1/auth.rb

@@ -34,7 +34,6 @@ def generate_jwt_response(response_body: nil, refresh_cookie: nil, audience: nil
           end
 
           jwt_response = generate_auth_info(response_body, refresh_cookie, true, audience)
-          @logger.debug "jwt_response: #{jwt_response}"
           jwt_response['user'] = response_body.key?('user') ? response_body['user'] : {}
           jwt_response['firstSeen'] = response_body.key?('firstSeen') ? response_body['firstSeen'] : true
 
@@ -62,6 +61,8 @@ def select_tenant(tenant_id: nil, refresh_token: nil)
           validate_refresh_token_not_nil(refresh_token)
           res = post(SELECT_TENANT_PATH, { tenantId: tenant_id }, {}, refresh_token)
           @logger.debug "select_tenant response: #{res}"
+          cookies = res.fetch('cookies')
+          generate_jwt_response(response_body: res, refresh_cookie: cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil))
           generate_jwt_response(
             response_body: res,
             refresh_cookie: res['refreshJwt']
@@ -231,24 +232,41 @@ def validate_token(token, _audience = nil)
         private
 
         def generate_auth_info(response_body, refresh_token, user_jwt, audience = nil)
-          @logger.debug "generating auth info: #{response_body}, #{refresh_token}, #{user_jwt}, #{audience}"
+          @logger.debug "generating auth info: response_body: #{response_body}, refresh_token: #{refresh_token}, user_jwt: #{user_jwt}, audience: #{audience}"
           jwt_response = {}
 
           # validate the session token if sessionJwt is not empty
           st_jwt = response_body.fetch('sessionJwt', '')
           unless st_jwt.empty?
-            @logger.debug "validating session token with refresh_token: #{refresh_token}" if st_jwt
-            jwt_response[SESSION_TOKEN_NAME] = validate_token(st_jwt, audience) if st_jwt
+            @logger.debug 'found sessionJwt in response body, adding to jwt_response'
+            jwt_response[SESSION_TOKEN_NAME] = validate_token(st_jwt, audience)
           end
 
           # validate refresh token if refresh_token was passed or if refreshJwt is not empty
           rt_jwt = response_body.fetch('refreshJwt', '')
 
-          if !refresh_token.nil? || !refresh_token.to_s.empty?
-            @logger.debug "validating refresh token: #{refresh_token}" if refresh_token
-            jwt_response[REFRESH_SESSION_TOKEN_NAME] = validate_token(refresh_token, audience)
-          elsif !rt_jwt.empty?
+          if !rt_jwt.empty?
+            @logger.debug 'found refreshJwt in response body, adding to jwt_response'
+            @logger.debug 'validating refreshJwt token...'
             jwt_response[REFRESH_SESSION_TOKEN_NAME] = validate_token(rt_jwt, audience)
+          elsif refresh_token && !refresh_token.empty?
+            # if refresh_token is in response body (local storage)
+            @logger.debug 'refreshJwt is empty, but refresh_token was passed, adding to jwt_response'
+            @logger.debug 'validating passed-in refresh token...'
+            jwt_response[REFRESH_SESSION_TOKEN_NAME] = validate_token(refresh_token, audience)
+          else
+            cookies = response_body.fetch('cookies', {})
+            # else if refresh token is in response cookie
+            cookies.each do |cookie_name, cookie_value|
+              if cookie_name == REFRESH_SESSION_COOKIE_NAME
+                jwt_response[REFRESH_SESSION_TOKEN_NAME] = validate_token(cookie_value, audience)
+              end
+            end
+          end
+
+          if jwt_response[REFRESH_SESSION_TOKEN_NAME].nil?
+            @logger.debug "Error: Could not find refreshJwt in response body: #{response_body} / cookies: #{cookies} / passed in refresh_token ->#{refresh_token}<-"
+            raise Descope::AuthException.new('Could not find refreshJwt in response body / cookies / passed in refresh_token', code: 500)
           end
 
           jwt_response = adjust_properties(jwt_response, user_jwt)

lib/descope/api/v1/auth/enchantedlink.rb

@@ -67,7 +67,9 @@ def enchanted_link_verify_token(token = nil)
           def enchanted_link_get_session(pending_ref = nil)
             # @see https://docs.descope.com/api/openapi/enchantedlink/operation/GetEnchantedLinkSession/
             res = post(GET_SESSION_ENCHANTEDLINK_AUTH_PATH, { pendingRef: pending_ref })
-            generate_jwt_response(response_body: res, refresh_cookie: res['refreshJwt'])
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           private

lib/descope/api/v1/auth/magiclink.rb

@@ -42,7 +42,9 @@ def magiclink_sign_up_or_in(method: nil, login_id: nil, uri: nil, login_options:
           def magiclink_verify_token(token = nil)
             validate_token_not_empty(token)
             res = post(VERIFY_MAGICLINK_AUTH_PATH, { token: })
-            generate_jwt_response(response_body: res)
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def magiclink_update_user_email(login_id: nil, email: nil, uri: nil, add_to_login_ids: nil, on_merge_use_existing: nil, provider_id: nil, template_id: nil, template_options: nil, refresh_token: nil)

lib/descope/api/v1/auth/otp.rb

@@ -63,7 +63,9 @@ def otp_verify_code(method: nil, login_id: nil, code: nil)
               code:
             }
             res = post(uri, request_params)
-            generate_jwt_response(response_body: res, refresh_cookie: res.fetch('refreshJwt', {}))
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def otp_update_user_email(login_id: nil, email: nil, refresh_token: nil, add_to_login_ids: false,

lib/descope/api/v1/auth/password.rb

@@ -22,7 +22,9 @@ def password_sign_up(login_id: nil, password: nil, user: nil)
 
             request_params[:user] = password_user_compose_update_body(**user) unless user.nil?
             res = post(SIGN_UP_PASSWORD_PATH, request_params)
-            generate_jwt_response(response_body: res)
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def password_sign_in(login_id: nil, password: nil, sso_app_id: nil)
@@ -38,7 +40,9 @@ def password_sign_in(login_id: nil, password: nil, sso_app_id: nil)
               ssoAppId: sso_app_id
             }
             res = post(SIGN_IN_PASSWORD_PATH, request_params)
-            generate_jwt_response(response_body: res)
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def password_replace(login_id: nil, old_password: nil, new_password: nil)

lib/descope/api/v1/auth/totp.rb

@@ -17,7 +17,9 @@ def totp_sign_in_code(login_id: nil, login_options: nil, code: nil)
             uri = VERIFY_TOTP_PATH
             body = totp_compose_signin_body(login_id, code, login_options)
             res = post(uri, body, {}, nil)
-            generate_jwt_response(response_body: res, refresh_cookie: res.fetch('refreshJwt', {}))
+            cookies = res.fetch(COOKIE_DATA_NAME, {})
+            refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+            generate_jwt_response(response_body: res, refresh_cookie:)
           end
 
           def totp_sign_up(login_id: nil, user: nil, sso_app_id: nil)

lib/descope/api/v1/session.rb

@@ -20,11 +20,13 @@ def refresh_session(refresh_token: nil, audience: nil)
           #  [amr, drn, exp, iss, rexp, sub, jwt] in the top level of the response dict, please use
           #  them from the sessionToken key instead, as these claims will soon be deprecated from the top level
           #  of the response dict.
-
+          #  Make sure you set Enable refresh token rotation in the Project Settings before using this.
           validate_refresh_token_not_nil(refresh_token)
           validate_token(refresh_token, audience)
           res = post(REFRESH_TOKEN_PATH, {}, {}, refresh_token)
-          generate_jwt_response(response_body: res, refresh_cookie: refresh_token, audience:)
+          cookies = res.fetch(COOKIE_DATA_NAME, {})
+          refresh_cookie = cookies.fetch(REFRESH_SESSION_COOKIE_NAME, nil) || res.fetch('refreshJwt', nil)
+          generate_jwt_response(response_body: res, refresh_cookie:, audience:)
         end
 
         def me(refresh_token = nil)

lib/descope/mixins/http.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-require "addressable/uri"
+require 'addressable/uri'
 require 'retryable'
 require_relative '../exception'
 
@@ -44,9 +44,17 @@ def retry_options
         }
       end
 
-      def safe_parse_json(body)
+      def safe_parse_json(body, cookies: {})
         @logger.debug "response => #{JSON.parse(body.to_s)}"
-        JSON.parse(body.to_s)
+        res = JSON.parse(body.to_s)
+
+        # Handle DSR cookie in response.
+        if cookies.key?(REFRESH_SESSION_COOKIE_NAME)
+          res['cookies'] = {}
+          res['cookies'][REFRESH_SESSION_COOKIE_NAME] = cookies[REFRESH_SESSION_COOKIE_NAME]
+        end
+
+        res
       rescue JSON::ParserError
         body
       end
@@ -94,11 +102,12 @@ def request(method, uri, body = {}, extra_headers = {})
                    call(method, encode_uri(uri), timeout, @headers, body.to_json)
                  end
 
-        raise Descope::Unsupported.new("No response from server", code: 400) unless result && result.respond_to?(:code)
+        raise Descope::Unsupported.new('No response from server', code: 400) unless result.respond_to?(:code)
 
         @logger.info("API Request: [#{method}] #{uri} - Response Code: #{result.code}")
+
         case result.code
-        when 200...226 then safe_parse_json(result.body)
+        when 200...226 then safe_parse_json(result.body, cookies: result.cookies)
         when 400       then raise Descope::BadRequest.new(result.body, code: result.code, headers: result.headers)
         when 401       then raise Descope::Unauthorized.new(result.body, code: result.code, headers: result.headers)
         when 403       then raise Descope::AccessDenied.new(result.body, code: result.code, headers: result.headers)

spec/integration/lib.descope/api/v1/auth/session_spec.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Descope::Api::V1::Session do
+  before(:all) do
+    @client = DescopeClient.new(Configuration.config)
+  end
+
+  after(:all) do
+    @client.logger.info('Cleaning up test users...')
+    all_users = @client.search_all_users
+    all_users['users'].each do |user|
+      if user['middleName'] == 'Ruby SDK User'
+        @client.logger.info("Deleting ruby spec test user #{user['loginIds'][0]}")
+        @client.delete_user(user['loginIds'][0])
+      end
+    end
+  end
+
+  context 'test session methods' do
+    it 'should refresh session with refresh token' do
+      @password = SpecUtils.generate_password
+      user = build(:user)
+
+      @client.logger.info('1. Sign up with password')
+      res = @client.password_sign_up(login_id: user[:login_id], password: @password, user:)
+      @client.logger.info("sign up with password res: #{res}")
+      original_refresh_token = res[REFRESH_SESSION_TOKEN_NAME]['jwt']
+
+      @client.logger.info('2. Sign in with password')
+      login_res = @client.password_sign_in(login_id: user[:login_id], password: @password)
+      @client.logger.info("sign_in res: #{login_res}")
+
+      @client.logger.info('3. sleep 1 second before calling refresh_session')
+      sleep(1)
+
+      @client.logger.info('4. Refresh session')
+      refresh_session_res = @client.refresh_session(refresh_token: login_res[REFRESH_SESSION_TOKEN_NAME]['jwt'])
+      @client.logger.info("refresh_session_res: #{refresh_session_res}")
+
+      new_refresh_token = refresh_session_res[REFRESH_SESSION_TOKEN_NAME]['jwt']
+      @client.logger.info("new_refresh_token: #{new_refresh_token}")
+
+      @client.logger.info('5. Check new refresh token is not the same as the original one')
+      expect(original_refresh_token).not_to eq(new_refresh_token)
+    end
+  end
+end

spec/integration/lib.descope/api/v1/management/access_key_spec.rb

@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::AccessKey do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
   end
 
@@ -28,6 +30,7 @@
       @tenant_id = @client.create_tenant(name: 'some-new-tenant')['id']
       @client.logger.info('creating access key')
       @access_key = @client.create_access_key(name: @key_name, key_tenants: [{ tenant_id: @tenant_id }])
+      @client.logger.info("waiting for access key #{@access_key['key']['id']} to be active 60 seconds")
       sleep 60
     end
 

spec/integration/lib.descope/api/v1/management/audit_spec.rb

@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Audit do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     @client.logger.info('Deleting all tenants for Ruby SDK...')
     @client.search_all_tenants(names: ['Ruby-SDK-test'])['tenants'].each do |tenant|
@@ -39,15 +41,14 @@
     created_user = @client.create_user(**user)['user']
 
     expect do
-      res = @client.audit_create_event(
+      @client.audit_create_event(
         user_id: created_user['loginId'],
         action: 'pencil.created',
         type: 'info',
         tenant_id:,
         actor_id: created_user['loginIds'][0],
         data: { 'key' => 'value' }
       )
-      expect(res).to eq({})
     end.not_to raise_error
   end
 end

spec/integration/lib.descope/api/v1/management/authz_spec.rb

@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Authz do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     puts 'authz schema delete'
   end

spec/integration/lib.descope/api/v1/management/flow_spec.rb

@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Flow do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
   end
 

spec/integration/lib.descope/api/v1/management/permissions_spec.rb

@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Permission do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     @client.load_all_permissions['permissions'].each do |perm|
       if perm['description'] == 'Ruby SDK'

spec/integration/lib.descope/api/v1/management/project_spec.rb

@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Project do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     @export_output = @client.export_project
   end

spec/integration/lib.descope/api/v1/management/roles_spec.rb

@@ -4,6 +4,8 @@
 
 describe Descope::Api::V1::Management::Role do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
     @client = DescopeClient.new(Configuration.config)
     @client.logger.info('Staring cleanup before tests...')
     @client.logger.info('Deleting all permissions for Ruby SDK...')

spec/integration/lib.descope/api/v1/management/user_spec.rb

@@ -4,10 +4,14 @@
 
 describe Descope::Api::V1::Management::User do
   before(:all) do
+    raise 'DESCOPE_MANAGEMENT_KEY is not set' if ENV['DESCOPE_MANAGEMENT_KEY'].nil?
+
+    @client = DescopeClient.new(Configuration.config)
+
     @password = SpecUtils.generate_password
     @new_password = SpecUtils.generate_password
     @user = build(:user)
-    @client = DescopeClient.new(Configuration.config)
+
     include Descope::Mixins::Common::DeliveryMethod
   end
 
@@ -226,8 +230,8 @@
       new_password = SpecUtils.generate_password
       @client.set_password(login_id: user['loginIds'][0], password: new_password)
       @client.password_sign_in(login_id: user['loginIds'][0], password:)
-    rescue Descope::Unauthorized => e
-      expect(e.message).to match(/"errorDescription":"Invalid signin credentials"/)
+    rescue Descope::ServerError => e
+      expect(e.message).to match(/"Password expired"/)
     end
   end
 

spec/lib.descope/api/v1/auth/enchantedlink_spec.rb

@@ -51,7 +51,7 @@
 
     it 'is expected to validate refresh token and not raise an error with refresh token and valid login options' do
       expect do
-        @instance.send(:validate_refresh_token_provided, { mfa: true, stepup: true },  'some-token')
+        @instance.send(:validate_refresh_token_provided, { mfa: true, stepup: true }, 'some-token')
       end.not_to raise_error
     end
 
@@ -148,7 +148,16 @@
     end
 
     it 'is expected to get session by pending ref with enchanted link' do
-      jwt_response = { 'fake': 'response' }
+      jwt_response = {
+        'sessionJwt' => 'fake_session_jwt',
+        'refreshJwt' => 'fake_refresh_jwt',
+        'cookies' => {
+          'refresh_token' => 'fake_refresh_cookie'
+        }
+      }
+      allow(@instance).to receive(:post).with(
+        GET_SESSION_ENCHANTEDLINK_AUTH_PATH, { pendingRef: 'pendingRef' }
+      ).and_return(jwt_response)
       allow(@instance).to receive(:generate_jwt_response).and_return(jwt_response)
 
       expect do