From 567fb070e3fbd7a0eb5a936cb52addc5de9298bf Mon Sep 17 00:00:00 2001
From: Yeuoly <admin@srmxy.cn>
Date: Thu, 29 Aug 2024 00:38:13 +0800
Subject: [PATCH] refactor: storage declaration

---
 .../backwards_invocation/task.go              |  6 ++++
 .../remote_manager/server_test.go             |  1 -
 .../plugin_entities/plugin_declaration.go     | 12 +++++--
 .../plugin_declaration_test.go                | 32 +++++++++++++++++--
 4 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/internal/core/plugin_daemon/backwards_invocation/task.go b/internal/core/plugin_daemon/backwards_invocation/task.go
index 1074e9c..40edab9 100644
--- a/internal/core/plugin_daemon/backwards_invocation/task.go
+++ b/internal/core/plugin_daemon/backwards_invocation/task.go
@@ -116,6 +116,12 @@ var (
 			},
 			"error": "permission denied, you need to enable app access in plugin manifest",
 		},
+		dify_invocation.INVOKE_TYPE_STORAGE: {
+			"func": func(declaration *plugin_entities.PluginDeclaration) bool {
+				return declaration.Resource.Permission.AllowInvokeStorage()
+			},
+			"error": "permission denied, you need to enable storage access in plugin manifest",
+		},
 	}
 )
 
diff --git a/internal/core/plugin_manager/remote_manager/server_test.go b/internal/core/plugin_manager/remote_manager/server_test.go
index 77b1d67..c8482cc 100644
--- a/internal/core/plugin_manager/remote_manager/server_test.go
+++ b/internal/core/plugin_manager/remote_manager/server_test.go
@@ -142,7 +142,6 @@ func TestAcceptConnection(t *testing.T) {
 		CreatedAt: time.Now(),
 		Resource: plugin_entities.PluginResourceRequirement{
 			Memory:     1,
-			Storage:    1,
 			Permission: nil,
 		},
 		Plugins: []string{
diff --git a/internal/types/entities/plugin_entities/plugin_declaration.go b/internal/types/entities/plugin_entities/plugin_declaration.go
index 70f7476..99ff2c0 100644
--- a/internal/types/entities/plugin_entities/plugin_declaration.go
+++ b/internal/types/entities/plugin_entities/plugin_declaration.go
@@ -22,6 +22,7 @@ type PluginPermissionRequirement struct {
 	Node     *PluginPermissionNodeRequirement     `json:"node" yaml:"node" validate:"omitempty"`
 	Endpoint *PluginPermissionEndpointRequirement `json:"endpoint" yaml:"endpoint" validate:"omitempty"`
 	App      *PluginPermissionAppRequirement      `json:"app" yaml:"app" validate:"omitempty"`
+	Storage  *PluginPermissionStorageRequirement  `json:"storage" yaml:"storage" validate:"omitempty"`
 }
 
 func (p *PluginPermissionRequirement) AllowInvokeTool() bool {
@@ -64,6 +65,10 @@ func (p *PluginPermissionRequirement) AllowRegistryEndpoint() bool {
 	return p != nil && p.Endpoint != nil && p.Endpoint.Enabled
 }
 
+func (p *PluginPermissionRequirement) AllowInvokeStorage() bool {
+	return p != nil && p.Storage != nil && p.Storage.Enabled
+}
+
 type PluginPermissionToolRequirement struct {
 	Enabled bool `json:"enabled" yaml:"enabled"`
 }
@@ -90,11 +95,14 @@ type PluginPermissionAppRequirement struct {
 	Enabled bool `json:"enabled" yaml:"enabled"`
 }
 
+type PluginPermissionStorageRequirement struct {
+	Enabled bool   `json:"enabled" yaml:"enabled"`
+	Size    uint64 `json:"size" yaml:"size" validate:"min=1024,max=1073741824"` // min 1024 bytes, max 1G
+}
+
 type PluginResourceRequirement struct {
 	// Memory in bytes
 	Memory int64 `json:"memory" yaml:"memory" validate:"required"`
-	// Storage in bytes
-	Storage int64 `json:"storage" yaml:"storage" validate:"required"`
 	// Permission requirements
 	Permission *PluginPermissionRequirement `json:"permission" yaml:"permission" validate:"omitempty"`
 }
diff --git a/internal/types/entities/plugin_entities/plugin_declaration_test.go b/internal/types/entities/plugin_entities/plugin_declaration_test.go
index db4e81e..827e3cd 100644
--- a/internal/types/entities/plugin_entities/plugin_declaration_test.go
+++ b/internal/types/entities/plugin_entities/plugin_declaration_test.go
@@ -16,8 +16,7 @@ func preparePluginDeclaration() PluginDeclaration {
 		Author:    "test",
 		CreatedAt: time.Now(),
 		Resource: PluginResourceRequirement{
-			Memory:  1,
-			Storage: 1,
+			Memory: 1,
 			Permission: &PluginPermissionRequirement{
 				Tool: &PluginPermissionToolRequirement{
 					Enabled: true,
@@ -28,6 +27,10 @@ func preparePluginDeclaration() PluginDeclaration {
 				Node: &PluginPermissionNodeRequirement{
 					Enabled: true,
 				},
+				Storage: &PluginPermissionStorageRequirement{
+					Enabled: true,
+					Size:    1024,
+				},
 			},
 		},
 		Plugins: []string{},
@@ -96,6 +99,10 @@ func TestPluginDeclarationFullTest(t *testing.T) {
 		return
 	}
 
+	if new_declaration.Resource.Permission.Storage == nil {
+		t.Errorf("storage permission is nil")
+		return
+	}
 }
 
 func TestPluginDeclarationIncorrectVersion(t *testing.T) {
@@ -134,6 +141,27 @@ func TestPluginUnsupportedArch(t *testing.T) {
 	}
 }
 
+func TestPluginStorageSizeTooSmallOrTooLarge(t *testing.T) {
+	declaration := preparePluginDeclaration()
+	declaration.Resource.Permission.Storage.Size = 1023
+	declaration_bytes := parser.MarshalJsonBytes(declaration)
+
+	_, err := parser.UnmarshalJsonBytes[PluginDeclaration](declaration_bytes)
+	if err == nil {
+		t.Errorf("failed to validate storage size")
+		return
+	}
+
+	declaration.Resource.Permission.Storage.Size = 1073741825
+	declaration_bytes = parser.MarshalJsonBytes(declaration)
+
+	_, err = parser.UnmarshalJsonBytes[PluginDeclaration](declaration_bytes)
+	if err == nil {
+		t.Errorf("failed to validate storage size")
+		return
+	}
+}
+
 func TestPluginDeclarationIncorrectType(t *testing.T) {
 	declaration := preparePluginDeclaration()
 	declaration.Type = "test"