feat(lang/py): add trust for pip_index (#1395)

* feat(lang/py): add trust for pip_index, fix #1387 Signed-off-by: weixiao-huang <hwx.simle@gmail.com> * feat(lang/py): make trust optional and use bool for trust Signed-off-by: weixiao-huang <hwx.simle@gmail.com> Signed-off-by: weixiao-huang <hwx.simle@gmail.com>
tensorchord · Jan 11, 2023 · e7cabe6 · e7cabe6
1 parent 31464f7
commit e7cabe6
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 8 deletions.
diff --git a/pkg/lang/frontend/starlark/v0/config/config.go b/pkg/lang/frontend/starlark/v0/config/config.go
@@ -95,18 +95,19 @@ func ruleFuncJupyter(thread *starlark.Thread, _ *starlark.Builtin,
 func ruleFuncPyPIIndex(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var url, extraURL starlark.String
+	var trust bool
 
 	if err := starlark.UnpackArgs(rulePyPIIndex, args, kwargs,
-		"url?", &url, "extra_url?", &extraURL); err != nil {
+		"url?", &url, "extra_url?", &extraURL, "trust?", &trust); err != nil {
 		return nil, err
 	}
 
 	indexStr := url.GoString()
 	extraIndexStr := extraURL.GoString()
 
-	logger.Debugf("rule `%s` is invoked, index=%s, extraIndex=%s",
-		rulePyPIIndex, indexStr, extraIndexStr)
-	if err := ir.PyPIIndex(indexStr, extraIndexStr); err != nil {
+	logger.Debugf("rule `%s` is invoked, index=%s, extraIndex=%s, trust=%t",
+		rulePyPIIndex, indexStr, extraIndexStr, trust)
+	if err := ir.PyPIIndex(indexStr, extraIndexStr, trust); err != nil {
 		return nil, err
 	}
 

diff --git a/pkg/lang/ir/v0/interface.go b/pkg/lang/ir/v0/interface.go
@@ -114,14 +114,15 @@ func UbuntuAPT(source string) error {
 	return nil
 }
 
-func PyPIIndex(url, extraURL string) error {
+func PyPIIndex(url, extraURL string, trust bool) error {
 	if url == "" {
 		return errors.New("url is required")
 	}
 	g := DefaultGraph.(*generalGraph)
 
 	g.PyPIIndexURL = &url
 	g.PyPIExtraIndexURL = &extraURL
+	g.PyPITrust = trust
 	return nil
 }
 

diff --git a/pkg/lang/ir/v0/python.go b/pkg/lang/ir/v0/python.go
@@ -16,6 +16,7 @@ package v0
 
 import (
 	"fmt"
+	"net/url"
 	"path/filepath"
 	"strings"
 
@@ -199,12 +200,26 @@ func (g generalGraph) compilePyPIPackages(root llb.State) llb.State {
 func (g generalGraph) compilePyPIIndex(root llb.State) llb.State {
 	if g.PyPIIndexURL != nil {
 		logrus.WithField("index", *g.PyPIIndexURL).Debug("using custom PyPI index")
-		var extraIndex string
+		var extra string
 		if g.PyPIExtraIndexURL != nil {
 			logrus.WithField("index", *g.PyPIIndexURL).Debug("using extra PyPI index")
-			extraIndex = "extra-index-url=" + *g.PyPIExtraIndexURL
+			extra = "extra-index-url=" + *g.PyPIExtraIndexURL
 		}
-		content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extraIndex)
+		if g.PyPITrust {
+			var hosts []string
+			for _, p := range []*string{g.PyPIIndexURL, g.PyPIExtraIndexURL} {
+				if p != nil {
+					u, err := url.Parse(*p)
+					if err == nil && u != nil && u.Hostname() != "" {
+						hosts = append(hosts, u.Hostname())
+					}
+				}
+			}
+			if len(hosts) > 0 {
+				extra += "\ntrusted-host=" + strings.Join(hosts, " ")
+			}
+		}
+		content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extra)
 		dir := filepath.Dir(pypiIndexFilePath)
 		pypiMirror := root.
 			File(llb.Mkdir(dir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)),

diff --git a/pkg/lang/ir/v0/types.go b/pkg/lang/ir/v0/types.go
@@ -42,6 +42,7 @@ type generalGraph struct {
 	JuliaPackageServer *string
 	PyPIIndexURL       *string
 	PyPIExtraIndexURL  *string
+	PyPITrust          bool
 
 	PublicKeyPath string