netbird/management/server/http/middleware/auth_middleware.go

package middleware

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
	log "github.com/sirupsen/logrus"
	"go.opentelemetry.io/otel/metric"

	serverauth "github.com/netbirdio/netbird/management/server/auth"
	nbcontext "github.com/netbirdio/netbird/management/server/context"
	"github.com/netbirdio/netbird/management/server/http/middleware/bypass"
	"github.com/netbirdio/netbird/management/server/types"
	"github.com/netbirdio/netbird/shared/auth"
	"github.com/netbirdio/netbird/shared/management/http/util"
	"github.com/netbirdio/netbird/shared/management/status"
)

type EnsureAccountFunc func(ctx context.Context, userAuth auth.UserAuth) (string, string, error)
type SyncUserJWTGroupsFunc func(ctx context.Context, userAuth auth.UserAuth) error

type GetUserFromUserAuthFunc func(ctx context.Context, userAuth auth.UserAuth) (*types.User, error)

// jwtTokenCtxKey carries the parsed JWT token.
type jwtTokenCtxKey struct{}

// AuthMiddleware middleware to verify personal access tokens (PAT) and JWT tokens
type AuthMiddleware struct {
	authManager         serverauth.Manager
	ensureAccount       EnsureAccountFunc
	getUserFromUserAuth GetUserFromUserAuthFunc
	syncUserJWTGroups   SyncUserJWTGroupsFunc
	rateLimiter         *APIRateLimiter
	patUsageTracker     *PATUsageTracker
}

// NewAuthMiddleware instance constructor
func NewAuthMiddleware(
	authManager serverauth.Manager,
	ensureAccount EnsureAccountFunc,
	syncUserJWTGroups SyncUserJWTGroupsFunc,
	getUserFromUserAuth GetUserFromUserAuthFunc,
	rateLimiter *APIRateLimiter,
	meter metric.Meter,
) *AuthMiddleware {
	var patUsageTracker *PATUsageTracker
	if meter != nil {
		var err error
		patUsageTracker, err = NewPATUsageTracker(context.Background(), meter)
		if err != nil {
			log.Errorf("Failed to create PAT usage tracker: %s", err)
		}
	}

	return &AuthMiddleware{
		authManager:         authManager,
		ensureAccount:       ensureAccount,
		syncUserJWTGroups:   syncUserJWTGroups,
		getUserFromUserAuth: getUserFromUserAuth,
		rateLimiter:         rateLimiter,
		patUsageTracker:     patUsageTracker,
	}
}

// Handler composes the full authentication chain by wrapping the given
// handler with ValidationHandler followed by AccountAccessHandler.
func (m *AuthMiddleware) Handler(h http.Handler) http.Handler {
	return m.ValidationHandler(m.AccountAccessHandler(h))
}

// ValidationHandler authenticates the caller via JWT or PAT and stores the
// resulting UserAuth in the request context. It performs no account-level work.
func (m *AuthMiddleware) ValidationHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if bypass.ShouldBypass(r.URL.Path, h, w, r) {
			return
		}

		authHeader := strings.Split(r.Header.Get("Authorization"), " ")
		authType := strings.ToLower(authHeader[0])

		// fallback to token when receive pat as bearer
		if len(authHeader) >= 2 && authType == "bearer" && strings.HasPrefix(authHeader[1], "nbp_") {
			authType = "token"
			authHeader[0] = authType
		}

		switch authType {
		case "bearer":
			if err := m.validateJWT(r, authHeader); err != nil {
				log.WithContext(r.Context()).Errorf("Error when validating JWT: %s", err.Error())
				util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "token invalid"), w)
				return
			}
			h.ServeHTTP(w, r)
		case "token":
			if err := m.validatePAT(r, authHeader); err != nil {
				log.WithContext(r.Context()).Debugf("Error when validating PAT: %s", err.Error())
				// Check if it's a status error, otherwise default to Unauthorized
				if _, ok := status.FromError(err); !ok {
					err = status.Errorf(status.Unauthorized, "token invalid")
				}
				util.WriteError(r.Context(), err, w)
				return
			}
			h.ServeHTTP(w, r)
		default:
			util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "no valid authentication provided"), w)
			return
		}
	})
}

// AccountAccessHandler runs post-validation access checks for JWT-authenticated
// requests. PAT requests pass through unchanged.
func (m *AuthMiddleware) AccountAccessHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if bypass.ShouldBypass(r.URL.Path, h, w, r) {
			return
		}

		userAuth, err := nbcontext.GetUserAuthFromRequest(r)
		if err != nil {
			util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "no valid authentication provided"), w)
			return
		}

		if userAuth.IsPAT {
			h.ServeHTTP(w, r)
			return
		}

		validatedToken, _ := r.Context().Value(jwtTokenCtxKey{}).(*jwt.Token)

		if err := m.applyAccountAccess(r, userAuth, validatedToken); err != nil {
			log.WithContext(r.Context()).Errorf("Error applying JWT account access: %s", err.Error())
			util.WriteError(r.Context(), status.Errorf(status.Unauthorized, "token invalid"), w)
			return
		}

		h.ServeHTTP(w, r)
	})
}

