feat: Add listening_addr to expose fun (#1110)

* Add listening_addr to expose fun Signed-off-by: nullday <aseaday@hotmail.com> * Rename parameter Signed-off-by: nullday <aseaday@hotmail.com> Signed-off-by: nullday <aseaday@hotmail.com>
tensorchord · Oct 29, 2022 · dedd731 · dedd731
1 parent f672c8f
commit dedd731
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 13 deletions.
diff --git a/pkg/envd/docker.go b/pkg/envd/docker.go
@@ -419,7 +419,7 @@ func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartRes
 			natPort := nat.Port(fmt.Sprintf("%d/tcp", item.EnvdPort))
 			hostConfig.PortBindings[natPort] = []nat.PortBinding{
 				{
-					HostIP:   localhost,
+					HostIP:   item.ListeningAddr,
 					HostPort: strconv.Itoa(item.HostPort),
 				},
 			}

diff --git a/pkg/lang/frontend/starlark/runtime/runtime.go b/pkg/lang/frontend/starlark/runtime/runtime.go
@@ -15,6 +15,7 @@
 package runtime
 
 import (
+	"net"
 	"os/user"
 	"path/filepath"
 	"strings"
@@ -103,13 +104,14 @@ func ruleFuncDaemon(thread *starlark.Thread, _ *starlark.Builtin,
 func ruleFuncExpose(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var (
-		envdPort    starlark.Int
-		hostPort    = starlark.MakeInt(0) // 0 means envd can randomly choose a free port
-		serviceName = starlark.String("")
+		envdPort      starlark.Int
+		hostPort      = starlark.MakeInt(0) // 0 means envd can randomly choose a free port
+		serviceName   = starlark.String("")
+		listeningAddr = starlark.String("127.0.0.1") // default to lisen only on local loopback interface
 	)
 
 	if err := starlark.UnpackArgs(ruleExpose,
-		args, kwargs, "envd_port", &envdPort, "host_port?", &hostPort, "service?", &serviceName); err != nil {
+		args, kwargs, "envd_port", &envdPort, "host_port?", &hostPort, "service?", &serviceName, "listen_addr?", &listeningAddr); err != nil {
 		return nil, err
 	}
 	envdPortInt, ok := envdPort.Int64()
@@ -121,9 +123,13 @@ func ruleFuncExpose(thread *starlark.Thread, _ *starlark.Builtin,
 		return nil, errors.New("host_port must be a positive integer less than 65535")
 	}
 	serviceNameStr := serviceName.GoString()
+	listeningAddrStr := listeningAddr.GoString()
+	if net.ParseIP(listeningAddrStr) == nil {
+		return nil, errors.New("listening_addr must be a valid IP address")
+	}
 
 	logger.Debugf("rule `%s` is invoked, envd_port=%d, host_port=%d, service=%s", ruleExpose, envdPortInt, hostPortInt, serviceNameStr)
-	err := ir.RuntimeExpose(int(envdPortInt), int(hostPortInt), serviceNameStr)
+	err := ir.RuntimeExpose(int(envdPortInt), int(hostPortInt), serviceNameStr, listeningAddrStr)
 	return starlark.None, err
 }
 

diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -214,11 +214,12 @@ func RuntimeDaemon(commands [][]string) {
 	DefaultGraph.RuntimeDaemon = append(DefaultGraph.RuntimeDaemon, commands...)
 }
 
-func RuntimeExpose(envdPort, hostPort int, serviceName string) error {
+func RuntimeExpose(envdPort, hostPort int, serviceName string, listeningAddr string) error {
 	DefaultGraph.RuntimeExpose = append(DefaultGraph.RuntimeExpose, ExposeItem{
-		EnvdPort:    envdPort,
-		HostPort:    hostPort,
-		ServiceName: serviceName,
+		EnvdPort:      envdPort,
+		HostPort:      hostPort,
+		ServiceName:   serviceName,
+		ListeningAddr: listeningAddr,
 	})
 	return nil
 }

diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
@@ -126,9 +126,10 @@ type GitConfig struct {
 }
 
 type ExposeItem struct {
-	EnvdPort    int
-	HostPort    int
-	ServiceName string
+	EnvdPort      int
+	HostPort      int
+	ServiceName   string
+	ListeningAddr string
 }
 
 type JupyterConfig struct {