Преглед изворни кода

improve: generalize plugin packager

Yeuoly пре 1 година
родитељ
комит
faee533c9c

+ 7 - 1
cmd/packager/main.go

@@ -4,6 +4,7 @@ import (
 	"flag"
 	"os"
 
+	"github.com/langgenius/dify-plugin-daemon/internal/core/plugin_packager/decoder"
 	"github.com/langgenius/dify-plugin-daemon/internal/core/plugin_packager/packager"
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/log"
 )
@@ -25,7 +26,12 @@ func main() {
 		os.Exit(0)
 	}
 
-	packager := packager.NewPackager(in_path)
+	decoder, err := decoder.NewFSPluginDecoder(in_path)
+	if err != nil {
+		log.Panic("failed to create plugin decoder , plugin path: %s, error: %v", in_path, err)
+	}
+
+	packager := packager.NewPackager(decoder)
 	zip_file, err := packager.Pack()
 
 	if err != nil {

+ 8 - 0
internal/core/plugin_packager/decoder/decoder.go

@@ -0,0 +1,8 @@
+package decoder
+
+type PluginDecoder interface {
+	Open() error
+	Walk(fn func(filename string, dir string) error) error
+	ReadFile(filename string) ([]byte, error)
+	Close() error
+}

+ 72 - 0
internal/core/plugin_packager/decoder/fs.go

@@ -0,0 +1,72 @@
+package decoder
+
+import (
+	"errors"
+	"io/fs"
+	"os"
+	"path"
+)
+
+var (
+	ErrNotDir = errors.New("not a directory")
+)
+
+type FSPluginDecoder struct {
+	PluginDecoder
+
+	// root directory of the plugin
+	root string
+
+	fs fs.FS
+}
+
+func NewFSPluginDecoder(root string) (*FSPluginDecoder, error) {
+	decoder := &FSPluginDecoder{
+		root: root,
+	}
+
+	err := decoder.Open()
+	if err != nil {
+		return nil, err
+	}
+
+	return decoder, nil
+}
+
+func (d *FSPluginDecoder) Open() error {
+	d.fs = os.DirFS(d.root)
+
+	// try to stat the root directory
+	s, err := os.Stat(d.root)
+	if err != nil {
+		return err
+	}
+
+	if !s.IsDir() {
+		return ErrNotDir
+	}
+
+	return nil
+}
+
+func (d *FSPluginDecoder) Walk(fn func(filename string, dir string) error) error {
+	return fs.WalkDir(d.fs, ".", func(root_path string, d fs.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+
+		if d.IsDir() {
+			return nil
+		}
+
+		return fn(root_path, d.Name())
+	})
+}
+
+func (d *FSPluginDecoder) Close() error {
+	return nil
+}
+
+func (d *FSPluginDecoder) ReadFile(filename string) ([]byte, error) {
+	return os.ReadFile(path.Join(d.root, filename))
+}

+ 78 - 0
internal/core/plugin_packager/decoder/zip.go

@@ -0,0 +1,78 @@
+package decoder
+
+import (
+	"archive/zip"
+	"bytes"
+	"path"
+)
+
+type ZipPluginDecoder struct {
+	PluginDecoder
+
+	reader *zip.Reader
+	err    error
+}
+
+func NewZipPluginDecoder(binary []byte) (*ZipPluginDecoder, error) {
+	reader, err := zip.NewReader(bytes.NewReader(binary), int64(len(binary)))
+
+	decoder := &ZipPluginDecoder{
+		reader: reader,
+		err:    err,
+	}
+
+	err = decoder.Open()
+	if err != nil {
+		return nil, err
+	}
+
+	return decoder, nil
+}
+
+func (z *ZipPluginDecoder) Open() error {
+	if z.reader == nil {
+		return z.err
+	}
+
+	return nil
+}
+
+func (z *ZipPluginDecoder) Walk(fn func(filename string, dir string) error) error {
+	if z.reader == nil {
+		return z.err
+	}
+
+	for _, file := range z.reader.File {
+		// split the path into directory and filename
+		dir, filename := path.Split(file.Name)
+		if err := fn(filename, dir); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (z *ZipPluginDecoder) Close() error {
+	return nil
+}
+
+func (z *ZipPluginDecoder) ReadFile(filename string) ([]byte, error) {
+	if z.reader == nil {
+		return nil, z.err
+	}
+
+	file, err := z.reader.Open(filename)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	data := new(bytes.Buffer)
+	_, err = data.ReadFrom(file)
+	if err != nil {
+		return nil, err
+	}
+
+	return data.Bytes(), nil
+}

+ 1 - 4
internal/core/plugin_packager/packager/file_scanner.go

@@ -1,7 +1,6 @@
 package packager
 
 import (
-	"os"
 	"path"
 
 	"github.com/langgenius/dify-plugin-daemon/internal/types/entities/plugin_entities"
@@ -12,9 +11,7 @@ func (p *Packager) ScanProvider() error {
 }
 
 func (p *Packager) fetchManifest() (*plugin_entities.PluginDeclaration, error) {
-
-	file_path := path.Join(p.wp, p.manifest)
-	file, err := os.ReadFile(file_path)
+	file, err := p.decoder.ReadFile(path.Clean(p.manifest))
 	if err != nil {
 		return nil, err
 	}

+ 9 - 19
internal/core/plugin_packager/packager/packager.go

@@ -3,48 +3,38 @@ package packager
 import (
 	"archive/zip"
 	"bytes"
-	"io/fs"
-	"os"
-	"path"
+
+	"github.com/langgenius/dify-plugin-daemon/internal/core/plugin_packager/decoder"
 )
 
 type Packager struct {
-	wp string // working path
-
+	decoder  decoder.PluginDecoder
 	manifest string // manifest file path
 }
 
-func NewPackager(plugin_path string) *Packager {
+func NewPackager(decoder decoder.PluginDecoder) *Packager {
 	return &Packager{
-		wp:       plugin_path,
+		decoder:  decoder,
 		manifest: "manifest.yaml",
 	}
 }
 
 func (p *Packager) Pack() ([]byte, error) {
-	// read manifest
-	_, err := p.fetchManifest()
+	err := p.Validate()
 	if err != nil {
 		return nil, err
 	}
 
 	zip_buffer := new(bytes.Buffer)
 	zip_writer := zip.NewWriter(zip_buffer)
-	err = fs.WalkDir(os.DirFS(p.wp), ".", func(root_path string, d fs.DirEntry, err error) error {
-		if err != nil {
-			return err
-		}
-
-		if d.IsDir() {
-			return nil
-		}
 
-		file, err := os.ReadFile(path.Join(p.wp, root_path))
+	p.decoder.Walk(func(filename, dir string) error {
+		file, err := p.decoder.ReadFile(filename)
 		if err != nil {
 			return err
 		}
 
-		zip_file, err := zip_writer.Create(root_path)
+		zip_file, err := zip_writer.Create(filename)
 		if err != nil {
 			return err
 		}

+ 11 - 0
internal/core/plugin_packager/packager/validator.go

@@ -0,0 +1,11 @@
+package packager
+
+func (p *Packager) Validate() error {
+	// read manifest
+	_, err := p.fetchManifest()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

+ 30 - 34
internal/core/plugin_packager/signer/sign.go

@@ -6,10 +6,12 @@ import (
 	"crypto/sha256"
 	"encoding/base64"
 	"io"
+	"path"
 	"strconv"
 	"time"
 
 	"github.com/langgenius/dify-plugin-daemon/internal/core/license/private_key"
+	"github.com/langgenius/dify-plugin-daemon/internal/core/plugin_packager/decoder"
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/encryption"
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/parser"
 )
@@ -29,35 +31,49 @@ func SignPlugin(plugin []byte) ([]byte, error) {
 		return nil, err
 	}
 
-	// construct zip
-	zip_reader, err := zip.NewReader(bytes.NewReader(plugin), int64(len(plugin)))
+	decoder, err := decoder.NewZipPluginDecoder(plugin)
 	if err != nil {
 		return nil, err
 	}
 
+	// create a new zip writer
+	zip_buffer := new(bytes.Buffer)
+	zip_writer := zip.NewWriter(zip_buffer)
+
+	defer zip_writer.Close()
+	// store temporary hash
 	data := new(bytes.Buffer)
 	// read one by one
-	for _, file := range zip_reader.File {
-		// read file bytes
-		file_reader, err := file.Open()
-		if err != nil {
-			return nil, err
-		}
-		defer file_reader.Close()
-
-		temp_data := new(bytes.Buffer)
-		_, err = temp_data.ReadFrom(file_reader)
+	err = decoder.Walk(func(filename, dir string) error {
+		file, err := decoder.ReadFile(path.Join(dir, filename))
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		// calculate sha256 hash of the file
 		hash := sha256.New()
-		hash.Write(temp_data.Bytes())
+		hash.Write(file)
 		hashed := hash.Sum(nil)
 
 		// write the hash into data
 		data.Write(hashed)
+
+		// create a new file in the zip writer
+		file_writer, err := zip_writer.Create(path.Join(dir, filename))
+		if err != nil {
+			return err
+		}
+
+		_, err = io.Copy(file_writer, bytes.NewReader(file))
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		return nil, err
 	}
 
 	// get current time
@@ -76,26 +92,6 @@ func SignPlugin(plugin []byte) ([]byte, error) {
 	}
 
 	// write the signature into the comment field of the zip file
-	zip_buffer := new(bytes.Buffer)
-	zip_writer := zip.NewWriter(zip_buffer)
-	for _, file := range zip_reader.File {
-		file_writer, err := zip_writer.Create(file.Name)
-		if err != nil {
-			return nil, err
-		}
-
-		file_reader, err := file.Open()
-		if err != nil {
-			return nil, err
-		}
-		defer file_reader.Close()
-
-		_, err = io.Copy(file_writer, file_reader)
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	comments := parser.MarshalJson(map[string]any{
 		"signature": base64.StdEncoding.EncodeToString(signature),
 		"time":      ct,