hscontrol/oidc: fix ACL policy not applied to new OIDC nodes (#2890)

Fixes #2888
Fixes #2896

hscontrol/auth.go

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/juanfont/headscale/hscontrol/types"
-	"github.com/juanfont/headscale/hscontrol/types/change"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/rs/zerolog/log"
 	"gorm.io/gorm"
@@ -364,16 +363,13 @@ func (h *Headscale) handleRegisterWithAuthKey(
 	// eventbus.
 	// TODO(kradalby): This needs to be ran as part of the batcher maybe?
 	// now since we dont update the node/pol here anymore
-	routeChange := h.state.AutoApproveRoutes(node)
-
-	if _, _, err := h.state.SaveNode(node); err != nil {
-		return nil, fmt.Errorf("saving auto approved routes to node: %w", err)
+	routesChange, err := h.state.AutoApproveRoutes(node)
+	if err != nil {
+		return nil, fmt.Errorf("auto approving routes: %w", err)
 	}
 
-	if routeChange && changed.Empty() {
-		changed = change.NodeAdded(node.ID())
-	}
-	h.Change(changed)
+	// Send both changes. Empty changes are ignored by Change().
+	h.Change(changed, routesChange)
 
 	// TODO(kradalby): I think this is covered above, but we need to validate that.
 	// // If policy changed due to node registration, send a separate policy change

hscontrol/grpcv1.go

@@ -312,13 +312,13 @@ func (api headscaleV1APIServer) RegisterNode(
 	// ensure we send an update.
 	// This works, but might be another good candidate for doing some sort of
 	// eventbus.
-	_ = api.h.state.AutoApproveRoutes(node)
-	_, _, err = api.h.state.SaveNode(node)
+	routeChange, err := api.h.state.AutoApproveRoutes(node)
 	if err != nil {
-		return nil, fmt.Errorf("saving auto approved routes to node: %w", err)
+		return nil, fmt.Errorf("auto approving routes: %w", err)
 	}
 
-	api.h.Change(nodeChange)
+	// Send both changes. Empty changes are ignored by Change().
+	api.h.Change(nodeChange, routeChange)
 
 	return &v1.RegisterNodeResponse{Node: node.Proto()}, nil
 }

hscontrol/oidc.go

@@ -41,10 +41,6 @@ var (
 	errOIDCAllowedUsers  = errors.New(
 		"authenticated principal does not match any allowed user",
 	)
-	errOIDCInvalidNodeState = errors.New(
-		"requested node state key expired before authorisation completed",
-	)
-	errOIDCNodeKeyMissing = errors.New("could not get node key from cache")
 )
 
 // RegistrationInfo contains both machine key and verifier information for OIDC validation.
@@ -107,16 +103,8 @@ func (a *AuthProviderOIDC) AuthURL(registrationID types.RegistrationID) string {
 		registrationID.String())
 }
 
-func (a *AuthProviderOIDC) determineNodeExpiry(idTokenExpiration time.Time) time.Time {
-	if a.cfg.UseExpiryFromToken {
-		return idTokenExpiration
-	}
-
-	return time.Now().Add(a.cfg.Expiry)
-}
-
-// RegisterOIDC redirects to the OIDC provider for authentication
-// Puts NodeKey in cache so the callback can retrieve it using the oidc state param
+// RegisterHandler registers the OIDC callback handler with the given router.
+// It puts NodeKey in cache so the callback can retrieve it using the oidc state param.
 // Listens in /register/:registration_id.
 func (a *AuthProviderOIDC) RegisterHandler(
 	writer http.ResponseWriter,
@@ -296,7 +284,7 @@ func (a *AuthProviderOIDC) OIDCCallbackHandler(
 		return
 	}
 
-	user, c, err := a.createOrUpdateUserFromClaim(&claims)
+	user, _, err := a.createOrUpdateUserFromClaim(&claims)
 	if err != nil {
 		log.Error().
 			Err(err).
@@ -315,9 +303,6 @@ func (a *AuthProviderOIDC) OIDCCallbackHandler(
 		return
 	}
 
-	// Send policy update notifications if needed
-	a.h.Change(c)
-
 	// 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
@@ -364,6 +349,14 @@ func (a *AuthProviderOIDC) OIDCCallbackHandler(
 	httpError(writer, NewHTTPError(http.StatusGone, "login session expired, try again", nil))
 }
 
+func (a *AuthProviderOIDC) determineNodeExpiry(idTokenExpiration time.Time) time.Time {
+	if a.cfg.UseExpiryFromToken {
+		return idTokenExpiration
+	}
+
+	return time.Now().Add(a.cfg.Expiry)
+}
+
 func extractCodeAndStateParamFromRequest(
 	req *http.Request,
 ) (string, string, error) {
@@ -496,8 +489,8 @@ func (a *AuthProviderOIDC) createOrUpdateUserFromClaim(
 	}
 
 	// if the user is still not found, create a new empty user.
-	// TODO(kradalby): This might cause us to not have an ID below which
-	// is a problem.
+	// TODO(kradalby): This context is not inherited from the request, which is probably not ideal.
+	// However, we need a context to use the OIDC provider.
 	if user == nil {
 		newUser = true
 		user = &types.User{}
@@ -549,18 +542,13 @@ func (a *AuthProviderOIDC) handleRegistration(
 	// ensure we send an update.
 	// This works, but might be another good candidate for doing some sort of
 	// eventbus.
-	_ = a.h.state.AutoApproveRoutes(node)
-	_, policyChange, err := a.h.state.SaveNode(node)
+	routesChange, err := a.h.state.AutoApproveRoutes(node)
 	if err != nil {
-		return false, fmt.Errorf("saving auto approved routes to node: %w", err)
+		return false, fmt.Errorf("auto approving routes: %w", err)
 	}
 
-	// Policy updates are full and take precedence over node changes.
-	if !policyChange.Empty() {
-		a.h.Change(policyChange)
-	} else {
-		a.h.Change(nodeChange)
-	}
+	// Send both changes. Empty changes are ignored by Change().
+	a.h.Change(nodeChange, routesChange)
 
 	return !nodeChange.Empty(), nil
 }

hscontrol/state/state.go

@@ -826,7 +826,7 @@ func (s *State) SetPolicy(pol []byte) (bool, error) {
 
 // AutoApproveRoutes checks if a node's routes should be auto-approved.
 // AutoApproveRoutes checks if any routes should be auto-approved for a node and updates them.
-func (s *State) AutoApproveRoutes(nv types.NodeView) bool {
+func (s *State) AutoApproveRoutes(nv types.NodeView) (change.ChangeSet, error) {
 	approved, changed := policy.ApproveRoutesWithPolicy(s.polMan, nv, nv.ApprovedRoutes().AsSlice(), nv.AnnouncedRoutes())
 	if changed {
 		log.Debug().
@@ -839,7 +839,7 @@ func (s *State) AutoApproveRoutes(nv types.NodeView) bool {
 
 		// Persist the auto-approved routes to database and NodeStore via SetApprovedRoutes
 		// This ensures consistency between database and NodeStore
-		_, _, err := s.SetApprovedRoutes(nv.ID(), approved)
+		_, c, err := s.SetApprovedRoutes(nv.ID(), approved)
 		if err != nil {
 			log.Error().
 				Uint64("node.id", nv.ID().Uint64()).
@@ -847,13 +847,15 @@ func (s *State) AutoApproveRoutes(nv types.NodeView) bool {
 				Err(err).
 				Msg("Failed to persist auto-approved routes")
 
-			return false
+			return change.EmptySet, err
 		}
 
 		log.Info().Uint64("node.id", nv.ID().Uint64()).Str("node.name", nv.Hostname()).Strs("routes.approved", util.PrefixesToString(approved)).Msg("Routes approved")
+
+		return c, nil
 	}
 
-	return changed
+	return change.EmptySet, nil
 }
 
 // GetPolicy retrieves the current policy from the database.
@@ -1643,6 +1645,7 @@ func (s *State) UpdateNodeFromMapRequest(id types.NodeID, req tailcfg.MapRequest
 		routeChange        bool
 		hostinfoChanged    bool
 		needsRouteApproval bool
+		autoApprovedRoutes []netip.Prefix
 		endpointChanged    bool
 		derpChanged        bool
 	)
@@ -1674,7 +1677,6 @@ func (s *State) UpdateNodeFromMapRequest(id types.NodeID, req tailcfg.MapRequest
 		}
 
 		// Calculate route approval before NodeStore update to avoid calling View() inside callback
-		var autoApprovedRoutes []netip.Prefix
 		var hasNewRoutes bool
 		if hi := req.Hostinfo; hi != nil {
 			hasNewRoutes = len(hi.RoutableIPs) > 0
@@ -1740,7 +1742,6 @@ func (s *State) UpdateNodeFromMapRequest(id types.NodeID, req tailcfg.MapRequest
 					Strs("newApprovedRoutes", util.PrefixesToString(autoApprovedRoutes)).
 					Bool("routeChanged", routeChange).
 					Msg("applying route approval results")
-				currentNode.ApprovedRoutes = autoApprovedRoutes
 			}
 		}
 	})
@@ -1749,6 +1750,24 @@ func (s *State) UpdateNodeFromMapRequest(id types.NodeID, req tailcfg.MapRequest
 		return change.EmptySet, fmt.Errorf("node not found in NodeStore: %d", id)
 	}
 
+	if routeChange {
+		log.Debug().
+			Uint64("node.id", id.Uint64()).
+			Strs("autoApprovedRoutes", util.PrefixesToString(autoApprovedRoutes)).
+			Msg("Persisting auto-approved routes from MapRequest")
+
+		// SetApprovedRoutes will update both database and PrimaryRoutes table
+		_, c, err := s.SetApprovedRoutes(id, autoApprovedRoutes)
+		if err != nil {
+			return change.EmptySet, fmt.Errorf("persisting auto-approved routes: %w", err)
+		}
+
+		// If SetApprovedRoutes resulted in a policy change, return it
+		if !c.Empty() {
+			return c, nil
+		}
+	} // Continue with the rest of the processing using the updated node
+
 	nodeRouteChange := change.EmptySet
 
 	// Handle route changes after NodeStore update
@@ -1763,13 +1782,8 @@ func (s *State) UpdateNodeFromMapRequest(id types.NodeID, req tailcfg.MapRequest
 			routesChangedButNotApproved = true
 		}
 	}
-	if routeChange {
-		needsRouteUpdate = true
-		log.Debug().
-			Caller().
-			Uint64("node.id", id.Uint64()).
-			Msg("updating routes because approved routes changed")
-	} else if routesChangedButNotApproved {
+
+	if routesChangedButNotApproved {
 		needsRouteUpdate = true
 		log.Debug().
 			Caller().
@@ -1824,25 +1838,26 @@ func (s *State) UpdateNodeFromMapRequest(id types.NodeID, req tailcfg.MapRequest
 	return change.NodeAdded(id), nil
 }
 
-func hostinfoEqual(oldNode types.NodeView, new *tailcfg.Hostinfo) bool {
-	if !oldNode.Valid() && new == nil {
+func hostinfoEqual(oldNode types.NodeView, newHI *tailcfg.Hostinfo) bool {
+	if !oldNode.Valid() && newHI == nil {
 		return true
 	}
-	if !oldNode.Valid() || new == nil {
+
+	if !oldNode.Valid() || newHI == nil {
 		return false
 	}
 	old := oldNode.AsStruct().Hostinfo
 
-	return old.Equal(new)
+	return old.Equal(newHI)
 }
 
-func routesChanged(oldNode types.NodeView, new *tailcfg.Hostinfo) bool {
+func routesChanged(oldNode types.NodeView, newHI *tailcfg.Hostinfo) bool {
 	var oldRoutes []netip.Prefix
 	if oldNode.Valid() && oldNode.AsStruct().Hostinfo != nil {
 		oldRoutes = oldNode.AsStruct().Hostinfo.RoutableIPs
 	}
 
-	newRoutes := new.RoutableIPs
+	newRoutes := newHI.RoutableIPs
 	if newRoutes == nil {
 		newRoutes = []netip.Prefix{}
 	}