-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[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
- Loading branch information
Showing
6 changed files
with
641 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
} | ||
}) | ||
} | ||
} |
Oops, something went wrong.