From 35be17ab3043c4d2b9d8ee4f60f0394dc6d71183 Mon Sep 17 00:00:00 2001 From: Andrii Chubatiuk Date: Tue, 1 Oct 2024 10:47:55 +0300 Subject: [PATCH] cluster: obtain tenant information from headers --- app/vmagent/main.go | 12 ++++++------ app/vminsert/main.go | 2 +- app/vminsert/opentsdbhttp/request_handler.go | 2 +- app/vmselect/main.go | 6 +++--- lib/auth/auth.go | 8 ++++++-- lib/auth/auth_test.go | 15 ++++++++++----- 6 files changed, 27 insertions(+), 18 deletions(-) diff --git a/app/vmagent/main.go b/app/vmagent/main.go index c20b9f53f..22791444d 100644 --- a/app/vmagent/main.go +++ b/app/vmagent/main.go @@ -197,16 +197,16 @@ func getOpenTSDBHTTPInsertHandler() func(req *http.Request) error { } } return func(req *http.Request) error { - path := strings.Replace(req.URL.Path, "//", "/", -1) - at, err := getAuthTokenFromPath(path) + at, err := getAuthTokenFromReq(req) if err != nil { - return fmt.Errorf("cannot obtain auth token from path %q: %w", path, err) + return fmt.Errorf("cannot obtain auth token: %w", err) } return opentsdbhttp.InsertHandler(at, req) } } -func getAuthTokenFromPath(path string) (*auth.Token, error) { +func getAuthTokenFromReq(req *http.Request) (*auth.Token, error) { + path := strings.Replace(req.URL.Path, "//", "/", -1) p, err := httpserver.ParsePath(path) if err != nil { return nil, fmt.Errorf("cannot parse multitenant path: %w", err) @@ -217,7 +217,7 @@ func getAuthTokenFromPath(path string) (*auth.Token, error) { if p.Suffix != "opentsdb/api/put" { return nil, fmt.Errorf("unsupported path requested: %q; expecting 'opentsdb/api/put'", p.Suffix) } - return auth.NewTokenPossibleMultitenant(p.AuthToken) + return auth.NewTokenPossibleMultitenant(p.AuthToken, req.Header) } func requestHandler(w http.ResponseWriter, r *http.Request) bool { @@ -498,7 +498,7 @@ func processMultitenantRequest(w http.ResponseWriter, r *http.Request, path stri httpserver.Errorf(w, r, `unsupported multitenant prefix: %q; expected "insert"`, p.Prefix) return true } - at, err := auth.NewToken(p.AuthToken) + at, err := auth.NewTokenPossibleMultitenant(p.AuthToken, r.Header) if err != nil { httpserver.Errorf(w, r, "cannot obtain auth token: %s", err) return true diff --git a/app/vminsert/main.go b/app/vminsert/main.go index d8fc4be22..8eea0fb39 100644 --- a/app/vminsert/main.go +++ b/app/vminsert/main.go @@ -217,7 +217,7 @@ func requestHandler(w http.ResponseWriter, r *http.Request) bool { // This is not our link. return false } - at, err := auth.NewTokenPossibleMultitenant(p.AuthToken) + at, err := auth.NewTokenPossibleMultitenant(p.AuthToken, r.Header) if err != nil { httpserver.Errorf(w, r, "auth error: %s", err) return true diff --git a/app/vminsert/opentsdbhttp/request_handler.go b/app/vminsert/opentsdbhttp/request_handler.go index fe40fd11c..3fd4852e6 100644 --- a/app/vminsert/opentsdbhttp/request_handler.go +++ b/app/vminsert/opentsdbhttp/request_handler.go @@ -34,7 +34,7 @@ func InsertHandler(req *http.Request) error { // This is not our link. return fmt.Errorf("unexpected path requested on HTTP OpenTSDB server: %q", path) } - at, err := auth.NewTokenPossibleMultitenant(p.AuthToken) + at, err := auth.NewTokenPossibleMultitenant(p.AuthToken, req.Header) if err != nil { return fmt.Errorf("auth error: %w", err) } diff --git a/app/vmselect/main.go b/app/vmselect/main.go index 0b7db09bd..ed1a29b27 100644 --- a/app/vmselect/main.go +++ b/app/vmselect/main.go @@ -269,7 +269,7 @@ func requestHandler(w http.ResponseWriter, r *http.Request) bool { httpserver.Errorf(w, r, "cannot parse path %q: %s", path, err) return true } - at, err := auth.NewTokenPossibleMultitenant(p.AuthToken) + at, err := auth.NewTokenPossibleMultitenant(p.AuthToken, nil) if err != nil { httpserver.Errorf(w, r, "auth error: %s", err) return true @@ -656,7 +656,7 @@ func handleStaticAndSimpleRequests(w http.ResponseWriter, r *http.Request, path } switch p.Suffix { case "prometheus/api/v1/status/active_queries": - at, err := auth.NewTokenPossibleMultitenant(p.AuthToken) + at, err := auth.NewTokenPossibleMultitenant(p.AuthToken, nil) if err != nil { return false } @@ -665,7 +665,7 @@ func handleStaticAndSimpleRequests(w http.ResponseWriter, r *http.Request, path promql.ActiveQueriesHandler(at, w, r) return true case "prometheus/api/v1/status/top_queries": - at, err := auth.NewTokenPossibleMultitenant(p.AuthToken) + at, err := auth.NewTokenPossibleMultitenant(p.AuthToken, nil) if err != nil { return false } diff --git a/lib/auth/auth.go b/lib/auth/auth.go index 893515d59..334e2b439 100644 --- a/lib/auth/auth.go +++ b/lib/auth/auth.go @@ -2,6 +2,7 @@ package auth import ( "fmt" + "net/http" "strconv" "strings" ) @@ -34,9 +35,12 @@ func NewToken(authToken string) (*Token, error) { // NewTokenPossibleMultitenant returns new Token for the given authToken. // -// If authToken == "multitenant", then nil Token is returned. -func NewTokenPossibleMultitenant(authToken string) (*Token, error) { +// If authToken == "multitenant" and TenantID header is not set, then nil Token is returned. +func NewTokenPossibleMultitenant(authToken string, headers http.Header) (*Token, error) { if authToken == "multitenant" { + if tenantID := headers.Get("TenantID"); tenantID != "" { + return NewToken(tenantID) + } return nil, nil } return NewToken(authToken) diff --git a/lib/auth/auth_test.go b/lib/auth/auth_test.go index 49f8ca765..9edaa8f36 100644 --- a/lib/auth/auth_test.go +++ b/lib/auth/auth_test.go @@ -1,6 +1,7 @@ package auth import ( + "net/http" "testing" ) @@ -29,9 +30,11 @@ func TestNewTokenSuccess(t *testing.T) { } func TestNewTokenPossibleMultitenantSuccess(t *testing.T) { - f := func(token string, want string) { + f := func(token string, tenantIDValue string, want string) { t.Helper() - newToken, err := NewTokenPossibleMultitenant(token) + headers := http.Header{} + headers.Set("TenantID", tenantIDValue) + newToken, err := NewTokenPossibleMultitenant(token, headers) if err != nil { t.Fatalf("unexpected error: %s", err) } @@ -41,11 +44,13 @@ func TestNewTokenPossibleMultitenantSuccess(t *testing.T) { } } // token with accountID only - f("1", "1") + f("1", "", "1") // token with accountID and projecTID - f("1:2", "1:2") + f("1:2", "", "1:2") // multitenant - f("multitenant", "multitenant") + f("multitenant", "", "multitenant") + // multitenant with tenantID in headers + f("multitenant", "1:2", "1:2") } func TestNewTokenFailure(t *testing.T) {