diff --git a/oidc.go b/oidc.go index 5509bd4..8b5f024 100644 --- a/oidc.go +++ b/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( - writer http.ResponseWriter, - idToken *oidc.IDToken, -) (*IDTokenClaims, bool) { + // Extract custom claims 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 -} - -// validateOIDCAllowedDomains checks that if AllowedDomains is provided, -// that the authenticated principal ends with @. -func validateOIDCAllowedDomains( - writer http.ResponseWriter, - allowedDomains []string, - claims *IDTokenClaims, -) bool { - if len(allowedDomains) > 0 { + // If AllowedDomains is provided, check that the authenticated principal ends with @. + if len(h.cfg.OIDC.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 -} - -// 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) { + // 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) { 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( - writer http.ResponseWriter, - namespace *Namespace, - machineKey *key.MachinePublic, -) bool { - machineKeyStr := MachinePublicKeyStripPrefix(*machineKey) - - if _, err := h.RegisterMachineFromAuthCallback( + _, 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") + } } diff --git a/poll.go b/poll.go index b9a757a..9c17b5c 100644 --- a/poll.go +++ b/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 - // when an outdated machine object is kept alive, e.g. db is update from - // command line, but then overwritten. + // 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 + // 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 - // when an outdated machine object is kept alive, e.g. db is update from - // command line, but then overwritten. + // 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 + // 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 - // when an outdated machine object is kept alive, e.g. db is update from - // command line, but then overwritten. + // 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 + // command line, but then overwritten. err := h.UpdateMachineFromDatabase(machine) if err != nil { log.Error().