Update
This commit is contained in:
1 parent
3530afb854
commit
6903fcff27
2146 files changed
+112256
-190703
No files matched your search
+1
@@ -0,0 +1 @@
|
||||
PLAN.md
|
||||
+159
@@ -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.
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
Copyright (c) 2023 Anmol Sethi <hi@nhooyr.io>
|
||||
Copyright (c) 2025 Coder
|
||||
|
||||
Permission to use, copy, modify, and distribute this software for any
|
||||
purpose with or without fee is hereby granted, provided that the above
|
||||
|
||||
+18
@@ -0,0 +1,18 @@
|
||||
.PHONY: all
|
||||
all: fmt lint test
|
||||
|
||||
.PHONY: fmt
|
||||
fmt:
|
||||
./ci/fmt.sh
|
||||
|
||||
.PHONY: lint
|
||||
lint:
|
||||
./ci/lint.sh
|
||||
|
||||
.PHONY: test
|
||||
test:
|
||||
./ci/test.sh
|
||||
|
||||
.PHONY: bench
|
||||
bench:
|
||||
./ci/bench.sh
|
||||
+16
-16
@@ -1,7 +1,7 @@
|
||||
# websocket
|
||||
|
||||
[](https://pkg.go.dev/github.com/coder/websocket)
|
||||
[](https://github.com/coder/websocket/coverage.html)
|
||||
[](https://coder.github.io/websocket/coverage.html)
|
||||
|
||||
websocket is a minimal and idiomatic WebSocket library for Go.
|
||||
|
||||
@@ -36,15 +36,15 @@ go get github.com/coder/websocket
|
||||
|
||||
See GitHub issues for minor issues but the major future enhancements are:
|
||||
|
||||
- [ ] Perfect examples [#217](https://github.com/nhooyr/websocket/issues/217)
|
||||
- [ ] wstest.Pipe for in memory testing [#340](https://github.com/nhooyr/websocket/issues/340)
|
||||
- [ ] Ping pong heartbeat helper [#267](https://github.com/nhooyr/websocket/issues/267)
|
||||
- [ ] Ping pong instrumentation callbacks [#246](https://github.com/nhooyr/websocket/issues/246)
|
||||
- [ ] Graceful shutdown helpers [#209](https://github.com/nhooyr/websocket/issues/209)
|
||||
- [ ] Assembly for WebSocket masking [#16](https://github.com/nhooyr/websocket/issues/16)
|
||||
- WIP at [#326](https://github.com/nhooyr/websocket/pull/326), about 3x faster
|
||||
- [ ] HTTP/2 [#4](https://github.com/nhooyr/websocket/issues/4)
|
||||
- [ ] The holy grail [#402](https://github.com/nhooyr/websocket/issues/402)
|
||||
- [ ] Perfect examples [#217](https://github.com/coder/websocket/issues/217)
|
||||
- [ ] wstest.Pipe for in memory testing [#340](https://github.com/coder/websocket/issues/340)
|
||||
- [ ] Ping pong heartbeat helper [#267](https://github.com/coder/websocket/issues/267)
|
||||
- [ ] Ping pong instrumentation callbacks [#246](https://github.com/coder/websocket/issues/246)
|
||||
- [ ] Graceful shutdown helpers [#209](https://github.com/coder/websocket/issues/209)
|
||||
- [ ] Assembly for WebSocket masking [#16](https://github.com/coder/websocket/issues/16)
|
||||
- WIP at [#326](https://github.com/coder/websocket/pull/326), about 3x faster
|
||||
- [ ] HTTP/2 [#4](https://github.com/coder/websocket/issues/4)
|
||||
- [ ] The holy grail [#402](https://github.com/coder/websocket/issues/402)
|
||||
|
||||
## Examples
|
||||
|
||||
@@ -63,10 +63,12 @@ http.HandlerFunc(func (w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
defer c.CloseNow()
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), time.Second*10)
|
||||
// Set the context as needed. Use of r.Context() is not recommended
|
||||
// to avoid surprising behavior (see http.Hijacker).
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
|
||||
defer cancel()
|
||||
|
||||
var v interface{}
|
||||
var v any
|
||||
err = wsjson.Read(ctx, c, &v)
|
||||
if err != nil {
|
||||
// ...
|
||||
@@ -107,8 +109,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:
|
||||
|
||||
@@ -126,9 +126,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))
|
||||
|
||||
+37
-11
@@ -1,10 +1,10 @@
|
||||
//go:build !js
|
||||
// +build !js
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha1"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
@@ -14,7 +14,7 @@ import (
|
||||
"net/http"
|
||||
"net/textproto"
|
||||
"net/url"
|
||||
"path/filepath"
|
||||
"path"
|
||||
"strings"
|
||||
|
||||
"github.com/coder/websocket/internal/errd"
|
||||
@@ -40,9 +40,10 @@ type AcceptOptions struct {
|
||||
// In such a case, example.com is the origin and chat.example.com is the request host.
|
||||
// One would set this field to []string{"example.com"} to authorize example.com to connect.
|
||||
//
|
||||
// Each pattern is matched case insensitively against the request origin host
|
||||
// with filepath.Match.
|
||||
// See https://golang.org/pkg/path/filepath/#Match
|
||||
// Each pattern is matched case insensitively with path.Match (see
|
||||
// https://golang.org/pkg/path/#Match). By default, it is matched
|
||||
// against the request origin host. If the pattern contains a URI
|
||||
// scheme ("://"), it will be matched against "scheme://host".
|
||||
//
|
||||
// Please ensure you understand the ramifications of enabling this.
|
||||
// If used incorrectly your WebSocket server will be open to CSRF attacks.
|
||||
@@ -62,6 +63,22 @@ type AcceptOptions struct {
|
||||
// Defaults to 512 bytes for CompressionNoContextTakeover and 128 bytes
|
||||
// for CompressionContextTakeover.
|
||||
CompressionThreshold int
|
||||
|
||||
// OnPingReceived is an optional callback invoked synchronously when a ping frame is received.
|
||||
//
|
||||
// The payload contains the application data of the ping frame.
|
||||
// If the callback returns false, the subsequent pong frame will not be sent.
|
||||
// To avoid blocking, any expensive processing should be performed asynchronously using a goroutine.
|
||||
OnPingReceived func(ctx context.Context, payload []byte) bool
|
||||
|
||||
// OnPongReceived is an optional callback invoked synchronously when a pong frame is received.
|
||||
//
|
||||
// The payload contains the application data of the pong frame.
|
||||
// To avoid blocking, any expensive processing should be performed asynchronously using a goroutine.
|
||||
//
|
||||
// Unlike OnPingReceived, this callback does not return a value because a pong frame
|
||||
// is a response to a ping and does not trigger any further frame transmission.
|
||||
OnPongReceived func(ctx context.Context, payload []byte)
|
||||
}
|
||||
|
||||
func (opts *AcceptOptions) cloneWithDefaults() *AcceptOptions {
|
||||
@@ -73,12 +90,15 @@ func (opts *AcceptOptions) cloneWithDefaults() *AcceptOptions {
|
||||
}
|
||||
|
||||
// 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.
|
||||
//
|
||||
// Accept will write a response to w on all errors.
|
||||
//
|
||||
// Note that using the http.Request Context after Accept returns may lead to
|
||||
// unexpected behavior (see http.Hijacker).
|
||||
func Accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (*Conn, error) {
|
||||
return accept(w, r, opts)
|
||||
}
|
||||
@@ -96,7 +116,7 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con
|
||||
if !opts.InsecureSkipVerify {
|
||||
err = authenticateOrigin(r, opts.OriginPatterns)
|
||||
if err != nil {
|
||||
if errors.Is(err, filepath.ErrBadPattern) {
|
||||
if errors.Is(err, path.ErrBadPattern) {
|
||||
log.Printf("websocket: %v", err)
|
||||
err = errors.New(http.StatusText(http.StatusForbidden))
|
||||
}
|
||||
@@ -105,7 +125,7 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con
|
||||
}
|
||||
}
|
||||
|
||||
hj, ok := w.(http.Hijacker)
|
||||
hj, ok := hijacker(w)
|
||||
if !ok {
|
||||
err = errors.New("http.ResponseWriter does not implement http.Hijacker")
|
||||
http.Error(w, http.StatusText(http.StatusNotImplemented), http.StatusNotImplemented)
|
||||
@@ -153,6 +173,8 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con
|
||||
client: false,
|
||||
copts: copts,
|
||||
flateThreshold: opts.CompressionThreshold,
|
||||
onPingReceived: opts.OnPingReceived,
|
||||
onPongReceived: opts.OnPongReceived,
|
||||
|
||||
br: brw.Reader,
|
||||
bw: brw.Writer,
|
||||
@@ -219,9 +241,13 @@ func authenticateOrigin(r *http.Request, originHosts []string) error {
|
||||
}
|
||||
|
||||
for _, hostPattern := range originHosts {
|
||||
matched, err := match(hostPattern, u.Host)
|
||||
target := u.Host
|
||||
if strings.Contains(hostPattern, "://") {
|
||||
target = u.Scheme + "://" + u.Host
|
||||
}
|
||||
matched, err := match(hostPattern, target)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to parse filepath pattern %q: %w", hostPattern, err)
|
||||
return fmt.Errorf("failed to parse path pattern %q: %w", hostPattern, err)
|
||||
}
|
||||
if matched {
|
||||
return nil
|
||||
@@ -234,7 +260,7 @@ func authenticateOrigin(r *http.Request, originHosts []string) error {
|
||||
}
|
||||
|
||||
func match(pattern, s string) (bool, error) {
|
||||
return filepath.Match(strings.ToLower(pattern), strings.ToLower(s))
|
||||
return path.Match(strings.ToLower(pattern), strings.ToLower(s))
|
||||
}
|
||||
|
||||
func selectSubprotocol(r *http.Request, subprotocols []string) string {
|
||||
|
||||
+3
-16
@@ -1,5 +1,4 @@
|
||||
//go:build !js
|
||||
// +build !js
|
||||
|
||||
package websocket
|
||||
|
||||
@@ -100,7 +99,7 @@ func CloseStatus(err error) StatusCode {
|
||||
func (c *Conn) Close(code StatusCode, reason string) (err error) {
|
||||
defer errd.Wrap(&err, "failed to close WebSocket")
|
||||
|
||||
if !c.casClosing() {
|
||||
if c.casClosing() {
|
||||
err = c.waitGoroutines()
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -133,7 +132,7 @@ func (c *Conn) Close(code StatusCode, reason string) (err error) {
|
||||
func (c *Conn) CloseNow() (err error) {
|
||||
defer errd.Wrap(&err, "failed to immediately close WebSocket")
|
||||
|
||||
if !c.casClosing() {
|
||||
if c.casClosing() {
|
||||
err = c.waitGoroutines()
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -232,12 +231,6 @@ func (c *Conn) waitGoroutines() error {
|
||||
t := time.NewTimer(time.Second * 15)
|
||||
defer t.Stop()
|
||||
|
||||
select {
|
||||
case <-c.timeoutLoopDone:
|
||||
case <-t.C:
|
||||
return errors.New("failed to wait for timeoutLoop goroutine to exit")
|
||||
}
|
||||
|
||||
c.closeReadMu.Lock()
|
||||
closeRead := c.closeReadCtx != nil
|
||||
c.closeReadMu.Unlock()
|
||||
@@ -329,13 +322,7 @@ func (ce CloseError) bytesErr() ([]byte, error) {
|
||||
}
|
||||
|
||||
func (c *Conn) casClosing() bool {
|
||||
c.closeMu.Lock()
|
||||
defer c.closeMu.Unlock()
|
||||
if !c.closing {
|
||||
c.closing = true
|
||||
return true
|
||||
}
|
||||
return false
|
||||
return c.closing.Swap(true)
|
||||
}
|
||||
|
||||
func (c *Conn) isClosed() bool {
|
||||
|
||||
+4
-3
@@ -1,5 +1,4 @@
|
||||
//go:build !js
|
||||
// +build !js
|
||||
|
||||
package websocket
|
||||
|
||||
@@ -168,8 +167,10 @@ type slidingWindow struct {
|
||||
buf []byte
|
||||
}
|
||||
|
||||
var swPoolMu sync.RWMutex
|
||||
var swPool = map[int]*sync.Pool{}
|
||||
var (
|
||||
swPoolMu sync.RWMutex
|
||||
swPool = map[int]*sync.Pool{}
|
||||
)
|
||||
|
||||
func slidingWindowPool(n int) *sync.Pool {
|
||||
swPoolMu.RLock()
|
||||
|
||||
+57
-36
@@ -1,5 +1,4 @@
|
||||
//go:build !js
|
||||
// +build !js
|
||||
|
||||
package websocket
|
||||
|
||||
@@ -52,9 +51,8 @@ type Conn struct {
|
||||
br *bufio.Reader
|
||||
bw *bufio.Writer
|
||||
|
||||
readTimeout chan context.Context
|
||||
writeTimeout chan context.Context
|
||||
timeoutLoopDone chan struct{}
|
||||
readTimeoutStop atomic.Pointer[func() bool]
|
||||
writeTimeoutStop atomic.Pointer[func() bool]
|
||||
|
||||
// Read state.
|
||||
readMu *mu
|
||||
@@ -69,17 +67,25 @@ type Conn struct {
|
||||
writeHeaderBuf [8]byte
|
||||
writeHeader header
|
||||
|
||||
// Close handshake state.
|
||||
closeStateMu sync.RWMutex
|
||||
closeReceivedErr error
|
||||
closeSentErr error
|
||||
|
||||
// CloseRead state.
|
||||
closeReadMu sync.Mutex
|
||||
closeReadCtx context.Context
|
||||
closeReadDone chan struct{}
|
||||
|
||||
closing atomic.Bool
|
||||
closeMu sync.Mutex // Protects following.
|
||||
closed chan struct{}
|
||||
closeMu sync.Mutex
|
||||
closing bool
|
||||
|
||||
pingCounter int32
|
||||
activePingsMu sync.Mutex
|
||||
activePings map[string]chan<- struct{}
|
||||
pingCounter atomic.Int64
|
||||
activePingsMu sync.Mutex
|
||||
activePings map[string]chan<- struct{}
|
||||
onPingReceived func(context.Context, []byte) bool
|
||||
onPongReceived func(context.Context, []byte)
|
||||
}
|
||||
|
||||
type connConfig struct {
|
||||
@@ -88,6 +94,8 @@ type connConfig struct {
|
||||
client bool
|
||||
copts *compressionOptions
|
||||
flateThreshold int
|
||||
onPingReceived func(context.Context, []byte) bool
|
||||
onPongReceived func(context.Context, []byte)
|
||||
|
||||
br *bufio.Reader
|
||||
bw *bufio.Writer
|
||||
@@ -104,12 +112,10 @@ func newConn(cfg connConfig) *Conn {
|
||||
br: cfg.br,
|
||||
bw: cfg.bw,
|
||||
|
||||
readTimeout: make(chan context.Context),
|
||||
writeTimeout: make(chan context.Context),
|
||||
timeoutLoopDone: make(chan struct{}),
|
||||
|
||||
closed: make(chan struct{}),
|
||||
activePings: make(map[string]chan<- struct{}),
|
||||
closed: make(chan struct{}),
|
||||
activePings: make(map[string]chan<- struct{}),
|
||||
onPingReceived: cfg.onPingReceived,
|
||||
onPongReceived: cfg.onPongReceived,
|
||||
}
|
||||
|
||||
c.readMu = newMu(c)
|
||||
@@ -133,8 +139,6 @@ func newConn(cfg connConfig) *Conn {
|
||||
c.close()
|
||||
})
|
||||
|
||||
go c.timeoutLoop()
|
||||
|
||||
return c
|
||||
}
|
||||
|
||||
@@ -164,27 +168,44 @@ func (c *Conn) close() error {
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *Conn) timeoutLoop() {
|
||||
defer close(c.timeoutLoopDone)
|
||||
func (c *Conn) setupWriteTimeout(ctx context.Context) bool {
|
||||
if ctx.Done() == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
readCtx := context.Background()
|
||||
writeCtx := context.Background()
|
||||
stop := context.AfterFunc(ctx, func() {
|
||||
c.clearWriteTimeout()
|
||||
c.close()
|
||||
})
|
||||
swapTimeoutStop(&c.writeTimeoutStop, &stop)
|
||||
return true
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return
|
||||
func (c *Conn) clearWriteTimeout() {
|
||||
swapTimeoutStop(&c.writeTimeoutStop, nil)
|
||||
}
|
||||
|
||||
case writeCtx = <-c.writeTimeout:
|
||||
case readCtx = <-c.readTimeout:
|
||||
func (c *Conn) setupReadTimeout(ctx context.Context) bool {
|
||||
if ctx.Done() == nil {
|
||||
return false
|
||||
}
|
||||
|
||||
case <-readCtx.Done():
|
||||
c.close()
|
||||
return
|
||||
case <-writeCtx.Done():
|
||||
c.close()
|
||||
return
|
||||
}
|
||||
stop := context.AfterFunc(ctx, func() {
|
||||
c.clearReadTimeout()
|
||||
c.close()
|
||||
})
|
||||
swapTimeoutStop(&c.readTimeoutStop, &stop)
|
||||
return true
|
||||
}
|
||||
|
||||
func (c *Conn) clearReadTimeout() {
|
||||
swapTimeoutStop(&c.readTimeoutStop, nil)
|
||||
}
|
||||
|
||||
func swapTimeoutStop(p *atomic.Pointer[func() bool], newStop *func() bool) {
|
||||
oldStop := p.Swap(newStop)
|
||||
if oldStop != nil {
|
||||
(*oldStop)()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -200,9 +221,9 @@ func (c *Conn) flate() bool {
|
||||
//
|
||||
// TCP Keepalives should suffice for most use cases.
|
||||
func (c *Conn) Ping(ctx context.Context) error {
|
||||
p := atomic.AddInt32(&c.pingCounter, 1)
|
||||
p := c.pingCounter.Add(1)
|
||||
|
||||
err := c.ping(ctx, strconv.Itoa(int(p)))
|
||||
err := c.ping(ctx, strconv.FormatInt(p, 10))
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to ping: %w", err)
|
||||
}
|
||||
|
||||
+18
-1
@@ -1,5 +1,4 @@
|
||||
//go:build !js
|
||||
// +build !js
|
||||
|
||||
package websocket
|
||||
|
||||
@@ -48,6 +47,22 @@ type DialOptions struct {
|
||||
// Defaults to 512 bytes for CompressionNoContextTakeover and 128 bytes
|
||||
// for CompressionContextTakeover.
|
||||
CompressionThreshold int
|
||||
|
||||
// OnPingReceived is an optional callback invoked synchronously when a ping frame is received.
|
||||
//
|
||||
// The payload contains the application data of the ping frame.
|
||||
// If the callback returns false, the subsequent pong frame will not be sent.
|
||||
// To avoid blocking, any expensive processing should be performed asynchronously using a goroutine.
|
||||
OnPingReceived func(ctx context.Context, payload []byte) bool
|
||||
|
||||
// OnPongReceived is an optional callback invoked synchronously when a pong frame is received.
|
||||
//
|
||||
// The payload contains the application data of the pong frame.
|
||||
// To avoid blocking, any expensive processing should be performed asynchronously using a goroutine.
|
||||
//
|
||||
// Unlike OnPingReceived, this callback does not return a value because a pong frame
|
||||
// is a response to a ping and does not trigger any further frame transmission.
|
||||
OnPongReceived func(ctx context.Context, payload []byte)
|
||||
}
|
||||
|
||||
func (opts *DialOptions) cloneWithDefaults(ctx context.Context) (context.Context, context.CancelFunc, *DialOptions) {
|
||||
@@ -163,6 +178,8 @@ func dial(ctx context.Context, urls string, opts *DialOptions, rand io.Reader) (
|
||||
client: true,
|
||||
copts: copts,
|
||||
flateThreshold: opts.CompressionThreshold,
|
||||
onPingReceived: opts.OnPingReceived,
|
||||
onPongReceived: opts.OnPongReceived,
|
||||
br: getBufioReader(rwc),
|
||||
bw: getBufioWriter(rwc),
|
||||
}), resp, nil
|
||||
|
||||
-1
@@ -1,5 +1,4 @@
|
||||
//go:build !js
|
||||
// +build !js
|
||||
|
||||
// Package websocket implements the RFC 6455 WebSocket protocol.
|
||||
//
|
||||
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"errors"
|
||||
)
|
||||
|
||||
// ErrMessageTooBig is returned when a message exceeds the read limit.
|
||||
var ErrMessageTooBig = errors.New("websocket: message too big")
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
//go:build !js
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
)
|
||||
|
||||
type rwUnwrapper interface {
|
||||
Unwrap() http.ResponseWriter
|
||||
}
|
||||
|
||||
// hijacker returns the Hijacker interface of the http.ResponseWriter.
|
||||
// It follows the Unwrap method of the http.ResponseWriter if available,
|
||||
// matching the behavior of http.ResponseController. If the Hijacker
|
||||
// interface is not found, it returns false.
|
||||
//
|
||||
// Since the http.ResponseController is not available in Go 1.19, and
|
||||
// does not support checking the presence of the Hijacker interface,
|
||||
// this function is used to provide a consistent way to check for the
|
||||
// Hijacker interface across Go versions.
|
||||
func hijacker(rw http.ResponseWriter) (http.Hijacker, bool) {
|
||||
for {
|
||||
switch t := rw.(type) {
|
||||
case http.Hijacker:
|
||||
return t, true
|
||||
case rwUnwrapper:
|
||||
rw = t.Unwrap()
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
}
|
||||
+5
-4
@@ -5,15 +5,16 @@ import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
var bpool sync.Pool
|
||||
var bpool = sync.Pool{
|
||||
New: func() any {
|
||||
return &bytes.Buffer{}
|
||||
},
|
||||
}
|
||||
|
||||
// Get returns a buffer from the pool or creates a new one if
|
||||
// the pool is empty.
|
||||
func Get() *bytes.Buffer {
|
||||
b := bpool.Get()
|
||||
if b == nil {
|
||||
return &bytes.Buffer{}
|
||||
}
|
||||
return b.(*bytes.Buffer)
|
||||
}
|
||||
|
||||
|
||||
+1
-1
@@ -7,7 +7,7 @@ import (
|
||||
// Wrap wraps err with fmt.Errorf if err is non nil.
|
||||
// Intended for use with defer and a named error return.
|
||||
// Inspired by https://github.com/golang/go/issues/32676.
|
||||
func Wrap(err *error, f string, v ...interface{}) {
|
||||
func Wrap(err *error, f string, v ...any) {
|
||||
if *err != nil {
|
||||
*err = fmt.Errorf(f+": %w", append(v, *err)...)
|
||||
}
|
||||
|
||||
+4
-4
@@ -33,7 +33,7 @@ func New(url string, protocols []string) (c WebSocket, err error) {
|
||||
c = WebSocket{}
|
||||
})
|
||||
|
||||
jsProtocols := make([]interface{}, len(protocols))
|
||||
jsProtocols := make([]any, len(protocols))
|
||||
for i, p := range protocols {
|
||||
jsProtocols[i] = p
|
||||
}
|
||||
@@ -57,7 +57,7 @@ func (c WebSocket) setBinaryType(typ string) {
|
||||
}
|
||||
|
||||
func (c WebSocket) addEventListener(eventType string, fn func(e js.Value)) func() {
|
||||
f := js.FuncOf(func(this js.Value, args []js.Value) interface{} {
|
||||
f := js.FuncOf(func(this js.Value, args []js.Value) any {
|
||||
fn(args[0])
|
||||
return nil
|
||||
})
|
||||
@@ -97,7 +97,7 @@ func (c WebSocket) OnError(fn func(e js.Value)) (remove func()) {
|
||||
// MessageEvent is the type passed to a message handler.
|
||||
type MessageEvent struct {
|
||||
// string or []byte.
|
||||
Data interface{}
|
||||
Data any
|
||||
|
||||
// There are more fields to the interface but we don't use them.
|
||||
// See https://developer.mozilla.org/en-US/docs/Web/API/MessageEvent
|
||||
@@ -106,7 +106,7 @@ type MessageEvent struct {
|
||||
// OnMessage registers a function to be called when the WebSocket receives a message.
|
||||
func (c WebSocket) OnMessage(fn func(m MessageEvent)) (remove func()) {
|
||||
return c.addEventListener("message", func(e js.Value) {
|
||||
var data interface{}
|
||||
var data any
|
||||
|
||||
arrayBuffer := e.Get("data")
|
||||
if arrayBuffer.Type() == js.TypeString {
|
||||
|
||||
-26
@@ -1,26 +0,0 @@
|
||||
package xsync
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"runtime/debug"
|
||||
)
|
||||
|
||||
// Go allows running a function in another goroutine
|
||||
// and waiting for its error.
|
||||
func Go(fn func() error) <-chan error {
|
||||
errs := make(chan error, 1)
|
||||
go func() {
|
||||
defer func() {
|
||||
r := recover()
|
||||
if r != nil {
|
||||
select {
|
||||
case errs <- fmt.Errorf("panic in go fn: %v, %s", r, debug.Stack()):
|
||||
default:
|
||||
}
|
||||
}
|
||||
}()
|
||||
errs <- fn()
|
||||
}()
|
||||
|
||||
return errs
|
||||
}
|
||||
-23
@@ -1,23 +0,0 @@
|
||||
package xsync
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// Int64 represents an atomic int64.
|
||||
type Int64 struct {
|
||||
// We do not use atomic.Load/StoreInt64 since it does not
|
||||
// work on 32 bit computers but we need 64 bit integers.
|
||||
i atomic.Value
|
||||
}
|
||||
|
||||
// Load loads the int64.
|
||||
func (v *Int64) Load() int64 {
|
||||
i, _ := v.i.Load().(int64)
|
||||
return i
|
||||
}
|
||||
|
||||
// Store stores the int64.
|
||||
func (v *Int64) Store(i int64) {
|
||||
v.i.Store(i)
|
||||
}
|
||||
-12
@@ -1,12 +0,0 @@
|
||||
#!/bin/sh
|
||||
set -eu
|
||||
cd -- "$(dirname "$0")"
|
||||
|
||||
echo "=== fmt.sh"
|
||||
./ci/fmt.sh
|
||||
echo "=== lint.sh"
|
||||
./ci/lint.sh
|
||||
echo "=== test.sh"
|
||||
./ci/test.sh "$@"
|
||||
echo "=== bench.sh"
|
||||
./ci/bench.sh
|
||||
+19
-23
@@ -68,7 +68,7 @@ func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
|
||||
defer nc.writeMu.unlock()
|
||||
|
||||
// Prevents future writes from writing until the deadline is reset.
|
||||
atomic.StoreInt64(&nc.writeExpired, 1)
|
||||
nc.writeExpired.Store(1)
|
||||
})
|
||||
if !nc.writeTimer.Stop() {
|
||||
<-nc.writeTimer.C
|
||||
@@ -84,7 +84,7 @@ func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
|
||||
defer nc.readMu.unlock()
|
||||
|
||||
// Prevents future reads from reading until the deadline is reset.
|
||||
atomic.StoreInt64(&nc.readExpired, 1)
|
||||
nc.readExpired.Store(1)
|
||||
})
|
||||
if !nc.readTimer.Stop() {
|
||||
<-nc.readTimer.C
|
||||
@@ -94,25 +94,22 @@ func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
|
||||
}
|
||||
|
||||
type netConn struct {
|
||||
// These must be first to be aligned on 32 bit platforms.
|
||||
// https://github.com/nhooyr/websocket/pull/438
|
||||
readExpired int64
|
||||
writeExpired int64
|
||||
|
||||
c *Conn
|
||||
msgType MessageType
|
||||
|
||||
writeTimer *time.Timer
|
||||
writeMu *mu
|
||||
writeCtx context.Context
|
||||
writeCancel context.CancelFunc
|
||||
writeTimer *time.Timer
|
||||
writeMu *mu
|
||||
writeExpired atomic.Int64
|
||||
writeCtx context.Context
|
||||
writeCancel context.CancelFunc
|
||||
|
||||
readTimer *time.Timer
|
||||
readMu *mu
|
||||
readCtx context.Context
|
||||
readCancel context.CancelFunc
|
||||
readEOFed bool
|
||||
reader io.Reader
|
||||
readTimer *time.Timer
|
||||
readMu *mu
|
||||
readExpired atomic.Int64
|
||||
readCtx context.Context
|
||||
readCancel context.CancelFunc
|
||||
readEOFed bool
|
||||
reader io.Reader
|
||||
}
|
||||
|
||||
var _ net.Conn = &netConn{}
|
||||
@@ -129,7 +126,7 @@ func (nc *netConn) Write(p []byte) (int, error) {
|
||||
nc.writeMu.forceLock()
|
||||
defer nc.writeMu.unlock()
|
||||
|
||||
if atomic.LoadInt64(&nc.writeExpired) == 1 {
|
||||
if nc.writeExpired.Load() == 1 {
|
||||
return 0, fmt.Errorf("failed to write: %w", context.DeadlineExceeded)
|
||||
}
|
||||
|
||||
@@ -157,7 +154,7 @@ func (nc *netConn) Read(p []byte) (int, error) {
|
||||
}
|
||||
|
||||
func (nc *netConn) read(p []byte) (int, error) {
|
||||
if atomic.LoadInt64(&nc.readExpired) == 1 {
|
||||
if nc.readExpired.Load() == 1 {
|
||||
return 0, fmt.Errorf("failed to read: %w", context.DeadlineExceeded)
|
||||
}
|
||||
|
||||
@@ -191,8 +188,7 @@ func (nc *netConn) read(p []byte) (int, error) {
|
||||
return n, err
|
||||
}
|
||||
|
||||
type websocketAddr struct {
|
||||
}
|
||||
type websocketAddr struct{}
|
||||
|
||||
func (a websocketAddr) Network() string {
|
||||
return "websocket"
|
||||
@@ -209,7 +205,7 @@ func (nc *netConn) SetDeadline(t time.Time) error {
|
||||
}
|
||||
|
||||
func (nc *netConn) SetWriteDeadline(t time.Time) error {
|
||||
atomic.StoreInt64(&nc.writeExpired, 0)
|
||||
nc.writeExpired.Store(0)
|
||||
if t.IsZero() {
|
||||
nc.writeTimer.Stop()
|
||||
} else {
|
||||
@@ -223,7 +219,7 @@ func (nc *netConn) SetWriteDeadline(t time.Time) error {
|
||||
}
|
||||
|
||||
func (nc *netConn) SetReadDeadline(t time.Time) error {
|
||||
atomic.StoreInt64(&nc.readExpired, 0)
|
||||
nc.readExpired.Store(0)
|
||||
if t.IsZero() {
|
||||
nc.readTimer.Stop()
|
||||
} else {
|
||||
|
||||
-1
@@ -1,5 +1,4 @@
|
||||
//go:build !js
|
||||
// +build !js
|
||||
|
||||
package websocket
|
||||
|
||||
|
||||
+82
-47
@@ -1,5 +1,4 @@
|
||||
//go:build !js
|
||||
// +build !js
|
||||
|
||||
package websocket
|
||||
|
||||
@@ -11,11 +10,11 @@ import (
|
||||
"io"
|
||||
"net"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/coder/websocket/internal/errd"
|
||||
"github.com/coder/websocket/internal/util"
|
||||
"github.com/coder/websocket/internal/xsync"
|
||||
)
|
||||
|
||||
// Reader reads from the connection until there is a WebSocket
|
||||
@@ -91,7 +90,8 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context {
|
||||
//
|
||||
// By default, the connection has a message read limit of 32768 bytes.
|
||||
//
|
||||
// When the limit is hit, the connection will be closed with StatusMessageTooBig.
|
||||
// When the limit is hit, reads return an error wrapping ErrMessageTooBig and
|
||||
// the connection is closed with StatusMessageTooBig.
|
||||
//
|
||||
// Set to -1 to disable.
|
||||
func (c *Conn) SetReadLimit(n int64) {
|
||||
@@ -217,60 +217,74 @@ func (c *Conn) readLoop(ctx context.Context) (header, error) {
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) readFrameHeader(ctx context.Context) (header, 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 header{}, net.ErrClosed
|
||||
case c.readTimeout <- ctx:
|
||||
return false, net.ErrClosed
|
||||
default:
|
||||
}
|
||||
timeoutSet := c.setupReadTimeout(ctx)
|
||||
|
||||
c.closeStateMu.Lock()
|
||||
closeReceivedErr := c.closeReceivedErr
|
||||
c.closeStateMu.Unlock()
|
||||
if closeReceivedErr != nil {
|
||||
if timeoutSet {
|
||||
c.clearReadTimeout()
|
||||
}
|
||||
return false, closeReceivedErr
|
||||
}
|
||||
|
||||
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) {
|
||||
timeoutSet, err := c.prepareRead(ctx)
|
||||
if err != nil {
|
||||
return header{}, err
|
||||
}
|
||||
defer c.finishRead(ctx, &err, timeoutSet)
|
||||
|
||||
h, err := readFrameHeader(c.br, c.readHeaderBuf[:])
|
||||
if err != nil {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return header{}, net.ErrClosed
|
||||
case <-ctx.Done():
|
||||
return header{}, ctx.Err()
|
||||
default:
|
||||
return header{}, err
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case <-c.closed:
|
||||
return header{}, net.ErrClosed
|
||||
case c.readTimeout <- context.Background():
|
||||
return header{}, err
|
||||
}
|
||||
|
||||
return h, nil
|
||||
}
|
||||
|
||||
func (c *Conn) readFramePayload(ctx context.Context, p []byte) (int, error) {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return 0, net.ErrClosed
|
||||
case c.readTimeout <- ctx:
|
||||
func (c *Conn) readFramePayload(ctx context.Context, p []byte) (_ int, err error) {
|
||||
timeoutSet, err := c.prepareRead(ctx)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer c.finishRead(ctx, &err, timeoutSet)
|
||||
|
||||
n, err := io.ReadFull(c.br, p)
|
||||
if err != nil {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return n, net.ErrClosed
|
||||
case <-ctx.Done():
|
||||
return n, ctx.Err()
|
||||
default:
|
||||
return n, fmt.Errorf("failed to read frame payload: %w", err)
|
||||
}
|
||||
return n, fmt.Errorf("failed to read frame payload: %w", err)
|
||||
}
|
||||
|
||||
select {
|
||||
case <-c.closed:
|
||||
return n, net.ErrClosed
|
||||
case c.readTimeout <- context.Background():
|
||||
}
|
||||
|
||||
return n, err
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
|
||||
@@ -301,8 +315,16 @@ func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
|
||||
|
||||
switch h.opcode {
|
||||
case opPing:
|
||||
if c.onPingReceived != nil {
|
||||
if !c.onPingReceived(ctx, b) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return c.writeControl(ctx, opPong, b)
|
||||
case opPong:
|
||||
if c.onPongReceived != nil {
|
||||
c.onPongReceived(ctx, b)
|
||||
}
|
||||
c.activePingsMu.Lock()
|
||||
pong, ok := c.activePings[string(b)]
|
||||
c.activePingsMu.Unlock()
|
||||
@@ -325,9 +347,22 @@ func (c *Conn) handleControl(ctx context.Context, h header) (err error) {
|
||||
}
|
||||
|
||||
err = fmt.Errorf("received close frame: %w", ce)
|
||||
c.writeClose(ce.Code, ce.Reason)
|
||||
c.readMu.unlock()
|
||||
c.close()
|
||||
c.closeStateMu.Lock()
|
||||
c.closeReceivedErr = err
|
||||
closeSent := c.closeSentErr != nil
|
||||
c.closeStateMu.Unlock()
|
||||
|
||||
// Only unlock readMu if this connection is being closed becaue
|
||||
// c.close will try to acquire the readMu lock. We unlock for
|
||||
// writeClose as well because it may also call c.close.
|
||||
if !closeSent {
|
||||
c.readMu.unlock()
|
||||
_ = c.writeClose(ce.Code, ce.Reason)
|
||||
}
|
||||
if !c.casClosing() {
|
||||
c.readMu.unlock()
|
||||
_ = c.close()
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -465,7 +500,7 @@ func (mr *msgReader) read(p []byte) (int, error) {
|
||||
type limitReader struct {
|
||||
c *Conn
|
||||
r io.Reader
|
||||
limit xsync.Int64
|
||||
limit atomic.Int64
|
||||
n int64
|
||||
}
|
||||
|
||||
@@ -489,9 +524,9 @@ func (lr *limitReader) Read(p []byte) (int, error) {
|
||||
}
|
||||
|
||||
if lr.n == 0 {
|
||||
err := fmt.Errorf("read limited at %v bytes", lr.limit.Load())
|
||||
lr.c.writeError(StatusMessageTooBig, err)
|
||||
return 0, err
|
||||
reason := fmt.Errorf("read limited at %d bytes", lr.limit.Load())
|
||||
lr.c.writeError(StatusMessageTooBig, reason)
|
||||
return 0, fmt.Errorf("%w: %v", ErrMessageTooBig, reason)
|
||||
}
|
||||
|
||||
if int64(len(p)) > lr.n {
|
||||
|
||||
+88
-34
@@ -1,10 +1,10 @@
|
||||
//go:build !js
|
||||
// +build !js
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"compress/flate"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/binary"
|
||||
@@ -14,8 +14,7 @@ import (
|
||||
"net"
|
||||
"time"
|
||||
|
||||
"compress/flate"
|
||||
|
||||
"github.com/coder/websocket/internal/bpool"
|
||||
"github.com/coder/websocket/internal/errd"
|
||||
"github.com/coder/websocket/internal/util"
|
||||
)
|
||||
@@ -102,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 {
|
||||
@@ -144,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)
|
||||
@@ -249,25 +292,36 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
|
||||
}
|
||||
defer c.writeFrameMu.unlock()
|
||||
|
||||
select {
|
||||
case <-c.closed:
|
||||
return 0, net.ErrClosed
|
||||
case c.writeTimeout <- ctx:
|
||||
}
|
||||
|
||||
defer func() {
|
||||
if c.isClosed() && opcode == opClose {
|
||||
err = nil
|
||||
}
|
||||
if err != nil {
|
||||
select {
|
||||
case <-c.closed:
|
||||
err = net.ErrClosed
|
||||
case <-ctx.Done():
|
||||
if ctx.Err() != nil {
|
||||
err = ctx.Err()
|
||||
default:
|
||||
} else if c.isClosed() {
|
||||
err = net.ErrClosed
|
||||
}
|
||||
err = fmt.Errorf("failed to write frame: %w", err)
|
||||
}
|
||||
}()
|
||||
|
||||
c.closeStateMu.Lock()
|
||||
closeSentErr := c.closeSentErr
|
||||
c.closeStateMu.Unlock()
|
||||
if closeSentErr != nil {
|
||||
return 0, net.ErrClosed
|
||||
}
|
||||
|
||||
select {
|
||||
case <-c.closed:
|
||||
return 0, net.ErrClosed
|
||||
default:
|
||||
}
|
||||
if c.setupWriteTimeout(ctx) {
|
||||
defer c.clearWriteTimeout()
|
||||
}
|
||||
|
||||
c.writeHeader.fin = fin
|
||||
c.writeHeader.opcode = opcode
|
||||
c.writeHeader.payloadLength = int64(len(p))
|
||||
@@ -303,13 +357,16 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case <-c.closed:
|
||||
if opcode == opClose {
|
||||
return n, nil
|
||||
if opcode == opClose {
|
||||
c.closeStateMu.Lock()
|
||||
c.closeSentErr = fmt.Errorf("sent close frame: %w", net.ErrClosed)
|
||||
closeReceived := c.closeReceivedErr != nil
|
||||
c.closeStateMu.Unlock()
|
||||
|
||||
if closeReceived && !c.casClosing() {
|
||||
c.writeFrameMu.unlock()
|
||||
_ = c.close()
|
||||
}
|
||||
return n, net.ErrClosed
|
||||
case c.writeTimeout <- context.Background():
|
||||
}
|
||||
|
||||
return n, nil
|
||||
@@ -335,10 +392,7 @@ func (c *Conn) writeFramePayload(p []byte) (n int, err error) {
|
||||
// Start of next write in the buffer.
|
||||
i := c.bw.Buffered()
|
||||
|
||||
j := len(p)
|
||||
if j > c.bw.Available() {
|
||||
j = c.bw.Available()
|
||||
}
|
||||
j := min(len(p), c.bw.Available())
|
||||
|
||||
_, err := c.bw.Write(p[:j])
|
||||
if err != nil {
|
||||
|
||||
+7
-7
@@ -12,11 +12,11 @@ import (
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall/js"
|
||||
|
||||
"github.com/coder/websocket/internal/bpool"
|
||||
"github.com/coder/websocket/internal/wsjs"
|
||||
"github.com/coder/websocket/internal/xsync"
|
||||
)
|
||||
|
||||
// opcode represents a WebSocket opcode.
|
||||
@@ -45,7 +45,7 @@ type Conn struct {
|
||||
ws wsjs.WebSocket
|
||||
|
||||
// read limit for a message in bytes.
|
||||
msgReadLimit xsync.Int64
|
||||
msgReadLimit atomic.Int64
|
||||
|
||||
closeReadMu sync.Mutex
|
||||
closeReadCtx context.Context
|
||||
@@ -144,9 +144,9 @@ func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
|
||||
}
|
||||
readLimit := c.msgReadLimit.Load()
|
||||
if readLimit >= 0 && int64(len(p)) > readLimit {
|
||||
err := fmt.Errorf("read limited at %v bytes", c.msgReadLimit.Load())
|
||||
c.Close(StatusMessageTooBig, err.Error())
|
||||
return 0, nil, err
|
||||
reason := fmt.Errorf("read limited at %d bytes", c.msgReadLimit.Load())
|
||||
c.Close(StatusMessageTooBig, reason.Error())
|
||||
return 0, nil, fmt.Errorf("%w: %v", ErrMessageTooBig, reason)
|
||||
}
|
||||
return typ, p, nil
|
||||
}
|
||||
@@ -196,7 +196,7 @@ func (c *Conn) Ping(ctx context.Context) error {
|
||||
// Write writes a message of the given type to the connection.
|
||||
// Always non blocking.
|
||||
func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error {
|
||||
err := c.write(ctx, typ, p)
|
||||
err := c.write(typ, p)
|
||||
if err != nil {
|
||||
// Have to ensure the WebSocket is closed after a write error
|
||||
// to match the Go API. It can only error if the message type
|
||||
@@ -210,7 +210,7 @@ func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Conn) write(ctx context.Context, typ MessageType, p []byte) error {
|
||||
func (c *Conn) write(typ MessageType, p []byte) error {
|
||||
if c.isClosed() {
|
||||
return net.ErrClosed
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user