// Copyright (c) 2026 Petr BalvĂ­n (https://petrbalvin.org) // SPDX-License-Identifier: MIT package server import ( "context" "errors" "io" "net" "testing" "time" ) func dial(t *testing.T, addr string) net.Conn { t.Helper() conn, err := net.Dial("tcp", addr) if err != nil { t.Fatalf("dial %s: %v", addr, err) } t.Cleanup(func() { conn.Close() }) return conn } func TestServeHandsConnectionsToHandle(t *testing.T) { ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatalf("listen: %v", err) } accepted := make(chan net.Conn, 2) s := &Server{Handle: func(conn net.Conn) { accepted <- conn conn.Close() }} ctx, cancel := context.WithCancel(context.Background()) serveErr := make(chan error, 1) go func() { serveErr <- s.Serve(ctx, ln) }() for range 2 { conn := dial(t, ln.Addr().String()) select { case got := <-accepted: if got.RemoteAddr() == nil { t.Fatal("a connection without a remote address arrived") } case <-time.After(2 * time.Second): t.Fatal("the connection did not reach Handle") } // Handle closed the connection, so the next read reports it. if _, err := conn.Read(make([]byte, 1)); !errors.Is(err, io.EOF) { t.Fatalf("read after Handle returned: %v", err) } } cancel() select { case err := <-serveErr: if err != nil { t.Fatalf("Serve returned %v after cancellation", err) } case <-time.After(2 * time.Second): t.Fatal("Serve did not return after cancellation") } } func TestServeClosesWithoutHandle(t *testing.T) { ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatalf("listen: %v", err) } s := &Server{} ctx, cancel := context.WithCancel(context.Background()) serveErr := make(chan error, 1) go func() { serveErr <- s.Serve(ctx, ln) }() conn := dial(t, ln.Addr().String()) if _, err := conn.Read(make([]byte, 1)); !errors.Is(err, io.EOF) { t.Fatalf("a connection without Handle read %v, want EOF", err) } cancel() select { case err := <-serveErr: if err != nil { t.Fatalf("Serve returned %v after cancellation", err) } case <-time.After(2 * time.Second): t.Fatal("Serve did not return after cancellation") } } func TestServeListenerFailure(t *testing.T) { ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatalf("listen: %v", err) } if err := ln.Close(); err != nil { t.Fatalf("close: %v", err) } // A closed listener is a stop, not a failure: the shutdown raced the // accept loop, and the daemon exits cleanly either way. s := &Server{} if err := s.Serve(context.Background(), ln); err != nil { t.Fatalf("Serve returned %v for a closed listener, want nil", err) } } func TestServeCapsConnections(t *testing.T) { ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatalf("listen: %v", err) } // The first connection holds its slot until released, so the state of // the cap is deterministic for the second. release := make(chan struct{}) s := &Server{MaxConns: 1, Handle: func(conn net.Conn) { <-release conn.Close() }} ctx, cancel := context.WithCancel(context.Background()) serveErr := make(chan error, 1) go func() { serveErr <- s.Serve(ctx, ln) }() first := dial(t, ln.Addr().String()) // The refused connection ends at once: the read sees the close, not a // timeout. refused := dial(t, ln.Addr().String()) refused.SetReadDeadline(time.Now().Add(2 * time.Second)) if _, err := refused.Read(make([]byte, 1)); !errors.Is(err, io.EOF) { t.Fatalf("a connection above the cap read %v, want the immediate end", err) } close(release) first.Close() cancel() select { case err := <-serveErr: if err != nil { t.Fatalf("Serve returned %v after cancellation", err) } case <-time.After(2 * time.Second): t.Fatal("Serve did not return after cancellation") } }