From 889eff265fdcaff348c406270666fa3067194559 Mon Sep 17 00:00:00 2001 From: Grigoriy Mikhalkin Date: Thu, 30 Jun 2022 23:35:22 +0200 Subject: [PATCH 1/4] graceful shutdown fix --- app.go | 23 +++++++++++++++++------ poll.go | 21 ++++++++++++--------- 2 files changed, 29 insertions(+), 15 deletions(-) diff --git a/app.go b/app.go index 11c8d685..de6ef669 100644 --- a/app.go +++ b/app.go @@ -95,6 +95,7 @@ type Headscale struct { ipAllocationMutex sync.Mutex shutdownChan chan struct{} + wg sync.WaitGroup } // Look up the TLS constant relative to user-supplied TLS client @@ -153,6 +154,7 @@ func NewHeadscale(cfg *Config) (*Headscale, error) { privateKey: privKey, aclRules: tailcfg.FilterAllowAll, // default allowall registrationCache: registrationCache, + wg: sync.WaitGroup{}, } err = app.initDB() @@ -567,6 +569,8 @@ func (h *Headscale) Serve() error { // https://github.com/soheilhy/cmux/issues/68 // https://github.com/soheilhy/cmux/issues/91 + var grpcServer *grpc.Server + var grpcListener net.Listener if tlsConfig != nil || h.cfg.GRPCAllowInsecure { log.Info().Msgf("Enabling remote gRPC at %s", h.cfg.GRPCAddr) @@ -587,12 +591,12 @@ func (h *Headscale) Serve() error { log.Warn().Msg("gRPC is running without security") } - grpcServer := grpc.NewServer(grpcOptions...) + grpcServer = grpc.NewServer(grpcOptions...) v1.RegisterHeadscaleServiceServer(grpcServer, newHeadscaleV1APIServer(h)) reflection.Register(grpcServer) - grpcListener, err := net.Listen("tcp", h.cfg.GRPCAddr) + grpcListener, err = net.Listen("tcp", h.cfg.GRPCAddr) if err != nil { return fmt.Errorf("failed to bind to TCP address: %w", err) } @@ -668,7 +672,7 @@ func (h *Headscale) Serve() error { syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGHUP) - go func(c chan os.Signal) { + sig_func := func(c chan os.Signal) { // Wait for a SIGINT or SIGKILL: for { sig := <-c @@ -678,7 +682,7 @@ func (h *Headscale) Serve() error { Str("signal", sig.String()). Msg("Received SIGHUP, reloading ACL and Config") - // TODO(kradalby): Reload config on SIGHUP + // TODO(kradalby): Reload config on SIGHUP if h.cfg.ACL.PolicyPath != "" { aclPath := AbsolutePathFromConfigPath(h.cfg.ACL.PolicyPath) @@ -698,7 +702,8 @@ func (h *Headscale) Serve() error { Str("signal", sig.String()). Msg("Received signal to stop, shutting down gracefully") - h.shutdownChan <- struct{}{} + close(h.shutdownChan) + h.wg.Wait() // Gracefully shut down servers ctx, cancel := context.WithTimeout(context.Background(), HTTPShutdownTimeout) @@ -710,6 +715,11 @@ func (h *Headscale) Serve() error { } grpcSocket.GracefulStop() + if grpcServer != nil { + grpcServer.GracefulStop() + grpcListener.Close() + } + // Close network listeners promHTTPListener.Close() httpListener.Close() @@ -736,7 +746,8 @@ func (h *Headscale) Serve() error { os.Exit(0) } } - }(sigc) + } + errorGroup.Go(func() error { sig_func(sigc); return nil }) return errorGroup.Wait() } diff --git a/poll.go b/poll.go index 6628a179..94941aa3 100644 --- a/poll.go +++ b/poll.go @@ -290,6 +290,9 @@ func (h *Headscale) PollNetMapStream( keepAliveChan chan []byte, updateChan chan struct{}, ) { + h.wg.Add(1) + defer h.wg.Done() + ctx := context.WithValue(req.Context(), machineNameContextKey, machine.Hostname) ctx, cancel := context.WithCancel(ctx) @@ -353,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(). @@ -431,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(). @@ -588,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(). From 3f0639c87ddbb86fd9af3d210eafa98b80301ad1 Mon Sep 17 00:00:00 2001 From: Grigoriy Mikhalkin Date: Mon, 11 Jul 2022 20:33:24 +0200 Subject: [PATCH 2/4] graceful shutdown lint fixes --- app.go | 32 ++++++++++++++++++-------------- config.go | 13 +++++++++---- poll.go | 4 ++-- 3 files changed, 29 insertions(+), 20 deletions(-) diff --git a/app.go b/app.go index de6ef669..5f5e2611 100644 --- a/app.go +++ b/app.go @@ -94,8 +94,8 @@ type Headscale struct { ipAllocationMutex sync.Mutex - shutdownChan chan struct{} - wg sync.WaitGroup + shutdownChan chan struct{} + pollNetMapStreamWG sync.WaitGroup } // Look up the TLS constant relative to user-supplied TLS client @@ -148,13 +148,13 @@ func NewHeadscale(cfg *Config) (*Headscale, error) { ) app := Headscale{ - cfg: cfg, - dbType: cfg.DBtype, - dbString: dbString, - privateKey: privKey, - aclRules: tailcfg.FilterAllowAll, // default allowall - registrationCache: registrationCache, - wg: sync.WaitGroup{}, + cfg: cfg, + dbType: cfg.DBtype, + dbString: dbString, + privateKey: privKey, + aclRules: tailcfg.FilterAllowAll, // default allowall + registrationCache: registrationCache, + pollNetMapStreamWG: sync.WaitGroup{}, } err = app.initDB() @@ -672,7 +672,7 @@ func (h *Headscale) Serve() error { syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGHUP) - sig_func := func(c chan os.Signal) { + sigFunc := func(c chan os.Signal) { // Wait for a SIGINT or SIGKILL: for { sig := <-c @@ -703,7 +703,7 @@ func (h *Headscale) Serve() error { Msg("Received signal to stop, shutting down gracefully") close(h.shutdownChan) - h.wg.Wait() + h.pollNetMapStreamWG.Wait() // Gracefully shut down servers ctx, cancel := context.WithTimeout(context.Background(), HTTPShutdownTimeout) @@ -747,7 +747,11 @@ func (h *Headscale) Serve() error { } } } - errorGroup.Go(func() error { sig_func(sigc); return nil }) + errorGroup.Go(func() error { + sigFunc(sigc) + + return nil + }) return errorGroup.Wait() } @@ -771,13 +775,13 @@ func (h *Headscale) getTLSSettings() (*tls.Config, error) { } switch h.cfg.TLS.LetsEncrypt.ChallengeType { - case "TLS-ALPN-01": + case tlsALPN01ChallengeType: // Configuration via autocert with TLS-ALPN-01 (https://tools.ietf.org/html/rfc8737) // The RFC requires that the validation is done on port 443; in other words, headscale // must be reachable on port 443. return certManager.TLSConfig(), nil - case "HTTP-01": + case http01ChallengeType: // Configuration via autocert with HTTP-01. This requires listening on // port 80 for the certificate validation in addition to the headscale // service, which can be configured to run on any other port. diff --git a/config.go b/config.go index 6789f6f0..69358401 100644 --- a/config.go +++ b/config.go @@ -18,6 +18,11 @@ import ( "tailscale.com/types/dnstype" ) +const ( + tlsALPN01ChallengeType = "TLS-ALPN-01" + http01ChallengeType = "HTTP-01" +) + // Config contains the initial Headscale configuration. type Config struct { ServerURL string @@ -136,7 +141,7 @@ func LoadConfig(path string, isFile bool) error { viper.AutomaticEnv() viper.SetDefault("tls_letsencrypt_cache_dir", "/var/www/.cache") - viper.SetDefault("tls_letsencrypt_challenge_type", "HTTP-01") + viper.SetDefault("tls_letsencrypt_challenge_type", http01ChallengeType) viper.SetDefault("tls_client_auth_mode", "relaxed") viper.SetDefault("log_level", "info") @@ -179,15 +184,15 @@ func LoadConfig(path string, isFile bool) error { } if (viper.GetString("tls_letsencrypt_hostname") != "") && - (viper.GetString("tls_letsencrypt_challenge_type") == "TLS-ALPN-01") && + (viper.GetString("tls_letsencrypt_challenge_type") == tlsALPN01ChallengeType) && (!strings.HasSuffix(viper.GetString("listen_addr"), ":443")) { // this is only a warning because there could be something sitting in front of headscale that redirects the traffic (e.g. an iptables rule) log.Warn(). Msg("Warning: when using tls_letsencrypt_hostname with TLS-ALPN-01 as challenge type, headscale must be reachable on port 443, i.e. listen_addr should probably end in :443") } - if (viper.GetString("tls_letsencrypt_challenge_type") != "HTTP-01") && - (viper.GetString("tls_letsencrypt_challenge_type") != "TLS-ALPN-01") { + if (viper.GetString("tls_letsencrypt_challenge_type") != http01ChallengeType) && + (viper.GetString("tls_letsencrypt_challenge_type") != tlsALPN01ChallengeType) { errorText += "Fatal config error: the only supported values for tls_letsencrypt_challenge_type are HTTP-01 and TLS-ALPN-01\n" } diff --git a/poll.go b/poll.go index 94941aa3..b9a757aa 100644 --- a/poll.go +++ b/poll.go @@ -290,8 +290,8 @@ func (h *Headscale) PollNetMapStream( keepAliveChan chan []byte, updateChan chan struct{}, ) { - h.wg.Add(1) - defer h.wg.Done() + h.pollNetMapStreamWG.Add(1) + defer h.pollNetMapStreamWG.Done() ctx := context.WithValue(req.Context(), machineNameContextKey, machine.Hostname) From 395caaad421797ac7fba6a0ca0e7df87d13262f2 Mon Sep 17 00:00:00 2001 From: Grigoriy Mikhalkin Date: Mon, 11 Jul 2022 23:25:13 +0200 Subject: [PATCH 3/4] decompose OIDCCallback method --- oidc.go | 240 +++++++++++++++++++++++++++++++++++++++++++------------- 1 file changed, 187 insertions(+), 53 deletions(-) diff --git a/oidc.go b/oidc.go index 8b5f0242..5509bd47 100644 --- a/oidc.go +++ b/oidc.go @@ -136,6 +136,82 @@ 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") @@ -150,9 +226,16 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return "", "", false } + 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(). @@ -169,7 +252,7 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return "", false } log.Trace(). @@ -190,11 +273,17 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return "", false } - verifier := h.oidcProvider.Verifier(&oidc.Config{ClientID: h.cfg.OIDC.ClientID}) + 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(). @@ -211,19 +300,18 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } - // 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 - // } + return idToken, true +} - // Extract custom claims +func extractIDTokenClaims( + 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(). @@ -238,13 +326,22 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } - // If AllowedDomains is provided, check that the authenticated principal ends with @. - if len(h.cfg.OIDC.AllowedDomains) > 0 { + 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 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 +353,22 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return false } } - // 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) { + 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) { 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 +380,23 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return false } + 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") @@ -293,13 +410,12 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } - machineKeyFromCache, machineKeyOK := machineKeyIf.(string) - var machineKey key.MachinePublic - err = machineKey.UnmarshalText( + machineKeyFromCache, machineKeyOK := machineKeyIf.(string) + err := machineKey.UnmarshalText( []byte(MachinePublicKeyEnsurePrefix(machineKeyFromCache)), ) if err != nil { @@ -315,7 +431,7 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } if !machineKeyOK { @@ -330,7 +446,7 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } // retrieve machine information if it exist @@ -353,7 +469,7 @@ func (h *Headscale) OIDCCallback( Msg("Failed to refresh machine") http.Error(writer, "Failed to refresh machine", http.StatusInternalServerError) - return + return nil, false } var content bytes.Buffer @@ -377,7 +493,7 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } writer.Header().Set("Content-Type", "text/html; charset=utf-8") @@ -390,12 +506,20 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } + return &machineKey, true +} + +func getNamespaceName( + writer http.ResponseWriter, + claims *IDTokenClaims, + stripEmaildomain bool, +) (string, bool) { namespaceName, err := NormalizeToFQDNRules( claims.Email, - h.cfg.OIDC.StripEmaildomain, + stripEmaildomain, ) if err != nil { log.Error().Err(err).Caller().Msgf("couldn't normalize email") @@ -409,12 +533,16 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return "", false } - // register the machine if it's new - log.Debug().Msg("Registering new machine after successful callback") + return namespaceName, true +} +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) @@ -434,7 +562,7 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } } else if err != nil { log.Error(). @@ -452,17 +580,24 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } - machineKeyStr := MachinePublicKeyStripPrefix(machineKey) + return namespace, true +} - _, err = h.RegisterMachineFromAuthCallback( +func (h *Headscale) registerMachineForOIDCCallback( + writer http.ResponseWriter, + namespace *Namespace, + machineKey *key.MachinePublic, +) bool { + machineKeyStr := MachinePublicKeyStripPrefix(*machineKey) + + if _, err := h.RegisterMachineFromAuthCallback( machineKeyStr, namespace.Name, RegisterMethodOIDC, - ) - if err != nil { + ); err != nil { log.Error(). Caller(). Err(err). @@ -477,9 +612,16 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return false } + 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, @@ -501,16 +643,8 @@ func (h *Headscale) OIDCCallback( Msg("Failed to write response") } - return + return nil, false } - 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") - } + return &content, true } From 56858a56dbd50cc1c0547542a0b78948318e0f30 Mon Sep 17 00:00:00 2001 From: Grigoriy Mikhalkin Date: Thu, 21 Jul 2022 23:47:59 +0200 Subject: [PATCH 4/4] Revert "decompose OIDCCallback method" This reverts commit 395caaad421797ac7fba6a0ca0e7df87d13262f2. --- oidc.go | 240 +++++++++++++------------------------------------------- poll.go | 18 ++--- 2 files changed, 62 insertions(+), 196 deletions(-) diff --git a/oidc.go b/oidc.go index 5509bd47..8b5f0242 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 b9a757aa..9c17b5cb 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().