mc-proxy/internal/firewall/firewall_test.go

package firewall

import (
	"net/netip"
	"testing"
	"time"
)

func TestEmptyFirewall(t *testing.T) {
	fw, err := New("", nil, nil, nil, 0, 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer fw.Close()

	addrs := []string{"192.168.1.1", "10.0.0.1", "::1", "2001:db8::1"}
	for _, a := range addrs {
		addr := netip.MustParseAddr(a)
		if fw.Blocked(addr) {
			t.Fatalf("empty firewall blocked %s", addr)
		}
	}
}

func TestIPBlocking(t *testing.T) {
	fw, err := New("", []string{"192.0.2.1", "2001:db8::dead"}, nil, nil, 0, 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer fw.Close()

	tests := []struct {
		addr    string
		blocked bool
	}{
		{"192.0.2.1", true},
		{"192.0.2.2", false},
		{"2001:db8::dead", true},
		{"2001:db8::beef", false},
	}

	for _, tt := range tests {
		addr := netip.MustParseAddr(tt.addr)
		if got := fw.Blocked(addr); got != tt.blocked {
			t.Fatalf("Blocked(%s) = %v, want %v", tt.addr, got, tt.blocked)
		}
	}
}

func TestCIDRBlocking(t *testing.T) {
	fw, err := New("", nil, []string{"198.51.100.0/24", "2001:db8::/32"}, nil, 0, 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer fw.Close()

	tests := []struct {
		addr    string
		blocked bool
	}{
		{"198.51.100.1", true},
		{"198.51.100.254", true},
		{"198.51.101.1", false},
		{"2001:db8::1", true},
		{"2001:db9::1", false},
	}

	for _, tt := range tests {
		addr := netip.MustParseAddr(tt.addr)
		if got := fw.Blocked(addr); got != tt.blocked {
			t.Fatalf("Blocked(%s) = %v, want %v", tt.addr, got, tt.blocked)
		}
	}
}

func TestIPv4MappedIPv6(t *testing.T) {
	fw, err := New("", []string{"192.0.2.1"}, nil, nil, 0, 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer fw.Close()

	addr := netip.MustParseAddr("::ffff:192.0.2.1")
	if !fw.Blocked(addr) {
		t.Fatal("expected IPv4-mapped IPv6 address to be blocked")
	}
}

func TestInvalidIP(t *testing.T) {
	_, err := New("", []string{"not-an-ip"}, nil, nil, 0, 0)
	if err == nil {
		t.Fatal("expected error for invalid IP")
	}
}

func TestInvalidCIDR(t *testing.T) {
	_, err := New("", nil, []string{"not-a-cidr"}, nil, 0, 0)
	if err == nil {
		t.Fatal("expected error for invalid CIDR")
	}
}

func TestCombinedRules(t *testing.T) {
	fw, err := New("", []string{"10.0.0.1"}, []string{"192.168.0.0/16"}, nil, 0, 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer fw.Close()

	tests := []struct {
		addr    string
		blocked bool
	}{
		{"10.0.0.1", true},
		{"10.0.0.2", false},
		{"192.168.1.1", true},
		{"172.16.0.1", false},
	}

	for _, tt := range tests {
		addr := netip.MustParseAddr(tt.addr)
		if got := fw.Blocked(addr); got != tt.blocked {
			t.Fatalf("Blocked(%s) = %v, want %v", tt.addr, got, tt.blocked)
		}
	}
}

func TestRateLimitBlocking(t *testing.T) {
	fw, err := New("", nil, nil, nil, 2, time.Minute)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer fw.Close()

	addr := netip.MustParseAddr("10.0.0.1")

	if fw.Blocked(addr) {
		t.Fatal("first request should be allowed")
	}
	if fw.Blocked(addr) {
		t.Fatal("second request should be allowed")
	}
	if !fw.Blocked(addr) {
		t.Fatal("third request should be blocked (limit=2)")
	}
}

func TestRateLimitBlocklistFirst(t *testing.T) {
	// A blocklisted IP should be blocked without consuming rate limit quota.
	fw, err := New("", []string{"10.0.0.1"}, nil, nil, 1, time.Minute)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer fw.Close()

	blockedAddr := netip.MustParseAddr("10.0.0.1")
	otherAddr := netip.MustParseAddr("10.0.0.2")

	// Blocked by blocklist — should not touch the rate limiter.
	if !fw.Blocked(blockedAddr) {
		t.Fatal("blocklisted IP should be blocked")
	}

	// Other address should still have its full rate limit quota.
	if fw.Blocked(otherAddr) {
		t.Fatal("other IP should be allowed (within rate limit)")
	}
	if !fw.Blocked(otherAddr) {
		t.Fatal("other IP should be blocked after exceeding rate limit")
	}
}

func TestBlockedWithReason(t *testing.T) {
	fw, err := New("", []string{"10.0.0.1"}, []string{"192.168.0.0/16"}, nil, 2, time.Minute)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer fw.Close()

	tests := []struct {
		addr       string
		wantBlock  bool
		wantReason string
	}{
		{"10.0.0.1", true, "ip"},
		{"192.168.1.1", true, "cidr"},
		{"172.16.0.1", false, ""},
	}

	for _, tt := range tests {
		addr := netip.MustParseAddr(tt.addr)
		blocked, reason := fw.BlockedWithReason(addr)
		if blocked != tt.wantBlock {
			t.Fatalf("BlockedWithReason(%s) blocked = %v, want %v", tt.addr, blocked, tt.wantBlock)
		}
		if reason != tt.wantReason {
			t.Fatalf("BlockedWithReason(%s) reason = %q, want %q", tt.addr, reason, tt.wantReason)
		}
	}

	// Test rate limit reason: use a fresh IP that will exceed the limit.
	rlAddr := netip.MustParseAddr("10.10.10.10")
	fw.BlockedWithReason(rlAddr)                    // 1
	fw.BlockedWithReason(rlAddr)                    // 2
	blocked, reason := fw.BlockedWithReason(rlAddr) // 3 — should be blocked
	if !blocked {
		t.Fatal("expected rate limit block")
	}
	if reason != "rate_limit" {
		t.Fatalf("reason = %q, want %q", reason, "rate_limit")
	}
}

func TestRuntimeMutation(t *testing.T) {
	fw, err := New("", nil, nil, nil, 0, 0)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	defer fw.Close()

	addr := netip.MustParseAddr("10.0.0.1")
	if fw.Blocked(addr) {
		t.Fatal("should not be blocked initially")
	}

	if err := fw.AddIP("10.0.0.1"); err != nil {
		t.Fatalf("add IP: %v", err)
	}
	if !fw.Blocked(addr) {
		t.Fatal("should be blocked after AddIP")
	}

	if err := fw.RemoveIP("10.0.0.1"); err != nil {
		t.Fatalf("remove IP: %v", err)
	}
	if fw.Blocked(addr) {
		t.Fatal("should not be blocked after RemoveIP")
	}
}