Merge pull request #685 from GrigoriyMikhalkin/oidc-refactoring

Decompose OIDCCallback method

oidc.go

@@ -21,6 +21,13 @@ import (
 
 const (
 	randomByteSize = 16
+
+	errEmptyOIDCCallbackParams = Error("empty OIDC callback params")
+	errNoOIDCIDToken           = Error("could not extract ID Token for OIDC callback")
+	errOIDCAllowedDomains      = Error("authenticated principal does not match any allowed domain")
+	errOIDCAllowedUsers        = Error("authenticated principal does not match any allowed user")
+	errOIDCInvalidMachineState = Error("requested machine state key expired before authorisation completed")
+	errOIDCMachineKeyMissing   = Error("could not get machine key from cache")
 )
 
 type IDTokenClaims struct {
@@ -136,6 +143,82 @@ func (h *Headscale) OIDCCallback(
 	writer http.ResponseWriter,
 	req *http.Request,
 ) {
+	code, state, err := validateOIDCCallbackParams(writer, req)
+	if err != nil {
+		return
+	}
+
+	rawIDToken, err := h.getIDTokenForOIDCCallback(writer, code, state)
+	if err != nil {
+		return
+	}
+
+	idToken, err := h.verifyIDTokenForOIDCCallback(writer, rawIDToken)
+	if err != nil {
+		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, err := extractIDTokenClaims(writer, idToken)
+	if err != nil {
+		return
+	}
+
+	if err := validateOIDCAllowedDomains(writer, h.cfg.OIDC.AllowedDomains, claims); err != nil {
+		return
+	}
+
+	if err := validateOIDCAllowedUsers(writer, h.cfg.OIDC.AllowedUsers, claims); err != nil {
+		return
+	}
+
+	machineKey, machineExists, err := h.validateMachineForOIDCCallback(writer, state, claims)
+	if err != nil || machineExists {
+		return
+	}
+
+	namespaceName, err := getNamespaceName(writer, claims, h.cfg.OIDC.StripEmaildomain)
+	if err != nil {
+		return
+	}
+
+	// register the machine if it's new
+	log.Debug().Msg("Registering new machine after successful callback")
+
+	namespace, err := h.findOrCreateNewNamespaceForOIDCCallback(writer, namespaceName)
+	if err != nil {
+		return
+	}
+
+	if err := h.registerMachineForOIDCCallback(writer, namespace, machineKey); err != nil {
+		return
+	}
+
+	content, err := renderOIDCCallbackTemplate(writer, claims)
+	if err != nil {
+		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, error) {
 	code := req.URL.Query().Get("code")
 	state := req.URL.Query().Get("state")
 
@@ -150,9 +233,16 @@ func (h *Headscale) OIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return
+		return "", "", errEmptyOIDCCallbackParams
 	}
 
+	return code, state, nil
+}
+
+func (h *Headscale) getIDTokenForOIDCCallback(
+	writer http.ResponseWriter,
+	code, state string,
+) (string, error) {
 	oauth2Token, err := h.oauth2Config.Exchange(context.Background(), code)
 	if err != nil {
 		log.Error().
@@ -161,15 +251,15 @@ func (h *Headscale) OIDCCallback(
 			Msg("Could not exchange code for token")
 		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		writer.WriteHeader(http.StatusBadRequest)
-		_, err := writer.Write([]byte("Could not exchange code for token"))
+		_, werr := writer.Write([]byte("Could not exchange code for token"))
-		if err != nil {
+		if werr != nil {
 			log.Error().
 				Caller().
-				Err(err).
+				Err(werr).
 				Msg("Failed to write response")
 		}
 
-		return
+		return "", err
 	}
 
 	log.Trace().
@@ -190,11 +280,17 @@ func (h *Headscale) OIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return
+		return "", errNoOIDCIDToken
 	}
 
-	verifier := h.oidcProvider.Verifier(&oidc.Config{ClientID: h.cfg.OIDC.ClientID})
+	return rawIDToken, nil
+}
 
+func (h *Headscale) verifyIDTokenForOIDCCallback(
+	writer http.ResponseWriter,
+	rawIDToken string,
+) (*oidc.IDToken, error) {
+	verifier := h.oidcProvider.Verifier(&oidc.Config{ClientID: h.cfg.OIDC.ClientID})
 	idToken, err := verifier.Verify(context.Background(), rawIDToken)
 	if err != nil {
 		log.Error().
@@ -203,48 +299,56 @@ func (h *Headscale) OIDCCallback(
 			Msg("failed to verify id token")
 		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		writer.WriteHeader(http.StatusBadRequest)
-		_, err := writer.Write([]byte("Failed to verify id token"))
+		_, werr := writer.Write([]byte("Failed to verify id token"))
-		if err != nil {
+		if werr != nil {
 			log.Error().
 				Caller().
-				Err(err).
+				Err(werr).
 				Msg("Failed to write response")
 		}
 
-		return
+		return nil, err
 	}
 
-	// TODO: we can use userinfo at some point to grab additional information about the user (groups membership, etc)
+	return idToken, nil
-	// userInfo, err := oidcProvider.UserInfo(context.Background(), oauth2.StaticTokenSource(oauth2Token))
+}
-	// if err != nil {
-	// 	c.String(http.StatusBadRequest, fmt.Sprintf("Failed to retrieve userinfo"))
-	// 	return
-	// }
 
-	// Extract custom claims
+func extractIDTokenClaims(
+	writer http.ResponseWriter,
+	idToken *oidc.IDToken,
+) (*IDTokenClaims, error) {
 	var claims IDTokenClaims
-	if err = idToken.Claims(&claims); err != nil {
+	if err := idToken.Claims(claims); err != nil {
 		log.Error().
 			Err(err).
 			Caller().
 			Msg("Failed to decode id token claims")
 		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		writer.WriteHeader(http.StatusBadRequest)
-		_, err := writer.Write([]byte("Failed to decode id token claims"))
+		_, werr := writer.Write([]byte("Failed to decode id token claims"))
-		if err != nil {
+		if werr != nil {
 			log.Error().
 				Caller().
-				Err(err).
+				Err(werr).
 				Msg("Failed to write response")
 		}
 
-		return
+		return nil, err
 	}
 
-	// If AllowedDomains is provided, check that the authenticated principal ends with @<alloweddomain>.
+	return &claims, nil
-	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,
+) error {
+	if len(allowedDomains) > 0 {
 		if at := strings.LastIndex(claims.Email, "@"); at < 0 ||
-			!IsStringInSlice(h.cfg.OIDC.AllowedDomains, claims.Email[at+1:]) {
+			!IsStringInSlice(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)
@@ -256,13 +360,22 @@ func (h *Headscale) OIDCCallback(
 					Msg("Failed to write response")
 			}
 
-			return
+			return errOIDCAllowedDomains
 		}
 	}
 
-	// If AllowedUsers is provided, check that the authenticated princial is part of that list.
+	return nil
-	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,
+) error {
+	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)
@@ -274,12 +387,23 @@ func (h *Headscale) OIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return
+		return errOIDCAllowedUsers
 	}
 
+	return nil
+}
+
+// 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, error) {
 	// retrieve machinekey from state cache
 	machineKeyIf, machineKeyFound := h.registrationCache.Get(state)
-
 	if !machineKeyFound {
 		log.Error().
 			Msg("requested machine state key expired before authorisation completed")
@@ -293,13 +417,12 @@ func (h *Headscale) OIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return
+		return nil, false, errOIDCInvalidMachineState
 	}
 
-	machineKeyFromCache, machineKeyOK := machineKeyIf.(string)
-
 	var machineKey key.MachinePublic
-	err = machineKey.UnmarshalText(
+	machineKeyFromCache, machineKeyOK := machineKeyIf.(string)
+	err := machineKey.UnmarshalText(
 		[]byte(MachinePublicKeyEnsurePrefix(machineKeyFromCache)),
 	)
 	if err != nil {
@@ -307,15 +430,15 @@ func (h *Headscale) OIDCCallback(
 			Msg("could not parse machine public key")
 		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		writer.WriteHeader(http.StatusBadRequest)
-		_, err := writer.Write([]byte("could not parse public key"))
+		_, werr := writer.Write([]byte("could not parse public key"))
-		if err != nil {
+		if werr != nil {
 			log.Error().
 				Caller().
-				Err(err).
+				Err(werr).
 				Msg("Failed to write response")
 		}
 
-		return
+		return nil, false, err
 	}
 
 	if !machineKeyOK {
@@ -330,7 +453,7 @@ func (h *Headscale) OIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return
+		return nil, false, errOIDCMachineKeyMissing
 	}
 
 	// retrieve machine information if it exist
@@ -353,7 +476,7 @@ func (h *Headscale) OIDCCallback(
 				Msg("Failed to refresh machine")
 			http.Error(writer, "Failed to refresh machine", http.StatusInternalServerError)
 
-			return
+			return nil, true, err
 		}
 
 		var content bytes.Buffer
@@ -369,15 +492,15 @@ func (h *Headscale) OIDCCallback(
 
 			writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 			writer.WriteHeader(http.StatusInternalServerError)
-			_, err := writer.Write([]byte("Could not render OIDC callback template"))
+			_, werr := writer.Write([]byte("Could not render OIDC callback template"))
-			if err != nil {
+			if werr != nil {
 				log.Error().
 					Caller().
-					Err(err).
+					Err(werr).
 					Msg("Failed to write response")
 			}
 
-			return
+			return nil, true, err
 		}
 
 		writer.Header().Set("Content-Type", "text/html; charset=utf-8")
@@ -390,31 +513,43 @@ func (h *Headscale) OIDCCallback(
 				Msg("Failed to write response")
 		}
 
-		return
+		return nil, true, nil
 	}
 
+	return &machineKey, false, nil
+}
+
+func getNamespaceName(
+	writer http.ResponseWriter,
+	claims *IDTokenClaims,
+	stripEmaildomain bool,
+) (string, error) {
 	namespaceName, err := NormalizeToFQDNRules(
 		claims.Email,
-		h.cfg.OIDC.StripEmaildomain,
+		stripEmaildomain,
 	)
 	if err != nil {
 		log.Error().Err(err).Caller().Msgf("couldn't normalize email")
 		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		writer.WriteHeader(http.StatusInternalServerError)
-		_, err := writer.Write([]byte("couldn't normalize email"))
+		_, werr := writer.Write([]byte("couldn't normalize email"))
-		if err != nil {
+		if werr != nil {
 			log.Error().
 				Caller().
-				Err(err).
+				Err(werr).
 				Msg("Failed to write response")
 		}
 
-		return
+		return "", err
 	}
 
-	// register the machine if it's new
+	return namespaceName, nil
-	log.Debug().Msg("Registering new machine after successful callback")
+}
 
+func (h *Headscale) findOrCreateNewNamespaceForOIDCCallback(
+	writer http.ResponseWriter,
+	namespaceName string,
+) (*Namespace, error) {
 	namespace, err := h.GetNamespace(namespaceName)
 	if errors.Is(err, ErrNamespaceNotFound) {
 		namespace, err = h.CreateNamespace(namespaceName)
@@ -426,15 +561,15 @@ func (h *Headscale) OIDCCallback(
 				Msgf("could not create new namespace '%s'", namespaceName)
 			writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 			writer.WriteHeader(http.StatusInternalServerError)
-			_, err := writer.Write([]byte("could not create namespace"))
+			_, werr := writer.Write([]byte("could not create namespace"))
-			if err != nil {
+			if werr != nil {
 				log.Error().
 					Caller().
-					Err(err).
+					Err(werr).
 					Msg("Failed to write response")
 			}
 
-			return
+			return nil, err
 		}
 	} else if err != nil {
 		log.Error().
@@ -444,42 +579,56 @@ func (h *Headscale) OIDCCallback(
 			Msg("could not find or create namespace")
 		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		writer.WriteHeader(http.StatusInternalServerError)
-		_, err := writer.Write([]byte("could not find or create namespace"))
+		_, werr := writer.Write([]byte("could not find or create namespace"))
-		if err != nil {
+		if werr != nil {
 			log.Error().
 				Caller().
-				Err(err).
+				Err(werr).
 				Msg("Failed to write response")
 		}
 
-		return
+		return nil, err
 	}
 
-	machineKeyStr := MachinePublicKeyStripPrefix(machineKey)
+	return namespace, nil
+}
 
-	_, err = h.RegisterMachineFromAuthCallback(
+func (h *Headscale) registerMachineForOIDCCallback(
+	writer http.ResponseWriter,
+	namespace *Namespace,
+	machineKey *key.MachinePublic,
+) error {
+	machineKeyStr := MachinePublicKeyStripPrefix(*machineKey)
+
+	if _, err := h.RegisterMachineFromAuthCallback(
 		machineKeyStr,
 		namespace.Name,
 		RegisterMethodOIDC,
-	)
+	); err != nil {
-	if err != nil {
 		log.Error().
 			Caller().
 			Err(err).
 			Msg("could not register machine")
 		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		writer.WriteHeader(http.StatusInternalServerError)
-		_, err := writer.Write([]byte("could not register machine"))
+		_, werr := writer.Write([]byte("could not register machine"))
-		if err != nil {
+		if werr != nil {
 			log.Error().
 				Caller().
-				Err(err).
+				Err(werr).
 				Msg("Failed to write response")
 		}
 
-		return
+		return err
 	}
 
+	return nil
+}
+
+func renderOIDCCallbackTemplate(
+	writer http.ResponseWriter,
+	claims *IDTokenClaims,
+) (*bytes.Buffer, error) {
 	var content bytes.Buffer
 	if err := oidcCallbackTemplate.Execute(&content, oidcCallbackTemplateConfig{
 		User: claims.Email,
@@ -493,24 +642,16 @@ func (h *Headscale) OIDCCallback(
 
 		writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
 		writer.WriteHeader(http.StatusInternalServerError)
-		_, err := writer.Write([]byte("Could not render OIDC callback template"))
+		_, werr := writer.Write([]byte("Could not render OIDC callback template"))
-		if err != nil {
+		if werr != nil {
 			log.Error().
 				Caller().
-				Err(err).
+				Err(werr).
 				Msg("Failed to write response")
 		}
 
-		return
+		return nil, err
 	}
 
-	writer.Header().Set("Content-Type", "text/html; charset=utf-8")
+	return &content, nil
-	writer.WriteHeader(http.StatusOK)
-	_, err = writer.Write(content.Bytes())
-	if err != nil {
-		log.Error().
-			Caller().
-			Err(err).
-			Msg("Failed to write response")
-	}
 }