headscale/integration/oidc_groups_test.go

package integration

import (
	"encoding/json"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
	"github.com/juanfont/headscale/integration/hsic"
	"github.com/juanfont/headscale/integration/tsic"
	"github.com/oauth2-proxy/mockoidc"
	"github.com/stretchr/testify/assert"
)

// TestOIDCGroupsExtraction tests that OIDC groups are properly extracted and stored
func TestOIDCGroupsExtraction(t *testing.T) {
	IntegrationSkip(t)

	// Create mock users with different group memberships
	spec := ScenarioSpec{
		NodesPerUser: 1,
		Users:        []string{"admin", "user", "readonly"},
		OIDCUsers: []mockoidc.MockUser{
			// Admin user with multiple groups
			{
				Subject:           "admin@example.com",
				Email:             "admin@example.com",
				PreferredUsername: "admin",
				Groups:            []string{"admins", "users", "engineering"},
			},
			// Regular user with single group
			{
				Subject:           "user@example.com", 
				Email:             "user@example.com",
				PreferredUsername: "user",
				Groups:            []string{"users"},
			},
			// Readonly user with different groups
			{
				Subject:           "readonly@example.com",
				Email:             "readonly@example.com", 
				PreferredUsername: "readonly",
				Groups:            []string{"readonly", "auditors"},
			},
		},
	}

	scenario, err := NewScenario(spec)
	assertNoErr(t, err)
	defer scenario.ShutdownAssertNoPanics(t)

	oidcMap := map[string]string{
		"HEADSCALE_OIDC_ISSUER":             scenario.mockOIDC.Issuer(),
		"HEADSCALE_OIDC_CLIENT_ID":          scenario.mockOIDC.ClientID(),
		"CREDENTIALS_DIRECTORY_TEST":        "/tmp",
		"HEADSCALE_OIDC_CLIENT_SECRET_PATH": "${CREDENTIALS_DIRECTORY_TEST}/hs_client_oidc_secret",
		// Enable groups claim extraction
		"HEADSCALE_OIDC_EXTRA_PARAMS":       `{"groups_claim": "groups"}`,
	}

	err = scenario.CreateHeadscaleEnvWithLoginURL(
		nil,
		hsic.WithTestName("oidcgroups"),
		hsic.WithConfigEnv(oidcMap),
		hsic.WithTLS(),
		hsic.WithHostnameAsServerURL(),
	)
	assertNoErr(t, err)

	// Perform OIDC logins for all users
	allClients, err := scenario.ListTailscaleClients()
	assertNoErr(t, err)

	for _, client := range allClients {
		user, ok := scenario.usernames[client.Hostname()]
		assertOK(t, ok)

		_ = client.Login(scenario.loginWaitGroup, user)

		// Wait for login to complete
		scenario.loginWaitGroup.Wait()
		time.Sleep(5 * time.Second)
	}

	// Test groups were extracted correctly
	t.Run("verify-groups-extracted", func(t *testing.T) {
		// Get all users from Headscale
		users, err := scenario.ControlServer().ListUsers(&v1.ListUsersRequest{})
		assertNoErr(t, err)

		// Verify each user has correct groups
		expectedGroups := map[string][]string{
			"admin":    {"admins", "users", "engineering"},
			"user":     {"users"},
			"readonly": {"readonly", "auditors"},
		}

		for _, user := range users.GetUsers() {
			// Parse groups from user
			var userGroups []string
			if user.GetGroups() != "" {
				err := json.Unmarshal([]byte(user.GetGroups()), &userGroups)
				assertNoErr(t, err)
			}

			// Check expected groups
			expected, exists := expectedGroups[user.GetName()]
			assert.True(t, exists, "Unexpected user: %s", user.GetName())
			
			// Sort both slices for comparison
			assert.ElementsMatch(t, expected, userGroups, 
				"User %s has incorrect groups. Expected: %v, Got: %v", 
				user.GetName(), expected, userGroups)
		}
	})

	// Test groups persist across logins
	t.Run("verify-groups-persistence", func(t *testing.T) {
		// Get a client and log it out then back in
		client := allClients[0]
		user := scenario.usernames[client.Hostname()]

		// Logout
		err := client.Logout()
		assertNoErr(t, err)

		// Login again
		_ = client.Login(scenario.loginWaitGroup, user)
		scenario.loginWaitGroup.Wait()
		time.Sleep(3 * time.Second)

		// Verify groups are still there
		users, err := scenario.ControlServer().ListUsers(&v1.ListUsersRequest{})
		assertNoErr(t, err)

		// Find our user
		var targetUser *v1.User
		for _, u := range users.GetUsers() {
			if u.GetName() == user {
				targetUser = u
				break
			}
		}
		assert.NotNil(t, targetUser, "User not found after re-login")

		// Verify groups are preserved
		var userGroups []string
		if targetUser.GetGroups() != "" {
			err := json.Unmarshal([]byte(targetUser.GetGroups()), &userGroups)
			assertNoErr(t, err)
		}

		assert.NotEmpty(t, userGroups, "Groups should persist after re-login")
	})
}

