diff --git a/CHANGELOG.md b/CHANGELOG.md index 7725303cf..04aabb779 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ ### Security - Metrics are no longer exposed to unregistered clients through the Noise protocol. They remain available only on the configured metrics listener; an empty `metrics_listen_addr` disables that listener. [#3571](https://github.com/juanfont/headscale/issues/3571) +- Preserve pending registrations and OIDC logins at capacity instead of evicting active approvals; refuse excess work with a retry response and reserve capacity for registered nodes' SSH checks. +- Retain pending registration data as owned serialized payloads capped at 64 KiB, preserving tags and expiry while bounding stored data. Oversized interactive registrations are rejected explicitly. +- Validate OIDC login session type before allocating state, consume callback state atomically, and release SSH authentication capacity after the client consumes its verdict. ### v1 REST API replaced; gRPC and Protobuf removed diff --git a/config-example.yaml b/config-example.yaml index d6a995d2e..9d3bcb4d1 100644 --- a/config-example.yaml +++ b/config-example.yaml @@ -510,7 +510,10 @@ auto_update: # # tuning: # # Maximum number of pending registration entries in the auth cache. -# # Oldest entries are evicted when the cap is reached. +# # At capacity, new registrations retry without displacing pending approvals. +# # SSH checks can use 64 additional reserved entries. Each registration +# # retains at most 64 KiB of serialized data, so the default registration +# # payload budget is 64 MiB plus record overhead. # # # # register_cache_max_entries: 1024 # diff --git a/docs/ref/registration.md b/docs/ref/registration.md index c2db42bdc..62162e043 100644 --- a/docs/ref/registration.md +++ b/docs/ref/registration.md @@ -91,6 +91,34 @@ headscale users create listed as "online" in the output of `headscale nodes list`. The "User" column displays `tagged-devices` as the owner of the node. See the "Tags" column for the list of assigned tags. +### Pending authentication limits + +Headscale keeps admitted authentication requests until they complete or expire. +When registration capacity is reached, new requests receive HTTP 503 with +`Retry-After` instead of invalidating someone else's pending approval. Existing +approvals and pre-authenticated key registration remain available. SSH checks +have 64 reserved entries beyond the registration allocation, including in the +OIDC state cache. + +The default registration allocation is 1024 entries. The advanced +`tuning.register_cache_max_entries` setting changes that allocation. Pending +registration retains only hostname, OS and requested tags from Hostinfo; other +metadata arrives with the node's first MapRequest. The full cached registration +must fit within 64 KiB when serialized; larger payloads are rejected without +truncating authentication metadata. Noise +registrations return HTTP 200 with `RegisterResponse.Error` so the Tailscale +client displays the permanent size error. The admin API reports this size error +as HTTP 413. +At the default allocation, retained registration payloads use at most 64 MiB +plus record overhead. This is a cache-storage bound, not a bound on total server +memory or concurrent request processing. + +OIDC login URLs must reference an existing pending authentication of the matching +type. OIDC callback state is single-use, including when the identity provider +exchange fails; restart the browser login flow after such an error. Capacity +limits protect admitted requests but cannot guarantee admission during sustained +overload. + ### Pre authenticated key Registration with a pre authenticated key (or auth key) is a non-interactive way to register a new node. The Headscale diff --git a/hscontrol/api/v1/errors.go b/hscontrol/api/v1/errors.go index 71ccfe035..2aa8e91f3 100644 --- a/hscontrol/api/v1/errors.go +++ b/hscontrol/api/v1/errors.go @@ -20,6 +20,13 @@ func mapError(msg string, err error) error { } switch { + case errors.Is(err, types.ErrRegistrationPayloadTooLarge): + return huma.Error413RequestEntityTooLarge(msg, err) + + case errors.Is(err, state.ErrAuthCacheFull), + errors.Is(err, state.ErrAuthCacheClosed): + return huma.Error503ServiceUnavailable(msg, err) + case errors.Is(err, gorm.ErrRecordNotFound), errors.Is(err, state.ErrNodeNotFound), errors.Is(err, state.ErrNodeNotInNodeStore), @@ -43,7 +50,8 @@ func mapError(msg string, err error) error { return huma.Error400BadRequest(msg, err) case errors.Is(err, state.ErrNodeKeyInUse), - errors.Is(err, state.ErrAmbiguousNodeOwnership): + errors.Is(err, state.ErrAmbiguousNodeOwnership), + errors.Is(err, state.ErrAuthCacheEntryExists): return huma.Error409Conflict(msg, err) default: diff --git a/hscontrol/api/v1/nodes.go b/hscontrol/api/v1/nodes.go index afe600397..b07cdd118 100644 --- a/hscontrol/api/v1/nodes.go +++ b/hscontrol/api/v1/nodes.go @@ -563,7 +563,11 @@ func registerNodeAdminOps(api huma.API, b Backend) { } authRegReq := types.NewRegisterAuthRequest(regData) - b.State.SetAuthCacheEntry(registrationID, authRegReq) + + err = b.State.SetAuthCacheEntry(registrationID, authRegReq) + if err != nil { + return nil, mapError("starting authentication", err) + } // Synthetic echo; the real node is created later via the auth path // from the cached registration data. diff --git a/hscontrol/apiv1_auth_test.go b/hscontrol/apiv1_auth_test.go index 1f157f414..1ff6fc194 100644 --- a/hscontrol/apiv1_auth_test.go +++ b/hscontrol/apiv1_auth_test.go @@ -37,7 +37,7 @@ func seedAuthRequest( Hostinfo: &tailcfg.Hostinfo{Hostname: hostname}, } - app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(regData)) + require.NoError(t, app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(regData))) } } @@ -124,7 +124,7 @@ func TestAPIV1AuthApprove(t *testing.T) { authID := types.MustAuthID() authReq := types.NewAuthRequest() - h.app.state.SetAuthCacheEntry(authID, authReq) + require.NoError(t, h.app.state.SetAuthCacheEntry(authID, authReq)) body := fmt.Appendf(nil, `{"authId":%q}`, authID.String()) res := h.callHuma(http.MethodPost, "/api/v1/auth/approve", body) @@ -159,7 +159,7 @@ func TestAPIV1AuthReject(t *testing.T) { authID := types.MustAuthID() authReq := types.NewAuthRequest() - h.app.state.SetAuthCacheEntry(authID, authReq) + require.NoError(t, h.app.state.SetAuthCacheEntry(authID, authReq)) body := fmt.Appendf(nil, `{"authId":%q}`, authID.String()) res := h.callHuma(http.MethodPost, "/api/v1/auth/reject", body) diff --git a/hscontrol/apiv1_nodes_test.go b/hscontrol/apiv1_nodes_test.go index d6197dc01..c483853ba 100644 --- a/hscontrol/apiv1_nodes_test.go +++ b/hscontrol/apiv1_nodes_test.go @@ -364,7 +364,7 @@ func TestAPIV1NodeRegister(t *testing.T) { MachineKey: mk.Public(), Hostname: "registered-node", } - app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(regData)) + require.NoError(t, app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(regData))) } assertParityIsolated(t, seed, http.MethodPost, @@ -453,3 +453,21 @@ func TestAPIV1NodeDebugCreate(t *testing.T) { assert.Equal(t, "0001-01-01T00:00:00Z", got.Node["lastSeen"]) }) } + +func TestAPIV1NodeDebugCreateAdmissionFailure(t *testing.T) { + app := createTestApp(t, func(cfg *types.Config) { + cfg.Tuning.RegisterCacheMaxEntries = 1 + }) + defer func() { require.NoError(t, app.state.Close()) }() + + app.state.CreateUserForTest("alice") + require.NoError(t, app.state.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest())) + + id := types.MustAuthID() + body := []byte(`{"user":"alice","key":"` + id.String() + `","name":"node","routes":[]}`) + res := callHandler(newHumaTestHandler(app), http.MethodPost, "/api/v1/debug/node", body) + require.Equal(t, http.StatusServiceUnavailable, res.status) + + _, ok := app.state.GetAuthCacheEntry(id) + require.False(t, ok) +} diff --git a/hscontrol/auth.go b/hscontrol/auth.go index 489b71e82..8f91f4af6 100644 --- a/hscontrol/auth.go +++ b/hscontrol/auth.go @@ -10,6 +10,7 @@ import ( "strings" "time" + "github.com/juanfont/headscale/hscontrol/state" "github.com/juanfont/headscale/hscontrol/types" "github.com/rs/zerolog/log" "gorm.io/gorm" @@ -24,6 +25,17 @@ type AuthProvider interface { AuthURL(authID types.AuthID) string } +func authAdmissionError(err error) HTTPError { + switch { + case errors.Is(err, types.ErrRegistrationPayloadTooLarge): + return NewHTTPError(http.StatusRequestEntityTooLarge, err.Error(), err) + case errors.Is(err, state.ErrAuthCacheFull), errors.Is(err, state.ErrAuthCacheClosed): + return NewHTTPError(http.StatusServiceUnavailable, "cannot start authentication; retry later", err) + default: + return NewHTTPError(http.StatusInternalServerError, "cannot start authentication", err) + } +} + // machineKeyMismatch fails closed when a node looked up by NodeKey was started // in a Noise session with a different machine key. Without this anyone holding a // target's NodeKey could open a session with a throwaway machine key and act on @@ -369,6 +381,8 @@ func (h *Headscale) waitForFollowup( } } + h.state.DeleteAuthCacheEntry(followupReg) + if verdict.Accept() { if !verdict.Node.Valid() { // registration is expired in the cache, instruct the client to try a new registration @@ -397,7 +411,11 @@ func (h *Headscale) reqToNewRegisterResponse( authRegReq := types.NewRegisterAuthRequest(regData) log.Info().Msgf("new followup node registration using auth id: %s", newAuthID) - h.state.SetAuthCacheEntry(newAuthID, authRegReq) + + err = h.state.SetAuthCacheEntry(newAuthID, authRegReq) + if err != nil { + return nil, authAdmissionError(err) + } return &tailcfg.RegisterResponse{ AuthURL: h.authProvider.RegisterURL(newAuthID), @@ -532,7 +550,10 @@ func (h *Headscale) handleRegisterInteractive( registrationDataFromRequest(req, machineKey), ) - h.state.SetAuthCacheEntry(authID, authRegReq) + err = h.state.SetAuthCacheEntry(authID, authRegReq) + if err != nil { + return nil, authAdmissionError(err) + } log.Info().Msgf("starting node registration using auth id: %s", authID) diff --git a/hscontrol/auth_tags_test.go b/hscontrol/auth_tags_test.go index da6ccf28d..f313a53ec 100644 --- a/hscontrol/auth_tags_test.go +++ b/hscontrol/auth_tags_test.go @@ -847,7 +847,7 @@ func TestExpiryDuringPersonalToTaggedConversion(t *testing.T) { }, Expiry: &clientExpiry, }) - app.state.SetAuthCacheEntry(registrationID1, regEntry1) + require.NoError(t, app.state.SetAuthCacheEntry(registrationID1, regEntry1)) node, _, err := app.state.HandleNodeFromAuthPath( registrationID1, types.UserID(user.ID), nil, "webauth", @@ -869,7 +869,7 @@ func TestExpiryDuringPersonalToTaggedConversion(t *testing.T) { }, Expiry: &clientExpiry, // Client still sends expiry }) - app.state.SetAuthCacheEntry(registrationID2, regEntry2) + require.NoError(t, app.state.SetAuthCacheEntry(registrationID2, regEntry2)) nodeAfter, _, err := app.state.HandleNodeFromAuthPath( registrationID2, types.UserID(user.ID), nil, "webauth", @@ -918,7 +918,7 @@ func TestExpiryDuringTaggedToPersonalConversion(t *testing.T) { RequestTags: []string{"tag:server"}, // Tagged node }, }) - app.state.SetAuthCacheEntry(registrationID1, regEntry1) + require.NoError(t, app.state.SetAuthCacheEntry(registrationID1, regEntry1)) node, _, err := app.state.HandleNodeFromAuthPath( registrationID1, types.UserID(user.ID), nil, "webauth", @@ -941,7 +941,7 @@ func TestExpiryDuringTaggedToPersonalConversion(t *testing.T) { }, Expiry: &clientExpiry, // Client requests expiry }) - app.state.SetAuthCacheEntry(registrationID2, regEntry2) + require.NoError(t, app.state.SetAuthCacheEntry(registrationID2, regEntry2)) nodeAfter, _, err := app.state.HandleNodeFromAuthPath( registrationID2, types.UserID(user.ID), nil, "webauth", diff --git a/hscontrol/auth_test.go b/hscontrol/auth_test.go index 60f6b8fcc..8f14abfc5 100644 --- a/hscontrol/auth_test.go +++ b/hscontrol/auth_test.go @@ -13,7 +13,9 @@ import ( "time" "github.com/juanfont/headscale/hscontrol/mapper" + "github.com/juanfont/headscale/hscontrol/state" "github.com/juanfont/headscale/hscontrol/types" + "github.com/rs/zerolog" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "tailscale.com/tailcfg" @@ -689,7 +691,7 @@ func TestAuthenticationFlows(t *testing.T) { MachineKey: machineKey1.Public(), Hostname: "followup-success-node", }) - app.state.SetAuthCacheEntry(regID, nodeToRegister) + require.NoError(t, app.state.SetAuthCacheEntry(regID, nodeToRegister)) // Simulate successful registration // [Headscale.handleRegister] will receive the value when it starts waiting @@ -737,7 +739,7 @@ func TestAuthenticationFlows(t *testing.T) { MachineKey: machineKey1.Public(), Hostname: "followup-timeout-node", }) - app.state.SetAuthCacheEntry(regID, nodeToRegister) + require.NoError(t, app.state.SetAuthCacheEntry(regID, nodeToRegister)) // Don't call FinishRegistration - will timeout return fmt.Sprintf("http://localhost:8080/register/%s", regID), nil @@ -1575,7 +1577,7 @@ func TestAuthenticationFlows(t *testing.T) { MachineKey: machineKey1.Public(), Hostname: "nil-response-node", }) - app.state.SetAuthCacheEntry(regID, nodeToRegister) + require.NoError(t, app.state.SetAuthCacheEntry(regID, nodeToRegister)) // Simulate registration that returns empty NodeView (cache expired during auth) go func() { @@ -3356,7 +3358,7 @@ func TestWebFlowReauthDifferentUser(t *testing.T) { } // Helper function to create test app. -func createTestApp(t *testing.T) *Headscale { +func createTestApp(t *testing.T, options ...func(*types.Config)) *Headscale { t.Helper() tmpDir := t.TempDir() @@ -3380,6 +3382,10 @@ func createTestApp(t *testing.T) *Headscale { }, } + for _, option := range options { + option(&cfg) + } + app, err := NewHeadscale(&cfg) require.NoError(t, err) @@ -3397,6 +3403,106 @@ func createTestApp(t *testing.T) *Headscale { return app } +func TestRegistrationFollowupReleasesRejectedSession(t *testing.T) { + app := createTestApp(t, func(cfg *types.Config) { + cfg.Tuning.RegisterCacheMaxEntries = 1 + }) + defer func() { require.NoError(t, app.state.Close()) }() + + machineKey := key.NewMachine().Public() + req := tailcfg.RegisterRequest{ + NodeKey: key.NewNode().Public(), + Hostinfo: &tailcfg.Hostinfo{Hostname: "rejected-registration"}, + } + initial, err := app.handleRegisterInteractive(req, machineKey) + require.NoError(t, err) + + req.Followup = initial.AuthURL + initialURL, err := url.Parse(initial.AuthURL) + require.NoError(t, err) + authID, err := types.AuthIDFromString(strings.TrimPrefix(initialURL.Path, "/register/")) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err = app.waitForFollowup(ctx, req.Followup, registrationDataFromRequest(req, machineKey)) + require.ErrorIs(t, err, context.Canceled) + + _, present := app.state.GetAuthCacheEntry(authID) + require.True(t, present, "a canceled poll must leave the pending session available") + + body := fmt.Appendf(nil, `{"authId":%q}`, authID.String()) + rejected := callHandler(newHumaTestHandler(app), http.MethodPost, "/api/v1/auth/reject", body) + require.Equal(t, http.StatusOK, rejected.status) + + _, present = app.state.GetAuthCacheEntry(authID) + require.True(t, present, "the verdict must remain available until follow-up consumes it") + + followup, err := app.waitForFollowup(t.Context(), req.Followup, registrationDataFromRequest(req, machineKey)) + require.NoError(t, err, "consuming the rejection must free capacity for the replacement") + require.NotEmpty(t, followup.AuthURL) + require.NotEqual(t, initial.AuthURL, followup.AuthURL) + + _, present = app.state.GetAuthCacheEntry(authID) + require.False(t, present, "the consumed session must release its slot") +} + +func TestAuthAdmissionErrors(t *testing.T) { + for _, flow := range []string{"interactive", "followup", "ssh"} { + t.Run(flow, func(t *testing.T) { + app := createTestApp(t, func(cfg *types.Config) { + cfg.Tuning.RegisterCacheMaxEntries = 1 + }) + defer func() { require.NoError(t, app.state.Close()) }() + + id := types.MustAuthID() + pending := types.NewAuthRequest() + require.NoError(t, app.state.SetAuthCacheEntry(id, pending)) + + if flow == "ssh" { + for range 64 { + require.NoError(t, app.state.SetAuthCacheEntry(types.MustAuthID(), types.NewSSHCheckAuthRequest(1, 2))) + } + } + + req := tailcfg.RegisterRequest{ + NodeKey: key.NewNode().Public(), + Hostinfo: &tailcfg.Hostinfo{Hostname: "pending-node"}, + } + machineKey := key.NewMachine().Public() + + var err error + + switch flow { + case "interactive": + var resp *tailcfg.RegisterResponse + + resp, err = app.handleRegisterInteractive(req, machineKey) + require.Nil(t, resp, "refused admission must not publish an auth URL") + case "followup": + var resp *tailcfg.RegisterResponse + + resp, err = app.reqToNewRegisterResponse(registrationDataFromRequest(req, machineKey)) + require.Nil(t, resp, "refused admission must not publish an auth URL") + case "ssh": + ns := &noiseServer{headscale: app} + + var action *tailcfg.SSHAction + + action, err = ns.sshActionHoldAndDelegate(zerolog.Nop(), &tailcfg.SSHAction{}, 1, 2, "alice") + require.Nil(t, action, "refused admission must not publish a delegation URL") + } + + require.ErrorIs(t, err, state.ErrAuthCacheFull) + + got, ok := app.state.GetAuthCacheEntry(id) + require.True(t, ok) + require.Same(t, pending, got) + }) + } +} + // TestGitHubIssue2830_NodeRestartWithUsedPreAuthKey tests the scenario reported in // https://github.com/juanfont/headscale/issues/2830 // @@ -3819,7 +3925,7 @@ func TestWebAuthRejectsUnauthorizedRequestTags(t *testing.T) { RequestTags: []string{"tag:unauthorized"}, // This tag is not in policy }, }) - app.state.SetAuthCacheEntry(registrationID, regEntry) + require.NoError(t, app.state.SetAuthCacheEntry(registrationID, regEntry)) // Complete the web auth - should fail because tag is unauthorized _, _, err := app.state.HandleNodeFromAuthPath( @@ -3882,7 +3988,7 @@ func TestWebAuthReauthWithEmptyTagsRemovesAllTags(t *testing.T) { RequestTags: []string{"tag:valid-owned", "tag:second"}, }, }) - app.state.SetAuthCacheEntry(registrationID1, regEntry1) + require.NoError(t, app.state.SetAuthCacheEntry(registrationID1, regEntry1)) // Complete initial registration with tags node, _, err := app.state.HandleNodeFromAuthPath( @@ -3909,7 +4015,7 @@ func TestWebAuthReauthWithEmptyTagsRemovesAllTags(t *testing.T) { RequestTags: []string{}, // EMPTY - should untag }, }) - app.state.SetAuthCacheEntry(registrationID2, regEntry2) + require.NoError(t, app.state.SetAuthCacheEntry(registrationID2, regEntry2)) // Complete reauth with empty tags nodeAfterReauth, _, err := app.state.HandleNodeFromAuthPath( @@ -3995,7 +4101,7 @@ func TestAuthKeyTaggedToUserOwnedViaReauth(t *testing.T) { RequestTags: []string{}, // EMPTY - should untag }, }) - app.state.SetAuthCacheEntry(registrationID, regEntry) + require.NoError(t, app.state.SetAuthCacheEntry(registrationID, regEntry)) // Complete reauth with empty tags nodeAfterReauth, _, err := app.state.HandleNodeFromAuthPath( @@ -4195,7 +4301,7 @@ func TestTaggedNodeWithoutUserToDifferentUser(t *testing.T) { RequestTags: []string{}, // Empty - transition to user-owned }, }) - app.state.SetAuthCacheEntry(registrationID, regEntry) + require.NoError(t, app.state.SetAuthCacheEntry(registrationID, regEntry)) // This should NOT panic - before the fix, this would panic with: // panic: runtime error: invalid memory address or nil pointer dereference @@ -4329,7 +4435,7 @@ func TestHandleNodeFromAuthPath_OldUserNil_NoPanic(t *testing.T) { Hostname: "authpath-orphan-newuser", }, }) - app.state.SetAuthCacheEntry(authID, regEntry) + require.NoError(t, app.state.SetAuthCacheEntry(authID, regEntry)) node, _, err := app.state.HandleNodeFromAuthPath( authID, @@ -4384,14 +4490,14 @@ func TestHandleNodeFromAuthPath_ReloginWholeNode(t *testing.T) { newNodeKey := key.NewNode() authID := types.MustAuthID() - app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: node.MachineKey, NodeKey: newNodeKey.Public(), Hostname: hostname, Hostinfo: &tailcfg.Hostinfo{ Hostname: hostname, }, - })) + }))) relogged, c, err := app.state.HandleNodeFromAuthPath( authID, @@ -4443,7 +4549,7 @@ func TestWaitForFollowupMachineKeyMismatch(t *testing.T) { NodeKey: key.NewNode().Public(), Hostname: hostname, }) - app.state.SetAuthCacheEntry(authID, regEntry) + require.NoError(t, app.state.SetAuthCacheEntry(authID, regEntry)) user := app.state.CreateUserForTest(hostname + "-user") node := app.state.CreateNodeForTest(user, hostname) @@ -4520,7 +4626,7 @@ func TestWaitForFollowupMachineKeyMismatch(t *testing.T) { NodeKey: key.NewNode().Public(), Hostname: "followup-pending", }) - app.state.SetAuthCacheEntry(authID, regEntry) + require.NoError(t, app.state.SetAuthCacheEntry(authID, regEntry)) url := fmt.Sprintf("http://localhost:8080/register/%s", authID) refusedAtOnce(t, url, foreignMachineKey.Public()) @@ -4539,7 +4645,7 @@ func TestWaitForFollowupMachineKeyMismatch(t *testing.T) { t.Run("ssh check session is not a registration", func(t *testing.T) { authID := types.MustAuthID() - app.state.SetAuthCacheEntry(authID, types.NewSSHCheckAuthRequest(1, 2)) + require.NoError(t, app.state.SetAuthCacheEntry(authID, types.NewSSHCheckAuthRequest(1, 2))) refusedAtOnce(t, fmt.Sprintf("http://localhost:8080/register/%s", authID), victimMachineKey.Public()) @@ -4583,7 +4689,7 @@ func TestFollowupWaitPrefersCompletedAuthOverExpiredContext(t *testing.T) { MachineKey: machineKey, Hostname: "followup-race-node", }) - app.state.SetAuthCacheEntry(regID, authReq) + require.NoError(t, app.state.SetAuthCacheEntry(regID, authReq)) // Registration completes BEFORE we wait: verdict is buffered. user := app.state.CreateUserForTest(fmt.Sprintf("followup-race-user-%d", i)) @@ -4711,7 +4817,7 @@ func TestFollowupParksOneRequestPerRegistration(t *testing.T) { MachineKey: machineKey, Hostname: "followup-park-node", }) - app.state.SetAuthCacheEntry(authID, reg) + require.NoError(t, app.state.SetAuthCacheEntry(authID, reg)) regData := &types.RegistrationData{ MachineKey: machineKey, @@ -4833,7 +4939,7 @@ func TestFollowupRefreshKeepsAuthMetadata(t *testing.T) { reg := types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: machineKey, Hostname: "previous-request", }) - app.state.SetAuthCacheEntry(authID, reg) + require.NoError(t, app.state.SetAuthCacheEntry(authID, reg)) reg.FinishAuth(types.AuthVerdict{}) } diff --git a/hscontrol/noise.go b/hscontrol/noise.go index 7d8b4c52b..d1b91e684 100644 --- a/hscontrol/noise.go +++ b/hscontrol/noise.go @@ -566,10 +566,13 @@ func (ns *noiseServer) sshActionHoldAndDelegate( ) } - ns.headscale.state.SetAuthCacheEntry( + err = ns.headscale.state.SetAuthCacheEntry( authID, types.NewSSHCheckAuthRequest(srcNodeID, dstNodeID), ) + if err != nil { + return nil, authAdmissionError(err) + } authURL := ns.headscale.authProvider.AuthURL(authID) @@ -691,6 +694,8 @@ func (ns *noiseServer) sshActionFollowUp( case verdict, verdictOK = <-auth.WaitForAuth(): } + ns.headscale.state.DeleteAuthCacheEntry(authID) + // FinishAuth buffers one verdict, then closes the channel. A closed // receive yields the zero verdict, which Accept() reports as success, // so a replayed follow-up would be approved even after a Reject. @@ -1004,6 +1009,7 @@ func (ns *noiseServer) RegistrationHandler( return } + statusCode := http.StatusOK registerResponse := func() *tailcfg.RegisterResponse { //nolint:contextcheck if decodeErr != nil { return regErr(decodeErr) @@ -1015,6 +1021,13 @@ func (ns *noiseServer) RegistrationHandler( resp, err := ns.headscale.handleRegister(req.Context(), registerRequest, ns.machineKey) if err != nil { if httpErr, ok := errors.AsType[HTTPError](err); ok { + // Keep permanent errors in HTTP 200 so controlclient surfaces RegisterResponse.Error. + if httpErr.Code == http.StatusServiceUnavailable { + statusCode = http.StatusServiceUnavailable + + writer.Header().Set("Retry-After", "10") + } + return &tailcfg.RegisterResponse{ Error: httpErr.Msg, } @@ -1027,7 +1040,7 @@ func (ns *noiseServer) RegistrationHandler( }() writer.Header().Set("Content-Type", "application/json; charset=utf-8") - writer.WriteHeader(http.StatusOK) + writer.WriteHeader(statusCode) err := json.NewEncoder(writer).Encode(registerResponse) if err != nil { diff --git a/hscontrol/noise_test.go b/hscontrol/noise_test.go index 788d398d9..ccad005d3 100644 --- a/hscontrol/noise_test.go +++ b/hscontrol/noise_test.go @@ -358,6 +358,69 @@ func serveRegister(t *testing.T, ns *noiseServer, body any) *httptest.ResponseRe return rec } +func TestRegistrationHandler_AdmissionRetry(t *testing.T) { + app := createTestApp(t, func(cfg *types.Config) { + cfg.Tuning.RegisterCacheMaxEntries = 1 + }) + defer func() { require.NoError(t, app.state.Close()) }() + + require.NoError(t, app.state.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest())) + + ns := &noiseServer{headscale: app, machineKey: key.NewMachine().Public()} + rec := serveRegister(t, ns, tailcfg.RegisterRequest{ + Version: capver.MinSupportedCapabilityVersion, + NodeKey: key.NewNode().Public(), + Hostinfo: &tailcfg.Hostinfo{Hostname: "node"}, + }) + require.Equal(t, http.StatusServiceUnavailable, rec.Code) + require.NotEmpty(t, rec.Header().Get("Retry-After")) + + var response tailcfg.RegisterResponse + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &response)) + require.Empty(t, response.AuthURL) + require.NotEmpty(t, response.Error) +} + +func TestRegistrationHandler_RegistrationPayloadLimit(t *testing.T) { + app := createTestApp(t) + defer func() { require.NoError(t, app.state.Close()) }() + + hostinfo := &tailcfg.Hostinfo{Hostname: "node", OS: strings.Repeat("x", types.MaxHostinfoSize-128)} + require.NoError(t, types.ValidateHostinfo(hostinfo), "the full registration must reach cache admission") + + ns := &noiseServer{headscale: app, machineKey: key.NewMachine().Public()} + rec := serveRegister(t, ns, tailcfg.RegisterRequest{ + Version: capver.MinSupportedCapabilityVersion, + NodeKey: key.NewNode().Public(), + Hostinfo: hostinfo, + }) + require.Equal(t, http.StatusOK, rec.Code, "controlclient only surfaces RegisterResponse.Error with HTTP 200") + require.Empty(t, rec.Header().Get("Retry-After")) + + var response tailcfg.RegisterResponse + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &response)) + require.Contains(t, response.Error, types.ErrRegistrationPayloadTooLarge.Error()) + require.Contains(t, response.Error, strconv.Itoa(types.MaxRegistrationPayloadBytes)) + require.Empty(t, response.AuthURL) + require.False(t, response.MachineAuthorized) + require.Zero(t, app.state.DebugRegistrationCache()["current_len"], "oversized registrations must not enter the auth cache") +} + +func TestSSHActionFollowUp_ReleasesConsumedSession(t *testing.T) { + f := newSSHVerdictFixture(t) + authID, auth := f.mint(t) + auth.FinishAuth(types.AuthVerdict{}) + + _, ok := f.ns.headscale.state.GetAuthCacheEntry(authID) + require.True(t, ok, "approval must remain available until the client consumes it") + + action := sshActionFromRecorder(t, f.followUp(t, authID)) + require.True(t, action.Accept) + + _, ok = f.ns.headscale.state.GetAuthCacheEntry(authID) + require.False(t, ok, "a consumed verdict must release its admission slot") +} + func requireBelowFloorRejected(t *testing.T, rec *httptest.ResponseRecorder) { t.Helper() @@ -721,10 +784,10 @@ func TestSSHActionFollowUp_RejectsBindingMismatch(t *testing.T) { // Mint an SSH-check auth request bound to (srcCached, dstCached). authID := types.MustAuthID() - app.state.SetAuthCacheEntry( + require.NoError(t, app.state.SetAuthCacheEntry( authID, types.NewSSHCheckAuthRequest(srcCached.ID, dstCached.ID), - ) + )) // Build a follow-up that claims to be for (srcOther, dstOther) but // reuses the bound auth_id. The Noise machineKey matches dstOther so @@ -881,7 +944,7 @@ func TestSSHActionFollowUp_RejectsVerdictAfterCheckRemoved(t *testing.T) { require.NoError(t, err) authID := types.MustAuthID() - app.state.SetAuthCacheEntry(authID, types.NewSSHCheckAuthRequest(src.ID, dst.ID)) + require.NoError(t, app.state.SetAuthCacheEntry(authID, types.NewSSHCheckAuthRequest(src.ID, dst.ID))) auth, ok := app.state.GetAuthCacheEntry(authID) require.True(t, ok) @@ -2261,7 +2324,7 @@ func TestSSHActionFollowUp_ParksOneRequestPerSession(t *testing.T) { authID := types.MustAuthID() auth := types.NewSSHCheckAuthRequest(src.ID, dst.ID) - app.state.SetAuthCacheEntry(authID, auth) + require.NoError(t, app.state.SetAuthCacheEntry(authID, auth)) ns := &noiseServer{headscale: app, machineKey: dst.MachineKey} diff --git a/hscontrol/oidc.go b/hscontrol/oidc.go index c237c6eb4..7aac502f9 100644 --- a/hscontrol/oidc.go +++ b/hscontrol/oidc.go @@ -10,6 +10,8 @@ import ( "net/url" "slices" "strings" + "sync" + "sync/atomic" "time" "github.com/coreos/go-oidc/v3/oidc" @@ -29,10 +31,10 @@ const ( defaultOAuthOptionsCount = 3 authCacheExpiration = time.Minute * 15 - // authCacheMaxEntries bounds the OIDC state→[AuthInfo] cache to prevent - // unauthenticated cache-fill DoS via repeated /register/{auth_id} or - // /auth/{auth_id} GETs that mint OIDC state cookies. + // Registrations may use authCacheMaxEntries slots; SSH checks may also use + // the reserve so new registrations cannot fill all pending OIDC capacity. authCacheMaxEntries = 1024 + authCacheSSHReserve = 64 // cookieNamePrefixLen is the number of leading characters from a // state/nonce value that [getCookieName] splices into the cookie name. @@ -47,7 +49,6 @@ var errOIDCStateTooShort = errors.New("oidc state parameter is too short") var ( errEmptyOIDCCallbackParams = errors.New("empty OIDC callback params") errNoOIDCIDToken = errors.New("extracting ID token") - errNoOIDCRegistrationInfo = errors.New("registration info not in cache") errOIDCAllowedDomains = errors.New( "authenticated principal does not match any allowed domain", ) @@ -71,10 +72,11 @@ type AuthProviderOIDC struct { serverURL string cfg *types.OIDCConfig - // authCache holds auth information between the auth and the callback - // steps. It is a bounded [expirable.LRU] keyed by OIDC state, evicting oldest - // entries to keep the cache footprint constant under attack. - authCache *expirable.LRU[string, AuthInfo] + // authCacheMu serializes non-evicting admission and single-use consumption. + // Reads use Peek to preserve insertion order, so the oldest entry expires first. + authCacheMu sync.Mutex + authCacheRegistrations atomic.Int64 + authCache *expirable.LRU[string, AuthInfo] oidcProvider *oidc.Provider oauth2Config *oauth2.Config @@ -101,21 +103,25 @@ func NewAuthProviderOIDC( Scopes: cfg.Scope, } - authCache := expirable.NewLRU[string, AuthInfo]( - authCacheMaxEntries, - nil, - authCacheExpiration, - ) - - return &AuthProviderOIDC{ + a := &AuthProviderOIDC{ h: h, serverURL: serverURL, cfg: cfg, - authCache: authCache, oidcProvider: oidcProvider, oauth2Config: oauth2Config, - }, nil + } + a.authCache = expirable.NewLRU[string, AuthInfo]( + authCacheMaxEntries+authCacheSSHReserve, + func(_ string, info AuthInfo) { + if info.Registration { + a.authCacheRegistrations.Add(-1) + } + }, + authCacheExpiration, + ) + + return a, nil } // cookiesSecure reports whether the OIDC cookies should carry the Secure flag. @@ -177,11 +183,24 @@ func (a *AuthProviderOIDC) authHandler( return } - // Set the state and nonce cookies to protect against CSRF attacks - state := a.setCSRFCookie(writer, req, "state") + authReq, ok := a.h.state.GetAuthCacheEntry(authID) + if !ok { + httpUserError(writer, NewHTTPError(http.StatusGone, "login session expired, try again", nil)) + return + } - // Set the state and nonce cookies to protect against CSRF attacks - nonce := a.setCSRFCookie(writer, req, "nonce") + if registration && !authReq.IsRegistration() { + httpUserError(writer, NewHTTPError(http.StatusBadRequest, "auth session is not for node registration", nil)) + return + } + + if !registration && !authReq.IsSSHCheck() { + httpUserError(writer, NewHTTPError(http.StatusBadRequest, "auth session is not for SSH check", nil)) + return + } + + state := rands.HexString(64) + nonce := rands.HexString(64) registrationInfo := AuthInfo{ AuthID: authID, @@ -218,8 +237,15 @@ func (a *AuthProviderOIDC) authHandler( extras = append(extras, oidc.Nonce(nonce)) - // Cache the registration info - a.authCache.Add(state, registrationInfo) + if !a.addAuthInfoToState(state, registrationInfo) { + writer.Header().Set("Retry-After", "10") + httpUserError(writer, NewHTTPError(http.StatusServiceUnavailable, "too many pending OIDC logins, try again later", nil)) + + return + } + + a.setCSRFCookie(writer, req, "state", state) + a.setCSRFCookie(writer, req, "nonce", nonce) authURL := a.oauth2Config.AuthCodeURL(state, extras...) log.Debug().Caller().Msgf("redirecting to %s for authentication", authURL) @@ -255,7 +281,20 @@ func (a *AuthProviderOIDC) OIDCCallbackHandler( return } - oauth2Token, err := a.getOauth2Token(req.Context(), code, state) + // A cookie-bound callback spends its state even if authentication fails; + // retrying then requires a fresh login flow rather than reusing the callback. + authInfo := a.getAuthInfoFromState(state) + if authInfo == nil { + httpUserError(writer, NewHTTPError(http.StatusGone, "login session expired, try again", nil)) + return + } + + if req.URL.Query().Get("error") != "" { + httpUserError(writer, NewHTTPError(http.StatusForbidden, "OIDC authentication failed", nil)) + return + } + + oauth2Token, err := a.getOauth2Token(req.Context(), code, authInfo) if err != nil { httpUserError(writer, err) return @@ -343,18 +382,6 @@ func (a *AuthProviderOIDC) OIDCCallbackHandler( return } - // TODO(kradalby): Is this comment right? - // If the node exists, then the node should be reauthenticated, - // if the node does not exist, and the machine key exists, then - // this is a new node that should be registered. - authInfo := a.getAuthInfoFromState(state) - if authInfo == nil { - log.Debug().Caller().Str("state", state).Msg("state not found in cache, login session may have expired") - httpUserError(writer, NewHTTPError(http.StatusGone, "login session expired, try again", nil)) - - return - } - // If this is a registration flow, render the confirmation // interstitial instead of finalising the registration immediately. // Without an explicit user click, a single GET to @@ -471,10 +498,11 @@ func (a *AuthProviderOIDC) determineNodeExpiry(idTokenExpiration time.Time) *tim func extractCodeAndStateParamFromRequest( req *http.Request, ) (string, string, error) { - code := req.URL.Query().Get("code") - state := req.URL.Query().Get("state") + query := req.URL.Query() + code := query.Get("code") + state := query.Get("state") - if code == "" || state == "" { + if (code == "" && query.Get("error") == "") || state == "" { return "", "", NewHTTPError(http.StatusBadRequest, "missing code or state parameter", errEmptyOIDCCallbackParams) } @@ -493,19 +521,12 @@ func extractCodeAndStateParamFromRequest( func (a *AuthProviderOIDC) getOauth2Token( ctx context.Context, code string, - state string, + authInfo *AuthInfo, ) (*oauth2.Token, error) { var exchangeOpts []oauth2.AuthCodeOption - if a.cfg.PKCE.Enabled { - regInfo, ok := a.authCache.Get(state) - if !ok { - return nil, NewHTTPError(http.StatusNotFound, "registration not found", errNoOIDCRegistrationInfo) - } - - if regInfo.Verifier != nil { - exchangeOpts = []oauth2.AuthCodeOption{oauth2.VerifierOption(*regInfo.Verifier)} - } + if a.cfg.PKCE.Enabled && authInfo.Verifier != nil { + exchangeOpts = []oauth2.AuthCodeOption{oauth2.VerifierOption(*authInfo.Verifier)} } oauth2Token, err := a.oauth2Config.Exchange(ctx, code, exchangeOpts...) @@ -631,17 +652,55 @@ func doOIDCAuthorization( return nil } +func (a *AuthProviderOIDC) addAuthInfoToState(state string, info AuthInfo) bool { + a.authCacheMu.Lock() + defer a.authCacheMu.Unlock() + + // Len includes expired entries until the sweeper removes their bucket. + for { + oldest, _, ok := a.authCache.GetOldest() + if !ok { + break + } + + if _, live := a.authCache.Peek(oldest); live { + break + } + + a.authCache.Remove(oldest) + } + + if a.authCache.Contains(state) || a.authCache.Len() >= a.authCache.Cap() { + return false + } + + if info.Registration && a.authCacheRegistrations.Load() >= int64(a.authCache.Cap()-authCacheSSHReserve) { + return false + } + + // Increment before Add: an expiry callback may run as soon as it is stored. + if info.Registration { + a.authCacheRegistrations.Add(1) + } + + a.authCache.Add(state, info) + + return true +} + // getAuthInfoFromState retrieves and consumes the auth info for a state. The // entry is removed on read so a state is single-use: a replayed callback cannot // resolve the same auth session twice, even within the cache TTL. func (a *AuthProviderOIDC) getAuthInfoFromState(state string) *AuthInfo { - authInfo, ok := a.authCache.Get(state) + a.authCacheMu.Lock() + defer a.authCacheMu.Unlock() + + authInfo, ok := a.authCache.Peek(state) + a.authCache.Remove(state) if !ok { return nil } - a.authCache.Remove(state) - return &authInfo } @@ -1101,9 +1160,8 @@ func (a *AuthProviderOIDC) setCSRFCookie( w http.ResponseWriter, r *http.Request, name string, -) string { - val := rands.HexString(64) - + val string, +) { //nolint:gosec // G124: Secure from server_url scheme or req.TLS; HttpOnly + SameSite set below c := &http.Cookie{ Path: a.oidcCallbackPath(), @@ -1120,6 +1178,4 @@ func (a *AuthProviderOIDC) setCSRFCookie( SameSite: http.SameSiteLaxMode, } http.SetCookie(w, c) - - return val } diff --git a/hscontrol/oidc_confirm_test.go b/hscontrol/oidc_confirm_test.go index bb4e4ee5a..e2cf58ee4 100644 --- a/hscontrol/oidc_confirm_test.go +++ b/hscontrol/oidc_confirm_test.go @@ -57,7 +57,7 @@ func TestRegisterConfirmHandler_RejectsCSRFMismatch(t *testing.T) { UserID: 1, CSRF: "expected-csrf", }) - app.state.SetAuthCacheEntry(authID, regReq) + require.NoError(t, app.state.SetAuthCacheEntry(authID, regReq)) rec := httptest.NewRecorder() provider.RegisterConfirmHandler(rec, @@ -89,9 +89,9 @@ func TestRegisterConfirmHandler_RejectsWithoutPending(t *testing.T) { authID := types.MustAuthID() // Cached registration with NO pending confirmation set — i.e. the // OIDC callback has not run yet. - app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest( + require.NoError(t, app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest( &types.RegistrationData{Hostname: "no-oidc-yet"}, - )) + ))) rec := httptest.NewRecorder() provider.RegisterConfirmHandler(rec, diff --git a/hscontrol/oidc_test.go b/hscontrol/oidc_test.go index a4325d006..3b004805e 100644 --- a/hscontrol/oidc_test.go +++ b/hscontrol/oidc_test.go @@ -3,6 +3,7 @@ package hscontrol import ( "context" "encoding/json" + "fmt" "io" "net/http" "net/http/cookiejar" @@ -10,14 +11,19 @@ import ( "net/url" "regexp" "strings" + "sync" + "sync/atomic" "testing" + "testing/synctest" "time" + "github.com/go-chi/chi/v5" "github.com/hashicorp/golang-lru/v2/expirable" "github.com/juanfont/headscale/hscontrol/types" "github.com/oauth2-proxy/mockoidc" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/oauth2" "tailscale.com/tailcfg" "tailscale.com/types/key" ) @@ -199,7 +205,7 @@ func TestSetCSRFCookieSameSite(t *testing.T) { w := httptest.NewRecorder() r := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/auth/abcdef0123456789", nil) - a.setCSRFCookie(w, r, "state") + a.setCSRFCookie(w, r, "state", "state-value") cookies := w.Result().Cookies() require.Len(t, cookies, 1) @@ -208,7 +214,7 @@ func TestSetCSRFCookieSameSite(t *testing.T) { } // TestExtractCodeAndStateParam covers the callback's first trust-boundary -// checks: both params required, and a too-short state is rejected before +// checks: a code or provider error is required, and a too-short state is rejected before // getCookieName can slice out of range. func TestExtractCodeAndStateParam(t *testing.T) { _, _, err := extractCodeAndStateParamFromRequest( @@ -242,6 +248,755 @@ func TestGetAuthInfoFromStateSingleUse(t *testing.T) { assert.Nil(t, a.getAuthInfoFromState("state-x"), "a consumed state must not resolve again") } +func newOIDCPendingProvider(t *testing.T, capacity int) *AuthProviderOIDC { + t.Helper() + + a := &AuthProviderOIDC{ + h: createTestApp(t), + serverURL: "https://headscale.example", + cfg: &types.OIDCConfig{}, + oauth2Config: &oauth2.Config{ + ClientID: "headscale-test", + Endpoint: oauth2.Endpoint{ + AuthURL: "https://issuer.example/authorize", + AuthStyle: oauth2.AuthStyleInParams, + }, + }, + } + t.Cleanup(func() { require.NoError(t, a.h.state.Close()) }) + + a.authCache = expirable.NewLRU[string, AuthInfo](capacity+authCacheSSHReserve, func(_ string, info AuthInfo) { + if info.Registration { + a.authCacheRegistrations.Add(-1) + } + }, time.Minute) + t.Cleanup(a.authCache.Close) + + return a +} + +func oidcPendingAuthRequest(t *testing.T, authID types.AuthID) *http.Request { + t.Helper() + + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/register/"+authID.String(), nil) + rctx := chi.NewRouteContext() + rctx.URLParams.Add("auth_id", authID.String()) + + return req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, rctx)) +} + +func startOIDCPendingAuth(t *testing.T, a *AuthProviderOIDC, registration bool) (*httptest.ResponseRecorder, types.AuthID) { + t.Helper() + + authID := types.MustAuthID() + + authReq := types.NewSSHCheckAuthRequest(1, 2) + if registration { + authReq = types.NewRegisterAuthRequest(&types.RegistrationData{Hostname: "pending-oidc"}) + } + + require.NoError(t, a.h.state.SetAuthCacheEntry(authID, authReq)) + + rec := httptest.NewRecorder() + a.authHandler(rec, oidcPendingAuthRequest(t, authID), registration) + + return rec, authID +} + +func TestOIDCPendingAdmissionPreservesSessions(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + a.cfg.PKCE.Enabled = true + a.cfg.PKCE.Method = types.PKCEMethodS256 + + states := make(map[string]AuthInfo) + + for _, registration := range []bool{true, false} { + rec, authID := startOIDCPendingAuth(t, a, registration) + require.Equal(t, http.StatusFound, rec.Code) + location, err := rec.Result().Location() + require.NoError(t, err) + + state := location.Query().Get("state") + info, ok := a.authCache.Peek(state) + require.True(t, ok) + require.NotNil(t, info.Verifier) + assert.Equal(t, authID, info.AuthID) + assert.Equal(t, registration, info.Registration) + assert.Equal(t, "S256", location.Query().Get("code_challenge_method")) + + states[state] = info + } + + rejected, _ := startOIDCPendingAuth(t, a, true) + assert.Equal(t, http.StatusServiceUnavailable, rejected.Code, + "a full pending-state cache must reject admission instead of evicting an in-flight login") + assert.Equal(t, "10", rejected.Header().Get("Retry-After")) + assert.Empty(t, rejected.Header().Get("Location")) + assert.Empty(t, rejected.Header().Values("Set-Cookie")) + + for state, want := range states { + assert.Equal(t, &want, a.getAuthInfoFromState(state), "an admitted session must survive a rejected admission") + } + + assert.Zero(t, a.authCache.Len()) +} + +func TestOIDCPendingAdmissionReservesSSHCapacity(t *testing.T) { + a := newOIDCPendingProvider(t, 2) + registrations := make(map[string]types.AuthID) + + for range 2 { + rec, authID := startOIDCPendingAuth(t, a, true) + require.Equal(t, http.StatusFound, rec.Code) + location, err := rec.Result().Location() + require.NoError(t, err) + + registrations[location.Query().Get("state")] = authID + } + + rejected, _ := startOIDCPendingAuth(t, a, true) + assert.Equal(t, http.StatusServiceUnavailable, rejected.Code) + assert.Empty(t, rejected.Header().Values("Set-Cookie")) + + for range 64 { + rec, authID := startOIDCPendingAuth(t, a, false) + require.Equal(t, http.StatusFound, rec.Code, "registration load must leave the SSH reserve available") + location, err := rec.Result().Location() + require.NoError(t, err) + + info, ok := a.authCache.Peek(location.Query().Get("state")) + require.True(t, ok) + assert.Equal(t, authID, info.AuthID) + assert.False(t, info.Registration) + } + + full, _ := startOIDCPendingAuth(t, a, false) + assert.Equal(t, http.StatusServiceUnavailable, full.Code, "the SSH reserve must also be bounded") + assert.Empty(t, full.Header().Get("Location")) + assert.Empty(t, full.Header().Values("Set-Cookie")) + assert.Equal(t, 66, a.authCache.Len()) + + for state, authID := range registrations { + info := a.getAuthInfoFromState(state) + require.NotNil(t, info, "admission at the total bound must not evict an earlier registration") + assert.Equal(t, authID, info.AuthID) + } +} + +func TestOIDCPendingAdmissionConcurrent(t *testing.T) { + a := newOIDCPendingProvider(t, 2) + + type result struct { + authID types.AuthID + rec *httptest.ResponseRecorder + } + + results := make(chan result, 4) + start := make(chan struct{}) + + var workers sync.WaitGroup + + for range 4 { + authID := types.MustAuthID() + require.NoError(t, a.h.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{Hostname: "concurrent-oidc"}))) + req := oidcPendingAuthRequest(t, authID) + + workers.Go(func() { + <-start + + rec := httptest.NewRecorder() + a.RegisterHandler(rec, req) + + results <- result{authID: authID, rec: rec} + }) + } + + close(start) + workers.Wait() + close(results) + + accepted := make(map[string]types.AuthID) + rejected := 0 + + for result := range results { + if result.rec.Code == http.StatusFound { + location, err := result.rec.Result().Location() + require.NoError(t, err) + + accepted[location.Query().Get("state")] = result.authID + } else { + assert.Equal(t, http.StatusServiceUnavailable, result.rec.Code) + assert.Empty(t, result.rec.Header().Get("Location")) + assert.Empty(t, result.rec.Header().Values("Set-Cookie")) + + rejected++ + } + } + + require.Len(t, accepted, 2, "concurrent admission must respect the pending-state bound") + assert.Equal(t, 2, rejected) + assert.Equal(t, 2, a.authCache.Len()) + + for state, authID := range accepted { + info := a.getAuthInfoFromState(state) + require.NotNil(t, info, "every successful admission must remain usable") + assert.Equal(t, authID, info.AuthID) + } +} + +func TestOIDCPendingAdmissionDuplicateState(t *testing.T) { + a := newOIDCPendingProvider(t, 2) + verifier := "original-verifier" + info := AuthInfo{AuthID: types.MustAuthID(), Verifier: &verifier, Registration: true} + require.True(t, a.addAuthInfoToState("same-state", info)) + assert.False(t, a.addAuthInfoToState("same-state", AuthInfo{AuthID: types.MustAuthID()}), + "a duplicate state must not replace the auth ID, flow type, or PKCE verifier") + assert.Equal(t, &info, a.getAuthInfoFromState("same-state")) + assert.Zero(t, a.authCacheRegistrations.Load()) +} + +func TestOIDCPendingAdmissionSSHUsesNormalSlots(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + for i := range 65 { + require.True(t, a.addAuthInfoToState(fmt.Sprintf("ssh-state-%d", i), AuthInfo{AuthID: types.MustAuthID()}), + "SSH checks may use the regular slots as well as their reserve") + } + + assert.False(t, a.addAuthInfoToState("one-too-many", AuthInfo{AuthID: types.MustAuthID()})) + assert.Equal(t, 65, a.authCache.Len()) + assert.Zero(t, a.authCacheRegistrations.Load()) +} + +func TestOIDCPendingAdmissionRequiresAuthRequest(t *testing.T) { + for _, registration := range []bool{true, false} { + for _, existing := range []bool{false, true} { + t.Run(fmt.Sprintf("registration=%t/existing=%t", registration, existing), func(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + authID := types.MustAuthID() + wantStatus := http.StatusGone + + if existing { + // A registration cannot authorize an SSH check, or vice versa. + authReq := types.NewRegisterAuthRequest(&types.RegistrationData{Hostname: "wrong-flow"}) + if registration { + authReq = types.NewSSHCheckAuthRequest(1, 2) + } + + require.NoError(t, a.h.state.SetAuthCacheEntry(authID, authReq)) + + wantStatus = http.StatusBadRequest + } + + rec := httptest.NewRecorder() + a.authHandler(rec, oidcPendingAuthRequest(t, authID), registration) + assert.Equal(t, wantStatus, rec.Code) + assert.Empty(t, rec.Header().Get("Location")) + assert.Empty(t, rec.Header().Values("Set-Cookie")) + assert.Zero(t, a.authCache.Len(), "invalid auth requests must not reserve OIDC capacity") + }) + } + } +} + +func TestOIDCPendingAdmissionInvalidPKCE(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + a.cfg.PKCE.Enabled = true + a.cfg.PKCE.Method = "invalid" + + rec, _ := startOIDCPendingAuth(t, a, true) + assert.Equal(t, http.StatusInternalServerError, rec.Code) + assert.Empty(t, rec.Header().Get("Location")) + assert.Empty(t, rec.Header().Values("Set-Cookie")) + assert.Zero(t, a.authCache.Len()) + + a.cfg.PKCE.Method = types.PKCEMethodPlain + rec, _ = startOIDCPendingAuth(t, a, true) + assert.Equal(t, http.StatusFound, rec.Code, "a preparation error must leave capacity available") +} + +func TestOIDCPendingAdmissionReclaimsExpired(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + a := newOIDCPendingProvider(t, 2) + a.authCache.Close() + + a.authCache = expirable.NewLRU[string, AuthInfo](66, func(_ string, info AuthInfo) { + if info.Registration { + a.authCacheRegistrations.Add(-1) + } + }, time.Second) + defer a.authCache.Close() + + old, _ := startOIDCPendingAuth(t, a, true) + require.Equal(t, http.StatusFound, old.Code) + oldURL, err := old.Result().Location() + require.NoError(t, err) + + <-time.After(time.Millisecond) + + live, liveID := startOIDCPendingAuth(t, a, true) + require.Equal(t, http.StatusFound, live.Code) + liveURL, err := live.Result().Location() + require.NoError(t, err) + + // Both entries share a cleanup bucket, which waits for the newer expiry. + <-time.After(time.Second - time.Millisecond + time.Nanosecond) + require.Equal(t, 2, a.authCache.Len(), "exercise expired entries that the sweeper has not removed") + _, ok := a.authCache.Peek(oldURL.Query().Get("state")) + require.False(t, ok) + + next, _ := startOIDCPendingAuth(t, a, true) + assert.Equal(t, http.StatusFound, next.Code, "expired entries must not cause a false capacity rejection") + assert.Nil(t, a.getAuthInfoFromState(oldURL.Query().Get("state"))) + info := a.getAuthInfoFromState(liveURL.Query().Get("state")) + require.NotNil(t, info, "reclaiming expiry must preserve the live session") + assert.Equal(t, liveID, info.AuthID) + }) +} + +func TestOIDCPendingAdmissionReclaimsSweptCapacity(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + a.authCache.Close() + + synctest.Test(t, func(t *testing.T) { + a.authCache = expirable.NewLRU[string, AuthInfo](65, func(_ string, info AuthInfo) { + if info.Registration { + a.authCacheRegistrations.Add(-1) + } + }, time.Second) + defer a.authCache.Close() + + require.True(t, a.addAuthInfoToState("expires", AuthInfo{Registration: true})) + + <-time.After(time.Second + 10*time.Millisecond) + synctest.Wait() + assert.Zero(t, a.authCache.Len()) + assert.Zero(t, a.authCacheRegistrations.Load(), "background expiry must return the registration budget") + assert.True(t, a.addAuthInfoToState("replacement", AuthInfo{Registration: true})) + }) +} + +func TestOIDCPendingCallbackReleasesOnError(t *testing.T) { + for _, pkce := range []bool{false, true} { + for _, missingToken := range []bool{false, true} { + t.Run(fmt.Sprintf("pkce=%t/missingToken=%t", pkce, missingToken), func(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + a.cfg.PKCE.Enabled = pkce + a.cfg.PKCE.Method = types.PKCEMethodS256 + + var exchanges atomic.Int32 + + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + exchanges.Add(1) + w.Header().Set("Content-Type", "application/json") + + if missingToken { + _, _ = io.WriteString(w, `{"access_token":"test-token","token_type":"Bearer"}`) + } else { + w.WriteHeader(http.StatusBadRequest) + _, _ = io.WriteString(w, `{"error":"invalid_grant"}`) + } + })) + t.Cleanup(tokenServer.Close) + a.oauth2Config.Endpoint.TokenURL = tokenServer.URL + + started, authID := startOIDCPendingAuth(t, a, true) + require.Equal(t, http.StatusFound, started.Code) + location, err := started.Result().Location() + require.NoError(t, err) + + state := location.Query().Get("state") + callbackURL := "/oidc/callback?code=test-code&state=" + state + callback := httptest.NewRequestWithContext(t.Context(), http.MethodGet, callbackURL, nil) + //nolint:gosec // G124: request cookies carry only name and value on the wire. + callback.AddCookie(&http.Cookie{Name: getCookieName("state", state), Value: state}) + + rec := httptest.NewRecorder() + a.OIDCCallbackHandler(rec, callback) + + wantStatus := http.StatusForbidden + if missingToken { + wantStatus = http.StatusBadRequest + } + + assert.Equal(t, wantStatus, rec.Code) + assert.Zero(t, a.authCache.Len(), "a callback error must release its pending-state slot") + _, pending := a.h.state.GetAuthCacheEntry(authID) + assert.True(t, pending, "consuming OIDC state must preserve the underlying auth request") + + replay := httptest.NewRecorder() + a.OIDCCallbackHandler(replay, callback) + assert.Equal(t, http.StatusGone, replay.Code, "a failed callback must still consume its state") + assert.EqualValues(t, 1, exchanges.Load(), "a consumed state must not reach token exchange again") + + next := httptest.NewRecorder() + a.RegisterHandler(next, oidcPendingAuthRequest(t, authID)) + assert.Equal(t, http.StatusFound, next.Code, "the original auth request must support a fresh browser flow") + nextURL, err := next.Result().Location() + require.NoError(t, err) + assert.NotEqual(t, state, nextURL.Query().Get("state")) + }) + } + } +} + +func TestOIDCPendingCallbackProviderError(t *testing.T) { + for _, tc := range []struct { + name string + providerError string + code string + }{ + {name: "cancelled", providerError: "access_denied"}, + {name: "error with code", providerError: "access_denied", code: "unused-code"}, + {name: "unknown error", providerError: strings.Repeat("untrusted-provider-error-", 2048)}, + } { + t.Run(tc.name, func(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + a.cfg.PKCE.Enabled = true + a.cfg.PKCE.Method = types.PKCEMethodS256 + + var exchanges atomic.Int32 + + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + exchanges.Add(1) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = io.WriteString(w, `{"error":"invalid_grant"}`) + })) + t.Cleanup(tokenServer.Close) + a.oauth2Config.Endpoint.TokenURL = tokenServer.URL + + started, authID := startOIDCPendingAuth(t, a, true) + require.Equal(t, http.StatusFound, started.Code) + location, err := started.Result().Location() + require.NoError(t, err) + + state := location.Query().Get("state") + + full := httptest.NewRecorder() + a.RegisterHandler(full, oidcPendingAuthRequest(t, authID)) + require.Equal(t, http.StatusServiceUnavailable, full.Code) + + query := url.Values{ + "error": {tc.providerError}, + "error_description": {strings.Repeat("untrusted-provider-description-", 2048)}, + "state": {state}, + "code": {tc.code}, + } + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + callback := httptest.NewRequestWithContext(ctx, http.MethodGet, "/oidc/callback?"+query.Encode(), nil) + //nolint:gosec // G124: request cookies carry only name and value on the wire. + callback.AddCookie(&http.Cookie{Name: getCookieName("state", state), Value: state}) + + rec := httptest.NewRecorder() + a.OIDCCallbackHandler(rec, callback) + assert.Equal(t, http.StatusForbidden, rec.Code) + assert.NotContains(t, rec.Body.String(), "untrusted-provider-description") + assert.NotContains(t, rec.Body.String(), "untrusted-provider-error") + assert.Less(t, rec.Body.Len(), 16*1024, "provider input must not inflate the error response") + assert.Zero(t, a.authCache.Len(), "provider cancellation must release its pending-state slot") + assert.Zero(t, a.authCacheRegistrations.Load()) + _, pending := a.h.state.GetAuthCacheEntry(authID) + assert.True(t, pending, "provider cancellation must preserve the underlying auth request") + + replay := httptest.NewRecorder() + a.OIDCCallbackHandler(replay, callback) + assert.Equal(t, http.StatusGone, replay.Code, "a provider-error callback must consume its state exactly once") + assert.Zero(t, exchanges.Load(), "provider-error callbacks must never exchange a code") + + next := httptest.NewRecorder() + a.RegisterHandler(next, oidcPendingAuthRequest(t, authID)) + require.Equal(t, http.StatusFound, next.Code, "cancellation must free capacity for a fresh login") + nextURL, err := next.Result().Location() + require.NoError(t, err) + assert.NotEqual(t, state, nextURL.Query().Get("state")) + }) + } +} + +func TestOIDCPendingCallbackProviderErrorValidation(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + started, authID := startOIDCPendingAuth(t, a, true) + require.Equal(t, http.StatusFound, started.Code) + location, err := started.Result().Location() + require.NoError(t, err) + + state := location.Query().Get("state") + cookieName := getCookieName("state", state) + cookie := cookieName + "=" + state + providerError := "error=access_denied&state=" + state + + for _, tc := range []struct { + name string + query string + cookie string + status int + }{ + {name: "missing state", query: "error=access_denied", cookie: cookie, status: http.StatusBadRequest}, + {name: "short state", query: "error=access_denied&state=abc", cookie: "state_abc=abc", status: http.StatusBadRequest}, + {name: "missing code and error", query: "state=" + state, cookie: cookie, status: http.StatusBadRequest}, + {name: "empty error", query: "error=&state=" + state, cookie: cookie, status: http.StatusBadRequest}, + {name: "description only", query: "error_description=cancelled&state=" + state, cookie: cookie, status: http.StatusBadRequest}, + {name: "missing cookie", query: providerError, status: http.StatusBadRequest}, + {name: "invalid cookie", query: providerError, cookie: cookieName + "=\"unterminated", status: http.StatusBadRequest}, + {name: "empty cookie", query: providerError, cookie: cookieName + "=", status: http.StatusForbidden}, + {name: "wrong cookie value", query: providerError, cookie: cookieName + "=different-state", status: http.StatusForbidden}, + {name: "wrong cookie name", query: providerError, cookie: "state_other=" + state, status: http.StatusBadRequest}, + {name: "unknown state", query: "error=access_denied&state=unknown-state", cookie: "state_unknow=unknown-state", status: http.StatusGone}, + } { + t.Run(tc.name, func(t *testing.T) { + callback := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/oidc/callback?"+tc.query, nil) + callback.Header.Set("Cookie", tc.cookie) + + rec := httptest.NewRecorder() + a.OIDCCallbackHandler(rec, callback) + assert.Equal(t, tc.status, rec.Code) + + info, ok := a.authCache.Peek(state) + require.True(t, ok, "invalid callbacks must preserve the pending state") + assert.Equal(t, authID, info.AuthID) + + full := httptest.NewRecorder() + a.RegisterHandler(full, oidcPendingAuthRequest(t, authID)) + assert.Equal(t, http.StatusServiceUnavailable, full.Code, "invalid callbacks must not free another browser's slot") + }) + } +} + +func TestOIDCPendingCallbackRequiresState(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + + var exchanges atomic.Int32 + + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + exchanges.Add(1) + w.WriteHeader(http.StatusBadRequest) + })) + t.Cleanup(tokenServer.Close) + a.oauth2Config.Endpoint.TokenURL = tokenServer.URL + + const state = "unknown-oidc-state" + + callback := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/oidc/callback?code=test-code&state="+state, nil) + //nolint:gosec // G124: request cookies carry only name and value on the wire. + callback.AddCookie(&http.Cookie{Name: getCookieName("state", state), Value: state}) + + rec := httptest.NewRecorder() + a.OIDCCallbackHandler(rec, callback) + assert.Equal(t, http.StatusGone, rec.Code) + assert.Zero(t, exchanges.Load(), "unknown state must be rejected before token exchange even without PKCE") +} + +func TestOIDCPendingCallbackConcurrent(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + a.cfg.PKCE.Enabled = true + a.cfg.PKCE.Method = types.PKCEMethodS256 + entered := make(chan struct{}, 2) + release := make(chan struct{}) + + var releaseOnce sync.Once + + unblock := func() { releaseOnce.Do(func() { close(release) }) } + + var exchanges atomic.Int32 + + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + exchanges.Add(1) + + entered <- struct{}{} + + <-release + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = io.WriteString(w, `{"error":"invalid_grant"}`) + })) + t.Cleanup(tokenServer.Close) + t.Cleanup(unblock) + + a.oauth2Config.Endpoint.TokenURL = tokenServer.URL + + started, _ := startOIDCPendingAuth(t, a, true) + require.Equal(t, http.StatusFound, started.Code) + location, err := started.Result().Location() + require.NoError(t, err) + + state := location.Query().Get("state") + callbackURL := "/oidc/callback?code=test-code&state=" + state + first := httptest.NewRequestWithContext(t.Context(), http.MethodGet, callbackURL, nil) + //nolint:gosec // G124: request cookies carry only name and value on the wire. + first.AddCookie(&http.Cookie{Name: getCookieName("state", state), Value: state}) + second := first.Clone(t.Context()) + firstRec, secondRec := httptest.NewRecorder(), httptest.NewRecorder() + firstDone, secondDone := make(chan struct{}), make(chan struct{}) + + go func() { + a.OIDCCallbackHandler(firstRec, first) + close(firstDone) + }() + + select { + case <-entered: + case <-firstDone: + t.Fatal("the first callback must reach token exchange") + case <-time.After(5 * time.Second): + t.Fatal("the first callback did not reach token exchange") + } + + go func() { + a.OIDCCallbackHandler(secondRec, second) + close(secondDone) + }() + + select { + case <-secondDone: + case <-entered: + // Release both requests before asserting the single-exchange invariant. + case <-time.After(5 * time.Second): + t.Fatal("the duplicate callback did not finish or reach token exchange") + } + + unblock() + <-firstDone + <-secondDone + assert.Equal(t, http.StatusForbidden, firstRec.Code) + assert.Equal(t, http.StatusGone, secondRec.Code) + assert.EqualValues(t, 1, exchanges.Load(), "only one callback may claim a pending state") + assert.Zero(t, a.authCache.Len()) +} + +func TestOIDCPendingCallbackCookieBinding(t *testing.T) { + a := newOIDCPendingProvider(t, 1) + started, _ := startOIDCPendingAuth(t, a, true) + require.Equal(t, http.StatusFound, started.Code) + location, err := started.Result().Location() + require.NoError(t, err) + + state := location.Query().Get("state") + + for _, value := range []string{"", "different-state"} { + callback := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/oidc/callback?code=test-code&state="+state, nil) + wantStatus := http.StatusBadRequest + + if value != "" { + //nolint:gosec // G124: request cookies carry only name and value on the wire. + callback.AddCookie(&http.Cookie{Name: getCookieName("state", state), Value: value}) + + wantStatus = http.StatusForbidden + } + + rec := httptest.NewRecorder() + a.OIDCCallbackHandler(rec, callback) + assert.Equal(t, wantStatus, rec.Code) + + _, ok := a.authCache.Peek(state) + assert.True(t, ok, "an unbound callback must not consume another browser's pending state") + } +} + +func TestOIDCPendingCallbackPKCEFlow(t *testing.T) { + for _, method := range []string{"", types.PKCEMethodPlain, types.PKCEMethodS256} { + t.Run("method="+method, func(t *testing.T) { + b := newOIDCBrowser(t) + a, ok := b.app.authProvider.(*AuthProviderOIDC) + require.True(t, ok) + + a.cfg.PKCE.Enabled = method != "" + a.cfg.PKCE.Method = method + a.authCache.Resize(65) + + authID, registerURL := b.pendingNode(t) + + status, _, body := b.get(t, registerURL) + require.Equal(t, http.StatusOK, status) + assert.Contains(t, body, "Confirm node registration") + assert.Zero(t, a.authCache.Len(), "successful callbacks must release pending OIDC capacity") + + authReq, ok := b.app.state.GetAuthCacheEntry(authID) + require.True(t, ok, "consuming OIDC state must preserve the registration for confirmation") + assert.NotNil(t, authReq.PendingConfirmation()) + }) + } +} + +func TestOIDCPendingCallbackPreservesSSHVerdict(t *testing.T) { + b := newOIDCBrowser(t) + a, ok := b.app.authProvider.(*AuthProviderOIDC) + require.True(t, ok) + + principal := mockoidc.DefaultUser() + user, _, err := a.createOrUpdateUserFromClaim(&types.OIDCClaims{ + Iss: b.idp.Issuer(), + Sub: principal.Subject, + Username: principal.PreferredUsername, + Email: principal.Email, + EmailVerified: true, + }) + require.NoError(t, err) + + src := b.app.state.CreateNodeForTest(user, "oidc-ssh-source") + b.app.state.PutNodeInStoreForTest(*src) + + authID := types.MustAuthID() + authReq := types.NewSSHCheckAuthRequest(src.ID, 2) + require.NoError(t, b.app.state.SetAuthCacheEntry(authID, authReq)) + + status, _, body := b.get(t, a.AuthURL(authID)) + require.Equal(t, http.StatusOK, status) + assert.Contains(t, body, "You may return to your terminal.") + assert.Zero(t, a.authCache.Len()) + + cached, ok := b.app.state.GetAuthCacheEntry(authID) + require.True(t, ok, "completing OIDC must preserve the SSH verdict for its polling client") + assert.Same(t, authReq, cached) + + select { + case verdict := <-cached.WaitForAuth(): + assert.True(t, verdict.Accept()) + default: + t.Fatal("successful SSH authentication must make its verdict available") + } +} + +func TestOIDCPendingCallbackNonceBinding(t *testing.T) { + b := newOIDCBrowser(t) + a, ok := b.app.authProvider.(*AuthProviderOIDC) + require.True(t, ok) + + b.client.CheckRedirect = func(req *http.Request, _ []*http.Request) error { + if req.URL.Path == "/oidc/callback" { + return http.ErrUseLastResponse + } + + return nil + } + _, registerURL := b.pendingNode(t) + //nolint:noctx,bodyclose // local test browser; closed below + response, err := b.client.Get(registerURL) + require.NoError(t, err) + + defer response.Body.Close() + + callbackURL, err := response.Location() + require.NoError(t, err) + require.Equal(t, "/oidc/callback", callbackURL.Path) + + callback := httptest.NewRequestWithContext(t.Context(), http.MethodGet, callbackURL.String(), nil) + for _, cookie := range b.client.Jar.Cookies(callbackURL) { + if strings.HasPrefix(cookie.Name, "nonce_") { + cookie.Value = "wrong-nonce" + } + + callback.AddCookie(cookie) + } + + rec := httptest.NewRecorder() + a.OIDCCallbackHandler(rec, callback) + assert.Equal(t, http.StatusForbidden, rec.Code) + assert.Zero(t, a.authCache.Len(), "nonce rejection must release the consumed pending state") +} + // TestClearOIDCCallbackCookie asserts the cookie is expired (negative MaxAge) on // the same path it was set with, so the browser drops it. func TestClearOIDCCallbackCookie(t *testing.T) { @@ -265,14 +1020,14 @@ func TestSetCSRFCookieSecure(t *testing.T) { secureRec := httptest.NewRecorder() secureProvider := &AuthProviderOIDC{serverURL: "https://hs.example.com"} - secureProvider.setCSRFCookie(secureRec, r, "state") + secureProvider.setCSRFCookie(secureRec, r, "state", "state-value") require.Len(t, secureRec.Result().Cookies(), 1) assert.True(t, secureRec.Result().Cookies()[0].Secure, "https server_url must set Secure even when req.TLS is nil (proxy case)") plainRec := httptest.NewRecorder() plainProvider := &AuthProviderOIDC{serverURL: "http://hs.example.com"} - plainProvider.setCSRFCookie(plainRec, r, "state") + plainProvider.setCSRFCookie(plainRec, r, "state", "state-value") require.Len(t, plainRec.Result().Cookies(), 1) assert.False(t, plainRec.Result().Cookies()[0].Secure, "plain-http server_url without req.TLS must not set Secure") @@ -379,11 +1134,11 @@ func (b *oidcBrowser) pendingNode(t *testing.T) (types.AuthID, string) { t.Helper() authID := types.MustAuthID() - b.app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, b.app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: key.NewMachine().Public(), NodeKey: key.NewNode().Public(), Hostname: "reload-victim", - })) + }))) return authID, b.publicURL + "/register/" + authID.String() } @@ -710,10 +1465,10 @@ func TestOIDCRegistrationRejectsOversizedMergedHostinfo(t *testing.T) { hostname := strings.Repeat("h", types.MaxHostnameLength) authID := types.MustAuthID() - app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, app.state.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: node.MachineKey, NodeKey: key.NewNode().Public(), Hostname: hostname, Hostinfo: &tailcfg.Hostinfo{Hostname: hostname}, - })) + }))) provider := &AuthProviderOIDC{h: app} changed, err := provider.handleRegistration(user, authID, nil) diff --git a/hscontrol/state/auth_cache_test.go b/hscontrol/state/auth_cache_test.go index 02301415f..5ae9b0de9 100644 --- a/hscontrol/state/auth_cache_test.go +++ b/hscontrol/state/auth_cache_test.go @@ -5,6 +5,7 @@ import ( "runtime" "runtime/pprof" "strings" + "sync" "testing" "time" @@ -45,10 +46,8 @@ func requireExpiredVerdict(t *testing.T, entry *types.AuthRequest) { } } -// TestAuthCacheBoundedLRU verifies that the registration auth cache is -// bounded by a maximum entry count, that exceeding the maxEntries evicts the -// oldest entry, and that the eviction callback resolves the parked -// AuthRequest with ErrRegistrationExpired so any waiting goroutine wakes. +// TestAuthCacheBoundedLRU checks the underlying library's eviction callback. +// State admission prevents this eviction behavior for pending auth requests. func TestAuthCacheBoundedLRU(t *testing.T) { const maxEntries = 4 @@ -97,6 +96,171 @@ func TestAuthCacheTTLExpiry(t *testing.T) { assert.False(t, ok, "expired entry must be gone from the cache") } +func TestAuthCacheAdmissionPreservesPending(t *testing.T) { + cfg := persistTestConfig(t.TempDir() + "/headscale.db") + cfg.Tuning.RegisterCacheMaxEntries = 1 + s, err := NewState(cfg) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + id := types.MustAuthID() + pending := types.NewAuthRequest() + require.NoError(t, s.SetAuthCacheEntry(id, pending)) + + otherID := types.MustAuthID() + require.ErrorIs(t, s.SetAuthCacheEntry(otherID, types.NewAuthRequest()), ErrAuthCacheFull) + + got, ok := s.GetAuthCacheEntry(id) + require.True(t, ok, "new work must not displace an admitted authentication") + require.Same(t, pending, got) + + _, ok = s.GetAuthCacheEntry(otherID) + require.False(t, ok, "new work must be refused at capacity") + + select { + case <-pending.WaitForAuth(): + t.Fatal("capacity refusal must not finish an existing authentication") + default: + } +} + +func TestAuthCacheAdmissionPreservesIdentity(t *testing.T) { + cfg := persistTestConfig(t.TempDir() + "/headscale.db") + s, err := NewState(cfg) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + id := types.MustAuthID() + pending := types.NewAuthRequest() + require.NoError(t, s.SetAuthCacheEntry(id, pending)) + require.NoError(t, s.SetAuthCacheEntry(id, pending)) + require.ErrorIs(t, s.SetAuthCacheEntry(id, types.NewAuthRequest()), ErrAuthCacheEntryExists) + + got, ok := s.GetAuthCacheEntry(id) + require.True(t, ok) + require.Same(t, pending, got, "an existing auth ID must keep its original request and waiter") +} + +func TestAuthCacheAdmissionConcurrent(t *testing.T) { + cfg := persistTestConfig(t.TempDir() + "/headscale.db") + cfg.Tuning.RegisterCacheMaxEntries = 2 + s, err := NewState(cfg) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + id := types.MustAuthID() + pending := types.NewAuthRequest() + require.NoError(t, s.SetAuthCacheEntry(id, pending)) + + var wg sync.WaitGroup + + results := make(chan error, 4) + + for range 4 { + wg.Go(func() { + results <- s.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest()) + }) + } + + wg.Wait() + close(results) + + var accepted int + + for err := range results { + if err == nil { + accepted++ + } else { + require.ErrorIs(t, err, ErrAuthCacheFull) + } + } + + require.Equal(t, 1, accepted) + + got, ok := s.GetAuthCacheEntry(id) + require.True(t, ok, "concurrent admission must preserve existing transactions") + require.Same(t, pending, got) + require.Equal(t, 2, s.authCache.Len()) +} + +func TestAuthCacheAdmissionReservesSSHCapacity(t *testing.T) { + cfg := persistTestConfig(t.TempDir() + "/headscale.db") + cfg.Tuning.RegisterCacheMaxEntries = 1 + s, err := NewState(cfg) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + registrationID := types.MustAuthID() + require.NoError(t, s.SetAuthCacheEntry(registrationID, types.NewAuthRequest())) + + sshID := types.MustAuthID() + require.NoError(t, s.SetAuthCacheEntry(sshID, types.NewSSHCheckAuthRequest(1, 2)), + "registration capacity must not consume the SSH reservation") + require.ErrorIs(t, s.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest()), ErrAuthCacheFull) + _, ok := s.GetAuthCacheEntry(registrationID) + require.True(t, ok) + _, ok = s.GetAuthCacheEntry(sshID) + require.True(t, ok) +} + +func TestAuthCacheAdmissionBoundsRegistrationPayload(t *testing.T) { + cfg := persistTestConfig(t.TempDir() + "/headscale.db") + s, err := NewState(cfg) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + id := types.MustAuthID() + request := types.NewRegisterAuthRequest(&types.RegistrationData{ + Hostname: strings.Repeat("x", types.MaxRegistrationPayloadBytes), + }) + require.ErrorIs(t, s.SetAuthCacheEntry(id, request), types.ErrRegistrationPayloadTooLarge) + _, ok := s.GetAuthCacheEntry(id) + require.False(t, ok) + require.Zero(t, s.authCache.Len()) + require.Zero(t, s.authCacheRegularEntries.Load()) +} + +func TestAuthCacheAdmissionDoesNotReuseExpiredRequest(t *testing.T) { + cfg := persistTestConfig(t.TempDir() + "/headscale.db") + cfg.Tuning.RegisterCacheExpiration = 50 * time.Millisecond + s, err := NewState(cfg) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + id := types.MustAuthID() + request := types.NewAuthRequest() + require.NoError(t, s.SetAuthCacheEntry(id, request)) + requireExpiredVerdict(t, request) + require.Error(t, s.SetAuthCacheEntry(id, request), "expired waiters cannot be reused as new pending requests") +} + +func TestAuthCacheAdmissionReleasesCapacity(t *testing.T) { + for _, removal := range []string{"consumption", "expiry"} { + t.Run(removal, func(t *testing.T) { + cfg := persistTestConfig(t.TempDir() + "/headscale.db") + cfg.Tuning.RegisterCacheMaxEntries = 1 + cfg.Tuning.RegisterCacheExpiration = 50 * time.Millisecond + s, err := NewState(cfg) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, s.Close()) }) + + id := types.MustAuthID() + pending := types.NewAuthRequest() + require.NoError(t, s.SetAuthCacheEntry(id, pending)) + + if removal == "consumption" { + pending.FinishAuth(types.AuthVerdict{}) + s.DeleteAuthCacheEntry(id) + } else { + requireExpiredVerdict(t, pending) + } + + require.NoError(t, s.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest())) + require.Equal(t, int64(1), s.authCacheRegularEntries.Load()) + }) + } +} + func goroutineID() string { buf := make([]byte, 64) n := runtime.Stack(buf, false) @@ -138,11 +302,12 @@ func TestStateCloseStopsAuthCacheSweeper(t *testing.T) { require.Equal(t, 1, authCacheSweepers(t, gid), "NewState must start a sweeper") pending := types.NewAuthRequest() - s.SetAuthCacheEntry(types.MustAuthID(), pending) + require.NoError(t, s.SetAuthCacheEntry(types.MustAuthID(), pending)) require.NoError(t, s.Close()) requireExpiredVerdict(t, pending) + require.ErrorIs(t, s.SetAuthCacheEntry(types.MustAuthID(), types.NewAuthRequest()), ErrAuthCacheClosed) assert.Eventually(t, func() bool { return authCacheSweepers(t, gid) == 0 diff --git a/hscontrol/state/auth_tagged_expiry_test.go b/hscontrol/state/auth_tagged_expiry_test.go index 18cc6d36e..77159c6c2 100644 --- a/hscontrol/state/auth_tagged_expiry_test.go +++ b/hscontrol/state/auth_tagged_expiry_test.go @@ -81,7 +81,7 @@ func TestTaggedReauthKeepsNilExpiry(t *testing.T) { } authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(regData)) + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(regData))) finalNode, _, err := s.HandleNodeFromAuthPath( authID, @@ -400,12 +400,12 @@ func TestAuthPathOnTaggedMachine(t *testing.T) { t.Helper() authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: mk, NodeKey: nodeKey, Hostname: "multi", Hostinfo: &tailcfg.Hostinfo{Hostname: "multi"}, - })) + }))) node, _, err := s.HandleNodeFromAuthPath(authID, types.UserID(user.ID), nil, util.RegisterMethodOIDC) @@ -443,12 +443,12 @@ func TestUserKeyOnLoginJoinedTaggedNode(t *testing.T) { nodeKey := key.NewNode().Public() authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: machineKey, NodeKey: nodeKey, Hostname: "login-node", Hostinfo: &tailcfg.Hostinfo{Hostname: "login-node"}, - })) + }))) joined, _, err := s.HandleNodeFromAuthPath(authID, types.UserID(owner.ID), nil, util.RegisterMethodOIDC) require.NoError(t, err) @@ -599,7 +599,7 @@ func (n seededTaggedNode) reauth(t *testing.T, authUser *types.User, requestTags rd.Expiry = clientExpiry authID := types.MustAuthID() - n.s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&rd)) + require.NoError(t, n.s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&rd))) node, _, err := n.s.HandleNodeFromAuthPath( authID, @@ -1001,7 +1001,7 @@ func TestIssue3371_TaggedNodeInteractiveReloginAfterLogout(t *testing.T) { }, } authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(regData)) + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(regData))) relogged, _, err := s.HandleNodeFromAuthPath( authID, types.UserID(user.ID), nil, util.RegisterMethodOIDC, diff --git a/hscontrol/state/debug.go b/hscontrol/state/debug.go index 2d26e2d11..706f17bd8 100644 --- a/hscontrol/state/debug.go +++ b/hscontrol/state/debug.go @@ -181,7 +181,8 @@ func (s *State) DebugRegistrationCache() map[string]any { return map[string]any{ "type": "expirable-lru", "expiration": registerCacheExpiration.String(), - "max_entries": defaultRegisterCacheMaxEntries, + "max_entries": s.authCacheMaxEntries, + "ssh_reserve": authCacheSSHReserve, "current_len": s.authCache.Len(), "status": "active", } diff --git a/hscontrol/state/persist_test.go b/hscontrol/state/persist_test.go index 9e0c38e7d..2c08dcce3 100644 --- a/hscontrol/state/persist_test.go +++ b/hscontrol/state/persist_test.go @@ -1387,11 +1387,11 @@ func TestAuthPathRejectsUnencodableExpiry(t *testing.T) { user := s.CreateUserForTest("auth-path-user") nodeKey := key.NewNode().Public() authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: key.NewMachine().Public(), NodeKey: nodeKey, Hostname: "auth-path-node", - })) + }))) // How go-oidc reads an exp claim in the year 10000. beyond := time.Unix(253402300800, 0) @@ -1419,11 +1419,11 @@ func TestAcceptedExpirySurvivesRestart(t *testing.T) { user := s.CreateUserForTest("oidc-user") authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: key.NewMachine().Public(), NodeKey: key.NewNode().Public(), Hostname: "oidc-node", - })) + }))) oidcNode, _, err := s.HandleNodeFromAuthPath( authID, types.UserID(user.ID), &east, util.RegisterMethodOIDC, @@ -1648,12 +1648,12 @@ func TestReauthKeepsOperationalMetadata(t *testing.T) { "precondition: the node serves its route") authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: node.MachineKey(), NodeKey: key.NewNode().Public(), Hostname: node.Hostname(), Hostinfo: &tailcfg.Hostinfo{Hostname: node.Hostname()}, - })) + }))) relogged, _, err := s.HandleNodeFromAuthPath( authID, node.TypedUserID(), nil, util.RegisterMethodOIDC, @@ -1696,10 +1696,10 @@ func TestReauthRejectsOversizedMergedHostinfo(t *testing.T) { require.NoError(t, types.ValidateHostinfo(&tailcfg.Hostinfo{Hostname: hostname})) authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: node.MachineKey(), NodeKey: key.NewNode().Public(), Hostname: hostname, Hostinfo: &tailcfg.Hostinfo{Hostname: hostname}, - })) + }))) _, _, err = s.HandleNodeFromAuthPath(authID, node.TypedUserID(), nil, util.RegisterMethodOIDC) require.ErrorIs(t, err, types.ErrHostinfoTooLarge) @@ -1734,11 +1734,11 @@ func TestReauthKeepsTaggedFutureExpiry(t *testing.T) { require.True(t, node.IsTagged()) authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: node.MachineKey(), NodeKey: key.NewNode().Public(), Hostname: node.Hostname(), Hostinfo: &tailcfg.Hostinfo{Hostname: node.Hostname(), RequestTags: []string{"tag:ci"}}, - })) + }))) relogged, _, err := s.HandleNodeFromAuthPath(authID, userID, nil, util.RegisterMethodOIDC) require.NoError(t, err) require.Equal(t, nodeID, relogged.ID()) @@ -1767,11 +1767,11 @@ func TestReauthReplacesReportedTags(t *testing.T) { for _, tags := range [][]string{{"tag:a"}, nil, {"tag:b"}} { want := slices.Clone(tags) authID := types.MustAuthID() - s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ + require.NoError(t, s.SetAuthCacheEntry(authID, types.NewRegisterAuthRequest(&types.RegistrationData{ MachineKey: node.MachineKey(), NodeKey: key.NewNode().Public(), Hostname: node.Hostname(), Hostinfo: &tailcfg.Hostinfo{Hostname: node.Hostname(), RequestTags: tags}, - })) + }))) relogged, _, err := s.HandleNodeFromAuthPath(authID, userID, nil, util.RegisterMethodOIDC) require.NoError(t, err) require.Equal(t, nodeID, relogged.ID()) diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index cf38bf8e8..3be9ce24e 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -50,11 +50,14 @@ const ( // registerCacheExpiration defines how long node registration entries remain in cache. registerCacheExpiration = time.Minute * 15 - // defaultRegisterCacheMaxEntries is the default upper bound on the number - // of pending registration entries the auth cache will hold. Entries keep - // only hostname, OS and requested tags, capped by [types.ValidateHostinfo]. + // defaultRegisterCacheMaxEntries bounds pending authentication transactions. + // Admission refuses new work at capacity instead of invalidating approvals. defaultRegisterCacheMaxEntries = 1024 + // authCacheSSHReserve keeps registered nodes' SSH checks available when + // unauthenticated registrations have filled their allocation. + authCacheSSHReserve = 64 + // defaultNodeStoreBatchSize is the default number of write operations to batch // before rebuilding the in-memory node snapshot. defaultNodeStoreBatchSize = 100 @@ -126,6 +129,18 @@ var nodeUpdateColumns = []string{ // ErrRegistrationExpired is returned when a registration has expired. var ErrRegistrationExpired = errors.New("registration expired") +// ErrAuthCacheFull is returned when pending authentication has reached capacity. +var ErrAuthCacheFull = errors.New("authentication capacity reached; retry later") + +// ErrAuthCacheEntryExists prevents replacing an existing authentication binding. +var ErrAuthCacheEntryExists = errors.New("authentication request already exists") + +// ErrAuthCacheClosed prevents admitting authentication after shutdown. +var ErrAuthCacheClosed = errors.New("authentication cache closed") + +// ErrAuthRequestFinished prevents reusing an already resolved verdict channel. +var ErrAuthRequestFinished = errors.New("authentication request already finished") + // ErrNodeKeyInUse is returned when a registration or re-auth claims a NodeKey // already bound to a different machine, enforcing the 1:1 NodeKey<->MachineKey // binding. @@ -167,12 +182,14 @@ type State struct { // polMan handles policy evaluation and management polMan policy.PolicyManager - // authCache holds any pending authentication requests from either auth - // type (Web and OIDC). It is a bounded LRU keyed by AuthID; oldest - // entries are evicted once the size cap is reached, and entries that - // time out have their auth verdict resolved with ErrRegistrationExpired - // via the eviction callback so any waiting goroutines wake. - authCache *expirable.LRU[types.AuthID, *types.AuthRequest] + // authCache holds pending Web, OIDC and SSH authentication requests. + // Admission preserves unexpired entries at capacity. Expiry resolves the + // verdict with ErrRegistrationExpired so waiting goroutines wake. + authCache *expirable.LRU[types.AuthID, *types.AuthRequest] + authCacheMu sync.Mutex + authCacheMaxEntries int + authCacheClosed bool + authCacheRegularEntries atomic.Int64 // pings tracks pending ping requests and their response channels. pings *pingTracker @@ -353,29 +370,33 @@ func NewState(cfg *types.Config) (*State, error) { nodeStore := NewNodeStore(nodes, policyPeersFunc(polMan), batchSize, batchTimeout) nodeStore.Start() - // Last: NewLRU starts the TTL sweeper goroutine, which only Close stops. - authCache := expirable.NewLRU[types.AuthID, *types.AuthRequest]( - cacheMaxEntries, - func(id types.AuthID, rn *types.AuthRequest) { - rn.FinishAuth(types.AuthVerdict{Err: ErrRegistrationExpired}) - }, - cacheExpiration, - ) - s := &State{ cfg: cfg, - db: db, - ipAlloc: ipAlloc, - polMan: polMan, - authCache: authCache, - nodeStore: nodeStore, - pings: newPingTracker(), + db: db, + ipAlloc: ipAlloc, + polMan: polMan, + authCacheMaxEntries: cacheMaxEntries, + nodeStore: nodeStore, + pings: newPingTracker(), sshCheckAuth: make(map[sshCheckPair]time.Time), registerLocks: xsync.NewMap[key.MachinePublic, *registerLock](), } + // Last: NewLRU starts the TTL sweeper goroutine, which only Close stops. + s.authCache = expirable.NewLRU[types.AuthID, *types.AuthRequest]( + cacheMaxEntries+authCacheSSHReserve, + func(_ types.AuthID, rn *types.AuthRequest) { + if !rn.IsSSHCheck() { + s.authCacheRegularEntries.Add(-1) + } + + rn.FinishAuth(types.AuthVerdict{Err: ErrRegistrationExpired}) + }, + cacheExpiration, + ) + // Surface nodes whose stored data would break map generation (e.g. an // invalid given name from a legacy row) so an operator can fix them. This // only logs; it never mutates a node's stored name at boot. @@ -411,6 +432,10 @@ func policyPeersFunc(pm policy.PolicyManager) PeersFunc { // Close gracefully shuts down the [State] instance and releases all resources. func (s *State) Close() error { + s.authCacheMu.Lock() + s.authCacheClosed = true + s.authCacheMu.Unlock() + s.pings.drain() s.nodeStore.Stop() // Purge before Close: Close stops the sweeper, so anything still cached @@ -1931,12 +1956,71 @@ func (s *State) DeleteExpiredAccessTokens(cutoff time.Time) (int64, error) { // GetAuthCacheEntry retrieves a pending auth request from the cache. func (s *State) GetAuthCacheEntry(id types.AuthID) (*types.AuthRequest, bool) { - return s.authCache.Get(id) + // Preserve insertion order so admission can reclaim expired entries from + // the oldest end without scanning every pending authentication. + return s.authCache.Peek(id) } -// SetAuthCacheEntry stores a pending auth request in the cache. -func (s *State) SetAuthCacheEntry(id types.AuthID, entry *types.AuthRequest) { +// SetAuthCacheEntry admits a pending authentication without displacing existing +// transactions. Re-admitting the same request does not extend its lifetime. +func (s *State) SetAuthCacheEntry(id types.AuthID, entry *types.AuthRequest) error { + s.authCacheMu.Lock() + defer s.authCacheMu.Unlock() + + if s.authCacheClosed { + return ErrAuthCacheClosed + } + + if existing, ok := s.authCache.Peek(id); ok { + if existing == entry { + return nil + } + + return ErrAuthCacheEntryExists + } + + // Remove an expired entry with this ID before Add can overwrite its waiter. + s.authCache.Remove(id) + + if entry.IsFinished() { + return ErrAuthRequestFinished + } + + for { + oldest, _, ok := s.authCache.GetOldest() + if !ok { + break + } + + if _, live := s.authCache.Peek(oldest); live { + break + } + + s.authCache.Remove(oldest) + } + + if (!entry.IsSSHCheck() && s.authCacheRegularEntries.Load() >= int64(s.authCacheMaxEntries)) || + s.authCache.Len() >= s.authCacheMaxEntries+authCacheSSHReserve { + return ErrAuthCacheFull + } + + err := entry.PrepareForCache() + if err != nil { + return err + } + + if !entry.IsSSHCheck() { + s.authCacheRegularEntries.Add(1) + } + s.authCache.Add(id, entry) + + return nil +} + +// DeleteAuthCacheEntry releases a consumed authentication transaction. +func (s *State) DeleteAuthCacheEntry(id types.AuthID) { + s.authCache.Remove(id) } // DeleteAuthCacheEntryForTest drops a pending auth request from the cache, @@ -1944,7 +2028,7 @@ func (s *State) SetAuthCacheEntry(id types.AuthID, entry *types.AuthRequest) { // (expired, evicted, or dropped on a control-plane restart) without faking an // auth_id. func (s *State) DeleteAuthCacheEntryForTest(id types.AuthID) { - s.authCache.Remove(id) + s.DeleteAuthCacheEntry(id) } // SetLastSSHAuth records a successful SSH check authentication diff --git a/hscontrol/types/common.go b/hscontrol/types/common.go index 05edd706f..d759fa466 100644 --- a/hscontrol/types/common.go +++ b/hscontrol/types/common.go @@ -4,6 +4,7 @@ package types //go:generate go run tailscale.com/cmd/viewer --type=User,Node,PreAuthKey,Credential import ( + "encoding/json" "errors" "fmt" "runtime" @@ -22,9 +23,10 @@ const ( // Common errors. var ( - ErrCannotParsePrefix = errors.New("cannot parse prefix") - ErrInvalidAuthIDLength = errors.New("auth ID has invalid length") - ErrInvalidAuthIDPrefix = errors.New("auth ID has invalid prefix") + ErrCannotParsePrefix = errors.New("cannot parse prefix") + ErrInvalidAuthIDLength = errors.New("auth ID has invalid length") + ErrInvalidAuthIDPrefix = errors.New("auth ID has invalid prefix") + ErrRegistrationPayloadTooLarge = errors.New("registration payload exceeds size limit") ) const ( @@ -109,7 +111,7 @@ type PendingRegistrationConfirmation struct { // AuthRequest represents a pending authentication request from a user or a // node. It carries the minimum data needed to either complete a node -// registration (regData populated) or an SSH check-mode auth (sshBinding +// registration (regData or regDataJSON populated) or an SSH check-mode auth (sshBinding // populated), and signals the verdict via the finished channel. The closed // flag guards [AuthRequest.FinishAuth] against double-close. // @@ -117,14 +119,14 @@ type PendingRegistrationConfirmation struct { // have a single canonical instance even when stored in caches that // internally copy values. type AuthRequest struct { - // regData is populated for node-registration flows (interactive web - // or OIDC). It carries the cached registration payload that the - // auth callback uses to promote this request into a real node. - // - // nil for non-registration flows. Use [AuthRequest.RegistrationData] to read it - // safely. + // regData holds the caller's registration payload until PrepareForCache + // replaces it with owned serialized data and releases the original pointer. regData *RegistrationData + // regDataJSON is immutable and bounded by MaxRegistrationPayloadBytes. + // RegistrationData decodes a fresh copy without retaining the decoded graph. + regDataJSON string + // sshBinding is populated for SSH check-mode flows. It captures the // (src, dst) node pair the request was minted for so the follow-up // and OIDC callback can refuse to record a verdict for any other @@ -160,7 +162,8 @@ func NewAuthRequest() *AuthRequest { // NewRegisterAuthRequest creates a pending auth request carrying the // minimal [RegistrationData] for a node-registration flow. The data is -// stored by pointer; callers must not mutate it after handing it off. +// stored by pointer until [AuthRequest.PrepareForCache] detaches it; callers +// must not mutate it until preparation finishes. func NewRegisterAuthRequest(data *RegistrationData) *AuthRequest { return &AuthRequest{ regData: data, @@ -169,6 +172,31 @@ func NewRegisterAuthRequest(data *RegistrationData) *AuthRequest { } } +// PrepareForCache replaces the registration payload with an owned serialized +// value of at most [MaxRegistrationPayloadBytes]. It must run before publishing +// the request to the cache. Serialization and size errors leave it unchanged; +// callers must reject admission on error. Repeated calls and non-registration +// requests are no-ops. +func (rn *AuthRequest) PrepareForCache() error { + if rn.regData == nil { + return nil + } + + data, err := json.Marshal(rn.regData) + if err != nil { + return fmt.Errorf("serializing pending registration: %w", err) + } + + if len(data) > MaxRegistrationPayloadBytes { + return fmt.Errorf("%w: limit %d bytes", ErrRegistrationPayloadTooLarge, MaxRegistrationPayloadBytes) + } + + rn.regDataJSON = string(data) + rn.regData = nil + + return nil +} + // NewSSHCheckAuthRequest creates a pending auth request bound to a // specific (src, dst) SSH check-mode pair. The follow-up handler and // OIDC callback must verify their incoming request matches this binding @@ -184,15 +212,27 @@ func NewSSHCheckAuthRequest(src, dst NodeID) *AuthRequest { } } -// RegistrationData returns the cached registration payload. It panics if +// RegistrationData returns the original registration payload before preparation, +// or a fresh decoded copy after [AuthRequest.PrepareForCache]. It panics if // called on an [AuthRequest] that was not created via // [NewRegisterAuthRequest]. func (rn *AuthRequest) RegistrationData() *RegistrationData { - if rn.regData == nil { + if rn.regData != nil { + return rn.regData + } + + if rn.regDataJSON == "" { panic("RegistrationData can only be used in registration requests") } - return rn.regData + var data RegistrationData + + err := json.Unmarshal([]byte(rn.regDataJSON), &data) + if err != nil { + panic(fmt.Sprintf("decoding prepared registration: %v", err)) + } + + return &data } // SSHCheckBinding returns the (src, dst) node pair an SSH check-mode @@ -209,7 +249,7 @@ func (rn *AuthRequest) SSHCheckBinding() *SSHCheckBinding { // IsRegistration reports whether this auth request carries registration // data (i.e. it was created via [NewRegisterAuthRequest]). func (rn *AuthRequest) IsRegistration() bool { - return rn.regData != nil + return rn.regData != nil || rn.regDataJSON != "" } // IsSSHCheck reports whether this auth request is bound to an SSH @@ -219,6 +259,11 @@ func (rn *AuthRequest) IsSSHCheck() bool { return rn.sshBinding != nil } +// IsFinished reports whether a verdict has already resolved this request. +func (rn *AuthRequest) IsFinished() bool { + return rn.closed.Load() +} + // SetPendingConfirmation marks this [AuthRequest] as having an // OIDC-resolved user that is waiting to confirm the registration on // the interstitial. The OIDC callback should call this and then render diff --git a/hscontrol/types/common_test.go b/hscontrol/types/common_test.go index d381d8723..67fe48417 100644 --- a/hscontrol/types/common_test.go +++ b/hscontrol/types/common_test.go @@ -1,10 +1,18 @@ package types import ( + "encoding/json" + "fmt" + "math" + "net/netip" + "strings" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "tailscale.com/tailcfg" + "tailscale.com/types/key" ) // TestNewSSHCheckAuthRequestBinding verifies that an SSH-check [AuthRequest] @@ -44,6 +52,195 @@ func TestNewRegisterAuthRequestPayload(t *testing.T) { }, "SSHCheckBinding() must panic on a registration AuthRequest") } +func TestAuthRequestPrepareForCacheOwnsRegistration(t *testing.T) { + expiry := time.Date(2027, time.January, 2, 3, 4, 5, 6, time.UTC) + data := &RegistrationData{ + MachineKey: key.NewMachine().Public(), + NodeKey: key.NewNode().Public(), + DiscoKey: key.NewDisco().Public(), + Hostname: "router-a", + Hostinfo: &tailcfg.Hostinfo{ + Hostname: "router-a", + IPNVersion: "1.104.0", + OS: "linux", + OSVersion: "6.18.0", + Distro: "nixos", + Container: "false", + Userspace: "true", + RequestTags: []string{"tag:router", "tag:production"}, + RoutableIPs: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/24"), netip.MustParsePrefix("2001:db8::/64")}, + WoLMACs: []string{"00:11:22:33:44:55"}, + Services: []tailcfg.Service{{Proto: tailcfg.TCP, Port: 22, Description: "ssh"}}, + SSH_HostKeys: []string{"ssh-ed25519 initial-host-key"}, + NetInfo: &tailcfg.NetInfo{ + WorkingUDP: "true", + WorkingIPv6: "false", + PreferredDERP: 1, + DERPLatency: map[string]float64{"1-v4": 0.01, "2-v6": 0.02}, + }, + Location: &tailcfg.Location{Country: "Norway", City: "Tromsø", Latitude: 69.65, Longitude: 18.96}, + TPM: &tailcfg.TPMInfo{Manufacturer: "MSFT", Model: 1}, + }, + Endpoints: []netip.AddrPort{netip.MustParseAddrPort("192.0.2.1:41641"), netip.MustParseAddrPort("[2001:db8::1]:41641")}, + Expiry: &expiry, + } + want := *data + want.Hostinfo = data.Hostinfo.Clone() + want.Endpoints = append([]netip.AddrPort(nil), data.Endpoints...) + wantExpiry := expiry + want.Expiry = &wantExpiry + + req := NewRegisterAuthRequest(data) + require.NoError(t, req.PrepareForCache()) + assert.True(t, req.IsRegistration()) + assert.False(t, req.IsSSHCheck()) + assert.Nil(t, req.regData, "prepared requests must release the decoded payload") + assert.Equal(t, &want, req.RegistrationData()) + + data.Hostname = "changed" + data.Hostinfo.RequestTags[0] = "tag:changed" + data.Hostinfo.RoutableIPs[0] = netip.MustParsePrefix("192.0.2.0/24") + data.Hostinfo.Services[0].Description = "changed" + data.Hostinfo.NetInfo.DERPLatency["1-v4"] = 1 + data.Hostinfo.Location.City = "changed" + data.Hostinfo.TPM.Manufacturer = "changed" + data.Endpoints[0] = netip.MustParseAddrPort("192.0.2.2:1") + *data.Expiry = time.Time{} + + assert.Equal(t, &want, req.RegistrationData(), "the producer must not retain ownership of cached data") + + got := req.RegistrationData() + got.Hostname = "changed by consumer" + got.Hostinfo.RequestTags[0] = "tag:consumer" + got.Hostinfo.Services[0].Description = "changed by consumer" + got.Hostinfo.NetInfo.DERPLatency["1-v4"] = 2 + got.Endpoints[0] = netip.MustParseAddrPort("192.0.2.3:1") + *got.Expiry = time.Time{} + + require.NoError(t, req.PrepareForCache()) + assert.Equal(t, &want, req.RegistrationData(), "access and repeated preparation must leave cached data unchanged") +} + +func TestAuthRequestPrepareForCachePayloadLimit(t *testing.T) { + const limit = MaxRegistrationPayloadBytes + + for _, test := range []struct { + name string + size int + }{ + {name: "below limit", size: limit - 1}, + {name: "exact limit", size: limit}, + {name: "over limit", size: limit + 1}, + } { + t.Run(test.name, func(t *testing.T) { + data := &RegistrationData{Hostinfo: &tailcfg.Hostinfo{OS: "linux"}} + base, err := json.Marshal(data) + require.NoError(t, err) + + data.Hostinfo.OS = strings.Repeat("a", test.size-len(base)+len(data.Hostinfo.OS)) + encoded, err := json.Marshal(data) + require.NoError(t, err) + require.Len(t, encoded, test.size) + + req := NewRegisterAuthRequest(data) + + err = req.PrepareForCache() + if test.size > limit { + require.ErrorIs(t, err, ErrRegistrationPayloadTooLarge, "oversize payloads must be rejected instead of truncated") + assert.Empty(t, req.regDataJSON) + assert.Same(t, data, req.RegistrationData(), "failed preparation must leave the original data intact") + + return + } + + require.NoError(t, err) + assert.Nil(t, req.regData) + assert.Len(t, req.regDataJSON, test.size) + assert.Equal(t, data, req.RegistrationData()) + }) + } +} + +func TestAuthRequestPrepareForCacheRouteAndServiceMetadata(t *testing.T) { + data := &RegistrationData{ + MachineKey: key.NewMachine().Public(), + NodeKey: key.NewNode().Public(), + Hostname: "regional-router", + Hostinfo: &tailcfg.Hostinfo{ + Hostname: "regional-router", + OS: "linux", + OSVersion: "6.18.0", + IPNVersion: "1.104.0", + RequestTags: []string{"tag:router", "tag:production"}, + }, + } + for i := range 128 { + data.Hostinfo.RoutableIPs = append(data.Hostinfo.RoutableIPs, + netip.MustParsePrefix(fmt.Sprintf("10.42.%d.0/24", i)), + netip.MustParsePrefix(fmt.Sprintf("2001:db8:%x::/64", i))) + } + + for port := uint16(8000); port < 8064; port++ { + data.Hostinfo.Services = append(data.Hostinfo.Services, tailcfg.Service{ + Proto: tailcfg.TCP, + Port: port, + Description: fmt.Sprintf("regional-service-%d", port), + }) + } + + req := NewRegisterAuthRequest(data) + require.NoError(t, req.PrepareForCache()) + assert.Nil(t, req.regData) + assert.LessOrEqual(t, len(req.regDataJSON), MaxRegistrationPayloadBytes) + assert.Equal(t, data, req.RegistrationData()) +} + +func TestAuthRequestPrepareForCacheInvalidPayload(t *testing.T) { + invalidExpiry := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + for _, test := range []struct { + name string + data *RegistrationData + }{ + {name: "invalid expiry", data: &RegistrationData{Expiry: &invalidExpiry}}, + {name: "invalid latency", data: &RegistrationData{Hostinfo: &tailcfg.Hostinfo{ + NetInfo: &tailcfg.NetInfo{DERPLatency: map[string]float64{"1-v4": math.NaN()}}, + }}}, + } { + t.Run(test.name, func(t *testing.T) { + req := NewRegisterAuthRequest(test.data) + require.Error(t, req.PrepareForCache()) + assert.Empty(t, req.regDataJSON) + assert.Same(t, test.data, req.RegistrationData()) + }) + } +} + +func TestAuthRequestPrepareForCacheNilValues(t *testing.T) { + for _, data := range []*RegistrationData{ + {}, + {Hostinfo: &tailcfg.Hostinfo{}}, + {Hostinfo: &tailcfg.Hostinfo{NetInfo: &tailcfg.NetInfo{}}}, + {Expiry: &time.Time{}}, + {Endpoints: []netip.AddrPort{}}, + } { + req := NewRegisterAuthRequest(data) + require.NoError(t, req.PrepareForCache()) + assert.True(t, req.IsRegistration()) + assert.Equal(t, data, req.RegistrationData()) + } + + for _, req := range []*AuthRequest{NewAuthRequest(), NewRegisterAuthRequest(nil), NewSSHCheckAuthRequest(7, 11)} { + require.NoError(t, req.PrepareForCache()) + require.NoError(t, req.PrepareForCache()) + assert.False(t, req.IsRegistration()) + assert.Panics(t, func() { _ = req.RegistrationData() }) + + if req.IsSSHCheck() { + assert.Equal(t, &SSHCheckBinding{SrcNodeID: 7, DstNodeID: 11}, req.SSHCheckBinding()) + } + } +} + // TestNewAuthRequestEmptyPayload verifies that a payload-less // [AuthRequest] reports both Is* helpers as false and panics on either // payload accessor. diff --git a/hscontrol/types/config.go b/hscontrol/types/config.go index 29fa7283d..ca7eaace7 100644 --- a/hscontrol/types/config.go +++ b/hscontrol/types/config.go @@ -462,9 +462,9 @@ type Tuning struct { // valid before being eligible for eviction. RegisterCacheExpiration time.Duration - // RegisterCacheMaxEntries bounds the number of pending registration - // entries the auth cache will hold. Older entries are evicted (LRU) - // when the cap is reached, preventing unauthenticated cache-fill DoS. + // RegisterCacheMaxEntries bounds pending registrations. At capacity, new + // registrations are refused without evicting admitted sessions. SSH checks + // can also use 64 reserved slots beyond this limit. // A value of 0 falls back to defaultRegisterCacheMaxEntries (1024). RegisterCacheMaxEntries int diff --git a/hscontrol/types/registration.go b/hscontrol/types/registration.go index be984b589..12bc749d0 100644 --- a/hscontrol/types/registration.go +++ b/hscontrol/types/registration.go @@ -8,14 +8,25 @@ import ( "tailscale.com/types/key" ) +// MaxRegistrationPayloadBytes limits the owned serialized registration payload +// retained by an [AuthRequest] after [AuthRequest.PrepareForCache]. It covers +// the authentication metadata and the minimal registration Hostinfo. +// It bounds retained serialized bytes, not transient encoding or decoded heap +// usage. A cache of 1024 entries retains at most 64 MiB of these payloads plus +// record overhead; a different entry limit scales that budget. +const MaxRegistrationPayloadBytes = 64 << 10 + // RegistrationData is the payload cached for a pending node registration. // It replaces the previous practice of caching a full *[Node] and carries // only the fields the registration callback path actually consumes when // promoting a pending registration to a real node. // -// Any client can open a pending registration. The bounded-LRU cache that -// holds these entries caps their number. Hostinfo retains only auth and +// Any client can open a pending registration. Cache admission caps the number +// of pending entries. Hostinfo retains only auth and // confirmation metadata, bounded by [ValidateHostinfo] before caching. +// +// [AuthRequest.PrepareForCache] retains an owned serialized representation +// bounded by [MaxRegistrationPayloadBytes], releasing this decoded payload. type RegistrationData struct { // MachineKey is the cryptographic identity of the machine being // registered. Required.