Browse Source

add backwards invocation permission check

Yeuoly 1 year ago
parent
commit
c2eed76a91

+ 78 - 0
internal/core/plugin_daemon/backwards_invocation/task.go

@@ -39,6 +39,13 @@ func InvokeDify(
 		return nil
 	}
 
+	// check permission
+	if err := checkPermission(runtime, request_handle); err != nil {
+		request_handle.WriteError(err)
+		request_handle.EndResponse()
+		return nil
+	}
+
 	// dispatch invocation task
 	routine.Submit(func() {
 		dispatchDifyInvocationTask(request_handle)
@@ -48,6 +55,77 @@ func InvokeDify(
 	return nil
 }
 
+var (
+	permissionMapping = map[dify_invocation.InvokeType]map[string]any{
+		dify_invocation.INVOKE_TYPE_TOOL: {
+			"func": func(runtime entities.PluginRuntimeInterface) bool {
+				return runtime.Configuration().Resource.Permission.AllowInvokeTool()
+			},
+			"error": "permission denied, you need to enable tool access in plugin manifest",
+		},
+		dify_invocation.INVOKE_TYPE_LLM: {
+			"func": func(runtime entities.PluginRuntimeInterface) bool {
+				return runtime.Configuration().Resource.Permission.AllowInvokeLLM()
+			},
+			"error": "permission denied, you need to enable llm access in plugin manifest",
+		},
+		dify_invocation.INVOKE_TYPE_TEXT_EMBEDDING: {
+			"func": func(runtime entities.PluginRuntimeInterface) bool {
+				return runtime.Configuration().Resource.Permission.AllowInvokeTextEmbedding()
+			},
+			"error": "permission denied, you need to enable text-embedding access in plugin manifest",
+		},
+		dify_invocation.INVOKE_TYPE_RERANK: {
+			"func": func(runtime entities.PluginRuntimeInterface) bool {
+				return runtime.Configuration().Resource.Permission.AllowInvokeRerank()
+			},
+			"error": "permission denied, you need to enable rerank access in plugin manifest",
+		},
+		dify_invocation.INVOKE_TYPE_TTS: {
+			"func": func(runtime entities.PluginRuntimeInterface) bool {
+				return runtime.Configuration().Resource.Permission.AllowInvokeTTS()
+			},
+			"error": "permission denied, you need to enable tts access in plugin manifest",
+		},
+		dify_invocation.INVOKE_TYPE_SPEECH2TEXT: {
+			"func": func(runtime entities.PluginRuntimeInterface) bool {
+				return runtime.Configuration().Resource.Permission.AllowInvokeSpeech2Text()
+			},
+			"error": "permission denied, you need to enable speech2text access in plugin manifest",
+		},
+		dify_invocation.INVOKE_TYPE_MODERATION: {
+			"func": func(runtime entities.PluginRuntimeInterface) bool {
+				return runtime.Configuration().Resource.Permission.AllowInvokeModeration()
+			},
+			"error": "permission denied, you need to enable moderation access in plugin manifest",
+		},
+		dify_invocation.INVOKE_TYPE_NODE: {
+			"func": func(runtime entities.PluginRuntimeInterface) bool {
+				return runtime.Configuration().Resource.Permission.AllowInvokeNode()
+			},
+			"error": "permission denied, you need to enable node access in plugin manifest",
+		},
+	}
+)
+
+func checkPermission(runtime entities.PluginRuntimeInterface, request_handle *BackwardsInvocation) error {
+	permission, ok := permissionMapping[request_handle.Type()]
+	if !ok {
+		return fmt.Errorf("unsupported invoke type: %s", request_handle.Type())
+	}
+
+	permission_func, ok := permission["func"].(func(runtime entities.PluginRuntimeInterface) bool)
+	if !ok {
+		return fmt.Errorf("permission function not found: %s", request_handle.Type())
+	}
+
+	if !permission_func(runtime) {
+		return fmt.Errorf(permission["error"].(string))
+	}
+
+	return nil
+}
+
 func prepareDifyInvocationArguments(session *session_manager.Session, request map[string]any) (*BackwardsInvocation, error) {
 	typ, ok := request["type"].(string)
 	if !ok {

+ 33 - 0
internal/types/entities/plugin_entities/plugin_declaration.go

@@ -21,6 +21,38 @@ type PluginPermissionRequirement struct {
 	Node  *PluginPermissionNodeRequirement  `json:"node" yaml:"node" validate:"omitempty"`
 }
 
+func (p *PluginPermissionRequirement) AllowInvokeTool() bool {
+	return p != nil && p.Tool != nil && p.Tool.Enabled
+}
+
+func (p *PluginPermissionRequirement) AllowInvokeLLM() bool {
+	return p != nil && p.Model != nil && p.Model.Enabled && p.Model.LLM
+}
+
+func (p *PluginPermissionRequirement) AllowInvokeTextEmbedding() bool {
+	return p != nil && p.Model != nil && p.Model.Enabled && p.Model.TextEmbedding
+}
+
+func (p *PluginPermissionRequirement) AllowInvokeRerank() bool {
+	return p != nil && p.Model != nil && p.Model.Enabled && p.Model.Rerank
+}
+
+func (p *PluginPermissionRequirement) AllowInvokeTTS() bool {
+	return p != nil && p.Model != nil && p.Model.Enabled && p.Model.TTS
+}
+
+func (p *PluginPermissionRequirement) AllowInvokeSpeech2Text() bool {
+	return p != nil && p.Model != nil && p.Model.Enabled && p.Model.Speech2text
+}
+
+func (p *PluginPermissionRequirement) AllowInvokeModeration() bool {
+	return p != nil && p.Model != nil && p.Model.Enabled && p.Model.Moderation
+}
+
+func (p *PluginPermissionRequirement) AllowInvokeNode() bool {
+	return p != nil && p.Node != nil && p.Node.Enabled
+}
+
 type PluginPermissionToolRequirement struct {
 	Enabled bool `json:"enabled" yaml:"enabled"`
 }
@@ -32,6 +64,7 @@ type PluginPermissionModelRequirement struct {
 	Rerank        bool `json:"rerank" yaml:"rerank"`
 	TTS           bool `json:"tts" yaml:"tts"`
 	Speech2text   bool `json:"speech2text" yaml:"speech2text"`
+	Moderation    bool `json:"moderation" yaml:"moderation"`
 }
 
 type PluginPermissionNodeRequirement struct {