实现分流器扩展

gost-geo-plugin/main_test.go

@@ -0,0 +1,91 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"net/netip"
+	"testing"
+
+	"github.com/go-gost/plugin/bypass/proto"
+)
+
+func TestBypass(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		addr string
+		ips  []netip.Addr
+		err  error
+		deny bool
+	}{
+		{name: "listed IPv4", addr: "203.0.113.1:443"},
+		{name: "foreign IPv4", addr: "198.51.100.1:443", deny: true},
+		{name: "listed IPv6", addr: "[2001:db8::1]:443"},
+		{name: "bare IPv6", addr: "2001:db8::1"},
+		{name: "foreign IPv6", addr: "[2001:db9::1]:443", deny: true},
+		{name: "mapped IPv4", addr: "[::ffff:203.0.113.1]:443"},
+		{name: "listed domain", addr: "allowed.test:443", ips: []netip.Addr{netip.MustParseAddr("203.0.113.2")}},
+		{name: "mixed answers", addr: "mixed.test:443", ips: []netip.Addr{netip.MustParseAddr("203.0.113.2"), netip.MustParseAddr("198.51.100.2")}, deny: true},
+		{name: "no addresses", addr: "empty.test:443", deny: true},
+		{name: "DNS failure", addr: "error.test:443", err: errors.New("DNS unavailable"), deny: true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			s := &server{
+				list: []netip.Prefix{netip.MustParsePrefix("203.0.113.0/24"), netip.MustParsePrefix("2001:db8::/32")},
+				lookup: func(ctx context.Context, _, _ string) ([]netip.Addr, error) {
+					if _, ok := ctx.Deadline(); !ok {
+						t.Error("DNS lookup has no deadline")
+					}
+					return tc.ips, tc.err
+				},
+			}
+			reply, err := s.Bypass(context.Background(), &proto.BypassRequest{Addr: tc.addr})
+			if err != nil || reply.Ok != tc.deny {
+				t.Fatalf("Bypass = %v, %v; want deny=%v", reply, err, tc.deny)
+			}
+		})
+	}
+}
+
+func TestRefresh(t *testing.T) {
+	status, body := 200, "# list\n203.0.113.7/24 # normalize\n2001:db8::/32\n"
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(status)
+		_, _ = w.Write([]byte(body))
+	}))
+	defer upstream.Close()
+	s := &server{}
+	if err := s.refresh(upstream.URL); err != nil {
+		t.Fatal(err)
+	}
+	for _, tc := range []struct {
+		name   string
+		status int
+		body   string
+	}{
+		{"HTTP failure", 503, "unavailable"},
+		{"empty list", 200, "# empty\n"},
+		{"malformed list", 200, "198.51.100.0/24\n<html>error</html>"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			status, body = tc.status, tc.body
+			if err := s.refresh(upstream.URL); err == nil {
+				t.Fatal("accepted invalid list")
+			}
+			reply, _ := s.Bypass(context.Background(), &proto.BypassRequest{Addr: "203.0.113.1:443"})
+			if reply.Ok || len(s.list) != 2 {
+				t.Fatal("failed refresh discarded previous rules")
+			}
+		})
+	}
+	status, body = 200, "198.51.100.0/24\n"
+	if err := s.refresh(upstream.URL); err != nil {
+		t.Fatal(err)
+	}
+	old, _ := s.Bypass(context.Background(), &proto.BypassRequest{Addr: "203.0.113.1:443"})
+	updated, _ := s.Bypass(context.Background(), &proto.BypassRequest{Addr: "198.51.100.1:443"})
+	if !old.Ok || updated.Ok {
+		t.Fatal("refresh did not replace rules")
+	}
+}