package tcpserver import ( "bufio" "context" "encoding/json" "errors" "io" "log/slog" "net" "strings" "syscall" "testing" "time" "git.avsdev.uk/Go/job-search/internal/config" "git.avsdev.uk/Go/job-search/internal/websocket" ) var discardLog = slog.New(slog.NewTextHandler(io.Discard, nil)) // testConf is a server configuration for tests, changed by mutate if it is // not nil. func testConf(mutate func(*config.JSCServer)) *config.JSCServer { c := &config.JSCServer{ Host: "127.0.0.1", Port: 0, WebSocket: config.JSCServerWebSocket{WriteTimeout: time.Second}, Timeouts: config.JSCServerTimeouts{ShutdownTimeout: time.Second}, } if mutate != nil { mutate(c) } return c } // newServer returns a server for conf with its own hub. func newServer(conf *config.JSCServer) (*Server, *websocket.Hub) { h := websocket.NewHub(conf.WebSocket, websocket.WithLogger(discardLog)) return New(conf, h, WithLogger(discardLog)), h } // startServer serves on a local port, returning the server's hub, its // address, a function to cancel it and a channel with serve's result. func startServer( t *testing.T, conf *config.JSCServer, ) (*websocket.Hub, string, context.CancelFunc, chan error) { t.Helper() s, h := newServer(conf) ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatal(err) } ctx, cancel := context.WithCancel(context.Background()) t.Cleanup(cancel) done := make(chan error, 1) go func() { done <- s.serve(ctx, ln) }() return h, ln.Addr().String(), cancel, done } // tcpClient is a plain TCP test client reading newline-delimited envelopes. type tcpClient struct { t *testing.T conn net.Conn r *bufio.Reader } func dial(t *testing.T, addr string) *tcpClient { t.Helper() conn, err := net.Dial("tcp", addr) if err != nil { t.Fatal(err) } t.Cleanup(func() { conn.Close() }) return &tcpClient{t: t, conn: conn, r: bufio.NewReader(conn)} } // read returns the next envelope, failing the test after a second. func (c *tcpClient) read() websocket.Envelope { c.t.Helper() c.conn.SetReadDeadline(time.Now().Add(time.Second)) line, err := c.r.ReadString('\n') if err != nil { c.t.Fatalf("read: %v (got %q)", err, line) } var env websocket.Envelope if err := json.Unmarshal([]byte(line), &env); err != nil { c.t.Fatalf("bad envelope %q: %v", line, err) } return env } // waitClosed checks the server closes the connection, skipping any lines // sent first. A reset counts as closed. func (c *tcpClient) waitClosed() { c.t.Helper() c.conn.SetReadDeadline(time.Now().Add(time.Second)) if _, err := io.ReadAll(c.r); err != nil && !errors.Is(err, syscall.ECONNRESET) { c.t.Fatalf("connection not closed: %v", err) } } // Clients get pushed messages; cancelling the server's context tells them it // is shutting down, closes them, stops accepting and returns nil. func TestServerShutdown(t *testing.T) { h, addr, cancel, done := startServer(t, testConf(nil)) c := dial(t, addr) if e := c.read(); e.Type != websocket.TypeWelcome { t.Fatalf("want welcome, got %+v", e) } for h.Count() != 1 { time.Sleep(time.Millisecond) } h.Push("item.created", map[string]string{"id": "x"}) if e := c.read(); e.Type != "item.created" { t.Fatalf("want item.created, got %+v", e) } cancel() if e := c.read(); e.Type != websocket.TypeError || e.Error != "server shutting down" { t.Fatalf("want shutdown error, got %+v", e) } c.waitClosed() select { case err := <-done: if err != nil { t.Fatalf("serve = %v", err) } case <-time.After(2 * time.Second): t.Fatal("serve did not return") } if conn, err := net.Dial("tcp", addr); err == nil { conn.Close() t.Error("listener still accepting") } } // Over websocket.connection_limit, a client is told it is rate limited and // closed; other clients are unaffected. func TestServerConnectionLimit(t *testing.T) { conf := testConf(func(c *config.JSCServer) { c.WebSocket.ConnectionLimit = config.JSCRateLimitRule{ Enabled: true, PerSecond: 1, Burst: 1, } c.RateLimits.IdleTTL = time.Minute }) _, addr, cancel, _ := startServer(t, conf) defer cancel() first := dial(t, addr) if e := first.read(); e.Type != websocket.TypeWelcome { t.Fatalf("first: want welcome, got %+v", e) } second := dial(t, addr) if e := second.read(); e.Type != websocket.TypeError || e.Error != "rate limited" { t.Fatalf("second: want rate limited, got %+v", e) } second.waitClosed() first.conn.Write([]byte(`{"type":"ping","id":"1"}` + "\n")) if e := first.read(); e.Type != websocket.TypePong { t.Fatalf("first: want pong, got %+v", e) } } // With the connection limit off (or no rate), there is no limiter. func TestServerConnectionLimitOff(t *testing.T) { for _, rule := range []config.JSCRateLimitRule{ {Enabled: false, PerSecond: 1, Burst: 1}, {Enabled: true, PerSecond: 0, Burst: 1}, } { s, _ := newServer(testConf(func(c *config.JSCServer) { c.WebSocket.ConnectionLimit = rule })) if l := s.connectionLimiter(context.Background()); l != nil { t.Errorf("rule %+v: got a limiter", rule) } } } func TestServerRun(t *testing.T) { s, _ := newServer(testConf(nil)) ctx, cancel := context.WithCancel(context.Background()) time.AfterFunc(20*time.Millisecond, cancel) if err := s.Run(ctx); err != nil { t.Fatalf("Run = %v", err) } } func TestServerRunListenError(t *testing.T) { s, _ := newServer(testConf(func(c *config.JSCServer) { c.Port = -1 })) if err := s.Run(context.Background()); err == nil || !strings.Contains(err.Error(), "listen") { t.Fatalf("Run = %v; want a listen error", err) } } // flakyListener fails Accept with a temporary error a number of times, then // as if closed. type flakyListener struct { net.Listener fails int } func (l *flakyListener) Accept() (net.Conn, error) { if l.fails > 0 { l.fails-- return nil, errors.New("too many open files") } return nil, net.ErrClosed } // Accept errors other than a closed listener are retried after a back off. func TestServerAcceptRetries(t *testing.T) { s, _ := newServer(testConf(nil)) ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatal(err) } defer ln.Close() fl := &flakyListener{Listener: ln, fails: 3} err = s.accept(context.Background(), context.Background(), fl) if !errors.Is(err, net.ErrClosed) || fl.fails != 0 { t.Fatalf("accept = %v with %d failures left; want net.ErrClosed", err, fl.fails) } } // Cancelling ctx stops accept while it is backing off. func TestServerAcceptCancelledInBackOff(t *testing.T) { s, _ := newServer(testConf(nil)) ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatal(err) } defer ln.Close() ctx, cancel := context.WithCancel(context.Background()) fl := &flakyListener{Listener: ln, fails: 1 << 30} time.AfterFunc(20*time.Millisecond, cancel) if err := s.accept(ctx, ctx, fl); !errors.Is(err, context.Canceled) { t.Fatalf("accept = %v; want context.Canceled", err) } } // A listener closed from elsewhere is an error, after a graceful shutdown. func TestServerListenerClosed(t *testing.T) { s, _ := newServer(testConf(nil)) ln, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatal(err) } ln.Close() if err := s.serve(context.Background(), ln); !errors.Is(err, net.ErrClosed) { t.Fatalf("serve = %v; want net.ErrClosed", err) } } func TestRemoteIP(t *testing.T) { a, b := net.Pipe() defer a.Close() defer b.Close() if ip := remoteIP(a); ip != "" { t.Errorf("remoteIP(pipe) = %q; want empty", ip) } }