Skip to content

Commit

Permalink
[feat] add lib/net: nslookup, tcping, httping (#118)
Browse files Browse the repository at this point in the history
* 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
hyorigo authored Jul 23, 2024
1 parent 5624d28 commit 97c1c1f
Show file tree
Hide file tree
Showing 6 changed files with 641 additions and 1 deletion.
2 changes: 2 additions & 0 deletions config.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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,
Expand Down
103 changes: 103 additions & 0 deletions lib/net/network.go
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
}
120 changes: 120 additions & 0 deletions lib/net/network_test.go
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
}
})
}
}
Loading

0 comments on commit 97c1c1f

Please sign in to comment.