From 3a2d03528334a4f907b978b6772eaf85e3a05d02 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Wed, 30 Aug 2023 14:08:47 +0200
Subject: [PATCH] lib/auth: add NewTokenPossibleMultitenant() for parsing auth
 token, which can be multitenant

Disallow parsing multitenant token at auth.NewToken().

Use auth.NewTokenPossibleMultitenant() at vminsert only. All the other callers should call auth.NewToken(),
since they do not support multitenant token.

This is a follow-up for f0c06b428e5c794705fb1f70a7392b3f2920f7e0

Updates https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4910
---
 app/vminsert/main.go  |  2 +-
 app/vmselect/main.go  |  5 -----
 lib/auth/auth.go      | 15 ++++++++++-----
 lib/auth/auth_test.go | 20 ++++++++++++++++++++
 4 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/app/vminsert/main.go b/app/vminsert/main.go
index 938dc3e0f4..b23827a02e 100644
--- a/app/vminsert/main.go
+++ b/app/vminsert/main.go
@@ -208,7 +208,7 @@ func requestHandler(w http.ResponseWriter, r *http.Request) bool {
 		// This is not our link.
 		return false
 	}
-	at, err := auth.NewToken(p.AuthToken)
+	at, err := auth.NewTokenPossibleMultitenant(p.AuthToken)
 	if err != nil {
 		httpserver.Errorf(w, r, "auth error: %s", err)
 		return true
diff --git a/app/vmselect/main.go b/app/vmselect/main.go
index 495f3d7348..9fe87822eb 100644
--- a/app/vmselect/main.go
+++ b/app/vmselect/main.go
@@ -292,11 +292,6 @@ func requestHandler(w http.ResponseWriter, r *http.Request) bool {
 		httpserver.Errorf(w, r, "auth error: %s", err)
 		return true
 	}
-	if at == nil {
-		// the only option for at to be nil is when p.AuthToken == "multitenant"
-		// vmselect does not have multitenant endpoint, so request must be rejected
-		return false
-	}
 	switch p.Prefix {
 	case "select":
 		return selectHandler(qt, startTime, w, r, p, at)
diff --git a/lib/auth/auth.go b/lib/auth/auth.go
index 38e2354482..9a7a4346fb 100644
--- a/lib/auth/auth.go
+++ b/lib/auth/auth.go
@@ -24,12 +24,7 @@ func (t *Token) String() string {
 }
 
 // NewToken returns new Token for the given authToken.
-//
-// If authToken == "multitenant", then nil Token is returned.
 func NewToken(authToken string) (*Token, error) {
-	if authToken == "multitenant" {
-		return nil, nil
-	}
 	var t Token
 	if err := t.Init(authToken); err != nil {
 		return nil, err
@@ -37,6 +32,16 @@ func NewToken(authToken string) (*Token, error) {
 	return &t, nil
 }
 
+// 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" {
+		return nil, nil
+	}
+	return NewToken(authToken)
+}
+
 // Init initializes t from authToken.
 func (t *Token) Init(authToken string) error {
 	tmp := strings.Split(authToken, ":")
diff --git a/lib/auth/auth_test.go b/lib/auth/auth_test.go
index eddf1c6fe0..49f8ca7652 100644
--- a/lib/auth/auth_test.go
+++ b/lib/auth/auth_test.go
@@ -26,6 +26,24 @@ func TestNewTokenSuccess(t *testing.T) {
 	f("1:4294967295", "1:4294967295")
 	// max uint32 accountID and projectID
 	f("4294967295:4294967295", "4294967295:4294967295")
+}
+
+func TestNewTokenPossibleMultitenantSuccess(t *testing.T) {
+	f := func(token string, want string) {
+		t.Helper()
+		newToken, err := NewTokenPossibleMultitenant(token)
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		got := newToken.String()
+		if got != want {
+			t.Fatalf("unexpected NewToken() result;got\n%s\nwant\n%s", got, want)
+		}
+	}
+	// token with accountID only
+	f("1", "1")
+	// token with accountID and projecTID
+	f("1:2", "1:2")
 	// multitenant
 	f("multitenant", "multitenant")
 }
@@ -75,4 +93,6 @@ func TestNewTokenFailure(t *testing.T) {
 	f("a:b:c")
 	// many int parts in the token"
 	f("1:2:3")
+	// multitenant
+	f("multitenant")
 }