Skip to content

Commit 97c1c1f

Browse files
hyorigohyorigo
authored
[feat] add lib/net: nslookup, tcping, httping (#118)
* draft for nslookup * for lookup * for invalid * for issues * for timeout * for test cases * draft for tcp ping * for checks * fix go sum * for ping and test * ping it now * include * fix test case * test coverage * fix lint * skip win * add notes for test case * refactor for httping * clean comments * use net/http/httptrace * split test files * add tests * fix var naming * for 200-400, and fix reuse * for mock server
1 parent 5624d28 commit 97c1c1f

File tree

6 files changed

+641
-1
lines changed

6 files changed

+641
-1
lines changed

‎config.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
libhttp "github.com/1set/starlet/lib/http"
1111
libjson "github.com/1set/starlet/lib/json"
1212
liblog "github.com/1set/starlet/lib/log"
13+
libnet "github.com/1set/starlet/lib/net"
1314
libpath "github.com/1set/starlet/lib/path"
1415
librand "github.com/1set/starlet/lib/random"
1516
libre "github.com/1set/starlet/lib/re"
@@ -49,6 +50,7 @@ var allBuiltinModules = ModuleLoaderMap{
4950
libfile.ModuleName: libfile.LoadModule,
5051
libhash.ModuleName: libhash.LoadModule,
5152
libhttp.ModuleName: libhttp.LoadModule,
53+
libnet.ModuleName: libnet.LoadModule,
5254
libjson.ModuleName: libjson.LoadModule,
5355
liblog.ModuleName: liblog.LoadModule,
5456
libpath.ModuleName: libpath.LoadModule,

‎lib/net/network.go

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
// Package net provides network-related functions for Starlark, inspired by Go's net package and Python's socket module.
2+
package net
3+
4+
import (
5+
"context"
6+
"fmt"
7+
"net"
8+
"strings"
9+
"sync"
10+
"time"
11+
12+
"github.com/1set/starlet/dataconv"
13+
tps "github.com/1set/starlet/dataconv/types"
14+
"go.starlark.net/starlark"
15+
"go.starlark.net/starlarkstruct"
16+
)
17+
18+
// ModuleName defines the expected name for this Module when used in starlark's load() function, eg: load('net', 'tcping')
19+
const ModuleName = "net"
20+
21+
var (
22+
none = starlark.None
23+
once sync.Once
24+
modFunc starlark.StringDict
25+
)
26+
27+
// LoadModule loads the net module. It is concurrency-safe and idempotent.
28+
func LoadModule() (starlark.StringDict, error) {
29+
once.Do(func() {
30+
modFunc = starlark.StringDict{
31+
ModuleName: &starlarkstruct.Module{
32+
Name: ModuleName,
33+
Members: starlark.StringDict{
34+
"nslookup": starlark.NewBuiltin(ModuleName+".nslookup", starLookup),
35+
"tcping": starlark.NewBuiltin(ModuleName+".tcping", starTCPPing),
36+
"httping": starlark.NewBuiltin(ModuleName+".httping", starHTTPing),
37+
},
38+
},
39+
}
40+
})
41+
return modFunc, nil
42+
}
43+
44+
func goLookup(ctx context.Context, domain, dnsServer string, timeout time.Duration) ([]string, error) {
45+
// create a custom resolver if a DNS server is specified
46+
var r *net.Resolver
47+
if dnsServer != "" {
48+
if !strings.Contains(dnsServer, ":") {
49+
// append default DNS port if not specified
50+
dnsServer = net.JoinHostPort(dnsServer, "53")
51+
}
52+
r = &net.Resolver{
53+
PreferGo: true,
54+
Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
55+
d := net.Dialer{
56+
Timeout: timeout,
57+
}
58+
return d.DialContext(ctx, "udp", dnsServer)
59+
},
60+
}
61+
} else {
62+
r = net.DefaultResolver
63+
}
64+
65+
// Create a new context with timeout
66+
ctxWithTimeout, cancel := context.WithTimeout(ctx, timeout)
67+
defer cancel()
68+
69+
// perform the DNS lookup
70+
return r.LookupHost(ctxWithTimeout, domain)
71+
}
72+
73+
func starLookup(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
74+
var (
75+
domain tps.StringOrBytes
76+
dnsServer tps.NullableStringOrBytes
77+
timeout tps.FloatOrInt = 10
78+
)
79+
if err := starlark.UnpackArgs(b.Name(), args, kwargs, "domain", &domain, "dns_server?", &dnsServer, "timeout?", &timeout); err != nil {
80+
return nil, err
81+
}
82+
83+
// correct timeout value
84+
if timeout <= 0 {
85+
timeout = 10
86+
}
87+
88+
// get the context
89+
ctx := dataconv.GetThreadContext(thread)
90+
91+
// perform the DNS lookup
92+
ips, err := goLookup(ctx, domain.GoString(), dnsServer.GoString(), time.Duration(timeout)*time.Second)
93+
94+
// return the result
95+
if err != nil {
96+
return none, fmt.Errorf("%s: %w", b.Name(), err)
97+
}
98+
var list []starlark.Value
99+
for _, ip := range ips {
100+
list = append(list, starlark.String(ip))
101+
}
102+
return starlark.NewList(list), nil
103+
}

‎lib/net/network_test.go

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
package net_test
2+
3+
import (
4+
"runtime"
5+
"testing"
6+
7+
itn "github.com/1set/starlet/internal"
8+
"github.com/1set/starlet/lib/net"
9+
)
10+
11+
func TestLoadModule_NSLookUp(t *testing.T) {
12+
isOnWindows := runtime.GOOS == "windows"
13+
tests := []struct {
14+
name string
15+
script string
16+
wantErr string
17+
skipWindows bool
18+
}{
19+
{
20+
name: `nslookup: normal`,
21+
script: itn.HereDoc(`
22+
load('net', 'nslookup')
23+
ips = nslookup('bing.com')
24+
print(ips)
25+
assert.true(len(ips) > 0)
26+
`),
27+
},
28+
{
29+
name: `nslookup: normal with timeout`,
30+
script: itn.HereDoc(`
31+
load('net', 'nslookup')
32+
ips = nslookup('bing.com', timeout=5)
33+
print(ips)
34+
assert.true(len(ips) > 0)
35+
`),
36+
},
37+
{
38+
name: `nslookup: normal with dns`,
39+
script: itn.HereDoc(`
40+
load('net', 'nslookup')
41+
ips = nslookup('bing.com', '8.8.8.8')
42+
print(ips)
43+
assert.true(len(ips) > 0)
44+
`),
45+
},
46+
{
47+
name: `nslookup: normal with dns:port`,
48+
script: itn.HereDoc(`
49+
load('net', 'nslookup')
50+
ips = nslookup('bing.com', '1.1.1.1:53')
51+
print(ips)
52+
assert.true(len(ips) > 0)
53+
`),
54+
},
55+
{
56+
name: `nslookup: ip`,
57+
script: itn.HereDoc(`
58+
load('net', 'nslookup')
59+
ips = nslookup('8.8.8.8', timeout=-1)
60+
print(ips)
61+
assert.true(len(ips) > 0)
62+
`),
63+
},
64+
{
65+
name: `nslookup: localhost`,
66+
script: itn.HereDoc(`
67+
load('net', 'nslookup')
68+
ips = nslookup('localhost')
69+
print(ips)
70+
assert.true(len(ips) > 0)
71+
`),
72+
},
73+
{
74+
name: `nslookup: not exists`,
75+
script: itn.HereDoc(`
76+
load('net', 'nslookup')
77+
ips = nslookup('missing.invalid')
78+
`),
79+
wantErr: `missing.invalid`, // mac/win: no such host, linux: server misbehaving
80+
},
81+
{
82+
name: `nslookup: wrong dns`,
83+
script: itn.HereDoc(`
84+
load('net', 'nslookup')
85+
ips = nslookup('bing.com', 'microsoft.com', timeout=1)
86+
`),
87+
wantErr: `i/o timeout`,
88+
skipWindows: true, // on Windows 2022 with Go 1.18.10, it returns results from the default DNS server
89+
},
90+
{
91+
name: `nslookup: no args`,
92+
script: itn.HereDoc(`
93+
load('net', 'nslookup')
94+
nslookup()
95+
`),
96+
wantErr: `net.nslookup: missing argument for domain`,
97+
},
98+
{
99+
name: `nslookup: invalid args`,
100+
script: itn.HereDoc(`
101+
load('net', 'nslookup')
102+
nslookup(1, 2, 3)
103+
`),
104+
wantErr: `net.nslookup: for parameter domain: got int, want string or bytes`,
105+
},
106+
}
107+
for _, tt := range tests {
108+
t.Run(tt.name, func(t *testing.T) {
109+
if isOnWindows && tt.skipWindows {
110+
t.Skipf("Skip test on Windows")
111+
return
112+
}
113+
res, err := itn.ExecModuleWithErrorTest(t, net.ModuleName, net.LoadModule, tt.script, tt.wantErr, nil)
114+
if (err != nil) != (tt.wantErr != "") {
115+
t.Errorf("net(%q) expects error = '%v', actual error = '%v', result = %v", tt.name, tt.wantErr, err, res)
116+
return
117+
}
118+
})
119+
}
120+
}

0 commit comments

Comments
 (0)