func (m *AuthMiddleware) validateJWT(r *http.Request, authHeaderParts []string) error {
	token, err := getTokenFromJWTRequest(authHeaderParts)
	if err != nil {
		return fmt.Errorf("error extracting token: %w", err)
	}

	userAuth, validatedToken, err := m.authManager.ValidateAndParseToken(r.Context(), token)
	if err != nil {
		return err
	}

	// propagates ctx change to upstream middleware
	*r = *nbcontext.SetUserAuthInRequest(r, userAuth)
	*r = *r.WithContext(context.WithValue(r.Context(), jwtTokenCtxKey{}, validatedToken))
	return nil
}

func (m *AuthMiddleware) validatePAT(r *http.Request, authHeaderParts []string) error {
	token, err := getTokenFromPATRequest(authHeaderParts)
	if err != nil {
		return fmt.Errorf("error extracting token: %w", err)
	}

	if m.patUsageTracker != nil {
		m.patUsageTracker.IncrementUsage(token)
	}

	if !isTerraformRequest(r) && !m.rateLimiter.Allow(token) {
		return status.Errorf(status.TooManyRequests, "too many requests")
	}

	ctx := r.Context()
	user, pat, accDomain, accCategory, err := m.authManager.GetPATInfo(ctx, token)
	if err != nil {
		return fmt.Errorf("invalid Token: %w", err)
	}
	if time.Now().After(pat.GetExpirationDate()) {
		return fmt.Errorf("token expired")
	}

	if err := m.authManager.MarkPATUsed(ctx, pat.ID); err != nil {
		return err
	}

	userAuth := auth.UserAuth{
		UserId:         user.Id,
		AccountId:      user.AccountID,
		Domain:         accDomain,
		DomainCategory: accCategory,
		IsPAT:          true,
	}

	// propagates ctx change to upstream middleware
	*r = *nbcontext.SetUserAuthInRequest(r, userAuth)
	return nil
}

// applyAccountAccess executes account-level checks for an authenticated JWT
// user: ensures the account exists, verifies access via JWT groups, syncs
// groups, and fetches the user record.
func (m *AuthMiddleware) applyAccountAccess(r *http.Request, userAuth auth.UserAuth, validatedToken *jwt.Token) error {
	ctx := r.Context()

	// we need to call this method because if user is new, we will automatically add it to existing or create a new account
	accountId, _, err := m.ensureAccount(ctx, userAuth)
	if err != nil {
		return err
	}

	if userAuth.AccountId != accountId {
		log.WithContext(ctx).Tracef("Auth middleware sets accountId from ensure, before %s, now %s", userAuth.AccountId, accountId)
		userAuth.AccountId = accountId
	}

	userAuth, err = m.authManager.EnsureUserAccessByJWTGroups(ctx, userAuth, validatedToken)
	if err != nil {
		return err
	}

	if err := m.syncUserJWTGroups(ctx, userAuth); err != nil {
		log.WithContext(ctx).Errorf("HTTP server failed to sync user JWT groups: %s", err)
	}

	if _, err := m.getUserFromUserAuth(ctx, userAuth); err != nil {
		log.WithContext(ctx).Errorf("HTTP server failed to update user from user auth: %s", err)
		return err
	}

	// propagates ctx change to upstream middleware
	*r = *nbcontext.SetUserAuthInRequest(r, userAuth)
	return nil
}

func isTerraformRequest(r *http.Request) bool {
	ua := strings.ToLower(r.Header.Get("User-Agent"))
	return strings.Contains(ua, "terraform")
}

// getTokenFromJWTRequest is a "TokenExtractor" that takes auth header parts and extracts
// the JWT token from the Authorization header.
func getTokenFromJWTRequest(authHeaderParts []string) (string, error) {
	if len(authHeaderParts) != 2 || strings.ToLower(authHeaderParts[0]) != "bearer" {
		return "", errors.New("authorization header format must be Bearer {token}")
	}

	return authHeaderParts[1], nil
}

// getTokenFromPATRequest is a "TokenExtractor" that takes auth header parts and extracts
// the PAT token from the Authorization header.
func getTokenFromPATRequest(authHeaderParts []string) (string, error) {
	if len(authHeaderParts) != 2 || strings.ToLower(authHeaderParts[0]) != "token" {
		return "", errors.New("authorization header format must be Token {token}")
	}

	return authHeaderParts[1], nil
}