From f5981c1447159220841ff6c6d83be1c2cc045dc6 Mon Sep 17 00:00:00 2001
From: Alexander Marshalov <_@marshalov.org>
Date: Thu, 20 Apr 2023 19:08:27 +0200
Subject: [PATCH] vmauth ip filters (refactoring) (#4059)

Added ip filters (allow_list and deny_list) for enterprise-version of vmauth (#3491)

---------

Signed-off-by: Alexander Marshalov <_@marshalov.org>
---
 app/vmauth/auth_config.go      | 48 ++++++++++++++++++++++------------
 app/vmauth/auth_config_test.go | 12 +++++++--
 app/vmauth/main.go             |  2 +-
 3 files changed, 43 insertions(+), 19 deletions(-)

diff --git a/app/vmauth/auth_config.go b/app/vmauth/auth_config.go
index 2310d5035a..3a4737c169 100644
--- a/app/vmauth/auth_config.go
+++ b/app/vmauth/auth_config.go
@@ -289,11 +289,11 @@ func initAuthConfig() {
 	// See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/1240
 	sighupCh := procutil.NewSighupChan()
 
-	m, err := readAuthConfig(*authConfigPath)
+	err := loadAuthConfig()
 	if err != nil {
 		logger.Fatalf("cannot load auth config from `-auth.config=%s`: %s", *authConfigPath, err)
 	}
-	authConfig.Store(m)
+
 	stopCh = make(chan struct{})
 	authConfigWG.Add(1)
 	go func() {
@@ -324,44 +324,60 @@ func authConfigReloader(sighupCh <-chan os.Signal) {
 			procutil.SelfSIGHUP()
 		case <-sighupCh:
 			logger.Infof("SIGHUP received; loading -auth.config=%q", *authConfigPath)
-			m, err := readAuthConfig(*authConfigPath)
+			err := loadAuthConfig()
 			if err != nil {
 				logger.Errorf("failed to load -auth.config=%q; using the last successfully loaded config; error: %s", *authConfigPath, err)
 				continue
 			}
-			authConfig.Store(m)
 			logger.Infof("Successfully reloaded -auth.config=%q", *authConfigPath)
 		}
 	}
 }
 
-var authConfig atomic.Value
+var authConfig atomic.Pointer[AuthConfig]
+var authUsers atomic.Pointer[map[string]*UserInfo]
 var authConfigWG sync.WaitGroup
 var stopCh chan struct{}
 
-func readAuthConfig(path string) (map[string]*UserInfo, error) {
+func loadAuthConfig() error {
+	ac, err := readAuthConfig(*authConfigPath)
+	if err != nil {
+		return fmt.Errorf("failed to load `-auth.config=%q`: %s", *authConfigPath, err)
+	}
+
+	m, err := parseAuthConfigUsers(ac)
+	if err != nil {
+		return fmt.Errorf("failed to parse users from `-auth.config=%q`: %s", *authConfigPath, err)
+	}
+	logger.Infof("Loaded information about %d users from %q", len(m), *authConfigPath)
+
+	authConfig.Store(ac)
+	authUsers.Store(&m)
+
+	return nil
+}
+
+func readAuthConfig(path string) (*AuthConfig, error) {
 	data, err := fs.ReadFileOrHTTP(path)
 	if err != nil {
 		return nil, err
 	}
-	m, err := parseAuthConfig(data)
-	if err != nil {
-		return nil, fmt.Errorf("cannot parse %q: %w", path, err)
-	}
-	logger.Infof("Loaded information about %d users from %q", len(m), path)
-	return m, nil
+	return parseAuthConfig(data)
 }
 
-func parseAuthConfig(data []byte) (map[string]*UserInfo, error) {
-	var err error
-	data, err = envtemplate.ReplaceBytes(data)
+func parseAuthConfig(data []byte) (*AuthConfig, error) {
+	data, err := envtemplate.ReplaceBytes(data)
 	if err != nil {
 		return nil, fmt.Errorf("cannot expand environment vars: %w", err)
 	}
 	var ac AuthConfig
-	if err := yaml.UnmarshalStrict(data, &ac); err != nil {
+	if err = yaml.UnmarshalStrict(data, &ac); err != nil {
 		return nil, fmt.Errorf("cannot unmarshal AuthConfig data: %w", err)
 	}
+	return &ac, nil
+}
+
+func parseAuthConfigUsers(ac *AuthConfig) (map[string]*UserInfo, error) {
 	uis := ac.Users
 	if len(uis) == 0 {
 		return nil, fmt.Errorf("`users` section cannot be empty in AuthConfig")
diff --git a/app/vmauth/auth_config_test.go b/app/vmauth/auth_config_test.go
index 6771513efe..8c66640593 100644
--- a/app/vmauth/auth_config_test.go
+++ b/app/vmauth/auth_config_test.go
@@ -13,7 +13,11 @@ import (
 func TestParseAuthConfigFailure(t *testing.T) {
 	f := func(s string) {
 		t.Helper()
-		_, err := parseAuthConfig([]byte(s))
+		ac, err := parseAuthConfig([]byte(s))
+		if err != nil {
+			return
+		}
+		_, err = parseAuthConfigUsers(ac)
 		if err == nil {
 			t.Fatalf("expecting non-nil error")
 		}
@@ -202,7 +206,11 @@ users:
 func TestParseAuthConfigSuccess(t *testing.T) {
 	f := func(s string, expectedAuthConfig map[string]*UserInfo) {
 		t.Helper()
-		m, err := parseAuthConfig([]byte(s))
+		ac, err := parseAuthConfig([]byte(s))
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		m, err := parseAuthConfigUsers(ac)
 		if err != nil {
 			t.Fatalf("unexpected error: %s", err)
 		}
diff --git a/app/vmauth/main.go b/app/vmauth/main.go
index 4b400c0631..7281527c82 100644
--- a/app/vmauth/main.go
+++ b/app/vmauth/main.go
@@ -94,7 +94,7 @@ func requestHandler(w http.ResponseWriter, r *http.Request) bool {
 		authToken = strings.Replace(authToken, "Token", "Bearer", 1)
 	}
 
-	ac := authConfig.Load().(map[string]*UserInfo)
+	ac := *authUsers.Load()
 	ui := ac[authToken]
 	if ui == nil {
 		invalidAuthTokenRequests.Inc()