diff --git a/ee/http/middleware/pat.go b/ee/http/middleware/pat.go index 292d209dd9..2bb188d9f2 100644 --- a/ee/http/middleware/pat.go +++ b/ee/http/middleware/pat.go @@ -29,47 +29,46 @@ func (p *Pat) Wrap(next http.Handler) http.Handler { for _, header := range p.headers { values = append(values, r.Header.Get(header)) - - if header == "SIGNOZ-API-KEY" { - patToken = values[0] - err := p.db.NewSelect().Model(&pat).Where("token = ?", patToken).Scan(r.Context()) - if err != nil { - next.ServeHTTP(w, r) - return - } - - if pat.ExpiresAt < time.Now().Unix() && pat.ExpiresAt != 0 { - next.ServeHTTP(w, r) - return - } - - // get user from db - user := types.User{} - err = p.db.NewSelect().Model(&user).Where("id = ?", pat.UserID).Scan(r.Context()) - if err != nil { - next.ServeHTTP(w, r) - return - } - - jwt := authtypes.Claims{ - UserID: user.ID, - GroupID: user.GroupID, - Email: user.Email, - OrgID: user.OrgID, - } - - ctx := authtypes.NewContextWithClaims(r.Context(), jwt) - r = r.WithContext(ctx) - - // Mark to update last used since SIGNOZ-API-KEY is present and successful - updateLastUsed = true - } } + ctx, err := p.uuid.ContextFromRequest(r.Context(), values...) if err != nil { next.ServeHTTP(w, r) return } + patToken, ok := authtypes.UUIDFromContext(ctx) + if !ok { + next.ServeHTTP(w, r) + return + } + + err = p.db.NewSelect().Model(&pat).Where("token = ?", patToken).Scan(r.Context()) + if err != nil { + next.ServeHTTP(w, r) + return + } + + if pat.ExpiresAt < time.Now().Unix() && pat.ExpiresAt != 0 { + next.ServeHTTP(w, r) + return + } + + // get user from db + user := types.User{} + err = p.db.NewSelect().Model(&user).Where("id = ?", pat.UserID).Scan(r.Context()) + if err != nil { + next.ServeHTTP(w, r) + return + } + + jwt := authtypes.Claims{ + UserID: user.ID, + GroupID: user.GroupID, + Email: user.Email, + OrgID: user.OrgID, + } + + ctx = authtypes.NewContextWithClaims(ctx, jwt) r = r.WithContext(ctx)