Revert "decompose OIDCCallback method"

This reverts commit 395caaad42.

oidc.go

@@ -136,82 +136,6 @@ func (h *Headscale) OIDCCallback(
 	writer http.ResponseWriter,
 	req *http.Request,
 ) {
-	code, state, ok := validateOIDCCallbackParams(writer, req)
-	if !ok {
-		return
-	}
-
-	rawIDToken, ok := h.getIDTokenForOIDCCallback(writer, code, state)
-	if !ok {
-		return
-	}
-
-	idToken, ok := h.verifyIDTokenForOIDCCallback(writer, rawIDToken)
-	if !ok {
-		return
-	}
-
-	// TODO: we can use userinfo at some point to grab additional information about the user (groups membership, etc)
-	// userInfo, err := oidcProvider.UserInfo(context.Background(), oauth2.StaticTokenSource(oauth2Token))
-	// if err != nil {
-	// 	c.String(http.StatusBadRequest, fmt.Sprintf("Failed to retrieve userinfo"))
-	// 	return
-	// }
-
-	claims, ok := extractIDTokenClaims(writer, idToken)
-	if !ok {
-		return
-	}
-
-	if ok := validateOIDCAllowedDomains(writer, h.cfg.OIDC.AllowedDomains, claims); !ok {
-		return
-	}
-
-	if ok := validateOIDCAllowedUsers(writer, h.cfg.OIDC.AllowedUsers, claims); !ok {
-		return
-	}
-
-	machineKey, ok := h.validateMachineForOIDCCallback(writer, state, claims)
-	if !ok {
-		return
-	}
-
-	namespaceName, ok := getNamespaceName(writer, claims, h.cfg.OIDC.StripEmaildomain)
-	if !ok {
-		return
-	}
-
-	// register the machine if it's new
-	log.Debug().Msg("Registering new machine after successful callback")
-
-	namespace, ok := h.findOrCreateNewNamespaceForOIDCCallback(writer, namespaceName)
-	if !ok {
-		return
-	}
-
-	if ok := h.registerMachineForOIDCCallback(writer, namespace, machineKey); !ok {
-		return
-	}
-
-	content, ok := renderOIDCCallbackTemplate(writer, claims)
-	if !ok {
-		return
-	}
-
-	writer.Header().Set("Content-Type", "text/html; charset=utf-8")
-	writer.WriteHeader(http.StatusOK)
-	if _, err := writer.Write(content.Bytes()); err != nil {
-		log.Error().
-			Caller().
-			Err(err).
-			Msg("Failed to write response")
-	}
-}
-
-func validateOIDCCallbackParams(
-	writer http.ResponseWriter,
-	req *http.Request,
-) (string, string, bool) {
 	code := req.URL.Query().Get("code")
 	state := req.URL.Query().Get("state")
 
@@ -226,16 +150,9 @@ func validateOIDCCallbackParams(
 				Msg("Failed to write response")
 		}
 
-		return "", "", false
+		return
 	}
 
-	return code, state, true
-}
-
-func (h *Headscale) getIDTokenForOIDCCallback(
-	writer http.ResponseWriter,
-	code, state string,
-) (string, bool) {
 	oauth2Token, err := h.oauth2Config.Exchange(context.Background(), code)
 	if err != nil {
 		log.Error().
@@ -252,7 +169,7 @@ func (h *Headscale) getIDTokenForOIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return "", false
+		return
 	}
 
 	log.Trace().
@@ -273,17 +190,11 @@ func (h *Headscale) getIDTokenForOIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return "", false
+		return
 	}
 
-	return rawIDToken, true
-}
-
-func (h *Headscale) verifyIDTokenForOIDCCallback(
-	writer http.ResponseWriter,
-	rawIDToken string,
-) (*oidc.IDToken, bool) {
 	verifier := h.oidcProvider.Verifier(&oidc.Config{ClientID: h.cfg.OIDC.ClientID})
+
 	idToken, err := verifier.Verify(context.Background(), rawIDToken)
 	if err != nil {
 		log.Error().
@@ -300,18 +211,19 @@ func (h *Headscale) verifyIDTokenForOIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return nil, false
+		return
 	}
 
-	return idToken, true
+	// TODO: we can use userinfo at some point to grab additional information about the user (groups membership, etc)
-}
+	// userInfo, err := oidcProvider.UserInfo(context.Background(), oauth2.StaticTokenSource(oauth2Token))
+	// if err != nil {
+	// 	c.String(http.StatusBadRequest, fmt.Sprintf("Failed to retrieve userinfo"))
+	// 	return
+	// }
 
