From ed7ceedef51d8aa137e64f94ec92ff4f47a69e06 Mon Sep 17 00:00:00 2001
From: Dan Milne <d@nmilne.com>
Date: Wed, 31 Dec 2025 14:45:38 +1100
Subject: [PATCH] Include the hash of the access token in the JWT / ID Token
 under the key at_hash as per the requirements.  Update the discovery endpoint
 to describe subject_type as 'pairwise', rather than 'public', since we do
 pairwise subject ids.

---
 app/controllers/oidc_controller.rb     | 21 ++++++++++++++++-----
 app/services/oidc_jwt_service.rb       | 10 +++++++++-
 test/services/oidc_jwt_service_test.rb | 19 +++++++++++++++++++
 3 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/app/controllers/oidc_controller.rb b/app/controllers/oidc_controller.rb
index 0b46ae0..a202b5b 100644
--- a/app/controllers/oidc_controller.rb
+++ b/app/controllers/oidc_controller.rb
@@ -26,7 +26,7 @@ class OidcController < ApplicationController
       response_types_supported: ["code"],
       response_modes_supported: ["query"],
       grant_types_supported: ["authorization_code", "refresh_token"],
-      subject_types_supported: ["public"],
+      subject_types_supported: ["pairwise"],
       id_token_signing_alg_values_supported: ["RS256"],
       scopes_supported: ["openid", "profile", "email", "groups", "offline_access"],
       token_endpoint_auth_methods_supported: ["client_secret_post", "client_secret_basic"],
@@ -422,8 +422,14 @@ class OidcController < ApplicationController
           return
         end
 
-        # Generate ID token (JWT) with pairwise SID
-        id_token = OidcJwtService.generate_id_token(user, application, consent: consent, nonce: auth_code.nonce)
+        # Generate ID token (JWT) with pairwise SID and at_hash
+        id_token = OidcJwtService.generate_id_token(
+          user,
+          application,
+          consent: consent,
+          nonce: auth_code.nonce,
+          access_token: access_token_record.plaintext_token
+        )
 
         # Return tokens
         render json: {
@@ -539,8 +545,13 @@ class OidcController < ApplicationController
       return
     end
 
-    # Generate new ID token (JWT with pairwise SID, no nonce for refresh grants)
-    id_token = OidcJwtService.generate_id_token(user, application, consent: consent)
+    # Generate new ID token (JWT with pairwise SID and at_hash, no nonce for refresh grants)
+    id_token = OidcJwtService.generate_id_token(
+      user,
+      application,
+      consent: consent,
+      access_token: new_access_token.plaintext_token
+    )
 
     # Return new tokens
     render json: {
diff --git a/app/services/oidc_jwt_service.rb b/app/services/oidc_jwt_service.rb
index 991da42..71eb8aa 100644
--- a/app/services/oidc_jwt_service.rb
+++ b/app/services/oidc_jwt_service.rb
@@ -3,7 +3,7 @@ class OidcJwtService
 
   class << self
     # Generate an ID token (JWT) for the user
-    def generate_id_token(user, application, consent: nil, nonce: nil)
+    def generate_id_token(user, application, consent: nil, nonce: nil, access_token: nil)
       now = Time.current.to_i
       # Use application's configured ID token TTL (defaults to 1 hour)
       ttl = application.id_token_expiry_seconds
@@ -26,6 +26,14 @@ class OidcJwtService
       # Add nonce if provided (OIDC requires this for implicit flow)
       payload[:nonce] = nonce if nonce.present?
 
+      # Add at_hash if access token is provided (OIDC Core spec §3.1.3.6)
+      # at_hash = left-most 128 bits of SHA-256 hash of access token, base64url encoded
+      if access_token.present?
+        sha256 = Digest::SHA256.digest(access_token)
+        at_hash = Base64.urlsafe_encode64(sha256[0..15], padding: false)
+        payload[:at_hash] = at_hash
+      end
+
       # Add groups if user has any
       if user.groups.any?
         payload[:groups] = user.groups.pluck(:name)
diff --git a/test/services/oidc_jwt_service_test.rb b/test/services/oidc_jwt_service_test.rb
index 32cc0b4..499a96b 100644
--- a/test/services/oidc_jwt_service_test.rb
+++ b/test/services/oidc_jwt_service_test.rb
@@ -476,4 +476,23 @@ class OidcJwtServiceTest < ActiveSupport::TestCase
     assert_includes decoded["roles"], "moderator"
     assert_includes decoded["roles"], "app_admin"
   end
+
+  test "should include at_hash when access token is provided" do
+    access_token = "test-access-token-abc123xyz"
+    token = @service.generate_id_token(@user, @application, access_token: access_token)
+
+    decoded = JWT.decode(token, nil, false).first
+    assert_includes decoded.keys, "at_hash", "Should include at_hash claim"
+
+    # Verify at_hash is correctly computed: base64url(sha256(access_token)[0:16])
+    expected_hash = Base64.urlsafe_encode64(Digest::SHA256.digest(access_token)[0..15], padding: false)
+    assert_equal expected_hash, decoded["at_hash"], "at_hash should match SHA-256 hash of access token"
+  end
+
+  test "should not include at_hash when access token is not provided" do
+    token = @service.generate_id_token(@user, @application)
+
+    decoded = JWT.decode(token, nil, false).first
+    refute_includes decoded.keys, "at_hash", "Should not include at_hash when no access token"
+  end
 end
\ No newline at end of file