From 98dcec0ee2023d24913941bfef7ecc5a7d2b78f8 Mon Sep 17 00:00:00 2001
From: Patrick Schmidt <patrick.schmidt@innogames.com>
Date: Fri, 1 Sep 2023 19:57:04 +0200
Subject: [PATCH] Clean up old signature hash pools

---
 weed/s3api/auth_credentials.go       |  6 ++-
 weed/s3api/auth_signature_v4.go      | 66 +++++++++++++++++++++-------
 weed/s3api/auto_signature_v4_test.go |  3 +-
 3 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/weed/s3api/auth_credentials.go b/weed/s3api/auth_credentials.go
index f2d057b90..234dc100b 100644
--- a/weed/s3api/auth_credentials.go
+++ b/weed/s3api/auth_credentials.go
@@ -33,6 +33,7 @@ type IdentityAccessManagement struct {
 	isAuthEnabled bool
 	domain        string
 	hashes        map[string]*sync.Pool
+	hashCounters  map[string]*int32
 	hashMu        sync.RWMutex
 }
 
@@ -79,8 +80,9 @@ func (action Action) getPermission() Permission {
 
 func NewIdentityAccessManagement(option *S3ApiServerOption) *IdentityAccessManagement {
 	iam := &IdentityAccessManagement{
-		domain: option.DomainName,
-		hashes: make(map[string]*sync.Pool),
+		domain:       option.DomainName,
+		hashes:       make(map[string]*sync.Pool),
+		hashCounters: make(map[string]*int32),
 	}
 	if option.Config != "" {
 		if err := iam.loadS3ApiConfigurationFromFile(option.Config); err != nil {
diff --git a/weed/s3api/auth_signature_v4.go b/weed/s3api/auth_signature_v4.go
index 06cdf67e4..04548cc6f 100644
--- a/weed/s3api/auth_signature_v4.go
+++ b/weed/s3api/auth_signature_v4.go
@@ -32,6 +32,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 	"unicode/utf8"
 
@@ -463,36 +464,67 @@ func (iam *IdentityAccessManagement) doesPresignedSignatureMatch(hashedPayload s
 	return identity, s3err.ErrNone
 }
 
-// getSignature
 func (iam *IdentityAccessManagement) getSignature(secretKey string, t time.Time, region string, service string, stringToSign string) string {
+	pool := iam.getSignatureHashPool(secretKey, t, region, service)
+	h := pool.Get().(hash.Hash)
+	defer pool.Put(h)
+
+	h.Reset()
+	h.Write([]byte(stringToSign))
+	sig := hex.EncodeToString(h.Sum(nil))
+
+	return sig
+}
+
+func (iam *IdentityAccessManagement) getSignatureHashPool(secretKey string, t time.Time, region string, service string) *sync.Pool {
+	// Build a caching key for the pool.
 	date := t.Format(yyyymmdd)
 	hashID := "AWS4" + secretKey + "/" + date + "/" + region + "/" + service + "/" + "aws4_request"
 
+	// Try to find an existing pool and return it.
 	iam.hashMu.RLock()
 	pool, ok := iam.hashes[hashID]
 	iam.hashMu.RUnlock()
 
 	if !ok {
 		iam.hashMu.Lock()
-		if pool, ok = iam.hashes[hashID]; !ok {
-			pool = &sync.Pool{
-				New: func() any {
-					signingKey := getSigningKey(secretKey, date, region, service)
-					return hmac.New(sha256.New, signingKey)
-				},
-			}
-			iam.hashes[hashID] = pool
-		}
-		iam.hashMu.Unlock()
+		defer iam.hashMu.Unlock()
+		pool, ok = iam.hashes[hashID]
 	}
 
-	h := pool.Get().(hash.Hash)
-	h.Reset()
-	h.Write([]byte(stringToSign))
-	sig := hex.EncodeToString(h.Sum(nil))
-	pool.Put(h)
+	if ok {
+		atomic.StoreInt32(iam.hashCounters[hashID], 1)
+		return pool
+	}
 
-	return sig
+	// Create a pool that returns HMAC hashers for the requested parameters to avoid expensive re-initializing
+	// of new instances on every request.
+	iam.hashes[hashID] = &sync.Pool{
+		New: func() any {
+			signingKey := getSigningKey(secretKey, date, region, service)
+			return hmac.New(sha256.New, signingKey)
+		},
+	}
+	iam.hashCounters[hashID] = new(int32)
+
+	// Clean up unused pools automatically after one hour of inactivity
+	ticker := time.NewTicker(time.Hour)
+	go func() {
+		for range ticker.C {
+			old := atomic.SwapInt32(iam.hashCounters[hashID], 0)
+			if old == 0 {
+				break
+			}
+		}
+
+		ticker.Stop()
+		iam.hashMu.Lock()
+		delete(iam.hashes, hashID)
+		delete(iam.hashCounters, hashID)
+		iam.hashMu.Unlock()
+	}()
+
+	return iam.hashes[hashID]
 }
 
 func contains(list []string, elem string) bool {
diff --git a/weed/s3api/auto_signature_v4_test.go b/weed/s3api/auto_signature_v4_test.go
index 15ca90b93..8d0b677f8 100644
--- a/weed/s3api/auto_signature_v4_test.go
+++ b/weed/s3api/auto_signature_v4_test.go
@@ -127,7 +127,8 @@ func TestCheckAdminRequestAuthType(t *testing.T) {
 func BenchmarkGetSignature(b *testing.B) {
 	t := time.Now()
 	iam := IdentityAccessManagement{
-		hashes: make(map[string]*sync.Pool),
+		hashes:       make(map[string]*sync.Pool),
+		hashCounters: make(map[string]*int32),
 	}
 
 	b.ReportAllocs()