From 38ed5925934e664841330065d17c44334fa6872a Mon Sep 17 00:00:00 2001
From: xuhaidong <haidong.xu@shanbay.com>
Date: Tue, 26 Sep 2023 10:20:53 +0800
Subject: [PATCH] fix: infinite recursion in memory cache

---
 extensions/cachext/backend/memory/memory.go | 24 +++++++++++++++------
 extensions/cachext/backend/redis/redis.go   |  2 +-
 extensions/cachext/ext.go                   | 10 ++++-----
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/extensions/cachext/backend/memory/memory.go b/extensions/cachext/backend/memory/memory.go
index 908ff073..5f613683 100644
--- a/extensions/cachext/backend/memory/memory.go
+++ b/extensions/cachext/backend/memory/memory.go
@@ -2,6 +2,7 @@ package memory
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	"github.com/spf13/viper"
@@ -10,7 +11,7 @@ import (
 )
 
 func init() {
-	if err := cachext.RegisteBackend("memory", func() cachext.CacheBackend { return &memoryBackend{} }); err != nil {
+	if err := cachext.RegisterBackend("memory", func() cachext.CacheBackend { return &memoryBackend{} }); err != nil {
 		panic("MemoryBackend Init error")
 	}
 }
@@ -21,6 +22,7 @@ type memoryBackendNode struct {
 }
 
 type memoryBackend struct {
+	lock   sync.Mutex
 	client map[string]*memoryBackendNode
 }
 
@@ -47,6 +49,8 @@ func (m *memoryBackend) Get(ctx context.Context, key string) ([]byte, error) {
 }
 
 func (m *memoryBackend) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error {
+	m.lock.Lock()
+	defer m.lock.Unlock()
 	node := &memoryBackendNode{Value: value, ExpiredAt: time.Now().Add(ttl)}
 	m.client[key] = node
 	return nil
@@ -70,16 +74,20 @@ func (m *memoryBackend) GetMany(ctx context.Context, keys []string) [][]byte {
 }
 
 func (m *memoryBackend) Delete(ctx context.Context, key string) bool {
-	exists := m.Exists(ctx, key)
-	delete(m.client, key)
-	return exists
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	if _, ok := m.client[key]; ok {
+		delete(m.client, key)
+		return true
+	}
+	return false
 }
 
 func (m *memoryBackend) DeleteMany(ctx context.Context, keys []string) bool {
-	var res bool
+	res := true
 	for _, key := range keys {
-		if m.Delete(ctx, key) {
-			res = true
+		if !m.Delete(ctx, key) {
+			res = false
 		}
 	}
 	return res
@@ -90,6 +98,8 @@ func (m *memoryBackend) Expire(ctx context.Context, key string, ttl time.Duratio
 	if val == nil {
 		return false
 	}
+	m.lock.Lock()
+	defer m.lock.Unlock()
 	m.client[key].ExpiredAt = time.Now().Add(ttl)
 	return true
 }
diff --git a/extensions/cachext/backend/redis/redis.go b/extensions/cachext/backend/redis/redis.go
index 16eccff7..b65a9cb2 100644
--- a/extensions/cachext/backend/redis/redis.go
+++ b/extensions/cachext/backend/redis/redis.go
@@ -12,7 +12,7 @@ import (
 )
 
 func init() {
-	if err := cachext.RegisteBackend("redis", func() cachext.CacheBackend { return &redisBackend{} }); err != nil {
+	if err := cachext.RegisterBackend("redis", func() cachext.CacheBackend { return &redisBackend{} }); err != nil {
 		panic("RedisBackend init error")
 	}
 }
diff --git a/extensions/cachext/ext.go b/extensions/cachext/ext.go
index aa155b49..3f8d78eb 100644
--- a/extensions/cachext/ext.go
+++ b/extensions/cachext/ext.go
@@ -8,11 +8,11 @@ import (
 	"sync"
 	"time"
 
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/shanbay/gobay"
 	"github.com/spf13/viper"
 	"github.com/vmihailenco/msgpack"
-	"github.com/prometheus/client_golang/prometheus"
-	"github.com/prometheus/client_golang/prometheus/promauto"
 )
 
 type void struct{}
@@ -119,16 +119,16 @@ func (c *CacheExt) CheckHealth(ctx context.Context) error {
 	return nil
 }
 
-// RegisteBackend if you want a new backend, use this func to registe your backend
+// RegisterBackend if you want a new backend, use this func to registe your backend
 // then load it by config
-func RegisteBackend(configBackend string, backendFunc func() CacheBackend) error {
+func RegisterBackend(configBackend string, backendFunc func() CacheBackend) error {
 	mu.Lock()
 	defer mu.Unlock()
 	if _, exist := backendMap[configBackend]; !exist {
 		backendMap[configBackend] = backendFunc
 		return nil
 	} else {
-		return errors.New("Backend already registered")
+		return errors.New("backend already registered")
 	}
 }