Refactor the method of packing function code into zip

shimmeris · Jan 3, 2023 · 2a79e10 · 2a79e10
1 parent a1a9837
commit 2a79e10
Show file tree

Hide file tree

Showing 6 changed files with 69 additions and 18 deletions.
diff --git a/function/code.go b/function/code.go
@@ -1,22 +1,79 @@
 package function
 
-import _ "embed"
+import (
+	"archive/zip"
+	"bytes"
+	_ "embed"
+	"encoding/base64"
 
-// compile socks code with `GOOS=linux GOARCH=amd64 go build main.go`
+	"github.com/sirupsen/logrus"
+)
+
+type File struct {
+	Name string
+	Content []byte
+	HighPriv bool
+}
 
+// compile socks code with `GOOS=linux GOARCH=amd64 go build main.go`
 var (
 	//go:embed http/tencent.py
-	TencentHttpCode []byte
+	tencentHttpCode []byte
+	TencentHttpCodeZip = CreateZipBase64([]File{{Name: "index.py", Content: tencentHttpCode}})
 
 	//go:embed http/alibaba.py
-	AlibabaHttpCode []byte
+	alibabaHttpCode []byte
+	AlibabaHttpCodeZip = CreateZipBase64([]File{{Name: "index.py", Content: alibabaHttpCode}})
 
 	//go:embed http/huawei.py
-	HuaweiHttpCode []byte
+	huaweiHttpCode []byte
+	HuaweiHttpCodeZip =CreateZipBase64([]File{{Name: "index.py", Content: huaweiHttpCode}})
 
 	//go:embed socks/tencent
-	TencentSocksCode []byte
+	tencentSocksCode []byte
+	TencentSocksCodeZip = CreateZipBase64([]File{{Name: "index.py", Content: tencentSocksCode, HighPriv: true}})
 
 	//go:embed socks/alibaba
-	AlibabaSocksCode []byte
+	alibabaSocksCode []byte
+	AlibabaSocksCodeZip = CreateZipBase64([]File{{Name: "index.py", Content: alibabaSocksCode, HighPriv: true}})
 )
+
+
+
+func CreateZipBase64(files []File) string {
+	buf := new(bytes.Buffer)
+
+	zw := zip.NewWriter(buf)
+
+	for _, f := range files{
+		if f.HighPriv {
+			fw, err := zw.CreateHeader(&zip.FileHeader{
+				CreatorVersion: 3 << 8,     // indicates Unix
+				ExternalAttrs:  0777 << 16, // -rwxrwxrwx file permissions
+				Name:           f.Name,
+				Method:         zip.Deflate,
+			})
+			if err != nil {
+				logrus.Error(err)
+			}
+
+			_, err = fw.Write(f.Content)
+			if err != nil {
+				logrus.Error(err)
+			}
+		} else {
+			fw, err := zw.Create(f.Name)
+			if err != nil {
+				logrus.Error(err)
+			}
+			_, err = fw.Write(f.Content)
+			if err != nil {
+				logrus.Error(err)
+			}
+		}
+	}
+
+	zw.Close()
+	return base64.StdEncoding.EncodeToString(buf.Bytes())
+
+}
diff --git a/sdk/provider/alibaba/http.go b/sdk/provider/alibaba/http.go
@@ -5,7 +5,6 @@ import (
 	"github.com/alibabacloud-go/tea/tea"
 	"github.com/sirupsen/logrus"
 
-	"github.com/shimmeris/SCFProxy/fileutil"
 	"github.com/shimmeris/SCFProxy/function"
 	"github.com/shimmeris/SCFProxy/sdk"
 )
@@ -57,7 +56,7 @@ func (p *Provider) createHttpFunction(serviceName, functionName string) error {
 		Timeout:      tea.Int32(30),
 		MemorySize:   tea.Int32(128),
 		Code: &fcopen.Code{
-			ZipFile: tea.String(fileutil.CreateZipBase64("index.py", function.AlibabaHttpCode)),
+			ZipFile: tea.String(function.AlibabaHttpCodeZip),
 		},
 	}
 

