owpengram-server/internal/mtprotoedge/same_port_mux_test.go

package mtprotoedge

import (
	"bytes"
	"context"
	"errors"
	"io"
	"net"
	"testing"
	"time"
)

func TestSamePortMuxListenerCloseWaitsAndReturnsBacklogAdmission(t *testing.T) {
	admission := newAdmissionController(4, 4, 1)
	listener := &samePortMuxListener{
		addr:   &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)},
		ch:     make(chan net.Conn, 1),
		closed: make(chan struct{}),
	}

	backlog, backlogPeer := trackedMuxPipe(t, admission, 1001)
	defer backlogPeer.Close()
	if !listener.deliver(context.Background(), backlog) {
		t.Fatal("initial backlog delivery was rejected")
	}

	// Deterministically model a producer that passed the delivery gate but has not yet
	// completed. Close must publish closed first, then wait before draining the backlog.
	if !listener.beginDelivery() {
		t.Fatal("in-flight delivery gate unexpectedly closed")
	}
	closeDone := make(chan struct{})
	go func() {
		_ = listener.Close()
		close(closeDone)
	}()
	select {
	case <-listener.closed:
	case <-time.After(time.Second):
		t.Fatal("Close did not publish listener closure")
	}
	select {
	case <-closeDone:
		t.Fatal("Close returned before in-flight delivery completed")
	default:
	}
	listener.deliveryWG.Done()
	select {
	case <-closeDone:
	case <-time.After(time.Second):
		t.Fatal("Close did not finish after delivery completed")
	}

	assertAdmissionConnections(t, admission, 0)

	late, latePeer := trackedMuxPipe(t, admission, 1002)
	defer latePeer.Close()
	if listener.deliver(context.Background(), late) {
		t.Fatal("delivery after Close unexpectedly succeeded")
	}
	_ = late.Close() // dispatch owns and closes a rejected delivery.
	assertAdmissionConnections(t, admission, 0)
}

func TestSamePortMuxPermanentAcceptErrorCancelsBlockedDeliveryBeforeWait(t *testing.T) {
	serverSide, clientSide := net.Pipe()
	defer clientSide.Close()
	wantErr := errors.New("same-port permanent accept failure")
	base := &connThenErrorListener{conn: serverSide, err: wantErr}
	closed := make(chan struct{})
	mux := &samePortMux{
		base:         base,
		addr:         base.Addr(),
		sniffTimeout: time.Hour,
		closed:       closed,
	}
	// An unbuffered child listener deterministically leaves dispatch blocked in deliver: no
	// consumer is running, and the base listener immediately returns a permanent second error.
	mux.tcp = &samePortMuxListener{addr: mux.addr, ch: make(chan net.Conn), closed: make(chan struct{})}
	mux.http = &samePortMuxListener{addr: mux.addr, ch: make(chan net.Conn), closed: make(chan struct{})}

	writeDone := make(chan error, 1)
	go func() {
		_, err := clientSide.Write([]byte{0xef, 0, 0, 0})
		writeDone <- err
	}()
	serveDone := make(chan error, 1)
	go func() {
		serveDone <- mux.Serve(context.Background())
	}()

	select {
	case err := <-serveDone:
		if !errors.Is(err, wantErr) {
			t.Fatalf("Serve error = %v, want %v", err, wantErr)
		}
	case <-time.After(time.Second):
		t.Fatal("same-port Serve waited for blocked delivery before canceling it")
	}
	select {
	case <-writeDone:
	case <-time.After(time.Second):
		t.Fatal("sniff writer remained blocked after same-port shutdown")
	}
}

