cluster: obtain tenant information from headers

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 := getAuthTokenFromReq(req)
-		at, err := getAuthTokenFromPath(path)
 		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

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

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)
 	}

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
 		}

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.
+// If authToken == "multitenant" and TenantID header is not set, then nil Token is returned.
-func NewTokenPossibleMultitenant(authToken string) (*Token, error) {
+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)

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) {