state: introduce state

this commit moves all of the read and write logic, and all different parts
of headscale that manages some sort of persistent and in memory state into
a separate package.

The goal of this is to clearly define the boundry between parts of the app
which accesses and modifies data, and where it happens. Previously, different
state (routes, policy, db and so on) was used directly, and sometime passed to
functions as pointers.

Now all access has to go through state. In the initial implementation,
most of the same functions exists and have just been moved. In the future
centralising this will allow us to optimise bottle necks with the database
(in memory state) and make the different parts talking to eachother do so
in the same way across headscale components.

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

hscontrol/oidc.go

@@ -17,7 +17,7 @@ import (
 	"github.com/gorilla/mux"
 	"github.com/juanfont/headscale/hscontrol/db"
 	"github.com/juanfont/headscale/hscontrol/notifier"
-	"github.com/juanfont/headscale/hscontrol/policy"
+	"github.com/juanfont/headscale/hscontrol/state"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/rs/zerolog/log"
@@ -28,6 +28,8 @@ import (
 const (
 	randomByteSize           = 16
 	defaultOAuthOptionsCount = 3
+	registerCacheExpiration  = time.Minute * 15
+	registerCacheCleanup     = time.Minute * 20
 )
 
 var (
@@ -56,11 +58,9 @@ type RegistrationInfo struct {
 type AuthProviderOIDC struct {
 	serverURL         string
 	cfg               *types.OIDCConfig
-	db                *db.HSDatabase
+	state             *state.State
 	registrationCache *zcache.Cache[string, RegistrationInfo]
 	notifier          *notifier.Notifier
-	ipAlloc           *db.IPAllocator
-	polMan            policy.PolicyManager
 
 	oidcProvider *oidc.Provider
 	oauth2Config *oauth2.Config
@@ -70,10 +70,8 @@ func NewAuthProviderOIDC(
 	ctx context.Context,
 	serverURL string,
 	cfg *types.OIDCConfig,
-	db *db.HSDatabase,
+	state *state.State,
 	notif *notifier.Notifier,
-	ipAlloc *db.IPAllocator,
-	polMan policy.PolicyManager,
 ) (*AuthProviderOIDC, error) {
 	var err error
 	// grab oidc config if it hasn't been already
@@ -101,11 +99,9 @@ func NewAuthProviderOIDC(
 	return &AuthProviderOIDC{
 		serverURL:         serverURL,
 		cfg:               cfg,
-		db:                db,
+		state:             state,
 		registrationCache: registrationCache,
 		notifier:          notif,
-		ipAlloc:           ipAlloc,
-		polMan:            polMan,
 
 		oidcProvider: oidcProvider,
 		oauth2Config: oauth2Config,
@@ -305,12 +301,31 @@ func (a *AuthProviderOIDC) OIDCCallbackHandler(
 		}
 	}
 
-	user, err := a.createOrUpdateUserFromClaim(&claims)
+	user, policyChanged, err := a.createOrUpdateUserFromClaim(&claims)
 	if err != nil {
-		httpError(writer, err)
+		log.Error().
+			Err(err).
+			Caller().
+			Msgf("could not create or update user")
+		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
+		writer.WriteHeader(http.StatusInternalServerError)
+		_, werr := writer.Write([]byte("Could not create or update user"))
+		if werr != nil {
+			log.Error().
+				Caller().
+				Err(werr).
+				Msg("Failed to write response")
+		}
+
 		return
 	}
 
+	// Send policy update notifications if needed
+	if policyChanged {
+		ctx := types.NotifyCtx(context.Background(), "oidc-user-created", user.Name)
+		a.notifier.NotifyAll(ctx, types.UpdateFull())
+	}
+
 	// 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
@@ -472,31 +487,40 @@ func (a *AuthProviderOIDC) getRegistrationIDFromState(state string) *types.Regis
 
 func (a *AuthProviderOIDC) createOrUpdateUserFromClaim(
 	claims *types.OIDCClaims,
-) (*types.User, error) {
+) (*types.User, bool, error) {
 	var user *types.User
 	var err error
-	user, err = a.db.GetUserByOIDCIdentifier(claims.Identifier())
+	var newUser bool
+	var policyChanged bool
+	user, err = a.state.GetUserByOIDCIdentifier(claims.Identifier())
 	if err != nil && !errors.Is(err, db.ErrUserNotFound) {
-		return nil, fmt.Errorf("creating or updating user: %w", err)
+		return nil, false, fmt.Errorf("creating or updating user: %w", err)
 	}
 
 	// if the user is still not found, create a new empty user.
 	if user == nil {
+		newUser = true
 		user = &types.User{}
 	}
 
 	user.FromClaim(claims)
-	err = a.db.DB.Save(user).Error
-	if err != nil {
-		return nil, fmt.Errorf("creating or updating user: %w", err)
+
+	if newUser {
+		user, policyChanged, err = a.state.CreateUser(*user)
+		if err != nil {
+			return nil, false, fmt.Errorf("creating user: %w", err)
+		}
+	} else {
+		_, policyChanged, err = a.state.UpdateUser(types.UserID(user.ID), func(u *types.User) error {
+			*u = *user
+			return nil
+		})
+		if err != nil {
+			return nil, false, fmt.Errorf("updating user: %w", err)
+		}
 	}
 
-	err = usersChangedHook(a.db, a.polMan, a.notifier)
-	if err != nil {
-		return nil, fmt.Errorf("updating resources using user: %w", err)
-	}
-
-	return user, nil
+	return user, policyChanged, nil
 }
 
 func (a *AuthProviderOIDC) handleRegistration(
@@ -504,47 +528,40 @@ func (a *AuthProviderOIDC) handleRegistration(
 	registrationID types.RegistrationID,
 	expiry time.Time,
 ) (bool, error) {
-	ipv4, ipv6, err := a.ipAlloc.Next()
-	if err != nil {
-		return false, err
-	}
-
-	node, newNode, err := a.db.HandleNodeFromAuthPath(
+	node, newNode, err := a.state.HandleNodeFromAuthPath(
 		registrationID,
 		types.UserID(user.ID),
 		&expiry,
 		util.RegisterMethodOIDC,
-		ipv4, ipv6,
 	)
 	if err != nil {
 		return false, fmt.Errorf("could not register node: %w", err)
 	}
 
-	// Send an update to all nodes if this is a new node that they need to know
-	// about.
-	// If this is a refresh, just send new expiry updates.
-	updateSent, err := nodesChangedHook(a.db, a.polMan, a.notifier)
-	if err != nil {
-		return false, fmt.Errorf("updating resources using node: %w", err)
-	}
-
 	// This is a bit of a back and forth, but we have a bit of a chicken and egg
 	// dependency here.
 	// Because the way the policy manager works, we need to have the node
 	// in the database, then add it to the policy manager and then we can
 	// approve the route. This means we get this dance where the node is
 	// first added to the database, then we add it to the policy manager via
-	// nodesChangedHook and then we can auto approve the routes.
+	// SaveNode (which automatically updates the policy manager) and then we can auto approve the routes.
 	// As that only approves the struct object, we need to save it again and
 	// ensure we send an update.
 	// This works, but might be another good candidate for doing some sort of
 	// eventbus.
-	routesChanged := policy.AutoApproveRoutes(a.polMan, node)
-	if err := a.db.DB.Save(node).Error; err != nil {
+	routesChanged := a.state.AutoApproveRoutes(node)
+	_, policyChanged, err := a.state.SaveNode(node)
+	if err != nil {
 		return false, fmt.Errorf("saving auto approved routes to node: %w", err)
 	}
 
-	if !updateSent || routesChanged {
+	// Send policy update notifications if needed (from SaveNode or route changes)
+	if policyChanged {
+		ctx := types.NotifyCtx(context.Background(), "oidc-nodes-change", "all")
+		a.notifier.NotifyAll(ctx, types.UpdateFull())
+	}
+
+	if routesChanged {
 		ctx := types.NotifyCtx(context.Background(), "oidc-expiry-self", node.Hostname)
 		a.notifier.NotifyByNodeID(
 			ctx,