func TestSamePortMuxShutdownInterruptsSlowSniffImmediately(t *testing.T) {
	tests := []struct {
		name     string
		shutdown func(context.CancelFunc, *samePortMux)
	}{
		{
			name: "context cancel",
			shutdown: func(cancel context.CancelFunc, _ *samePortMux) {
				cancel()
			},
		},
		{
			name: "mux close",
			shutdown: func(_ context.CancelFunc, mux *samePortMux) {
				_ = mux.Close()
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			base, err := net.Listen("tcp", "127.0.0.1:0")
			if err != nil {
				t.Fatalf("listen: %v", err)
			}
			mux := newSamePortMux(base, time.Minute)
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			serveDone := make(chan error, 1)
			go func() { serveDone <- mux.Serve(ctx) }()

			peer, err := net.Dial("tcp", base.Addr().String())
			if err != nil {
				t.Fatalf("dial: %v", err)
			}
			defer peer.Close()
			// No bytes are written: dispatch is blocked in the four-byte sniff with a one-minute
			// deadline. Shutdown must close this accepted socket instead of waiting for it.
			tt.shutdown(cancel, mux)

			select {
			case err := <-serveDone:
				if err != nil {
					t.Fatalf("Serve: %v", err)
				}
			case <-time.After(time.Second):
				t.Fatal("Serve waited for the sniff deadline after shutdown")
			}
			if err := peer.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
				t.Fatalf("set peer deadline: %v", err)
			}
			var one [1]byte
			if _, err := peer.Read(one[:]); err == nil {
				t.Fatal("slow sniff socket remained open after mux shutdown")
			}
		})
	}
}

func TestSamePortMuxSuccessfulHandoffReleasesSniffOwnership(t *testing.T) {
	base, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	mux := newSamePortMux(base, time.Minute)
	ctx, cancel := context.WithCancel(context.Background())
	serveDone := make(chan error, 1)
	go func() { serveDone <- mux.Serve(ctx) }()

	peer, err := net.Dial("tcp", base.Addr().String())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer peer.Close()
	if _, err := peer.Write([]byte{0xef, 0, 0, 0}); err != nil {
		t.Fatalf("write sniff prefix: %v", err)
	}
	accepted, err := mux.TCP().Accept()
	if err != nil {
		t.Fatalf("accept child: %v", err)
	}
	defer accepted.Close()

	// Once dispatch has delivered the Conn, canceling the mux may close listeners/backlog but
	// must not let the old sniff watcher close a socket now owned by the child consumer.
	cancel()
	select {
	case err := <-serveDone:
		if err != nil {
			t.Fatalf("Serve: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("Serve did not stop after cancel")
	}
	if _, err := peer.Write([]byte{1, 2, 3, 4}); err != nil {
		t.Fatalf("write after handoff/shutdown: %v", err)
	}
	if err := accepted.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
		t.Fatalf("set accepted deadline: %v", err)
	}
	got := make([]byte, 8)
	if _, err := io.ReadFull(accepted, got); err != nil {
		t.Fatalf("read handed-off connection: %v", err)
	}
	want := []byte{0xef, 0, 0, 0, 1, 2, 3, 4}
	if !bytes.Equal(got, want) {
		t.Fatalf("handed-off bytes = %x, want %x", got, want)
	}
}

func trackedMuxPipe(t *testing.T, admission *admissionController, port int) (net.Conn, net.Conn) {
	t.Helper()
	server, peer := net.Pipe()
	release, ok := admission.acquireConnection(&net.TCPAddr{
		IP:   net.ParseIP("203.0.113.20"),
		Port: port,
	})
	if !ok {
		_ = server.Close()
		_ = peer.Close()
		t.Fatal("test connection admission rejected")
	}
	return &admittedConn{Conn: server, release: release}, peer
}

func assertAdmissionConnections(t *testing.T, admission *admissionController, want int) {
	t.Helper()
	admission.mu.Lock()
	got := admission.connections
	byIP := len(admission.byIP)
	admission.mu.Unlock()
	if got != want || (want == 0 && byIP != 0) {
		t.Fatalf("admission state = connections:%d by_ip:%d, want connections:%d", got, byIP, want)
	}
}