SimpleRemoter/server/go/server/server_test.go

package server

import (
	"bytes"
	"net"
	"testing"
	"time"
)

func TestIPLimiterConcurrentCap(t *testing.T) {
	l := newIPLimiter(2, 0, 0) // 2 concurrent, no rate limit

	if !l.allow("10.0.0.1") {
		t.Fatal("first connection from an IP should be allowed")
	}
	if !l.allow("10.0.0.1") {
		t.Fatal("second connection from an IP should be allowed")
	}
	if l.allow("10.0.0.1") {
		t.Fatal("third concurrent connection from the same IP should be rejected")
	}

	// A different IP is unaffected.
	if !l.allow("10.0.0.2") {
		t.Fatal("a different IP should not share the first IP's cap")
	}

	// Releasing one slot frees capacity again.
	l.release("10.0.0.1")
	if !l.allow("10.0.0.1") {
		t.Fatal("after release, the IP should have a free slot")
	}
}

func TestIPLimiterRateLimit(t *testing.T) {
	l := newIPLimiter(0, 3, time.Minute) // no concurrent cap, 3/min per IP

	for range 3 {
		if !l.allow("10.0.0.1") {
			t.Fatal("an attempt within budget should be allowed")
		}
	}
	if l.allow("10.0.0.1") {
		t.Fatal("4th connection within the window should be rejected")
	}
	// A different IP has its own budget.
	if !l.allow("10.0.0.2") {
		t.Fatal("a different IP should have its own rate budget")
	}
}

func TestIPLimiterEmptyIPBypasses(t *testing.T) {
	l := newIPLimiter(1, 1, time.Minute)
	// An empty peer IP (shouldn't happen for real TCP) must not be throttled,
	// so a caller that loses the IP falls back to accepting.
	if !l.allow("") {
		t.Fatal("empty IP should bypass the limiter")
	}
}

func TestIPLimiterSweep(t *testing.T) {
	l := newIPLimiter(0, 1, time.Minute)

	l.mu.Lock()
	l.recent["stale"] = []time.Time{time.Now().Add(-time.Hour)}
	l.recent["fresh"] = []time.Time{time.Now()}
	l.mu.Unlock()

	l.sweep()

	l.mu.Lock()
	defer l.mu.Unlock()
	if _, ok := l.recent["stale"]; ok {
		t.Fatal("stale rate-limit entry should be swept")
	}
	if _, ok := l.recent["fresh"]; !ok {
		t.Fatal("fresh rate-limit entry should survive sweep")
	}
}

func TestPeerIP(t *testing.T) {
	tests := []struct {
		addr net.Addr
		want string
	}{
		{&net.TCPAddr{IP: net.ParseIP("1.2.3.4"), Port: 1234}, "1.2.3.4"},
		{&net.TCPAddr{IP: net.ParseIP("::1"), Port: 6543}, "::1"},
	}
	for _, tt := range tests {
		if got := peerIP(tt.addr); got != tt.want {
			t.Fatalf("peerIP(%v) = %q, want %q", tt.addr, got, tt.want)
		}
	}
}

// A peer that sends bytes which never assemble into a complete packet (here an
// HTTP-mask envelope with no terminating CRLF CRLF) must be disconnected once
// the per-connection buffer cap is exceeded, rather than allowed to grow it
// without bound.
func TestHandleConnectionBufferCap(t *testing.T) {
	cfg := DefaultConfig()
	cfg.MaxBufferedBytes = 1024 // small cap for a fast, in-memory test
	cfg.ReadTimeout = time.Second

	srv := New(cfg)
	srv.running.Store(true) // bypass acceptLoop; exercise handleConnection directly

	client, serverConn := net.Pipe()
	defer client.Close()

	done := make(chan struct{})
	go func() {
		srv.handleConnection(serverConn)
		close(done)
	}()

	// "POST " prefix with no terminator, plus filler to exceed the cap.
	payload := append([]byte("POST /x HTTP/1.1\r\n"), bytes.Repeat([]byte{'A'}, 4096)...)
	writeErr := make(chan error, 1)
	go func() {
		_, err := client.Write(payload)
		writeErr <- err
	}()

	select {
	case <-done:
		// handleConnection returned because the cap was exceeded and it closed
		// the connection.
	case <-time.After(2 * time.Second):
		t.Fatal("connection was not closed after exceeding the buffer cap")
	}
}