// TestOIDCGroupsWithoutClaim tests behavior when groups claim is not configured
func TestOIDCGroupsWithoutClaim(t *testing.T) {
	IntegrationSkip(t)

	spec := ScenarioSpec{
		NodesPerUser: 1,
		Users:        []string{"user1"},
		OIDCUsers: []mockoidc.MockUser{
			{
				Subject:           "user1@example.com",
				Email:             "user1@example.com",
				PreferredUsername: "user1",
				Groups:            []string{"admins", "users"}, // Groups present but not requested
			},
		},
	}

	scenario, err := NewScenario(spec)
	assertNoErr(t, err)
	defer scenario.ShutdownAssertNoPanics(t)

	// OIDC config WITHOUT groups claim
	oidcMap := map[string]string{
		"HEADSCALE_OIDC_ISSUER":             scenario.mockOIDC.Issuer(),
		"HEADSCALE_OIDC_CLIENT_ID":          scenario.mockOIDC.ClientID(),
		"CREDENTIALS_DIRECTORY_TEST":        "/tmp",
		"HEADSCALE_OIDC_CLIENT_SECRET_PATH": "${CREDENTIALS_DIRECTORY_TEST}/hs_client_oidc_secret",
		// No groups claim configured
	}

	err = scenario.CreateHeadscaleEnvWithLoginURL(
		nil,
		hsic.WithTestName("oidcnogroups"),
		hsic.WithConfigEnv(oidcMap),
		hsic.WithTLS(),
		hsic.WithHostnameAsServerURL(),
	)
	assertNoErr(t, err)

	// Login user
	allClients, err := scenario.ListTailscaleClients()
	assertNoErr(t, err)

	client := allClients[0]
	user := scenario.usernames[client.Hostname()]
	_ = client.Login(scenario.loginWaitGroup, user)
	scenario.loginWaitGroup.Wait()
	time.Sleep(3 * time.Second)

	// Verify user exists but has no groups
	users, err := scenario.ControlServer().ListUsers(&v1.ListUsersRequest{})
	assertNoErr(t, err)

	assert.Len(t, users.GetUsers(), 1, "Should have exactly one user")
	
	user1 := users.GetUsers()[0]
	assert.Empty(t, user1.GetGroups(), "User should have no groups when claim not configured")
}

// TestOIDCGroupsEmptyGroups tests behavior when user has no groups
func TestOIDCGroupsEmptyGroups(t *testing.T) {
	IntegrationSkip(t)

	spec := ScenarioSpec{
		NodesPerUser: 1,
		Users:        []string{"ungrouped"},
		OIDCUsers: []mockoidc.MockUser{
			{
				Subject:           "ungrouped@example.com",
				Email:             "ungrouped@example.com",
				PreferredUsername: "ungrouped",
				Groups:            []string{}, // User has no groups
			},
		},
	}

	scenario, err := NewScenario(spec)
	assertNoErr(t, err)
	defer scenario.ShutdownAssertNoPanics(t)

	oidcMap := map[string]string{
		"HEADSCALE_OIDC_ISSUER":             scenario.mockOIDC.Issuer(),
		"HEADSCALE_OIDC_CLIENT_ID":          scenario.mockOIDC.ClientID(),
		"CREDENTIALS_DIRECTORY_TEST":        "/tmp",
		"HEADSCALE_OIDC_CLIENT_SECRET_PATH": "${CREDENTIALS_DIRECTORY_TEST}/hs_client_oidc_secret",
		"HEADSCALE_OIDC_EXTRA_PARAMS":       `{"groups_claim": "groups"}`,
	}

	err = scenario.CreateHeadscaleEnvWithLoginURL(
		nil,
		hsic.WithTestName("oidcemptygroups"),
		hsic.WithConfigEnv(oidcMap),
		hsic.WithTLS(),
		hsic.WithHostnameAsServerURL(),
	)
	assertNoErr(t, err)

	// Login user with no groups
	allClients, err := scenario.ListTailscaleClients()
	assertNoErr(t, err)

	client := allClients[0]
	user := scenario.usernames[client.Hostname()]
	_ = client.Login(scenario.loginWaitGroup, user)
	scenario.loginWaitGroup.Wait()
	time.Sleep(3 * time.Second)

	// Verify user exists with empty groups
	users, err := scenario.ControlServer().ListUsers(&v1.ListUsersRequest{})
	assertNoErr(t, err)

	assert.Len(t, users.GetUsers(), 1, "Should have exactly one user")
	
	ungroupedUser := users.GetUsers()[0]
	assert.Empty(t, ungroupedUser.GetGroups(), "User with no groups should have empty groups field")
}

