diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..d1941b1c --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +PLAN.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..711bc96a --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,159 @@ +# Agents + +This document helps AI agents work effectively in this codebase. It explains the philosophy, patterns, and pitfalls behind the code, so you can make good decisions on any task, not just scenarios explicitly covered. + +## Philosophy + +This is a minimal, idiomatic WebSocket library. Simplicity is a feature. + +Before adding code, articulate why it's necessary in one sentence. If you cannot justify the addition, it probably isn't needed. + +Before adding a dependency, don't. Tests requiring external packages are isolated in `internal/thirdparty`. + +## Workflow + +Every task follows phases. Do not skip phases; if you realize you missed something, return to the appropriate phase. + +**Making changes:** + +1. **Research** - Understand the problem and codebase before acting +2. **Plan** - Articulate your approach before implementing (when asked, or for complex changes) +3. **Implement** - Make changes in small, verifiable steps +4. **Verify** - Confirm correctness using external tools + +For trivial changes (typo fixes, comment updates), an abbreviated workflow is acceptable. + +## Research + +Research in sequential passes. Each pass has one focus. Don't skip ahead to code until you've completed the earlier passes. + +**Pass 1: Read the issue.** If you were given a link, read it now. Do not explore code, do not pass go. Fetch the linked issue or document first. Summarize what it asks for. If you cannot restate the problem, you are not ready to proceed. + +**Pass 2: Read linked references.** Follow every link in the issue: related issues, RFCs, external docs. Document what you learn. Code comments reference RFC 6455 (WebSocket), RFC 7692 (compression), and RFC 8441 (HTTP/2). + +**Pass 3: Trace the code.** Start from public API inward: `Accept` (server) or `Dial` (client) → `Conn` → `Reader`/`Writer`. Read tests for intent; the autobahn-testsuite (`autobahn_test.go`) validates protocol compliance. + +**Pass 4: Check both platforms.** Native Go files have `//go:build !js`. WASM lives in `ws_js.go`. Same API, different implementations. WASM wraps browser APIs and cannot control framing, compression, or masking. + +**Pass 5: Search exhaustively.** If the change affects a pattern used in multiple places, grep for all instances. Missing one creates inconsistent behavior. + +**Pass 6: Document unknowns.** List what you still don't know. Unknown unknowns become known unknowns when you ask "what am I still unsure about?" + +**After all passes:** Can you restate the problem in your own words? If not, return to Pass 1. Gaps in earlier passes will cause problems later. + +## Plan + +When asked to plan, write to `PLAN.md`. Write for someone else, not yourself; don't skip context you already know. + +**Pass 1: Document research.** Summarize what you learned in Research. Follow every reference; document findings so the implementer can verify. If the research section is empty, you haven't researched enough. + +**Pass 2: Consider approaches.** For non-trivial problems, enumerate at least two approaches. For each, note: what would change, what could go wrong, what's the tradeoff. + +**Pass 3: Detail the chosen approach.** Explain what and why, not step-by-step how. Point to specific files, functions, line numbers. Make claims verifiable. Leave room for the implementer to find a better solution. + +**Pass 4: List open questions.** What's still unclear? What assumptions are you making? What would change your approach? + +**After all passes:** Review from the implementer's perspective. Could they start work with only this document? If not, add what's missing. + +## Implement + +Implement in sequential passes. Don't write code until you've completed the verification passes. + +**Pass 1: Verify understanding.** Did you do your research? Can you state your approach in one sentence? If requirements are ambiguous, stop and ask. A wrong assumption wastes more time than a quick question. + +**Pass 2: Check scope.** Does this need to exist? Check if it already exists in the API. Is this the library's job or the user's job? The library handles protocol correctness; application concerns (reconnection, auth, routing) belong in user code. + +**Pass 3: Check invariants.** Walk through Key Invariants before writing code: + +- Reads: Will something still read from the connection? +- Pools: Will pooled objects be returned on all paths? +- Locks: Are you using context-aware `mu`, not `sync.Mutex`? +- Independence: Are you coupling reads and writes unnecessarily? + +**Pass 4: Implement.** Make the change. Every change needs a reason; if you can't articulate why it improves things, don't make it. Preserve existing comments unless you can prove they're wrong. + +**Pass 5: Verify examples.** Trace through usage examples as if writing real code. If an example wouldn't compile, the design is wrong. Check edge cases: what happens on error? On cancellation? + +**After all passes:** If feedback identifies a problem, fix that specific problem. Don't pivot to a new approach without articulating what failed and why the new approach avoids it. + +## Verify + +Verify using external signals. Self-assessment is unreliable; these tools are the ground truth. + +**Pass 1: Run tests.** `go test ./...` must pass. If tests fail, return to Research to understand why, not to Implement to "fix" it. + +**Pass 2: Run vet.** `go vet ./...` must report no issues. + +**Pass 3: Check platforms.** Code must compile on both native Go and WASM. API changes need both implementations; test that method signatures match. + +**Pass 4: Protocol compliance.** For protocol changes, run `AUTOBAHN=1 go test`. Compare wire bytes when implementations disagree; check against the RFC. + +**After all passes:** If any pass fails, fix the issue and restart from Pass 1. Don't consider a change complete until all passes succeed. + +## When Uncertain + +If confidence is low or requirements are ambiguous: + +1. State what you're uncertain about +2. Identify what information would resolve the uncertainty +3. Gather that information or ask for clarification + +Do not proceed with low-confidence assumptions. + +For complex changes, approach from multiple angles. If two approaches give different answers, that discrepancy demands resolution before proceeding. + +## Code Style + +Follow Go conventions: [Effective Go](https://go.dev/doc/effective_go), [Go Code Review Comments](https://go.dev/wiki/CodeReviewComments), and [Go Proverbs](https://go-proverbs.github.io/). Be concise, declarative, and factual. + +Never use emdash. Use commas, semicolons, or separate sentences. + +**Doc comments** start with the name and state what it does. Structure: `// [Name] [verb]s [what]. [Optional: when/why to use it].` Put useful information where users make decisions (usually the constructor, not methods). + +**Inline comments** are terse. Prefer end-of-line when short enough. + +**Explain why, not what.** The code shows what it does; comments should explain reasoning, non-obvious decisions, or edge cases. + +**Wrap comments** at ~80 characters, continuing naturally at word boundaries. + +**Naming matters.** Before proposing a name, stop and review existing names in the file. Ask: what would someone assume from this name? Does it fit with how similar things are named? A good name is accurate on its own and consistent in context. + +**Comment content:** + +- Add information beyond what the code shows (not tautologies) +- State directly: "Returns X" (not "Note that this returns X") +- Drop filler: "basically", "actually", "really" add nothing + +## Key Invariants + +**Always read from connections.** Control frames (ping, pong, close) arrive on the read path. A connection that never reads will miss them and misbehave. `CloseRead` exists for write-only patterns. + +**Pooled objects must be returned.** `flate.Writer` is ~1.2MB. Leaking them causes memory growth. Follow `get*()` / `put*()` patterns; return on all paths including errors. + +**Locks must respect context.** The `mu` type in `conn.go` unblocks when context cancels or connection closes. Using `sync.Mutex` for user-facing operations would block forever on stuck connections. + +**Reads and writes are independent.** They have separate locks (`readMu`, `writeFrameMu`) and can happen concurrently. Keep them decoupled. + +**Masking is asymmetric.** Clients mask payloads; servers don't. The mask is applied into the bufio.Writer buffer, not in-place on source bytes. + +## Commits and PRs + +Use conventional commits: `type(scope): description`. Scope is optional but include it when changes are constrained to a single file or directory. + +**Choose precise verbs:** `add` (new functionality), `fix` (broken behavior), `prevent` (undesirable state), `allow` (enable action), `handle` (edge cases), `improve` (performance/quality), `use` (change approach), `skip` (bypass), `remove` (delete), `rewrite` (substantial rework). + +**Before writing the commit message**, articulate what was wrong with the previous behavior and why this change fixes it. The commit message captures this reasoning. + +**Commit messages explain why.** State what was wrong, why it mattered, and why this solution was chosen. One to three sentences. Keep the tone neutral. + +**PR descriptions are for humans.** Explain what reviewers can't see: the why, alternatives considered, and tradeoffs made. Show evidence (logs, benchmarks, screenshots). Be explicit about what the PR doesn't do. + +**Link issues:** `Fixes #123` (PR resolves it), `Refs #123` (related), `Updates #123` (partial progress). + +## RFC References + +- RFC 6455: The WebSocket Protocol +- RFC 7692: Compression Extensions for WebSocket (permessage-deflate) +- RFC 8441: Bootstrapping WebSockets with HTTP/2 + +Section numbers in code comments refer to these RFCs. diff --git a/README.md b/README.md index fb14ec4f..3475d871 100644 --- a/README.md +++ b/README.md @@ -116,8 +116,6 @@ 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) -- No extra goroutine per connection to support cancellation with context.Context. This costs github.com/coder/websocket 2 KB of memory per connection. - - Will be removed soon with [context.AfterFunc](https://github.com/golang/go/issues/57928). See [#411](https://github.com/nhooyr/websocket/issues/411) Advantages of github.com/coder/websocket: @@ -135,9 +133,9 @@ Advantages of github.com/coder/websocket: - Gorilla requires registering a pong callback before sending a Ping - Can target Wasm ([gorilla/websocket#432](https://github.com/gorilla/websocket/issues/432)) - Transparent message buffer reuse with [wsjson](https://pkg.go.dev/github.com/coder/websocket/wsjson) subpackage -- [1.75x](https://github.com/nhooyr/websocket/releases/tag/v1.7.4) faster WebSocket masking implementation in pure Go +- [1.75x](https://github.com/coder/websocket/releases/tag/v1.7.4) faster WebSocket masking implementation in pure Go - Gorilla's implementation is slower and uses [unsafe](https://golang.org/pkg/unsafe/). - Soon we'll have assembly and be 3x faster [#326](https://github.com/nhooyr/websocket/pull/326) + Soon we'll have assembly and be 3x faster [#326](https://github.com/coder/websocket/pull/326) - Full [permessage-deflate](https://tools.ietf.org/html/rfc7692) compression extension support - Gorilla only supports no context takeover mode - [CloseRead](https://pkg.go.dev/github.com/coder/websocket#Conn.CloseRead) helper for write only connections ([gorilla/websocket#492](https://github.com/gorilla/websocket/issues/492)) diff --git a/accept.go b/accept.go index 53f10334..0cf7708d 100644 --- a/accept.go +++ b/accept.go @@ -105,7 +105,7 @@ func (opts *AcceptOptions) cloneWithDefaults() (*AcceptOptions, error) { } // Accept accepts a WebSocket handshake from a client and upgrades the -// the connection to a WebSocket. +// connection to a WebSocket. // // Accept will not allow cross origin requests by default. // See the InsecureSkipVerify and OriginPatterns options to allow cross origin requests. diff --git a/compress_test.go b/compress_test.go index 1964c84f..81adce9a 100644 --- a/compress_test.go +++ b/compress_test.go @@ -3,11 +3,15 @@ package websocket import ( + "bufio" "bytes" "compress/flate" + "context" "io" + "net" "strings" "testing" + "time" "github.com/coder/websocket/internal/test/assert" "github.com/coder/websocket/internal/test/xrand" @@ -59,3 +63,249 @@ func BenchmarkFlateReader(b *testing.B) { io.ReadAll(r) } } + +// TestWriteSingleFrameCompressed verifies that Conn.Write sends compressed +// messages in a single frame instead of multiple frames, and that messages +// below the flateThreshold are sent uncompressed. +// This is a regression test for https://github.com/coder/websocket/issues/435 +func TestWriteSingleFrameCompressed(t *testing.T) { + t.Parallel() + + var ( + flateThreshold = 64 + + largeMsg = []byte(strings.Repeat("hello world ", 100)) + smallMsg = []byte("small message") + ) + + testCases := []struct { + name string + mode CompressionMode + msg []byte + wantRsv1 bool // true = compressed, false = uncompressed + }{ + {"ContextTakeover/AboveThreshold", CompressionContextTakeover, largeMsg, true}, + {"NoContextTakeover/AboveThreshold", CompressionNoContextTakeover, largeMsg, true}, + {"ContextTakeover/BelowThreshold", CompressionContextTakeover, smallMsg, false}, + {"NoContextTakeover/BelowThreshold", CompressionNoContextTakeover, smallMsg, false}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + c := newConn(connConfig{ + rwc: clientConn, + client: true, + copts: tc.mode.opts(), + flateThreshold: flateThreshold, + br: bufio.NewReader(clientConn), + bw: bufio.NewWriterSize(clientConn, 4096), + }) + + ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*100) + defer cancel() + + writeDone := make(chan error, 1) + go func() { + writeDone <- c.Write(ctx, MessageText, tc.msg) + }() + + reader := bufio.NewReader(serverConn) + readBuf := make([]byte, 8) + + h, err := readFrameHeader(reader, readBuf) + assert.Success(t, err) + + _, err = io.CopyN(io.Discard, reader, h.payloadLength) + assert.Success(t, err) + + assert.Equal(t, "opcode", opText, h.opcode) + assert.Equal(t, "rsv1 (compressed)", tc.wantRsv1, h.rsv1) + assert.Equal(t, "fin", true, h.fin) + + err = <-writeDone + assert.Success(t, err) + }) + } +} + +// TestWriteThenWriterContextTakeover verifies that using Conn.Write followed by +// Conn.Writer works correctly with context takeover enabled. This tests that +// the flateWriter destination is properly restored after Conn.Write redirects +// it to a temporary buffer. +func TestWriteThenWriterContextTakeover(t *testing.T) { + t.Parallel() + + clientConn, serverConn := net.Pipe() + defer clientConn.Close() + defer serverConn.Close() + + client := newConn(connConfig{ + rwc: clientConn, + client: true, + copts: CompressionContextTakeover.opts(), + flateThreshold: 64, + br: bufio.NewReader(clientConn), + bw: bufio.NewWriterSize(clientConn, 4096), + }) + + server := newConn(connConfig{ + rwc: serverConn, + client: false, + copts: CompressionContextTakeover.opts(), + flateThreshold: 64, + br: bufio.NewReader(serverConn), + bw: bufio.NewWriterSize(serverConn, 4096), + }) + + ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*500) + defer cancel() + + msg1 := []byte(strings.Repeat("first message ", 100)) + msg2 := []byte(strings.Repeat("second message ", 100)) + + type readResult struct { + typ MessageType + p []byte + err error + } + readCh := make(chan readResult, 3) + go func() { + for range 3 { + typ, p, err := server.Read(ctx) + readCh <- readResult{typ, p, err} + } + }() + + // We want to verify mixing `Write` and `Writer` usages still work. + // + // To this end, we call them in this order: + // - `Write` + // - `Writer` + // - `Write` + // + // This verifies that it works for a `Write` followed by a `Writer` + // as well as a `Writer` followed by a `Write`. + + // 1. `Write` API + err := client.Write(ctx, MessageText, msg1) + assert.Success(t, err) + + r := <-readCh + assert.Success(t, r.err) + assert.Equal(t, "Write type", MessageText, r.typ) + assert.Equal(t, "Write content", string(msg1), string(r.p)) + + // 2. `Writer` API + w, err := client.Writer(ctx, MessageBinary) + assert.Success(t, err) + _, err = w.Write(msg2) + assert.Success(t, err) + assert.Success(t, w.Close()) + + r = <-readCh + assert.Success(t, r.err) + assert.Equal(t, "Writer type", MessageBinary, r.typ) + assert.Equal(t, "Writer content", string(msg2), string(r.p)) + + // 3. `Write` API again + err = client.Write(ctx, MessageText, msg1) + assert.Success(t, err) + + r = <-readCh + assert.Success(t, r.err) + assert.Equal(t, "Write type", MessageText, r.typ) + assert.Equal(t, "Write content", string(msg1), string(r.p)) +} + +// TestCompressionDictionaryPreserved verifies that context takeover mode +// preserves the compression dictionary across Conn.Write calls, resulting +// in better compression for consecutive similar messages. +func TestCompressionDictionaryPreserved(t *testing.T) { + t.Parallel() + + msg := []byte(strings.Repeat(`{"type":"event","data":"value"}`, 50)) + + takeoverClient, takeoverServer := net.Pipe() + defer takeoverClient.Close() + defer takeoverServer.Close() + + withTakeover := newConn(connConfig{ + rwc: takeoverClient, + client: true, + copts: CompressionContextTakeover.opts(), + flateThreshold: 64, + br: bufio.NewReader(takeoverClient), + bw: bufio.NewWriterSize(takeoverClient, 4096), + }) + + noTakeoverClient, noTakeoverServer := net.Pipe() + defer noTakeoverClient.Close() + defer noTakeoverServer.Close() + + withoutTakeover := newConn(connConfig{ + rwc: noTakeoverClient, + client: true, + copts: CompressionNoContextTakeover.opts(), + flateThreshold: 64, + br: bufio.NewReader(noTakeoverClient), + bw: bufio.NewWriterSize(noTakeoverClient, 4096), + }) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + + // Capture compressed sizes for both modes + var withTakeoverSizes, withoutTakeoverSizes []int64 + + reader1 := bufio.NewReader(takeoverServer) + reader2 := bufio.NewReader(noTakeoverServer) + readBuf := make([]byte, 8) + + // Send 3 identical messages each + for range 3 { + // With context takeover + writeDone1 := make(chan error, 1) + go func() { + writeDone1 <- withTakeover.Write(ctx, MessageText, msg) + }() + + h1, err := readFrameHeader(reader1, readBuf) + assert.Success(t, err) + + _, err = io.CopyN(io.Discard, reader1, h1.payloadLength) + assert.Success(t, err) + + withTakeoverSizes = append(withTakeoverSizes, h1.payloadLength) + assert.Success(t, <-writeDone1) + + // Without context takeover + writeDone2 := make(chan error, 1) + go func() { + writeDone2 <- withoutTakeover.Write(ctx, MessageText, msg) + }() + + h2, err := readFrameHeader(reader2, readBuf) + assert.Success(t, err) + + _, err = io.CopyN(io.Discard, reader2, h2.payloadLength) + assert.Success(t, err) + + withoutTakeoverSizes = append(withoutTakeoverSizes, h2.payloadLength) + assert.Success(t, <-writeDone2) + } + + // With context takeover, the 2nd and 3rd messages should be smaller than + // without context takeover (dictionary helps compress repeated patterns). + // The first message will be similar size for both modes since there's no + // prior dictionary. But subsequent messages benefit from context takeover. + if withTakeoverSizes[2] >= withoutTakeoverSizes[2] { + t.Errorf("context takeover should compress better: with=%d, without=%d", + withTakeoverSizes[2], withoutTakeoverSizes[2]) + } +} diff --git a/conn.go b/conn.go index 09234871..bfc78aaa 100644 --- a/conn.go +++ b/conn.go @@ -168,24 +168,34 @@ func (c *Conn) close() error { return err } -func (c *Conn) setupWriteTimeout(ctx context.Context) { +func (c *Conn) setupWriteTimeout(ctx context.Context) bool { + if ctx.Done() == nil { + return false + } + stop := context.AfterFunc(ctx, func() { c.clearWriteTimeout() c.close() }) swapTimeoutStop(&c.writeTimeoutStop, &stop) + return true } func (c *Conn) clearWriteTimeout() { swapTimeoutStop(&c.writeTimeoutStop, nil) } -func (c *Conn) setupReadTimeout(ctx context.Context) { +func (c *Conn) setupReadTimeout(ctx context.Context) bool { + if ctx.Done() == nil { + return false + } + stop := context.AfterFunc(ctx, func() { c.clearReadTimeout() c.close() }) swapTimeoutStop(&c.readTimeoutStop, &stop) + return true } func (c *Conn) clearReadTimeout() { diff --git a/conn_test.go b/conn_test.go index c3ccc886..01af8cd4 100644 --- a/conn_test.go +++ b/conn_test.go @@ -592,10 +592,11 @@ func BenchmarkConn(b *testing.B) { msg := []byte(strings.Repeat("1234", 128)) readBuf := make([]byte, len(msg)) writes := make(chan struct{}) - defer close(writes) werrs := make(chan error) + writerDone := make(chan struct{}) go func() { + defer close(writerDone) for range writes { select { case werrs <- c1.Write(bb.ctx, websocket.MessageText, msg): @@ -650,6 +651,13 @@ func BenchmarkConn(b *testing.B) { } b.StopTimer() + close(writes) + select { + case <-writerDone: + case <-bb.ctx.Done(): + b.Fatal(bb.ctx.Err()) + } + b.ReportMetric(float64(*bytesWritten/b.N), "written/op") b.ReportMetric(float64(*bytesRead/b.N), "read/op") diff --git a/http2.go b/http2.go index fa5d4046..1dfae584 100644 --- a/http2.go +++ b/http2.go @@ -42,8 +42,15 @@ func (s *h2ServerStream) Flush() error { func (s *h2ServerStream) Close() error { // TODO(mafredri): Verify if the flush is necessary before closing the reader. - err := s.Flush() - return errors.Join(err, s.ReadCloser.Close()) + + //err := s.Flush() + //return errors.Join(err, s.ReadCloser.Close()) + // Write flushes after every write, so nothing is left buffered here. + // Flushing in Close is also unsafe: Close can run on a context.AfterFunc + // goroutine (e.g. after the peer sends RST_STREAM) while another + // goroutine is in Write, and http.ResponseWriter is not safe for + // concurrent use. + return s.ReadCloser.Close() } // h2ClientStream is a minimal io.ReadWriteCloser for an HTTP/2 extended CONNECT diff --git a/read.go b/read.go index 64822511..c19f41de 100644 --- a/read.go +++ b/read.go @@ -217,51 +217,52 @@ func (c *Conn) readLoop(ctx context.Context) (header, error) { } } -// prepareRead sets the readTimeout context and returns a done function -// to be called after the read is done. It also returns an error if the -// connection is closed. The reference to the error is used to assign -// an error depending on if the connection closed or the context timed -// out during use. Typically, the referenced error is a named return -// variable of the function calling this method. -func (c *Conn) prepareRead(ctx context.Context, err *error) (func(), error) { +// prepareRead sets the read timeout and checks whether the connection is closed. +func (c *Conn) prepareRead(ctx context.Context) (bool, error) { select { case <-c.closed: - return nil, net.ErrClosed + return false, net.ErrClosed default: } - c.setupReadTimeout(ctx) - - done := func() { - c.clearReadTimeout() - select { - case <-c.closed: - if *err != nil { - *err = net.ErrClosed - } - default: - } - if *err != nil && ctx.Err() != nil { - *err = ctx.Err() - } - } + timeoutSet := c.setupReadTimeout(ctx) c.closeStateMu.Lock() closeReceivedErr := c.closeReceivedErr c.closeStateMu.Unlock() if closeReceivedErr != nil { - defer done() - return nil, closeReceivedErr + if timeoutSet { + c.clearReadTimeout() + } + return false, closeReceivedErr } - return done, nil + return timeoutSet, nil +} + +// finishRead clears the read timeout and reports whether the connection or +// operation context ended while the read was in progress. +func (c *Conn) finishRead(ctx context.Context, err *error, timeoutSet bool) { + if timeoutSet { + c.clearReadTimeout() + } + select { + case <-c.closed: + if *err != nil { + *err = net.ErrClosed + } + default: + } + if *err != nil && ctx.Err() != nil { + *err = ctx.Err() + } } func (c *Conn) readFrameHeader(ctx context.Context) (_ header, err error) { - readDone, err := c.prepareRead(ctx, &err) + timeoutSet, err := c.prepareRead(ctx) if err != nil { return header{}, err } - defer readDone() + defer c.finishRead(ctx, &err, timeoutSet) h, err := readFrameHeader(c.br, c.readHeaderBuf[:]) if err != nil { @@ -272,11 +273,11 @@ func (c *Conn) readFrameHeader(ctx context.Context) (_ header, err error) { } func (c *Conn) readFramePayload(ctx context.Context, p []byte) (_ int, err error) { - readDone, err := c.prepareRead(ctx, &err) + timeoutSet, err := c.prepareRead(ctx) if err != nil { return 0, err } - defer readDone() + defer c.finishRead(ctx, &err, timeoutSet) n, err := io.ReadFull(c.br, p) if err != nil { diff --git a/write.go b/write.go index d7172a7b..81bab4fa 100644 --- a/write.go +++ b/write.go @@ -14,6 +14,7 @@ import ( "net" "time" + "github.com/coder/websocket/internal/bpool" "github.com/coder/websocket/internal/errd" "github.com/coder/websocket/internal/util" ) @@ -100,23 +101,17 @@ func (c *Conn) writer(ctx context.Context, typ MessageType) (io.WriteCloser, err } func (c *Conn) write(ctx context.Context, typ MessageType, p []byte) (int, error) { - mw, err := c.writer(ctx, typ) + err := c.msgWriter.reset(ctx, typ) if err != nil { return 0, err } + defer c.msgWriter.mu.unlock() - if !c.flate() { - defer c.msgWriter.mu.unlock() + if !c.flate() || len(p) < c.flateThreshold { return c.writeFrame(ctx, true, false, c.msgWriter.opcode, p) } - n, err := mw.Write(p) - if err != nil { - return n, err - } - - err = mw.Close() - return n, err + return c.msgWriter.writeCompressedFrame(ctx, p) } func (mw *msgWriter) reset(ctx context.Context, typ MessageType) error { @@ -142,6 +137,56 @@ func (mw *msgWriter) putFlateWriter() { } } +// writeCompressedFrame compresses and writes p as a single frame. +func (mw *msgWriter) writeCompressedFrame(ctx context.Context, p []byte) (int, error) { + err := mw.writeMu.lock(mw.ctx) + if err != nil { + return 0, fmt.Errorf("failed to write: %w", err) + } + defer mw.writeMu.unlock() + + if mw.closed { + return 0, errors.New("cannot use closed writer") + } + + mw.ensureFlate() + + buf := bpool.Get() + defer bpool.Put(buf) + + // Buffer compressed output so we can write as + // a single frame instead of chunked frames. + origWriter := mw.trimWriter.w + mw.trimWriter.w = buf + defer func() { + mw.trimWriter.w = origWriter + }() + + _, err = mw.flateWriter.Write(p) + if err != nil { + return 0, fmt.Errorf("failed to compress: %w", err) + } + + err = mw.flateWriter.Flush() + if err != nil { + return 0, fmt.Errorf("failed to flush compression: %w", err) + } + + mw.trimWriter.reset() + + if !mw.flateContextTakeover() { + mw.putFlateWriter() + } + + mw.closed = true + + _, err = mw.c.writeFrame(ctx, true, true, mw.opcode, buf.Bytes()) + if err != nil { + return 0, err + } + return len(p), nil +} + // Write writes the given bytes to the WebSocket connection. func (mw *msgWriter) Write(p []byte) (_ int, err error) { err = mw.writeMu.lock(mw.ctx) @@ -273,8 +318,9 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco return 0, net.ErrClosed default: } - c.setupWriteTimeout(ctx) - defer c.clearWriteTimeout() + if c.setupWriteTimeout(ctx) { + defer c.clearWriteTimeout() + } c.writeHeader.fin = fin c.writeHeader.opcode = opcode