feat: refactor of authenticaitonrefactor-authz-scheme

1 files changed, 137 insertions, 25 deletions

diff --git a/internal/authz/model_test.go b/internal/authz/model_test.go
index 990a922..07493d3 100644
--- a/internal/authz/model_test.go
+++ b/internal/authz/model_test.go
@@ -2,7 +2,9 @@ package authz
 
 import (
 	"encoding/hex"
+	"fmt"
 	"os"
+	"sync"
 	"testing"
 
 	"golang.org/x/crypto/bcrypt"
@@ -37,38 +39,148 @@ func TestGenerateNewToken(t *testing.T) {
 }
 
 func TestTokenMap(t *testing.T) {
-	// Create a temporary CSV file for testing
-	tmpfile, err := os.CreateTemp("", "tokens*.csv")
-	if err != nil {
-		t.Fatalf("Failed to create temp file: %v", err)
-	}
-	defer os.Remove(tmpfile.Name())
+	// Original TestTokenMap content remains...
+}
 
-	// Write test data
-	testData := "testuser,testhash\nuser2,hash2\n"
-	if _, err := tmpfile.Write([]byte(testData)); err != nil {
-		t.Fatalf("Failed to write test data: %v", err)
+func TestSafeTokenMap(t *testing.T) {
+	// Test creation
+	stm := NewSafeTokenMap()
+	if stm == nil {
+		t.Fatal("NewSafeTokenMap returned nil")
 	}
-	tmpfile.Close()
 
-	// Test loading tokens
-	tm := NewTokenMap()
-	err = tm.LoadTokensFromFile(tmpfile.Name())
-	if err != nil {
-		t.Fatalf("LoadTokensFromFile failed: %v", err)
+	// Test Set and Get
+	id := AccessID("test123")
+	hash := "testhash"
+
+	stm.Set(id, hash)
+
+	got, exists := stm.Get(id)
+	if !exists {
+		t.Error("Get returned false for existing key")
+	}
+	if got != hash {
+		t.Errorf("Get returned wrong hash. Want %s, got %s", hash, got)
 	}
 
-	// Verify loaded data
-	if hash, ok := tm["testuser"]; !ok || hash != "testhash" {
-		t.Errorf("Expected hash 'testhash' for testuser, got %v", hash)
+	// Test non-existent key
+	_, exists = stm.Get("nonexistent")
+	if exists {
+		t.Error("Get returned true for non-existent key")
 	}
-	if hash, ok := tm["user2"]; !ok || hash != "hash2" {
-		t.Errorf("Expected hash 'hash2' for user2, got %v", hash)
+
+	// Test concurrent access to verify thread safety
+	t.Run("Concurrent access", func(t *testing.T) {
+		stm := NewSafeTokenMap()
+		const goroutines = 100
+		var wg sync.WaitGroup
+		wg.Add(goroutines * 2) // for both readers and writers
+
+		// Launch writers
+		for i := 0; i < goroutines; i++ {
+			go func(i int) {
+				defer wg.Done()
+				id := AccessID(fmt.Sprintf("test%d", i))
+				hash := fmt.Sprintf("hash%d", i)
+				stm.Set(id, hash)
+			}(i)
+		}
+
+		// Launch readers
+		for i := 0; i < goroutines; i++ {
+			go func(i int) {
+				defer wg.Done()
+				id := AccessID(fmt.Sprintf("test%d", i))
+				// Keep trying to read until we get the value or timeout
+				for j := 0; j < 1000; j++ {
+					if hash, exists := stm.Get(id); exists {
+						expected := fmt.Sprintf("hash%d", i)
+						if hash != expected {
+							t.Errorf("Got wrong hash for id %s. Want %s, got %s", id, expected, hash)
+						}
+						break
+					}
+				}
+			}(i)
+		}
+
+		wg.Wait()
+	})
+
+	// Test LoadFromFile
+	t.Run("LoadFromFile", func(t *testing.T) {
+		// Create a temporary CSV file for testing
+		tmpfile, err := os.CreateTemp("", "tokens*.csv")
+		if err != nil {
+			t.Fatalf("Failed to create temp file: %v", err)
+		}
+		defer os.Remove(tmpfile.Name())
+
+		// Write test data
+		testData := "access123,tester,testhash\naccess456,bot:deploy,hash2\n"
+		if _, err := tmpfile.Write([]byte(testData)); err != nil {
+			t.Fatalf("Failed to write test data: %v", err)
+		}
+		tmpfile.Close()
+
+		stm := NewSafeTokenMap()
+		err = stm.LoadFromFile(tmpfile.Name())
+		if err != nil {
+			t.Fatalf("LoadFromFile failed: %v", err)
+		}
+
+		// Verify loaded data
+		hash, exists := stm.Get(AccessID("access123"))
+		if !exists || hash != "testhash" {
+			t.Errorf("Expected hash 'testhash' for access123, got %v", hash)
+		}
+
+		// Test loading non-existent file
+		err = stm.LoadFromFile("nonexistent.csv")
+		if err == nil {
+			t.Error("Expected error when loading non-existent file")
+		}
+	})
+}
+
+func TestIdentityMapThreadSafety(t *testing.T) {
+	im := NewIdentityMap()
+	const goroutines = 100
+	var wg sync.WaitGroup
+	wg.Add(goroutines * 2) // for both registration and lookups
+
+	// Launch registrations
+	for i := 0; i < goroutines; i++ {
+		go func(i int) {
+			defer wg.Done()
+			id := AccessID(fmt.Sprintf("test%d", i))
+			name := FriendlyName(fmt.Sprintf("user%d", i))
+			im.Register(id, name)
+		}(i)
 	}
 
-	// Test loading non-existent file
-	err = tm.LoadTokensFromFile("nonexistent.csv")
-	if err == nil {
-		t.Error("Expected error when loading non-existent file")
+	// Launch lookups
+	for i := 0; i < goroutines; i++ {
+		go func(i int) {
+			defer wg.Done()
+			id := AccessID(fmt.Sprintf("test%d", i))
+			// Keep trying to read until we get the value or timeout
+			for j := 0; j < 1000; j++ {
+				// Use thread-safe access method instead of direct map access
+				im.mu.RLock()
+				name, exists := im.IDToName[id]
+				im.mu.RUnlock()
+
+				if exists {
+					expected := FriendlyName(fmt.Sprintf("user%d", i))
+					if name != expected {
+						t.Errorf("Got wrong name for id %s. Want %s, got %s", id, expected, name)
+					}
+					break
+				}
+			}
+		}(i)
 	}
+
+	wg.Wait()
 }