From f19f4eb430828b9be6bced9e6c4b3078e3b47bb9 Mon Sep 17 00:00:00 2001
From: Shelikhoo <xiaokangwang@outlook.com>
Date: Sun, 6 Feb 2022 20:25:45 +0000
Subject: [PATCH] implement scoped transient storage

---
 .../transientstorageimpl/storage.go           | 78 +++++++++++++++++++
 features/extension/storage/storage.go         |  2 +-
 2 files changed, 79 insertions(+), 1 deletion(-)
 create mode 100644 common/environment/transientstorageimpl/storage.go

diff --git a/common/environment/transientstorageimpl/storage.go b/common/environment/transientstorageimpl/storage.go
new file mode 100644
index 000000000..e97c0c56b
--- /dev/null
+++ b/common/environment/transientstorageimpl/storage.go
@@ -0,0 +1,78 @@
+package transientstorageimpl
+
+//go:generate go run github.com/v2fly/v2ray-core/v5/common/errors/errorgen
+
+import (
+	"context"
+	"github.com/v2fly/v2ray-core/v5/features/extension/storage"
+	"strings"
+	"sync"
+)
+
+func NewScopedTransientStorageImpl() storage.ScopedTransientStorage {
+	return &scopedTransientStorageImpl{scopes: map[string]storage.ScopedTransientStorage{}, values: map[string]interface{}{}}
+}
+
+type scopedTransientStorageImpl struct {
+	access sync.Mutex
+	scopes map[string]storage.ScopedTransientStorage
+	values map[string]interface{}
+}
+
+func (s *scopedTransientStorageImpl) ScopedTransientStorage() {
+	panic("implement me")
+}
+
+func (s *scopedTransientStorageImpl) Put(ctx context.Context, key string, value interface{}) error {
+	s.access.Lock()
+	defer s.access.Unlock()
+	s.values[key] = value
+	return nil
+}
+
+func (s *scopedTransientStorageImpl) Get(ctx context.Context, key string) (interface{}, error) {
+	s.access.Lock()
+	defer s.access.Unlock()
+	sw, ok := s.values[key]
+	if !ok {
+		return nil, newError("unable to find ")
+	}
+	return sw, nil
+}
+
+func (s *scopedTransientStorageImpl) List(ctx context.Context, keyPrefix string) ([]string, error) {
+	s.access.Lock()
+	defer s.access.Unlock()
+	var ret []string
+	for key, _ := range s.values {
+		if strings.HasPrefix(key, keyPrefix) {
+			ret = append(ret, key)
+		}
+	}
+	return ret, nil
+}
+
+func (s *scopedTransientStorageImpl) Clear(ctx context.Context) {
+	s.access.Lock()
+	defer s.access.Unlock()
+	s.values = map[string]interface{}{}
+}
+
+func (s *scopedTransientStorageImpl) NarrowScope(ctx context.Context, key string) (storage.ScopedTransientStorage, error) {
+	s.access.Lock()
+	defer s.access.Unlock()
+	sw, ok := s.scopes[key]
+	if !ok {
+		scope := NewScopedTransientStorageImpl()
+		s.scopes[key] = scope
+		return scope, nil
+	}
+	return sw, nil
+}
+
+func (s *scopedTransientStorageImpl) DropScope(ctx context.Context, key string) error {
+	s.access.Lock()
+	defer s.access.Unlock()
+	delete(s.scopes, key)
+	return nil
+}
diff --git a/features/extension/storage/storage.go b/features/extension/storage/storage.go
index 4ed58671a..bcb407c69 100644
--- a/features/extension/storage/storage.go
+++ b/features/extension/storage/storage.go
@@ -20,6 +20,6 @@ type ScopedTransientStorage interface {
 	Get(ctx context.Context, key string) (interface{}, error)
 	List(ctx context.Context, keyPrefix string) ([]string, error)
 	Clear(ctx context.Context)
-	NarrowScope(ctx context.Context, key string) (ScopedPersistentStorage, error)
+	NarrowScope(ctx context.Context, key string) (ScopedTransientStorage, error)
 	DropScope(ctx context.Context, key string) error
 }