diff --git a/sdk/provider/alibaba/socks.go b/sdk/provider/alibaba/socks.go
@@ -6,7 +6,6 @@ import (
 	fcopen "github.com/alibabacloud-go/fc-open-20210406/client"
 	"github.com/alibabacloud-go/tea/tea"
 
-	"github.com/shimmeris/SCFProxy/fileutil"
 	"github.com/shimmeris/SCFProxy/function"
 	"github.com/shimmeris/SCFProxy/sdk"
 )
@@ -41,7 +40,7 @@ func (p *Provider) createSocksFunction(serviceName, functionName string) error {
 		Timeout:      tea.Int32(900),
 		MemorySize:   tea.Int32(128),
 		Code: &fcopen.Code{
-			ZipFile: tea.String(fileutil.CreateZipBase64("main", function.AlibabaSocksCode)),
+			ZipFile: tea.String(function.AlibabaSocksCodeZip),
 		},
 	}
 

diff --git a/sdk/provider/huawei/http.go b/sdk/provider/huawei/http.go
@@ -8,7 +8,6 @@ import (
 
 	"github.com/huaweicloud/huaweicloud-sdk-go-v3/services/functiongraph/v2/model"
 
-	"github.com/shimmeris/SCFProxy/fileutil"
 	"github.com/shimmeris/SCFProxy/function"
 	"github.com/shimmeris/SCFProxy/sdk"
 )
@@ -61,7 +60,6 @@ func (p *Provider) createGroup(groupName string) error {
 
 func (p *Provider) createFunction(functionName string) (string, error) {
 	r := &model.CreateFunctionRequest{}
-	code := fileutil.CreateZipBase64("index.py", function.HuaweiHttpCode)
 	r.Body = &model.CreateFunctionRequestBody{
 		Package:    "default",
 		FuncName:   functionName,
@@ -71,7 +69,7 @@ func (p *Provider) createFunction(functionName string) (string, error) {
 		CodeType:   model.GetCreateFunctionRequestBodyCodeTypeEnum().ZIP,
 		Runtime:    model.GetCreateFunctionRequestBodyRuntimeEnum().PYTHON3_9,
 		FuncCode: &model.FuncCode{
-			File: &code,
+			File: &function.HuaweiHttpCodeZip,
 		},
 	}
 

diff --git a/sdk/provider/tencent/http.go b/sdk/provider/tencent/http.go
@@ -8,7 +8,6 @@ import (
 	"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/errors"
 	scf "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/scf/v20180416"
 
-	"github.com/shimmeris/SCFProxy/fileutil"
 	"github.com/shimmeris/SCFProxy/function"
 	"github.com/shimmeris/SCFProxy/sdk"
 )
@@ -58,7 +57,7 @@ func (p *Provider) ClearHttpProxy(opts *sdk.HttpProxyOpts) error {
 func (p *Provider) createHttpFunction(functionName string) error {
 	r := scf.NewCreateFunctionRequest()
 	r.FunctionName = common.StringPtr(functionName)
-	r.Code = &scf.Code{ZipFile: common.StringPtr(fileutil.CreateZipBase64("index.py", function.TencentHttpCode))}
+	r.Code = &scf.Code{ZipFile: common.StringPtr(function.TencentHttpCodeZip)}
 	r.Handler = common.StringPtr("index.handler")
 	r.MemorySize = common.Int64Ptr(128)
 	r.Timeout = common.Int64Ptr(30)

diff --git a/sdk/provider/tencent/socks.go b/sdk/provider/tencent/socks.go
@@ -7,7 +7,6 @@ import (
 	"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/errors"
 	scf "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/scf/v20180416"
 
-	"github.com/shimmeris/SCFProxy/fileutil"
 	"github.com/shimmeris/SCFProxy/function"
 	"github.com/shimmeris/SCFProxy/sdk"
 )
@@ -45,7 +44,7 @@ func (p *Provider) createSocksFunction(functionName string) error {
 	r.FunctionName = common.StringPtr(functionName)
 	r.Handler = common.StringPtr("main")
 	r.Runtime = common.StringPtr("Go1")
-	r.Code = &scf.Code{ZipFile: common.StringPtr(fileutil.CreateZipBase64("main", function.TencentSocksCode))}
+	r.Code = &scf.Code{ZipFile: common.StringPtr(function.TencentSocksCodeZip)}
 	r.Timeout = common.Int64Ptr(900)
 	r.MemorySize = common.Int64Ptr(128)