Added websocket hub for handling websocket clients, with a wsserve command
which runs a plain TCP server for receiving websocket hub clients, with server config, CSRF and rate limit middleware. - internal/websocket: Hub pushing messages to WebSocket clients (ServeHTTP, example.com/ws upgrade with optional CSRF check) and plain TCP clients (ServeTCP, newline-delimited JSON envelopes); ping/pong keepalive, per-client message rate limit, slow consumer handling, graceful Shutdown, debug logging of messages sent and received - internal/websocket/tcpserver: standalone TCP server for a hub, with accept back-off, a per-IP connection limit (server.websocket.connection_limit) and graceful shutdown - internal/cli/wsserve: `wsserve` command running the TCP server, with --listen/--host/--port overrides - internal/middleware: keyed token bucket rate limiter, real client IP from trusted proxies - internal/httpx: RFC 9457 problem responses, request IDs in context - internal/config: server section (host, port, shutdown timeout, rate limiter idle TTL, websocket); util.ByteSize settings (e.g. "32KiB"); lists of sections as single keys; empty lists stored as nil; isExplicitlySet for printLoadedConfig; JobSearchConfigSource renamed to JSCSourcesSource; TestSections checks every section's methods - internal/util: ByteSize, parsed and formatted with go-humanize
This commit is contained in:
bwilliams committed
1 parent
803a8929e6
commit
c9f42b62ad
42 files changed
+3619
-57
No files matched your search
@@ -0,0 +1,263 @@
|
||||
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, discardLog)
|
||||
return New(conf, h, 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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user