143 lines
3.1 KiB
Go
143 lines
3.1 KiB
Go
package whois
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"net"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestWhoisReferralLoopGuard(t *testing.T) {
|
|
old, hadOld := defaultWhoisMap["loop"]
|
|
defer func() {
|
|
if hadOld {
|
|
defaultWhoisMap["loop"] = old
|
|
} else {
|
|
delete(defaultWhoisMap, "loop")
|
|
}
|
|
}()
|
|
|
|
var addrA, addrB string
|
|
var hitA, hitB atomic.Int32
|
|
|
|
addrB, stopB := startWhoisServerWithHandler(t, &hitB, func(_ string) string {
|
|
return strings.Join([]string{
|
|
"Domain Name: EXAMPLE.LOOP",
|
|
"Registrar: LOOP-REG-B",
|
|
"Whois Server: " + addrA,
|
|
"",
|
|
}, "\n")
|
|
})
|
|
defer stopB()
|
|
|
|
addrA, stopA := startWhoisServerWithHandler(t, &hitA, func(_ string) string {
|
|
return strings.Join([]string{
|
|
"Domain Name: EXAMPLE.LOOP",
|
|
"Registrar: LOOP-REG-A",
|
|
"Whois Server: " + addrB,
|
|
"",
|
|
}, "\n")
|
|
})
|
|
defer stopA()
|
|
|
|
defaultWhoisMap["loop"] = addrA
|
|
|
|
c := NewClient()
|
|
r, err := c.WhoisWithOptions("example.loop", QueryOptions{
|
|
Level: QueryBoth,
|
|
ReferralMaxDepth: 8,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("WhoisWithOptions() error: %v", err)
|
|
}
|
|
if !r.Exists() {
|
|
t.Fatal("expected exists=true")
|
|
}
|
|
if got := hitA.Load(); got != 1 {
|
|
t.Fatalf("expected A hit once (loop prevented), got=%d", got)
|
|
}
|
|
if got := hitB.Load(); got != 1 {
|
|
t.Fatalf("expected B hit once, got=%d", got)
|
|
}
|
|
}
|
|
|
|
func TestWhoisNegativeCacheNotFound(t *testing.T) {
|
|
clearNegativeCache()
|
|
old, hadOld := defaultWhoisMap["nfcache"]
|
|
defer func() {
|
|
if hadOld {
|
|
defaultWhoisMap["nfcache"] = old
|
|
} else {
|
|
delete(defaultWhoisMap, "nfcache")
|
|
}
|
|
clearNegativeCache()
|
|
}()
|
|
|
|
var hit atomic.Int32
|
|
addr, stop := startWhoisServerWithHandler(t, &hit, func(_ string) string {
|
|
return strings.Join([]string{
|
|
"No match for \"EXAMPLE.NFCACHE\"",
|
|
"",
|
|
}, "\n")
|
|
})
|
|
defer stop()
|
|
defaultWhoisMap["nfcache"] = addr
|
|
|
|
c := NewClient().SetNegativeCacheTTL(time.Minute)
|
|
r1, err := c.WhoisWithOptions("example.nfcache", QueryOptions{Level: QueryAuto})
|
|
if err != nil {
|
|
t.Fatalf("first WhoisWithOptions() error: %v", err)
|
|
}
|
|
r2, err := c.WhoisWithOptions("example.nfcache", QueryOptions{Level: QueryAuto})
|
|
if err != nil {
|
|
t.Fatalf("second WhoisWithOptions() error: %v", err)
|
|
}
|
|
if r1.Exists() || r2.Exists() {
|
|
t.Fatal("expected cached not-found results")
|
|
}
|
|
if got := hit.Load(); got != 1 {
|
|
t.Fatalf("expected one whois query hit with negative cache, got=%d", got)
|
|
}
|
|
}
|
|
|
|
func startWhoisServerWithHandler(t *testing.T, hit *atomic.Int32, handler func(query string) string) (addr string, shutdown func()) {
|
|
t.Helper()
|
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
|
if err != nil {
|
|
t.Fatalf("listen mock whois failed: %v", err)
|
|
}
|
|
done := make(chan struct{})
|
|
go func() {
|
|
for {
|
|
conn, err := ln.Accept()
|
|
if err != nil {
|
|
select {
|
|
case <-done:
|
|
return
|
|
default:
|
|
return
|
|
}
|
|
}
|
|
go func(c net.Conn) {
|
|
defer c.Close()
|
|
_ = c.SetDeadline(time.Now().Add(5 * time.Second))
|
|
line, _ := bufio.NewReader(c).ReadString('\n')
|
|
if hit != nil {
|
|
hit.Add(1)
|
|
}
|
|
resp := handler(strings.TrimSpace(line))
|
|
if resp != "" {
|
|
_, _ = fmt.Fprint(c, resp)
|
|
}
|
|
}(conn)
|
|
}
|
|
}()
|
|
return ln.Addr().String(), func() {
|
|
close(done)
|
|
_ = ln.Close()
|
|
}
|
|
}
|