From d323c0c150126d07d9f15f1deb176b6d369922f0 Mon Sep 17 00:00:00 2001
From: Dayuan <daniellee0929@gmail.com>
Date: Tue, 5 Mar 2024 14:26:14 +0800
Subject: [PATCH] feat: add kusion module plugin (#870)

---
 pkg/modules/plugin.go | 133 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 133 insertions(+)
 create mode 100644 pkg/modules/plugin.go

diff --git a/pkg/modules/plugin.go b/pkg/modules/plugin.go
new file mode 100644
index 00000000..3b0d7b39
--- /dev/null
+++ b/pkg/modules/plugin.go
@@ -0,0 +1,133 @@
+package modules
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path"
+	"runtime"
+	"strings"
+	"sync"
+
+	"github.com/hashicorp/go-plugin"
+
+	"kusionstack.io/kusion/pkg/util/kfile"
+)
+
+const (
+	DefaultModulePathEnv     = "KUSION_MODULE_PATH"
+	KusionModuleBinaryPrefix = "kusion-module-"
+	Dir                      = "modules"
+)
+
+var mu sync.Mutex
+
+// PluginMap is the map of plugins we can dispense.
+var PluginMap = map[string]plugin.Plugin{
+	PluginKey: &GRPCPlugin{},
+}
+
+type Plugin struct {
+	// key represents the module key, it consists of two parts: moduleName@version. e.g. "kusionstack/mysql@v0.1"
+	key    string
+	client *plugin.Client
+	// Module represents the real module impl
+	Module Module
+}
+
+func NewPlugin(key string) (*Plugin, error) {
+	if key == "" {
+		return nil, fmt.Errorf("module key can not be empty")
+	}
+	p := &Plugin{key: key}
+	err := p.initModule()
+	if err != nil {
+		return nil, err
+	}
+	return p, nil
+}
+
+func (p *Plugin) initModule() error {
+	key := p.key
+	split := strings.Split(key, "@")
+	msg := "invalid module key: %s. The correct format for a key should be as follows: namespace/resourceType@version. e.g. kusionstack/mysql@v0.1"
+	if len(split) != 2 {
+		return fmt.Errorf(msg, key)
+	}
+	prefix := strings.Split(split[0], "/")
+	if len(prefix) != 2 {
+		return fmt.Errorf(msg, key)
+	}
+
+	// build the plugin client
+	pluginPath, err := buildPluginPath(prefix[0], prefix[1], split[1])
+	if err != nil {
+		return err
+	}
+	client := newPluginClient(pluginPath)
+	p.client = client
+	rpcClient, err := client.Client()
+	if err != nil {
+		return err
+	}
+
+	// dispense the plugin to get the real module
+	raw, err := rpcClient.Dispense(PluginKey)
+	if err != nil {
+		return err
+	}
+	p.Module = raw.(Module)
+
+	return nil
+}
+
+func buildPluginPath(namespace, resourceType, version string) (string, error) {
+	mu.Lock()
+	defer mu.Unlock()
+
+	// validate the module path
+	prefixPath, err := PluginDir()
+	if err != nil {
+		return "", err
+	}
+	goOs := runtime.GOOS
+	goArch := runtime.GOARCH
+	p := path.Join(prefixPath, namespace, resourceType, version, goOs, goArch, KusionModuleBinaryPrefix+resourceType)
+	_, err = os.Stat(p)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return "", fmt.Errorf("module dir doesn't exist. %s", p)
+		} else {
+			return "", err
+		}
+	}
+	return p, nil
+}
+
+func newPluginClient(path string) *plugin.Client {
+	// We're a host! Start by launching the plugin process.
+	// need to defer kill
+	client := plugin.NewClient(&plugin.ClientConfig{
+		HandshakeConfig: HandshakeConfig,
+		Plugins:         PluginMap,
+		Cmd:             exec.Command(path),
+		AllowedProtocols: []plugin.Protocol{
+			plugin.ProtocolGRPC,
+		},
+	})
+	return client
+}
+
+func (p *Plugin) KillPluginClient() {
+	p.client.Kill()
+}
+
+func PluginDir() (string, error) {
+	if env, found := os.LookupEnv(DefaultModulePathEnv); found {
+		return env, nil
+	} else if dir, err := kfile.KusionDataFolder(); err == nil {
+		return path.Join(dir, Dir), nil
+	} else {
+		return "", err
+	}
+}