[feat] add lib/net: nslookup, tcping, httping

config.go

@@ -10,6 +10,7 @@ import (
 	libhttp "github.com/1set/starlet/lib/http"
 	libjson "github.com/1set/starlet/lib/json"
 	liblog "github.com/1set/starlet/lib/log"
+	libnet "github.com/1set/starlet/lib/net"
 	libpath "github.com/1set/starlet/lib/path"
 	librand "github.com/1set/starlet/lib/random"
 	libre "github.com/1set/starlet/lib/re"
@@ -49,6 +50,7 @@ var allBuiltinModules = ModuleLoaderMap{
 	libfile.ModuleName: libfile.LoadModule,
 	libhash.ModuleName: libhash.LoadModule,
 	libhttp.ModuleName: libhttp.LoadModule,
+	libnet.ModuleName:  libnet.LoadModule,
 	libjson.ModuleName: libjson.LoadModule,
 	liblog.ModuleName:  liblog.LoadModule,
 	libpath.ModuleName: libpath.LoadModule,

lib/net/network.go

@@ -0,0 +1,103 @@
+// Package net provides network-related functions for Starlark, inspired by Go's net package and Python's socket module.
+package net
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/1set/starlet/dataconv"
+	tps "github.com/1set/starlet/dataconv/types"
+	"go.starlark.net/starlark"
+	"go.starlark.net/starlarkstruct"
+)
+
+// ModuleName defines the expected name for this Module when used in starlark's load() function, eg: load('net', 'tcping')
+const ModuleName = "net"
+
+var (
+	none    = starlark.None
+	once    sync.Once
+	modFunc starlark.StringDict
+)
+
+// LoadModule loads the net module. It is concurrency-safe and idempotent.
+func LoadModule() (starlark.StringDict, error) {
+	once.Do(func() {
+		modFunc = starlark.StringDict{
+			ModuleName: &starlarkstruct.Module{
+				Name: ModuleName,
+				Members: starlark.StringDict{
+					"nslookup": starlark.NewBuiltin(ModuleName+".nslookup", starLookup),
+					"tcping":   starlark.NewBuiltin(ModuleName+".tcping", starTCPPing),
+					"httping":  starlark.NewBuiltin(ModuleName+".httping", starHTTPing),
+				},
+			},
+		}
+	})
+	return modFunc, nil
+}
+
+func goLookup(ctx context.Context, domain, dnsServer string, timeout time.Duration) ([]string, error) {
+	// create a custom resolver if a DNS server is specified
+	var r *net.Resolver
+	if dnsServer != "" {
+		if !strings.Contains(dnsServer, ":") {
+			// append default DNS port if not specified
+			dnsServer = net.JoinHostPort(dnsServer, "53")
+		}
+		r = &net.Resolver{
+			PreferGo: true,
+			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+				d := net.Dialer{
+					Timeout: timeout,
+				}
+				return d.DialContext(ctx, "udp", dnsServer)
+			},
+		}
+	} else {
+		r = net.DefaultResolver
+	}
+
+	// Create a new context with timeout
+	ctxWithTimeout, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	// perform the DNS lookup
+	return r.LookupHost(ctxWithTimeout, domain)
+}
+
+func starLookup(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var (
+		domain    tps.StringOrBytes
+		dnsServer tps.NullableStringOrBytes
+		timeout   tps.FloatOrInt = 10
+	)
+	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "domain", &domain, "dns_server?", &dnsServer, "timeout?", &timeout); err != nil {
+		return nil, err
+	}
+
+	// correct timeout value
+	if timeout <= 0 {
+		timeout = 10
+	}
+
+	// get the context
+	ctx := dataconv.GetThreadContext(thread)
+
+	// perform the DNS lookup
+	ips, err := goLookup(ctx, domain.GoString(), dnsServer.GoString(), time.Duration(timeout)*time.Second)
+
+	// return the result
+	if err != nil {
+		return none, fmt.Errorf("%s: %w", b.Name(), err)
+	}
+	var list []starlark.Value
+	for _, ip := range ips {
+		list = append(list, starlark.String(ip))
+	}
+	return starlark.NewList(list), nil
+}

lib/net/network_test.go

@@ -0,0 +1,120 @@
+package net_test
+
+import (
+	"runtime"
+	"testing"
+
+	itn "github.com/1set/starlet/internal"
+	"github.com/1set/starlet/lib/net"
+)
+
+func TestLoadModule_NSLookUp(t *testing.T) {
+	isOnWindows := runtime.GOOS == "windows"
+	tests := []struct {
+		name        string
+		script      string
+		wantErr     string
+		skipWindows bool
+	}{
+		{
+			name: `nslookup: normal`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				ips = nslookup('bing.com')
+				print(ips)
+				assert.true(len(ips) > 0)
+			`),
+		},
+		{
+			name: `nslookup: normal with timeout`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				ips = nslookup('bing.com', timeout=5)
+				print(ips)
+				assert.true(len(ips) > 0)
+			`),
+		},
+		{
+			name: `nslookup: normal with dns`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				ips = nslookup('bing.com', '8.8.8.8')
+				print(ips)
+				assert.true(len(ips) > 0)
+			`),
+		},
+		{
+			name: `nslookup: normal with dns:port`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				ips = nslookup('bing.com', '1.1.1.1:53')
+				print(ips)
+				assert.true(len(ips) > 0)
+			`),
+		},
+		{
+			name: `nslookup: ip`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				ips = nslookup('8.8.8.8', timeout=-1)
+				print(ips)
+				assert.true(len(ips) > 0)
+			`),
+		},
+		{
+			name: `nslookup: localhost`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				ips = nslookup('localhost')
+				print(ips)
+				assert.true(len(ips) > 0)
+			`),
+		},
+		{
+			name: `nslookup: not exists`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				ips = nslookup('missing.invalid')
+			`),
+			wantErr: `missing.invalid`, // mac/win: no such host, linux: server misbehaving
+		},
+		{
+			name: `nslookup: wrong dns`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				ips = nslookup('bing.com', 'microsoft.com', timeout=1)
+			`),
+			wantErr:     `i/o timeout`,
+			skipWindows: true, // on Windows 2022 with Go 1.18.10, it returns results from the default DNS server
+		},
+		{
+			name: `nslookup: no args`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				nslookup()
+			`),
+			wantErr: `net.nslookup: missing argument for domain`,
+		},
+		{
+			name: `nslookup: invalid args`,
+			script: itn.HereDoc(`
+				load('net', 'nslookup')
+				nslookup(1, 2, 3)
+			`),
+			wantErr: `net.nslookup: for parameter domain: got int, want string or bytes`,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if isOnWindows && tt.skipWindows {
+				t.Skipf("Skip test on Windows")
+				return
+			}
+			res, err := itn.ExecModuleWithErrorTest(t, net.ModuleName, net.LoadModule, tt.script, tt.wantErr, nil)
+			if (err != nil) != (tt.wantErr != "") {
+				t.Errorf("net(%q) expects error = '%v', actual error = '%v', result = %v", tt.name, tt.wantErr, err, res)
+				return
+			}
+		})
+	}
+}