Add entitlements to session (#349)

* Refactor session spec to mock less and split out the payload

Allowing more complete testing of the token with only the verification being mocked.

The separate payload enables creating more tests with differing payloads.

* Include entitlements in the authenticated session

If the users organisation has a stripe link the token includes any entitlements, this should be exposed to client code.

Fixes #336

Author: adam-h

lib/workos/session.rb

@@ -50,6 +50,7 @@ def authenticate
         organization_id: decoded['org_id'],
         role: decoded['role'],
         permissions: decoded['permissions'],
+        entitlements: decoded['entitlements'],
         user: session[:user],
         impersonator: session[:impersonator],
         reason: nil,

spec/lib/workos/session_spec.rb

@@ -103,24 +103,23 @@
 
   describe '.authenticate' do
     let(:user_management) { instance_double('UserManagement') }
-    let(:valid_access_token) do
-      payload = {
+    let(:payload) do
+      {
         sid: 'session_id',
         org_id: 'org_id',
         role: 'role',
         permissions: ['read'],
         exp: Time.now.to_i + 3600,
       }
-      headers = { kid: jwk[:kid] }
-      JWT.encode(payload, jwk.signing_key, jwk[:alg], headers)
     end
+    let(:valid_access_token) { JWT.encode(payload, jwk.signing_key, jwk[:alg], { kid: jwk[:kid] }) }
     let(:session_data) do
-  WorkOS::Session.seal_data({
-                              access_token: valid_access_token,
-                              user: 'user',
-                              impersonator: 'impersonator',
-                            }, cookie_password,)
-end
+      WorkOS::Session.seal_data({
+                                  access_token: valid_access_token,
+                                  user: 'user',
+                                  impersonator: 'impersonator',
+                                }, cookie_password,)
+    end
 
     before do
       allow(user_management).to receive(:get_jwks_url).with(client_id).and_return(jwks_url)
@@ -167,26 +166,55 @@
         session_data: session_data,
         cookie_password: cookie_password,
       )
-      allow(session).to receive(:is_valid_jwt).and_return(true)
-      allow(JWT).to receive(:decode).and_return([{
-                                                  'sid' => 'session_id',
-                                                  'org_id' => 'org_id',
-                                                  'role' => 'role',
-                                                  'permissions' => ['read'],
-                                                }])
-
+      allow_any_instance_of(JWT::Decode).to receive(:verify_signature).and_return(true)
       result = session.authenticate
       expect(result).to eq({
                              authenticated: true,
                              session_id: 'session_id',
                              organization_id: 'org_id',
                              role: 'role',
                              permissions: ['read'],
+                             entitlements: nil,
                              user: 'user',
                              impersonator: 'impersonator',
                              reason: nil,
                            })
     end
+
+    describe 'with entitlements' do
+      let(:payload) do
+        {
+          sid: 'session_id',
+          org_id: 'org_id',
+          role: 'role',
+          permissions: ['read'],
+          entitlements: ['billing'],
+          exp: Time.now.to_i + 3600,
+        }
+      end
+
+      it 'includes entitlements in the result' do
+        session = WorkOS::Session.new(
+          user_management: user_management,
+          client_id: client_id,
+          session_data: session_data,
+          cookie_password: cookie_password,
+        )
+        allow_any_instance_of(JWT::Decode).to receive(:verify_signature).and_return(true)
+        result = session.authenticate
+        expect(result).to eq({
+                               authenticated: true,
+                               session_id: 'session_id',
+                               organization_id: 'org_id',
+                               role: 'role',
+                               permissions: ['read'],
+                               entitlements: ['billing'],
+                               user: 'user',
+                               impersonator: 'impersonator',
+                               reason: nil,
+                             })
+      end
+    end
   end
 
   describe '.refresh' do