probo/pkg/iam/oauth2/id_token_test.go

// Copyright (c) 2026 Probo Inc <hello@probo.com>.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package oauth2_test

import (
	"crypto/sha256"
	"encoding/base64"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.probo.inc/probo/pkg/coredata"
	"go.probo.inc/probo/pkg/gid"
	"go.probo.inc/probo/pkg/iam/oauth2"
	"go.probo.inc/probo/pkg/uri"
)

var testIssuer = uri.URI("https://issuer.example.com")

func TestComputeAtHash(t *testing.T) {
	t.Parallel()

	t.Run(
		"returns left half of sha256 base64url encoded",
		func(t *testing.T) {
			t.Parallel()

			accessToken := "ya29.test-access-token"
			h := sha256.Sum256([]byte(accessToken))
			expected := base64.RawURLEncoding.EncodeToString(h[:16])

			result := oauth2.ComputeAtHash(accessToken)
			assert.Equal(t, expected, result)
		},
	)

	t.Run(
		"different tokens produce different hashes",
		func(t *testing.T) {
			t.Parallel()

			hash1 := oauth2.ComputeAtHash("token-a")
			hash2 := oauth2.ComputeAtHash("token-b")
			assert.NotEqual(t, hash1, hash2)
		},
	)

	t.Run(
		"empty token",
		func(t *testing.T) {
			t.Parallel()

			result := oauth2.ComputeAtHash("")
			assert.NotEmpty(t, result)
		},
	)

	t.Run(
		"deterministic",
		func(t *testing.T) {
			t.Parallel()

			hash1 := oauth2.ComputeAtHash("same-token")
			hash2 := oauth2.ComputeAtHash("same-token")
			assert.Equal(t, hash1, hash2)
		},
	)
}

func TestNewIDTokenClaims(t *testing.T) {
	t.Parallel()

	identityID := gid.Nil
	clientID := gid.Nil
	authTime := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)

	t.Run(
		"basic claims without optional scopes",
		func(t *testing.T) {
			t.Parallel()

			claims := oauth2.NewIDTokenClaims(
				testIssuer,
				identityID,
				clientID,
				authTime,
				coredata.OAuth2Scopes{oauth2.ScopeOpenID},
				"",
				"",
				"user@example.com",
				true,
				"John Doe",
				1*time.Hour,
			)

			assert.Equal(t, testIssuer, claims.Issuer)
			assert.Equal(t, identityID.String(), claims.Subject)
			assert.Equal(t, clientID.String(), claims.Audience)
			assert.Equal(t, authTime.Unix(), claims.AuthTime)
			assert.Empty(t, claims.Nonce)
			assert.Empty(t, claims.AtHash)
			assert.Empty(t, claims.Email)
			assert.Nil(t, claims.EmailVerified)
			assert.Empty(t, claims.Name)
		},
	)

	t.Run(
		"sets nonce when provided",
		func(t *testing.T) {
			t.Parallel()

			claims := oauth2.NewIDTokenClaims(
				testIssuer,
				identityID,
				clientID,
				authTime,
				coredata.OAuth2Scopes{oauth2.ScopeOpenID},
				"test-nonce",
				"",
				"",
				false,
				"",
				1*time.Hour,
			)

			assert.Equal(t, "test-nonce", claims.Nonce)
		},
	)

	t.Run(
		"computes at_hash when access token provided",
		func(t *testing.T) {
			t.Parallel()

			claims := oauth2.NewIDTokenClaims(
				testIssuer,
				identityID,
				clientID,
				authTime,
				coredata.OAuth2Scopes{oauth2.ScopeOpenID},
				"",
				"access-token-123",
				"",
				false,
				"",
				1*time.Hour,
			)

			expected := oauth2.ComputeAtHash("access-token-123")
			assert.Equal(t, expected, claims.AtHash)
		},
	)

	t.Run(
		"includes email claims with email scope",
		func(t *testing.T) {
			t.Parallel()

			claims := oauth2.NewIDTokenClaims(
				testIssuer,
				identityID,
				clientID,
				authTime,
				coredata.OAuth2Scopes{oauth2.ScopeOpenID, oauth2.ScopeEmail},
				"",
				"",
				"user@example.com",
				true,
				"",
				1*time.Hour,
			)

			assert.Equal(t, "user@example.com", claims.Email)
			require.NotNil(t, claims.EmailVerified)
			assert.True(t, *claims.EmailVerified)
		},
	)

	t.Run(
		"includes name with profile scope",
		func(t *testing.T) {
			t.Parallel()

			claims := oauth2.NewIDTokenClaims(
				testIssuer,
				identityID,
				clientID,
				authTime,
				coredata.OAuth2Scopes{oauth2.ScopeOpenID, oauth2.ScopeProfile},
				"",
				"",
				"",
				false,
				"Jane Doe",
				1*time.Hour,
			)

			assert.Equal(t, "Jane Doe", claims.Name)
		},
	)

	t.Run(
		"includes all claims with all scopes",
		func(t *testing.T) {
			t.Parallel()

			claims := oauth2.NewIDTokenClaims(
				testIssuer,
				identityID,
				clientID,
				authTime,
				coredata.OAuth2Scopes{
					oauth2.ScopeOpenID,
					oauth2.ScopeEmail,
					oauth2.ScopeProfile,
				},
				"nonce-val",
				"access-token",
				"user@example.com",
				false,
				"John Doe",
				1*time.Hour,
			)

			assert.Equal(t, "nonce-val", claims.Nonce)
			assert.NotEmpty(t, claims.AtHash)
			assert.Equal(t, "user@example.com", claims.Email)
			require.NotNil(t, claims.EmailVerified)
			assert.False(t, *claims.EmailVerified)
			assert.Equal(t, "John Doe", claims.Name)
		},
	)

	t.Run(
		"sets expiration based on ttl",
		func(t *testing.T) {
			t.Parallel()

			ttl := 2 * time.Hour
			before := time.Now()
			claims := oauth2.NewIDTokenClaims(
				testIssuer,
				identityID,
				clientID,
				authTime,
				coredata.OAuth2Scopes{oauth2.ScopeOpenID},
				"",
				"",
				"",
				false,
				"",
				ttl,
			)
			after := time.Now()

			assert.GreaterOrEqual(t, claims.ExpiresAt, before.Add(ttl).Unix())
			assert.LessOrEqual(t, claims.ExpiresAt, after.Add(ttl).Unix())
			assert.GreaterOrEqual(t, claims.IssuedAt, before.Unix())
			assert.LessOrEqual(t, claims.IssuedAt, after.Unix())
		},
	)

	t.Run(
		"email not verified",
		func(t *testing.T) {
			t.Parallel()

			claims := oauth2.NewIDTokenClaims(
				testIssuer,
				identityID,
				clientID,
				authTime,
				coredata.OAuth2Scopes{oauth2.ScopeOpenID, oauth2.ScopeEmail},
				"",
				"",
				"user@example.com",
				false,
				"",
				1*time.Hour,
			)

			require.NotNil(t, claims.EmailVerified)
			assert.False(t, *claims.EmailVerified)
		},
	)
}