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,258 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.avsdev.uk/Go/job-search/internal/config"
|
||||
ws "github.com/coder/websocket"
|
||||
)
|
||||
|
||||
// fakeConn is a msgConn that records how it was closed. Read blocks until its
|
||||
// context is cancelled.
|
||||
type fakeConn struct {
|
||||
closed chan ws.StatusCode
|
||||
}
|
||||
|
||||
func newFakeConn() *fakeConn { return &fakeConn{closed: make(chan ws.StatusCode, 1)} }
|
||||
|
||||
func (f *fakeConn) Read(ctx context.Context) (ws.MessageType, []byte, error) {
|
||||
<-ctx.Done()
|
||||
return 0, nil, ctx.Err()
|
||||
}
|
||||
|
||||
func (f *fakeConn) Write(context.Context, ws.MessageType, []byte) error { return nil }
|
||||
|
||||
func (f *fakeConn) Ping(context.Context) error { return nil }
|
||||
|
||||
func (f *fakeConn) Close(code ws.StatusCode, _ string) error {
|
||||
f.closed <- code
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeConn) CloseNow() error { return nil }
|
||||
|
||||
// newTestClient returns a client of a new hub on f, with no loops running,
|
||||
// and its context.
|
||||
func newTestClient(f *fakeConn) (*client, context.Context) {
|
||||
h := NewHub(config.JSCServerWebSocket{}, discardLog)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
return newClient(h, f, "192.0.2.1", cancel), ctx
|
||||
}
|
||||
|
||||
// An envelope that can't be marshalled is dropped, not queued.
|
||||
func TestSendEnvelopeBadPayload(t *testing.T) {
|
||||
c, _ := newTestClient(newFakeConn())
|
||||
c.sendEnvelope(Envelope{Type: TypeEcho, Payload: json.RawMessage("{")})
|
||||
if len(c.send) != 0 {
|
||||
t.Errorf("queued %d messages; want none", len(c.send))
|
||||
}
|
||||
}
|
||||
|
||||
// A client whose queue is full is closed as a slow consumer.
|
||||
func TestSendEnvelopeSlowConsumer(t *testing.T) {
|
||||
f := newFakeConn()
|
||||
c, ctx := newTestClient(f)
|
||||
for range clientSendBuffer {
|
||||
c.sendEnvelope(Envelope{Type: TypePong})
|
||||
}
|
||||
select {
|
||||
case code := <-f.closed:
|
||||
t.Fatalf("closed with %v before the queue was full", code)
|
||||
default:
|
||||
}
|
||||
|
||||
c.sendEnvelope(Envelope{Type: TypePong})
|
||||
select {
|
||||
case code := <-f.closed:
|
||||
if code != ws.StatusPolicyViolation {
|
||||
t.Errorf("closed with %v; want policy violation", code)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("slow consumer not closed")
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("slow consumer not cancelled")
|
||||
}
|
||||
}
|
||||
|
||||
// close closes the connection once, however often it is called.
|
||||
func TestClientCloseOnce(t *testing.T) {
|
||||
f := newFakeConn()
|
||||
c, _ := newTestClient(f)
|
||||
c.close(ws.StatusGoingAway, "bye")
|
||||
c.close(ws.StatusNormalClosure, "again")
|
||||
if code := <-f.closed; code != ws.StatusGoingAway {
|
||||
t.Errorf("closed with %v; want going away", code)
|
||||
}
|
||||
select {
|
||||
case code := <-f.closed:
|
||||
t.Errorf("closed a second time with %v", code)
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// The message rate limit applies only when it is enabled and has a rate.
|
||||
func TestClientMessageRateLimit(t *testing.T) {
|
||||
tests := []struct {
|
||||
rule config.JSCRateLimitRule
|
||||
limited bool
|
||||
}{
|
||||
{config.JSCRateLimitRule{Enabled: false, PerSecond: 1, Burst: 1}, false},
|
||||
{config.JSCRateLimitRule{Enabled: true, PerSecond: 0, Burst: 1}, false},
|
||||
{config.JSCRateLimitRule{Enabled: true, PerSecond: 1, Burst: 1}, true},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
h := NewHub(config.JSCServerWebSocket{MessageRateLimit: tt.rule}, discardLog)
|
||||
c := newClient(h, newFakeConn(), "", func() {})
|
||||
c.limiter.Allow()
|
||||
if limited := !c.limiter.Allow(); limited != tt.limited {
|
||||
t.Errorf("rule %+v: second message limited = %v; want %v", tt.rule, limited, tt.limited)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMustRaw(t *testing.T) {
|
||||
raw := json.RawMessage(`{"a":1}`)
|
||||
tests := []struct {
|
||||
name string
|
||||
v any
|
||||
want string
|
||||
}{
|
||||
{"nil", nil, ""},
|
||||
{"raw message as is", raw, `{"a":1}`},
|
||||
{"marshalled", map[string]int{"b": 2}, `{"b":2}`},
|
||||
{"unmarshallable", func() {}, ""},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
if got := string(mustRaw(tt.v)); got != tt.want {
|
||||
t.Errorf("%s: mustRaw = %q; want %q", tt.name, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithTimeout(t *testing.T) {
|
||||
ctx, cancel := withTimeout(context.Background(), 0)
|
||||
defer cancel()
|
||||
if _, ok := ctx.Deadline(); ok {
|
||||
t.Error("timeout 0 set a deadline")
|
||||
}
|
||||
ctx, cancel = withTimeout(context.Background(), time.Minute)
|
||||
defer cancel()
|
||||
if _, ok := ctx.Deadline(); !ok {
|
||||
t.Error("timeout 1m set no deadline")
|
||||
}
|
||||
}
|
||||
|
||||
// syncBuffer is a bytes.Buffer safe for the client's goroutines to log to.
|
||||
type syncBuffer struct {
|
||||
mu sync.Mutex
|
||||
buf bytes.Buffer
|
||||
}
|
||||
|
||||
func (b *syncBuffer) Write(p []byte) (int, error) {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
return b.buf.Write(p)
|
||||
}
|
||||
|
||||
// entries returns the JSON log entries written so far.
|
||||
func (b *syncBuffer) entries(t *testing.T) []map[string]any {
|
||||
t.Helper()
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
var entries []map[string]any
|
||||
for line := range strings.Lines(b.buf.String()) {
|
||||
var e map[string]any
|
||||
if err := json.Unmarshal([]byte(line), &e); err != nil {
|
||||
t.Fatalf("log line %q: %v", line, err)
|
||||
}
|
||||
entries = append(entries, e)
|
||||
}
|
||||
return entries
|
||||
}
|
||||
|
||||
// jsonLog returns a logger writing JSON entries at level and above to a
|
||||
// buffer.
|
||||
func jsonLog(level slog.Level) (*slog.Logger, *syncBuffer) {
|
||||
buf := &syncBuffer{}
|
||||
return slog.New(slog.NewJSONHandler(buf, &slog.HandlerOptions{Level: level})), buf
|
||||
}
|
||||
|
||||
// At debug level, messages received from and sent to a client are logged.
|
||||
func TestLogMessages(t *testing.T) {
|
||||
log, buf := jsonLog(slog.LevelDebug)
|
||||
h := NewHub(config.JSCServerWebSocket{WriteTimeout: time.Second}, log)
|
||||
peer, _ := pipeClient(t, h)
|
||||
|
||||
ping := `{"type":"ping","id":"1"}`
|
||||
go io.WriteString(peer, ping+"\n")
|
||||
peer.SetReadDeadline(time.Now().Add(time.Second))
|
||||
pong, err := bufio.NewReader(peer).ReadString('\n')
|
||||
if err != nil {
|
||||
t.Fatalf("pong: %v", err)
|
||||
}
|
||||
pong = strings.TrimSuffix(pong, "\n")
|
||||
|
||||
// The write is logged after it returns, so wait for its entry.
|
||||
want := map[string]string{
|
||||
"message received": ping,
|
||||
"message sent": pong,
|
||||
}
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for {
|
||||
found := map[string]bool{}
|
||||
for _, e := range buf.entries(t) {
|
||||
if msg, ok := e["msg"].(string); ok && want[msg] == e["message"] {
|
||||
found[msg] = true
|
||||
if e["bytes"] != float64(len(want[msg])) || e["truncated"] != false {
|
||||
t.Errorf("entry %v: bad bytes or truncated", e)
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(found) == len(want) {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("logged %v; want received ping and sent pong", buf.entries(t))
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
// A long message is logged cut to logMessageLimit bytes; nothing is logged
|
||||
// unless debug logging is on.
|
||||
func TestLogMessageLimit(t *testing.T) {
|
||||
log, buf := jsonLog(slog.LevelDebug)
|
||||
h := NewHub(config.JSCServerWebSocket{}, log)
|
||||
c := newClient(h, newFakeConn(), "", func() {})
|
||||
long := strings.Repeat("x", logMessageLimit+100)
|
||||
c.logMessage(context.Background(), "message sent", []byte(long))
|
||||
|
||||
entries := buf.entries(t)
|
||||
if len(entries) != 1 {
|
||||
t.Fatalf("logged %d entries; want 1", len(entries))
|
||||
}
|
||||
e := entries[0]
|
||||
if e["message"] != long[:logMessageLimit] || e["bytes"] != float64(len(long)) ||
|
||||
e["truncated"] != true || e["client"] != c.id {
|
||||
t.Errorf("entry = %v", e)
|
||||
}
|
||||
|
||||
log, buf = jsonLog(slog.LevelInfo)
|
||||
c.hub = NewHub(config.JSCServerWebSocket{}, log)
|
||||
c.logMessage(context.Background(), "message sent", []byte("x"))
|
||||
if entries := buf.entries(t); len(entries) != 0 {
|
||||
t.Errorf("logged %v at info level", entries)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user