// TestOIDCGroupsUpdatesOnLogin tests that groups are updated when user logs in again
func TestOIDCGroupsUpdatesOnLogin(t *testing.T) {
	IntegrationSkip(t)

	// Create scenario with user having initial groups
	spec := ScenarioSpec{
		NodesPerUser: 1,
		Users:        []string{"dynamic"},
		OIDCUsers: []mockoidc.MockUser{
			{
				Subject:           "dynamic@example.com",
				Email:             "dynamic@example.com",
				PreferredUsername: "dynamic",
				Groups:            []string{"initial-group"},
			},
		},
	}

	scenario, err := NewScenario(spec)
	assertNoErr(t, err)
	defer scenario.ShutdownAssertNoPanics(t)

	oidcMap := map[string]string{
		"HEADSCALE_OIDC_ISSUER":             scenario.mockOIDC.Issuer(),
		"HEADSCALE_OIDC_CLIENT_ID":          scenario.mockOIDC.ClientID(),
		"CREDENTIALS_DIRECTORY_TEST":        "/tmp",
		"HEADSCALE_OIDC_CLIENT_SECRET_PATH": "${CREDENTIALS_DIRECTORY_TEST}/hs_client_oidc_secret",
		"HEADSCALE_OIDC_EXTRA_PARAMS":       `{"groups_claim": "groups"}`,
	}

	err = scenario.CreateHeadscaleEnvWithLoginURL(
		nil,
		hsic.WithTestName("oidcdynamicgroups"),
		hsic.WithConfigEnv(oidcMap),
		hsic.WithTLS(),
		hsic.WithHostnameAsServerURL(),
	)
	assertNoErr(t, err)

	// Initial login
	allClients, err := scenario.ListTailscaleClients()
	assertNoErr(t, err)

	client := allClients[0]
	user := scenario.usernames[client.Hostname()]
	_ = client.Login(scenario.loginWaitGroup, user)
	scenario.loginWaitGroup.Wait()
	time.Sleep(3 * time.Second)

	// Verify initial groups
	users, err := scenario.ControlServer().ListUsers(&v1.ListUsersRequest{})
	assertNoErr(t, err)
	
	userObj := users.GetUsers()[0]
	var initialGroups []string
	if userObj.GetGroups() != "" {
		err := json.Unmarshal([]byte(userObj.GetGroups()), &initialGroups)
		assertNoErr(t, err)
	}
	assert.Equal(t, []string{"initial-group"}, initialGroups)

	// Update the mock user to have different groups
	// Note: In a real test, this would involve updating the OIDC provider
	// For this test, we'll simulate the scenario by modifying the mock

	// Logout and login again (simulating groups change in OIDC provider)
	err = client.Logout()
	assertNoErr(t, err)

	// Update mock user groups (this is test-specific, not production code)
	scenario.mockOIDC.SetUserGroups("dynamic@example.com", []string{"updated-group", "admin-group"})

	_ = client.Login(scenario.loginWaitGroup, user)
	scenario.loginWaitGroup.Wait()
	time.Sleep(3 * time.Second)

	// Verify groups were updated
	users, err = scenario.ControlServer().ListUsers(&v1.ListUsersRequest{})
	assertNoErr(t, err)
	
	updatedUserObj := users.GetUsers()[0]
	var updatedGroups []string
	if updatedUserObj.GetGroups() != "" {
		err := json.Unmarshal([]byte(updatedUserObj.GetGroups()), &updatedGroups)
		assertNoErr(t, err)
	}
	
	assert.ElementsMatch(t, []string{"updated-group", "admin-group"}, updatedGroups,
		"Groups should be updated on subsequent login")
}