diff --git a/README.md b/README.md index 3f174d5b..f9f00241 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,7 @@ go get github.com/coder/websocket - [net.Conn](https://pkg.go.dev/github.com/coder/websocket#NetConn) wrapper - [Ping pong](https://pkg.go.dev/github.com/coder/websocket#Conn.Ping) API - [RFC 7692](https://tools.ietf.org/html/rfc7692) permessage-deflate compression +- [Prepared message](https://pkg.go.dev/github.com/coder/websocket#PreparedMessage) for efficient broadcasting - [CloseRead](https://pkg.go.dev/github.com/coder/websocket#Conn.CloseRead) helper for write only connections - Compile to [Wasm](https://pkg.go.dev/github.com/coder/websocket#hdr-Wasm) @@ -48,10 +49,11 @@ See GitHub issues for minor issues but the major future enhancements are: ## Examples -For a production quality example that demonstrates the complete API, see the -[echo example](./internal/examples/echo). - -For a full stack example, see the [chat example](./internal/examples/chat). +- Production quality example that demonstrates the complete API, see the + [echo example](./internal/examples/echo). +- Full stack example, see the [chat example](./internal/examples/chat). +- Broadcasting a message to many connections, see the + [broadcast example](./internal/examples/broadcast). ### Server @@ -107,7 +109,6 @@ c.Close(websocket.StatusNormalClosure, "") Advantages of [gorilla/websocket](https://github.com/gorilla/websocket): - Mature and widely used -- [Prepared writes](https://pkg.go.dev/github.com/gorilla/websocket#PreparedMessage) - Configurable [buffer sizes](https://pkg.go.dev/github.com/gorilla/websocket#hdr-Buffers) Advantages of github.com/coder/websocket: diff --git a/compress.go b/compress.go index 41bd5bdb..c8b67469 100644 --- a/compress.go +++ b/compress.go @@ -41,10 +41,13 @@ const ( // // This means less efficient compression as the sliding window from previous messages will not be used but the // memory overhead will be lower as there will be no fixed cost for the flate.Writer nor the 32 KB sliding window. - // Especially if the connections are long lived and seldom written to. + // Especially if the connections are long-lived and seldom written to. // // Thus, it uses less memory than CompressionContextTakeover but compresses less efficiently. // + // A message written with Conn.Write or Conn.WritePrepared is sent uncompressed if compression does not make it + // smaller, such as already compressed data. + // // If the peer does not support CompressionNoContextTakeover then we will fall back to CompressionDisabled. CompressionNoContextTakeover ) diff --git a/compress_test.go b/compress_test.go index 81adce9a..c09f61a8 100644 --- a/compress_test.go +++ b/compress_test.go @@ -76,6 +76,8 @@ func TestWriteSingleFrameCompressed(t *testing.T) { largeMsg = []byte(strings.Repeat("hello world ", 100)) smallMsg = []byte("small message") + // Random bytes do not compress, so compression makes them larger. + incompressibleMsg = xrand.Bytes(1024) ) testCases := []struct { @@ -88,6 +90,10 @@ func TestWriteSingleFrameCompressed(t *testing.T) { {"NoContextTakeover/AboveThreshold", CompressionNoContextTakeover, largeMsg, true}, {"ContextTakeover/BelowThreshold", CompressionContextTakeover, smallMsg, false}, {"NoContextTakeover/BelowThreshold", CompressionNoContextTakeover, smallMsg, false}, + // With context takeover, the message is in the sliding window, so it + // stays compressed to keep the peer's window in sync. + {"ContextTakeover/Incompressible", CompressionContextTakeover, incompressibleMsg, true}, + {"NoContextTakeover/Incompressible", CompressionNoContextTakeover, incompressibleMsg, false}, } for _, tc := range testCases { @@ -126,6 +132,9 @@ func TestWriteSingleFrameCompressed(t *testing.T) { assert.Equal(t, "opcode", opText, h.opcode) assert.Equal(t, "rsv1 (compressed)", tc.wantRsv1, h.rsv1) + if !tc.wantRsv1 { + assert.Equal(t, "payload length", int64(len(tc.msg)), h.payloadLength) + } assert.Equal(t, "fin", true, h.fin) err = <-writeDone diff --git a/conn_test.go b/conn_test.go index 01af8cd4..8d5fd26c 100644 --- a/conn_test.go +++ b/conn_test.go @@ -50,9 +50,14 @@ func TestConn(t *testing.T) { c1.SetReadLimit(131072) + // Interleave prepared writes with regular writes, as they share + // the compression state of the connection. for range 5 { err := wstest.Echo(tt.ctx, c1, 131072) assert.Success(t, err) + + err = wstest.EchoPrepared(tt.ctx, c1, 131072) + assert.Success(t, err) } err := c1.Close(websocket.StatusNormalClosure, "") @@ -61,6 +66,69 @@ func TestConn(t *testing.T) { } }) + t.Run("writePreparedConcurrent", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), time.Second*30) + defer cancel() + + exp := strings.Repeat("prepared", 128) + pm := websocket.NewPreparedMessage(websocket.MessageText, []byte(exp)) + + modes := []websocket.CompressionMode{ + websocket.CompressionDisabled, + websocket.CompressionContextTakeover, + websocket.CompressionNoContextTakeover, + } + + // Write the same message with every compression mode from both sides, + // so connections with different settings share it and client masking + // runs concurrently on the shared payload. + var errs []<-chan error + for i := range 12 { + mode := modes[i%len(modes)] + client, server := wstest.Pipe(&websocket.DialOptions{ + CompressionMode: mode, + }, &websocket.AcceptOptions{ + CompressionMode: mode, + }) + t.Cleanup(func() { + client.CloseNow() + server.CloseNow() + }) + + w, r := client, server + if i%2 == 0 { + w, r = server, client + } + errs = append(errs, xsync.Go(func() error { + for range 5 { + err := w.WritePrepared(ctx, pm) + if err != nil { + return err + } + } + return nil + })) + errs = append(errs, xsync.Go(func() error { + for range 5 { + _, p, err := r.Read(ctx) + if err != nil { + return err + } + if string(p) != exp { + return fmt.Errorf("unexpected msg: %q", p) + } + } + return nil + })) + } + + for _, errc := range errs { + assert.Success(t, <-errc) + } + }) + t.Run("badClose", func(t *testing.T) { tt, c1, c2 := newConnTest(t, nil, nil) @@ -667,6 +735,85 @@ func BenchmarkConn(b *testing.B) { } } +func BenchmarkBroadcast(b *testing.B) { + const conns = 16 + + modes := []struct { + name string + mode websocket.CompressionMode + }{ + {"disabledCompress", websocket.CompressionDisabled}, + {"compressContextTakeover", websocket.CompressionContextTakeover}, + {"compressNoContext", websocket.CompressionNoContextTakeover}, + } + writes := []struct { + name string + write func(ctx context.Context, cs []*websocket.Conn, msg []byte) error + }{ + {"write", func(ctx context.Context, cs []*websocket.Conn, msg []byte) error { + for _, c := range cs { + err := c.Write(ctx, websocket.MessageText, msg) + if err != nil { + return err + } + } + return nil + }}, + {"writePrepared", func(ctx context.Context, cs []*websocket.Conn, msg []byte) error { + pm := websocket.NewPreparedMessage(websocket.MessageText, msg) + for _, c := range cs { + err := c.WritePrepared(ctx, pm) + if err != nil { + return err + } + } + return nil + }}, + } + + msg := []byte(strings.Repeat("1234", 256)) + + // Clients mask every write, so they are benchmarked separately. + for _, side := range []string{"server", "client"} { + for _, m := range modes { + for _, w := range writes { + b.Run(side+"/"+m.name+"/"+w.name, func(b *testing.B) { + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + writers := make([]*websocket.Conn, conns) + for i := range writers { + client, server := wstest.Pipe(&websocket.DialOptions{ + CompressionMode: m.mode, + }, &websocket.AcceptOptions{ + CompressionMode: m.mode, + }) + b.Cleanup(func() { + client.CloseNow() + server.CloseNow() + }) + writers[i] = server + if side == "client" { + writers[i] = client + } + writers[i].DiscardWrites() + } + + b.SetBytes(int64(len(msg) * conns)) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + err := w.write(ctx, writers, msg) + if err != nil { + b.Fatal(err) + } + } + }) + } + } + } +} + func echoServer(w http.ResponseWriter, r *http.Request, opts *websocket.AcceptOptions) (err error) { defer errd.Wrap(&err, "echo server failed") diff --git a/example_test.go b/example_test.go index 7026e311..e43c4810 100644 --- a/example_test.go +++ b/example_test.go @@ -118,7 +118,7 @@ func Example_writeOnly() { } func Example_crossOrigin() { - // This handler demonstrates how to safely accept cross origin WebSockets + // This handler demonstrates how to safely accept cross-origin WebSockets // from the origin example.com. fn := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { c, err := websocket.Accept(w, r, &websocket.AcceptOptions{ @@ -165,7 +165,13 @@ func Example_fullStackChat() { // https://github.com/nhooyr/websocket/tree/master/internal/examples/chat } -// This example demonstrates a echo server. +// This example demonstrates an echo server. func Example_echo() { // https://github.com/nhooyr/websocket/tree/master/internal/examples/echo } + +// This example demonstrates broadcasting messages to many subscribers with +// a PreparedMessage. +func ExampleConn_WritePrepared() { + // https://github.com/coder/websocket/tree/master/internal/examples/broadcast +} diff --git a/export_test.go b/export_test.go index e0071922..41ceb59c 100644 --- a/export_test.go +++ b/export_test.go @@ -3,6 +3,7 @@ package websocket import ( + "io" "net" "github.com/coder/websocket/internal/util" @@ -17,6 +18,10 @@ func (c *Conn) RecordBytesWritten() *int { return &bytesWritten } +func (c *Conn) DiscardWrites() { + c.bw.Reset(io.Discard) +} + func (c *Conn) RecordBytesRead() *int { var bytesRead int c.br.Reset(util.ReaderFunc(func(p []byte) (int, error) { diff --git a/internal/examples/broadcast/README.md b/internal/examples/broadcast/README.md new file mode 100644 index 00000000..e1e5d1d3 --- /dev/null +++ b/internal/examples/broadcast/README.md @@ -0,0 +1,38 @@ +# Broadcast Example + +This directory contains a broadcast server example using github.com/coder/websocket. + +```bash +$ cd internal/examples/broadcast +$ go run . localhost:8080 +listening on http://127.0.0.1:8080 +``` + +Subscribe with a WebSocket client like [websocat](https://github.com/vi/websocat) and publish +messages with curl: + +```bash +$ websocat ws://127.0.0.1:8080/subscribe +$ curl --data-binary 'hello' http://127.0.0.1:8080/publish +``` + +Every published message is delivered to all subscribers. Messages are sent as text, so +`/publish` rejects bodies that are not valid UTF-8. + +## Structure + +The server is in `server.go`. Subscribers connect to the WebSocket `/subscribe` endpoint and +messages are published via the HTTP POST `/publish` endpoint, so that you can easily publish +with curl. + +Each published message is wrapped in a single `PreparedMessage` and queued for every subscriber. +Connections are accepted with `CompressionNoContextTakeover`, so a message that reaches the +compression threshold is compressed once per broadcast instead of once per subscriber. +Subscribers that do not support compression receive the same message uncompressed. + +Publishing never blocks: a subscriber that cannot keep up with its queue is disconnected. + +`server_test.go` contains a test that subscribes clients with and without compression and +ensures every client receives every published message, while invalid UTF-8 is rejected. + +`main.go` brings it all together so that you can run it and play around with it. diff --git a/internal/examples/broadcast/main.go b/internal/examples/broadcast/main.go new file mode 100644 index 00000000..1edb37a5 --- /dev/null +++ b/internal/examples/broadcast/main.go @@ -0,0 +1,59 @@ +package main + +import ( + "context" + "errors" + "log" + "net" + "net/http" + "os" + "os/signal" + "time" +) + +func main() { + log.SetFlags(0) + + err := run() + if err != nil { + log.Fatal(err) + } +} + +// run starts a http.Server for the passed in address +// with all requests handled by broadcastServer. +func run() error { + if len(os.Args) < 2 { + return errors.New("please provide an address to listen on as the first argument") + } + + l, err := net.Listen("tcp", os.Args[1]) + if err != nil { + return err + } + log.Printf("listening on http://%v", l.Addr()) + + s := &http.Server{ + Handler: newBroadcastServer(log.Printf), + ReadTimeout: time.Second * 10, + WriteTimeout: time.Second * 10, + } + errc := make(chan error, 1) + go func() { + errc <- s.Serve(l) + }() + + sigs := make(chan os.Signal, 1) + signal.Notify(sigs, os.Interrupt) + select { + case err := <-errc: + log.Printf("failed to serve: %v", err) + case sig := <-sigs: + log.Printf("terminating: %v", sig) + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second*10) + defer cancel() + + return s.Shutdown(ctx) +} diff --git a/internal/examples/broadcast/server.go b/internal/examples/broadcast/server.go new file mode 100644 index 00000000..91af8330 --- /dev/null +++ b/internal/examples/broadcast/server.go @@ -0,0 +1,189 @@ +package main + +import ( + "context" + "errors" + "io" + "net/http" + "sync" + "time" + "unicode/utf8" + + "github.com/coder/websocket" +) + +// broadcastServer broadcasts every message published to /publish to all +// WebSocket clients subscribed to /subscribe. +// +// Each message is wrapped in a single websocket.PreparedMessage, so it is +// compressed once per broadcast instead of once per subscriber. +type broadcastServer struct { + // subscriberMessageBuffer controls the max number + // of messages that can be queued for a subscriber + // before it is kicked. + // + // Defaults to 16. + subscriberMessageBuffer int + + // logf controls where logs are sent. + logf func(f string, v ...any) + + // serveMux routes the various endpoints to the appropriate handler. + serveMux http.ServeMux + + subscribersMu sync.Mutex + subscribers map[*subscriber]struct{} +} + +// newBroadcastServer constructs a broadcastServer with the defaults. +func newBroadcastServer(logf func(f string, v ...any)) *broadcastServer { + bs := &broadcastServer{ + subscriberMessageBuffer: 16, + logf: logf, + subscribers: make(map[*subscriber]struct{}), + } + bs.serveMux.HandleFunc("/subscribe", bs.subscribeHandler) + bs.serveMux.HandleFunc("/publish", bs.publishHandler) + + return bs +} + +// subscriber represents a subscriber. +// Messages are sent on the msgs channel and if the client +// cannot keep up with the messages, the slow channel is closed. +type subscriber struct { + msgs chan *websocket.PreparedMessage + slow chan struct{} + slowOnce sync.Once +} + +func (s *subscriber) closeSlow() { + s.slowOnce.Do(func() { + close(s.slow) + }) +} + +func (bs *broadcastServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + bs.serveMux.ServeHTTP(w, r) +} + +// subscribeHandler accepts the WebSocket connection and then subscribes +// it to all future messages. +func (bs *broadcastServer) subscribeHandler(w http.ResponseWriter, r *http.Request) { + err := bs.subscribe(w, r) + if errors.Is(err, context.Canceled) { + return + } + if websocket.CloseStatus(err) == websocket.StatusNormalClosure || + websocket.CloseStatus(err) == websocket.StatusGoingAway { + return + } + if err != nil { + bs.logf("%v", err) + return + } +} + +// publishHandler reads the request body with a limit of 65536 bytes and then +// publishes the received message. +// +// Messages are sent as text, which must be valid UTF-8, so other bodies are +// rejected instead of being sent as invalid text messages. +func (bs *broadcastServer) publishHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed) + return + } + body := http.MaxBytesReader(w, r.Body, 65536) + msg, err := io.ReadAll(body) + if err != nil { + http.Error(w, http.StatusText(http.StatusRequestEntityTooLarge), http.StatusRequestEntityTooLarge) + return + } + if !utf8.Valid(msg) { + http.Error(w, "message must be valid UTF-8", http.StatusBadRequest) + return + } + + bs.publish(websocket.NewPreparedMessage(websocket.MessageText, msg)) + + w.WriteHeader(http.StatusAccepted) +} + +// subscribe subscribes the given WebSocket to all broadcast messages. +// The subscriber is registered before the WebSocket is accepted, so it +// receives every message published after the client's Dial returns. +// +// It uses CloseRead to keep reading from the connection to process control +// messages and cancel the context if the connection drops. +func (bs *broadcastServer) subscribe(w http.ResponseWriter, r *http.Request) error { + s := &subscriber{ + msgs: make(chan *websocket.PreparedMessage, bs.subscriberMessageBuffer), + slow: make(chan struct{}), + } + bs.addSubscriber(s) + defer bs.deleteSubscriber(s) + + // Without context takeover, connections compress messages independently + // of each other, which allows sharing the compressed payload. + c, err := websocket.Accept(w, r, &websocket.AcceptOptions{ + CompressionMode: websocket.CompressionNoContextTakeover, + }) + if err != nil { + return err + } + defer c.CloseNow() + + ctx := c.CloseRead(context.Background()) + + for { + select { + case pm := <-s.msgs: + err := writeTimeout(ctx, time.Second*5, c, pm) + if err != nil { + return err + } + case <-s.slow: + return c.Close(websocket.StatusPolicyViolation, "connection too slow to keep up with messages") + case <-ctx.Done(): + return ctx.Err() + } + } +} + +// publish publishes the pm to all subscribers. +// It never blocks and so messages to slow subscribers +// are dropped. +func (bs *broadcastServer) publish(pm *websocket.PreparedMessage) { + bs.subscribersMu.Lock() + defer bs.subscribersMu.Unlock() + + for s := range bs.subscribers { + select { + case s.msgs <- pm: + default: + s.closeSlow() + } + } +} + +// addSubscriber registers a subscriber. +func (bs *broadcastServer) addSubscriber(s *subscriber) { + bs.subscribersMu.Lock() + bs.subscribers[s] = struct{}{} + bs.subscribersMu.Unlock() +} + +// deleteSubscriber deletes the given subscriber. +func (bs *broadcastServer) deleteSubscriber(s *subscriber) { + bs.subscribersMu.Lock() + delete(bs.subscribers, s) + bs.subscribersMu.Unlock() +} + +func writeTimeout(ctx context.Context, timeout time.Duration, c *websocket.Conn, pm *websocket.PreparedMessage) error { + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + return c.WritePrepared(ctx, pm) +} diff --git a/internal/examples/broadcast/server_test.go b/internal/examples/broadcast/server_test.go new file mode 100644 index 00000000..82d07948 --- /dev/null +++ b/internal/examples/broadcast/server_test.go @@ -0,0 +1,76 @@ +package main + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/coder/websocket" +) + +// Test_broadcastServer subscribes clients with and without compression, +// publishes messages below and above the compression threshold and ensures +// every client receives all of them, while invalid UTF-8 is rejected. +func Test_broadcastServer(t *testing.T) { + t.Parallel() + + s := httptest.NewServer(newBroadcastServer(t.Logf)) + defer s.Close() + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + var clients []*websocket.Conn + for i := 0; i < 4; i++ { + mode := websocket.CompressionDisabled + if i%2 == 0 { + mode = websocket.CompressionNoContextTakeover + } + c, _, err := websocket.Dial(ctx, s.URL+"/subscribe", &websocket.DialOptions{ + CompressionMode: mode, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close(websocket.StatusInternalError, "the sky is falling") + clients = append(clients, c) + } + + msgs := []string{ + "hello", + strings.Repeat("hello world ", 512), + } + publish := func(msg string, expStatus int) { + resp, err := http.Post(s.URL+"/publish", "text/plain", strings.NewReader(msg)) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != expStatus { + t.Fatalf("expected status %v but got %v", expStatus, resp.StatusCode) + } + } + + // Published first, so reading the valid messages below fails if it was + // delivered. + publish("\xff", http.StatusBadRequest) + for _, msg := range msgs { + publish(msg, http.StatusAccepted) + } + + for i, c := range clients { + for _, exp := range msgs { + _, b, err := c.Read(ctx) + if err != nil { + t.Fatal(err) + } + if string(b) != exp { + t.Fatalf("client %v: expected %q but got %q", i, exp, b) + } + } + c.Close(websocket.StatusNormalClosure, "") + } +} diff --git a/internal/test/wstest/echo.go b/internal/test/wstest/echo.go index c0c8dcd7..18078433 100644 --- a/internal/test/wstest/echo.go +++ b/internal/test/wstest/echo.go @@ -49,6 +49,19 @@ func EchoLoop(ctx context.Context, c *websocket.Conn) error { // Echo writes a message and ensures the same is sent back on c. func Echo(ctx context.Context, c *websocket.Conn, max int) error { + return echo(ctx, c, max, func(typ websocket.MessageType, msg []byte) error { + return c.Write(ctx, typ, msg) + }) +} + +// EchoPrepared is like Echo but writes the message with Conn.WritePrepared. +func EchoPrepared(ctx context.Context, c *websocket.Conn, max int) error { + return echo(ctx, c, max, func(typ websocket.MessageType, msg []byte) error { + return c.WritePrepared(ctx, websocket.NewPreparedMessage(typ, msg)) + }) +} + +func echo(ctx context.Context, c *websocket.Conn, max int, write func(websocket.MessageType, []byte) error) error { expType := websocket.MessageBinary if xrand.Bool() { expType = websocket.MessageText @@ -57,7 +70,7 @@ func Echo(ctx context.Context, c *websocket.Conn, max int) error { msg := randMessage(expType, xrand.Int(max)) writeErr := xsync.Go(func() error { - return c.Write(ctx, expType, msg) + return write(expType, msg) }) actType, act, err := c.Read(ctx) diff --git a/prepared.go b/prepared.go new file mode 100644 index 00000000..8408b631 --- /dev/null +++ b/prepared.go @@ -0,0 +1,94 @@ +//go:build !js + +package websocket + +import ( + "bytes" + "fmt" + "sync" +) + +// PreparedMessage is a message that can be written to many connections +// efficiently with Conn.WritePrepared, such as when broadcasting the same +// update to every subscriber. +// +// With compression enabled, most of the cost of writing a message is +// compressing it. Conn.Write compresses the message once per connection, while +// Conn.WritePrepared compresses it once and shares the compressed payload between +// connections. Sharing is only possible when all the following hold: +// +// - The connection compresses without context takeover. For a server, this +// is the case with CompressionNoContextTakeover, or with +// CompressionContextTakeover when the client negotiates +// server_no_context_takeover. For a client, client_no_context_takeover +// plays the same role. +// - The message is at least as long as the compression threshold, 512 bytes +// by default when compressing without context takeover. Smaller messages +// are not compressed. +// +// Otherwise, Conn.WritePrepared costs the same as Conn.Write. With context takeover, +// each connection compresses against the sliding window of its own previous +// messages, so the compressed payload differs per connection and cannot be +// shared. Without compression there is no expensive work to share. Clients +// generate a new masking key for every write in all cases. +// +// The compressed payload is computed on the first write that needs it and +// kept for the lifetime of the PreparedMessage, so a PreparedMessage can be +// written any number of times, for example a greeting sent to every new +// connection. It is safe for concurrent use. +// +// In Wasm, the browser handles framing and compression, so Conn.WritePrepared is +// equivalent to Conn.Write. +type PreparedMessage struct { + typ MessageType + data []byte + + // Written to connections that compress without context takeover. + compressOnce sync.Once + deflated []byte // Compressed data, set only if smaller than data. + compressed bool // Whether deflated is written instead of data. + compressErr error +} + +// NewPreparedMessage returns a PreparedMessage of the specified type with cloned data. +func NewPreparedMessage(typ MessageType, data []byte) *PreparedMessage { + return &PreparedMessage{ + typ: typ, + data: bytes.Clone(data), + } +} + +// compress returns the bytes to write on connections that compress without +// context takeover and whether they are compressed. They are the original +// data if compression does not make it smaller. +func (pm *PreparedMessage) compress() ([]byte, bool, error) { + pm.compressOnce.Do(func() { + var buf bytes.Buffer + + tw := &trimLastFourBytesWriter{ + w: &buf, + } + fw := getFlateWriter(tw) + defer putFlateWriter(fw) + + _, err := fw.Write(pm.data) + if err != nil { + pm.compressErr = fmt.Errorf("failed to compress: %w", err) + return + } + + err = fw.Flush() + if err != nil { + pm.compressErr = fmt.Errorf("failed to flush compression: %w", err) + return + } + + if buf.Len() < len(pm.data) { + pm.deflated, pm.compressed = buf.Bytes(), true + } + }) + if !pm.compressed { + return pm.data, false, pm.compressErr + } + return pm.deflated, true, nil +} diff --git a/prepared_test.go b/prepared_test.go new file mode 100644 index 00000000..b9c19e5c --- /dev/null +++ b/prepared_test.go @@ -0,0 +1,97 @@ +//go:build !js + +package websocket + +import ( + "bufio" + "bytes" + "context" + "strings" + "testing" + + "github.com/coder/websocket/internal/test/assert" + "github.com/coder/websocket/internal/test/xrand" +) + +type bufConn struct { + bytes.Buffer +} + +func (*bufConn) Close() error { + return nil +} + +// TestWritePreparedWire verifies that WritePrepared puts the same bytes on the +// wire as Write, including for messages written after it, and that the +// shared compressed payload is written only when the connection compresses +// without context takeover and compression makes the message smaller. +func TestWritePreparedWire(t *testing.T) { + t.Parallel() + + largeMsg := strings.Repeat("prepared message ", 100) + smallMsg := "small" + nextMsg := []byte(strings.Repeat("next message ", 100)) + // Random bytes do not compress, so compression makes them larger. + incompressibleMsg := string(xrand.Bytes(1024)) + + testCases := []struct { + name string + copts *compressionOptions + msg string + // Whether the shared compressed payload is written. It is false both + // when the regular write path is used and when the message is written + // as is because compression does not make it smaller. + wantCompressed bool + }{ + {"Disabled", nil, largeMsg, false}, + {"ContextTakeover/AboveThreshold", CompressionContextTakeover.opts(), largeMsg, false}, + {"ContextTakeover/BelowThreshold", CompressionContextTakeover.opts(), smallMsg, false}, + {"NoContextTakeover/AboveThreshold", CompressionNoContextTakeover.opts(), largeMsg, true}, + {"NoContextTakeover/AtThreshold", CompressionNoContextTakeover.opts(), strings.Repeat("a", 512), true}, + {"NoContextTakeover/BelowThreshold", CompressionNoContextTakeover.opts(), strings.Repeat("a", 511), false}, + {"ContextTakeover/Incompressible", CompressionContextTakeover.opts(), incompressibleMsg, false}, + {"NoContextTakeover/Incompressible", CompressionNoContextTakeover.opts(), incompressibleMsg, false}, + // Only the server side of the negotiated options applies to server writes. + {"ServerNoContextTakeover", &compressionOptions{serverNoContextTakeover: true}, largeMsg, true}, + {"ClientNoContextTakeover", &compressionOptions{clientNoContextTakeover: true}, largeMsg, false}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + newServerConn := func() (*Conn, *bufConn) { + rwc := &bufConn{} + c := newConn(connConfig{ + rwc: rwc, + copts: tc.copts, + br: bufio.NewReader(rwc), + bw: bufio.NewWriter(rwc), + }) + t.Cleanup(func() { + _ = c.CloseNow() + }) + return c, rwc + } + + ctx := context.Background() + + data := []byte(tc.msg) + pm := NewPreparedMessage(MessageText, data) + copy(data, "modified") // NewPreparedMessage must copy data. + + want, wantBuf := newServerConn() + got, gotBuf := newServerConn() + for range 2 { + assert.Success(t, want.Write(ctx, MessageText, []byte(tc.msg))) + assert.Success(t, want.Write(ctx, MessageBinary, nextMsg)) + + assert.Success(t, got.WritePrepared(ctx, pm)) + assert.Success(t, got.Write(ctx, MessageBinary, nextMsg)) + } + + assert.Equal(t, "wire bytes", wantBuf.Bytes(), gotBuf.Bytes()) + assert.Equal(t, "compressed", tc.wantCompressed, pm.compressed) + }) + } +} diff --git a/write.go b/write.go index 81bab4fa..f9b52563 100644 --- a/write.go +++ b/write.go @@ -48,6 +48,39 @@ func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error { return nil } +// WritePrepared writes a prepared message to the connection. +// +// It behaves like Write but reuses the compressed payload of pm across +// connections where possible. See PreparedMessage. +func (c *Conn) WritePrepared(ctx context.Context, pm *PreparedMessage) error { + err := c.writePrepared(ctx, pm) + if err != nil { + return fmt.Errorf("failed to write prepared msg: %w", err) + } + return nil +} + +func (c *Conn) writePrepared(ctx context.Context, pm *PreparedMessage) error { + if !c.flate() || len(pm.data) < c.flateThreshold || c.msgWriter.flateContextTakeover() { + _, err := c.write(ctx, pm.typ, pm.data) + return err + } + + p, compressed, err := pm.compress() + if err != nil { + return err + } + + err = c.msgWriter.reset(ctx, pm.typ) + if err != nil { + return err + } + defer c.msgWriter.mu.unlock() + + _, err = c.writeFrame(ctx, true, compressed, c.msgWriter.opcode, p) + return err +} + type msgWriter struct { c *Conn @@ -180,7 +213,19 @@ func (mw *msgWriter) writeCompressedFrame(ctx context.Context, p []byte) (int, e mw.closed = true - _, err = mw.c.writeFrame(ctx, true, true, mw.opcode, buf.Bytes()) + // Without context takeover, the peer keeps no history between messages, + // so a message that compression does not make smaller is sent as is. + // With context takeover, the message is already in our sliding window + // and must be sent compressed to keep the peer's window in sync, as + // uncompressed messages do not affect it. + // See https://tools.ietf.org/html/rfc7692#section-7.2.3.2 + compressed := mw.flateContextTakeover() || buf.Len() < len(p) + b := p + if compressed { + b = buf.Bytes() + } + + _, err = mw.c.writeFrame(ctx, true, compressed, mw.opcode, b) if err != nil { return 0, err } diff --git a/ws_js.go b/ws_js.go index 026b75fc..cdf25a38 100644 --- a/ws_js.go +++ b/ws_js.go @@ -194,7 +194,7 @@ func (c *Conn) Ping(ctx context.Context) error { } // Write writes a message of the given type to the connection. -// Always non blocking. +// Always non-blocking. func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error { err := c.write(typ, p) if err != nil { @@ -210,6 +210,27 @@ func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error { return nil } +// PreparedMessage is a message that can be written to many connections with +// Conn.WritePrepared. +type PreparedMessage struct { + typ MessageType + data []byte +} + +// NewPreparedMessage returns a PreparedMessage of the specified type with cloned data. +func NewPreparedMessage(typ MessageType, data []byte) *PreparedMessage { + return &PreparedMessage{ + typ: typ, + data: bytes.Clone(data), + } +} + +// WritePrepared writes a prepared message to the connection. +// The browser handles framing and compression, so it is equivalent to Write. +func (c *Conn) WritePrepared(ctx context.Context, pm *PreparedMessage) error { + return c.Write(ctx, pm.typ, pm.data) +} + func (c *Conn) write(typ MessageType, p []byte) error { if c.isClosed() { return net.ErrClosed diff --git a/ws_js_test.go b/ws_js_test.go index 1fa242f4..f0e1e89b 100644 --- a/ws_js_test.go +++ b/ws_js_test.go @@ -37,6 +37,28 @@ func TestWasm(t *testing.T) { assert.Success(t, err) } +func TestWasmPrepared(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + + c, _, err := websocket.Dial(ctx, os.Getenv("WS_ECHO_SERVER_URL"), &websocket.DialOptions{ + Subprotocols: []string{"echo"}, + }) + assert.Success(t, err) + defer c.Close(websocket.StatusInternalError, "") + + c.SetReadLimit(65536) + for range 10 { + err = wstest.EchoPrepared(ctx, c, 65536) + assert.Success(t, err) + } + + err = c.Close(websocket.StatusNormalClosure, "") + assert.Success(t, err) +} + func TestWasmDialTimeout(t *testing.T) { t.Parallel()