Bamboo/discover/discover_test.go

package discover

import (
	"context"
	"net"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

func TestParseSubnetsRefusesPublicAndHuge(t *testing.T) {
	got, err := ParseSubnets("192.168.1.7/24, 10.0.0.0/22")
	if err != nil || len(got) != 2 || got[0].String() != "192.168.1.0/24" {
		t.Fatalf("got %v, %v", got, err)
	}
	for _, bad := range []string{"8.8.8.0/24", "10.0.0.0/8", "192.168.0.0/16", "nonsense", "1.2.3.4/32"} {
		if _, err := ParseSubnets(bad); err == nil {
			t.Errorf("%q should be refused", bad)
		}
	}
}

func TestParsePorts(t *testing.T) {
	got, err := ParsePorts("443, 80,8000-8002,80")
	if err != nil || len(got) != 5 || got[0] != 80 || got[4] != 8002 {
		t.Fatalf("got %v, %v", got, err)
	}
	for _, bad := range []string{"0", "70000", "9-3", "x", "1-5000"} {
		if _, err := ParsePorts(bad); err == nil {
			t.Errorf("%q should be refused", bad)
		}
	}
}

func TestParseProcNetTCP(t *testing.T) {
	v4 := `  sl  local_address rem_address   st tx_queue rx_queue
   0: 0100007F:1F90 00000000:0000 0A 00000000:00000000 00:00000000 00000000  1000 0 1 1
   1: 00000000:0016 00000000:0000 0A 00000000:00000000 00:00000000 00000000     0 0 2 1
   2: 0100007F:D431 0100007F:1F90 01 00000000:00000000 00:00000000 00000000  1000 0 3 1
`
	got := ParseProcNetTCP(v4, false)
	if len(got) != 2 || got[0].String() != "127.0.0.1:8080" || got[1].String() != "0.0.0.0:22" {
		t.Fatalf("v4: %v", got)
	}
	v6 := `  sl  local_address rem_address   st
   0: 00000000000000000000000001000000:0050 00000000000000000000000000000000:0000 0A 0:0 0:0 0 0 0 1
`
	got = ParseProcNetTCP(v6, true)
	if len(got) != 1 || got[0].String() != "[::1]:80" {
		t.Fatalf("v6: %v", got)
	}
}

// banner starts a listener that says line first (or stays silent when line is "").
func banner(t *testing.T, line string) netip.AddrPort {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			if line != "" {
				c.Write([]byte(line + "\r\n"))
			}
			time.Sleep(50 * time.Millisecond)
			c.Close()
		}
	}()
	return netip.MustParseAddrPort(ln.Addr().String())
}

func newTest(t *testing.T, cfg Config) *Discoverer {
	t.Helper()
	cfg.Rate = 10000
	cfg.timeout = time.Second
	d, err := New(cfg)
	if err != nil {
		t.Fatal(err)
	}
	return d
}

func TestIdentify(t *testing.T) {
	ssh := banner(t, "SSH-2.0-OpenSSH_9.6")
	smtp := banner(t, "220 mail.example ESMTP Postfix")
	web := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Server", "testd/1"); w.WriteHeader(200) }))
	defer web.Close()
	tlsWeb := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer tlsWeb.Close()
	silent := banner(t, "")

	d := newTest(t, Config{})
	ctx := context.Background()
	for _, c := range []struct {
		name string
		ap   netip.AddrPort
		want string
		det  string
	}{
		{"ssh", ssh, "ssh", "SSH-2.0"},
		{"smtp", smtp, "smtp", "ESMTP"},
		{"http", netip.MustParseAddrPort(strings.TrimPrefix(web.URL, "http://")), "http", "testd/1"},
		{"https", netip.MustParseAddrPort(strings.TrimPrefix(tlsWeb.URL, "https://")), "https", ""},
		{"silent", silent, "", ""},
	} {
		svc, detail := d.identify(ctx, c.ap)
		if svc != c.want || !strings.Contains(detail, c.det) {
			t.Errorf("%s: got %q %q, want %q containing %q", c.name, svc, detail, c.want, c.det)
		}
	}
}

func TestSweepLocalAndScan(t *testing.T) {
	ssh := banner(t, "SSH-2.0-test")
	d := newTest(t, Config{
		Local:   true,
		Subnets: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/30")},
		Ports:   []int{int(ssh.Port()), 1}, // port 1 is closed
		Exclude: []int{1},
		listening: func() ([]netip.AddrPort, error) {
			return []netip.AddrPort{netip.AddrPortFrom(netip.IPv4Unspecified(), ssh.Port())}, nil
		},
	})
	res := d.Sweep(context.Background())
	if len(res) != 1 { // the local hit and the scan hit are the same endpoint, reported once
		t.Fatalf("want one service, got %+v", res)
	}
	r := res[0]
	if r.Host != "127.0.0.1" || r.Port != int(ssh.Port()) || r.Service != "ssh" || r.Source != "local" || r.Proto != "tcp" {
		t.Errorf("unexpected: %+v", r)
	}
}

func TestExcludeLocalPort(t *testing.T) {
	ssh := banner(t, "SSH-2.0-test")
	d := newTest(t, Config{Local: true, Exclude: []int{int(ssh.Port())},
		listening: func() ([]netip.AddrPort, error) { return []netip.AddrPort{ssh}, nil }})
	if res := d.Sweep(context.Background()); len(res) != 0 {
		t.Fatalf("excluded port was reported: %+v", res)
	}
}

func TestScanSkipsNetworkAndBroadcast(t *testing.T) {
	var hosts = map[string]bool{}
	d := newTest(t, Config{Subnets: []netip.Prefix{netip.MustParsePrefix("192.168.9.0/30")}, Ports: []int{80}})
	d.cfg.dial = func(ctx context.Context, addr string) (net.Conn, error) {
		h, _, _ := net.SplitHostPort(addr)
		hosts[h] = true
		return nil, net.ErrClosed
	}
	d.Sweep(context.Background())
	if len(hosts) != 2 || hosts["192.168.9.0"] || hosts["192.168.9.3"] {
		t.Fatalf("probed %v", hosts)
	}
}

func TestRateLimit(t *testing.T) {
	var n atomic.Int32
	d, err := New(Config{Subnets: []netip.Prefix{netip.MustParsePrefix("192.168.9.0/29")}, Ports: []int{80, 81}, Rate: 50})
	if err != nil {
		t.Fatal(err)
	}
	d.cfg.dial = func(ctx context.Context, addr string) (net.Conn, error) { n.Add(1); return nil, net.ErrClosed }
	start := time.Now()
	d.Sweep(context.Background()) // 6 hosts x 2 ports = 12 probes at 50/s: at least ~220ms
	if n.Load() != 12 {
		t.Fatalf("probes: %d", n.Load())
	}
	if el := time.Since(start); el < 200*time.Millisecond {
		t.Fatalf("not rate limited: %v", el)
	}
}

func TestSweepStopsOnCancel(t *testing.T) {
	d, _ := New(Config{Subnets: []netip.Prefix{netip.MustParsePrefix("10.9.0.0/22")}, Ports: []int{80}, Rate: 1})
	d.cfg.dial = func(ctx context.Context, addr string) (net.Conn, error) { return nil, net.ErrClosed }
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	done := make(chan struct{})
	go func() { d.Sweep(ctx); close(done) }()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("sweep ignored cancellation")
	}
}

func TestNewRefusesPublicSubnet(t *testing.T) {
	if _, err := New(Config{Subnets: []netip.Prefix{netip.MustParsePrefix("8.8.8.0/24")}}); err == nil {
		t.Fatal("public subnet accepted")
	}
}