-func extractIDTokenClaims(
+	// Extract custom claims
-	writer http.ResponseWriter,
-	idToken *oidc.IDToken,
-) (*IDTokenClaims, bool) {
 	var claims IDTokenClaims
-	if err := idToken.Claims(claims); err != nil {
+	if err = idToken.Claims(&claims); err != nil {
 		log.Error().
 			Err(err).
 			Caller().
@@ -326,22 +238,13 @@ func extractIDTokenClaims(
 				Msg("Failed to write response")
 		}
 
-		return nil, false
+		return
 	}
 
-	return &claims, true
+	// If AllowedDomains is provided, check that the authenticated principal ends with @<alloweddomain>.
-}
+	if len(h.cfg.OIDC.AllowedDomains) > 0 {
-
-// validateOIDCAllowedDomains checks that if AllowedDomains is provided,
-// that the authenticated principal ends with @<alloweddomain>.
-func validateOIDCAllowedDomains(
-	writer http.ResponseWriter,
-	allowedDomains []string,
-	claims *IDTokenClaims,
-) bool {
-	if len(allowedDomains) > 0 {
 		if at := strings.LastIndex(claims.Email, "@"); at < 0 ||
-			!IsStringInSlice(allowedDomains, claims.Email[at+1:]) {
+			!IsStringInSlice(h.cfg.OIDC.AllowedDomains, claims.Email[at+1:]) {
 			log.Error().Msg("authenticated principal does not match any allowed domain")
 			writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 			writer.WriteHeader(http.StatusBadRequest)
@@ -353,22 +256,13 @@ func validateOIDCAllowedDomains(
 					Msg("Failed to write response")
 			}
 
-			return false
+			return
 		}
 	}
 
-	return true
+	// If AllowedUsers is provided, check that the authenticated princial is part of that list.
-}
+	if len(h.cfg.OIDC.AllowedUsers) > 0 &&
-
+		!IsStringInSlice(h.cfg.OIDC.AllowedUsers, claims.Email) {
-// validateOIDCAllowedUsers checks that if AllowedUsers is provided,
-// that the authenticated principal is part of that list.
-func validateOIDCAllowedUsers(
-	writer http.ResponseWriter,
-	allowedUsers []string,
-	claims *IDTokenClaims,
-) bool {
-	if len(allowedUsers) > 0 &&
-		!IsStringInSlice(allowedUsers, claims.Email) {
 		log.Error().Msg("authenticated principal does not match any allowed user")
 		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		writer.WriteHeader(http.StatusBadRequest)
@@ -380,23 +274,12 @@ func validateOIDCAllowedUsers(
 				Msg("Failed to write response")
 		}
 
-		return false
+		return
 	}
 
-	return true
-}
-
-// validateMachine retrieves machine information if it exist
-// The error is not important, because if it does not
-// exist, then this is a new machine and we will move
-// on to registration.
-func (h *Headscale) validateMachineForOIDCCallback(
-	writer http.ResponseWriter,
-	state string,
-	claims *IDTokenClaims,
-) (*key.MachinePublic, bool) {
 	// retrieve machinekey from state cache
 	machineKeyIf, machineKeyFound := h.registrationCache.Get(state)
+
 	if !machineKeyFound {
 		log.Error().
 			Msg("requested machine state key expired before authorisation completed")
@@ -410,12 +293,13 @@ func (h *Headscale) validateMachineForOIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return nil, false
+		return
 	}
 
-	var machineKey key.MachinePublic
 	machineKeyFromCache, machineKeyOK := machineKeyIf.(string)
-	err := machineKey.UnmarshalText(
+
+	var machineKey key.MachinePublic
+	err = machineKey.UnmarshalText(
 		[]byte(MachinePublicKeyEnsurePrefix(machineKeyFromCache)),
 	)
 	if err != nil {
@@ -431,7 +315,7 @@ func (h *Headscale) validateMachineForOIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return nil, false
+		return
 	}
 
 	if !machineKeyOK {
@@ -446,7 +330,7 @@ func (h *Headscale) validateMachineForOIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return nil, false
+		return
 	}
 
 	// retrieve machine information if it exist
@@ -469,7 +353,7 @@ func (h *Headscale) validateMachineForOIDCCallback(
 				Msg("Failed to refresh machine")
 			http.Error(writer, "Failed to refresh machine", http.StatusInternalServerError)
 
-			return nil, false
+			return
 		}
 
 		var content bytes.Buffer
@@ -493,7 +377,7 @@ func (h *Headscale) validateMachineForOIDCCallback(
 					Msg("Failed to write response")
 			}
 
-			return nil, false
+			return
 		}
 
 		writer.Header().Set("Content-Type", "text/html; charset=utf-8")
@@ -506,20 +390,12 @@ func (h *Headscale) validateMachineForOIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return nil, false
+		return
 	}
 
-	return &machineKey, true
-}
-
-func getNamespaceName(
-	writer http.ResponseWriter,
-	claims *IDTokenClaims,
-	stripEmaildomain bool,
-) (string, bool) {
 	namespaceName, err := NormalizeToFQDNRules(
 		claims.Email,
-		stripEmaildomain,
+		h.cfg.OIDC.StripEmaildomain,
 	)
 	if err != nil {
 		log.Error().Err(err).Caller().Msgf("couldn't normalize email")
@@ -533,16 +409,12 @@ func getNamespaceName(
 				Msg("Failed to write response")
 		}
 
-		return "", false
+		return
 	}
 
-	return namespaceName, true
+	// register the machine if it's new
-}
+	log.Debug().Msg("Registering new machine after successful callback")
 
-func (h *Headscale) findOrCreateNewNamespaceForOIDCCallback(
-	writer http.ResponseWriter,
-	namespaceName string,
-) (*Namespace, bool) {
 	namespace, err := h.GetNamespace(namespaceName)
 	if errors.Is(err, errNamespaceNotFound) {
 		namespace, err = h.CreateNamespace(namespaceName)
@@ -562,7 +434,7 @@ func (h *Headscale) findOrCreateNewNamespaceForOIDCCallback(
 					Msg("Failed to write response")
 			}
 
-			return nil, false
+			return
 		}
 	} else if err != nil {
 		log.Error().
@@ -580,24 +452,17 @@ func (h *Headscale) findOrCreateNewNamespaceForOIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return nil, false
+		return
 	}
 
-	return namespace, true
+	machineKeyStr := MachinePublicKeyStripPrefix(machineKey)
-}
 
-func (h *Headscale) registerMachineForOIDCCallback(
+	_, err = h.RegisterMachineFromAuthCallback(
-	writer http.ResponseWriter,
-	namespace *Namespace,
-	machineKey *key.MachinePublic,
-) bool {
-	machineKeyStr := MachinePublicKeyStripPrefix(*machineKey)
-
-	if _, err := h.RegisterMachineFromAuthCallback(
 		machineKeyStr,
 		namespace.Name,
 		RegisterMethodOIDC,
-	); err != nil {
+	)
+	if err != nil {
 		log.Error().
 			Caller().
 			Err(err).
@@ -612,16 +477,9 @@ func (h *Headscale) registerMachineForOIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return false
+		return
 	}
 
-	return true
-}
-
-func renderOIDCCallbackTemplate(
-	writer http.ResponseWriter,
-	claims *IDTokenClaims,
-) (*bytes.Buffer, bool) {
 	var content bytes.Buffer
 	if err := oidcCallbackTemplate.Execute(&content, oidcCallbackTemplateConfig{
 		User: claims.Email,
@@ -643,8 +501,16 @@ func renderOIDCCallbackTemplate(
 				Msg("Failed to write response")
 		}
 
-		return nil, false
+		return
 	}
 
-	return &content, true
+	writer.Header().Set("Content-Type", "text/html; charset=utf-8")
+	writer.WriteHeader(http.StatusOK)
+	_, err = writer.Write(content.Bytes())
+	if err != nil {
+		log.Error().
+			Caller().
+			Err(err).
+			Msg("Failed to write response")
+	}
 }

poll.go

@@ -356,9 +356,9 @@ func (h *Headscale) PollNetMapStream(
 				Str("channel", "pollData").
 				Int("bytes", len(data)).
 				Msg("Data from pollData channel written successfully")
-			// TODO(kradalby): Abstract away all the database calls, this can cause race conditions
+				// TODO(kradalby): Abstract away all the database calls, this can cause race conditions
-			// when an outdated machine object is kept alive, e.g. db is update from
+				// when an outdated machine object is kept alive, e.g. db is update from
-			// command line, but then overwritten.
+				// command line, but then overwritten.
 			err = h.UpdateMachineFromDatabase(machine)
 			if err != nil {
 				log.Error().
@@ -434,9 +434,9 @@ func (h *Headscale) PollNetMapStream(
 				Str("channel", "keepAlive").
 				Int("bytes", len(data)).
 				Msg("Keep alive sent successfully")
-			// TODO(kradalby): Abstract away all the database calls, this can cause race conditions
+				// TODO(kradalby): Abstract away all the database calls, this can cause race conditions
-			// when an outdated machine object is kept alive, e.g. db is update from
+				// when an outdated machine object is kept alive, e.g. db is update from
-			// command line, but then overwritten.
+				// command line, but then overwritten.
 			err = h.UpdateMachineFromDatabase(machine)
 			if err != nil {
 				log.Error().
@@ -591,9 +591,9 @@ func (h *Headscale) PollNetMapStream(
 				Str("handler", "PollNetMapStream").
 				Str("machine", machine.Hostname).
 				Msg("The client has closed the connection")
-			// TODO: Abstract away all the database calls, this can cause race conditions
+				// TODO: Abstract away all the database calls, this can cause race conditions
-			// when an outdated machine object is kept alive, e.g. db is update from
+				// when an outdated machine object is kept alive, e.g. db is update from
-			// command line, but then overwritten.
+				// command line, but then overwritten.
 			err := h.UpdateMachineFromDatabase(machine)
 			if err != nil {
 				log.Error().