Skip to content

Commit 0fb5b24

Browse files
nhooyrnhooyr
committed
Implement more of the API for WASM
Realized I can at least make the Reader/Writer/SetReadLimit methods work as expected even if they're not perfect.
1 parent 8c54bd9 commit 0fb5b24

File tree

8 files changed

+89
-147
lines changed

8 files changed

+89
-147
lines changed

‎conn.go

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ type Conn struct {
5959
msgReadLimit int64
6060

6161
// Used to ensure a previous writer is not used after being closed.
62-
activeWriter *messageWriter
62+
activeWriter atomic.Value
6363
// messageWriter state.
6464
writeMsgOpcode opcode
6565
writeMsgCtx context.Context
@@ -575,7 +575,7 @@ func (c *Conn) writer(ctx context.Context, typ MessageType) (io.WriteCloser, err
575575
w := &messageWriter{
576576
c: c,
577577
}
578-
c.activeWriter = w
578+
c.activeWriter.Store(w)
579579
return w, nil
580580
}
581581

@@ -607,7 +607,7 @@ type messageWriter struct {
607607
}
608608

609609
func (w *messageWriter) closed() bool {
610-
return w != w.c.activeWriter
610+
return w != w.c.activeWriter.Load()
611611
}
612612

613613
// Write writes the given bytes to the WebSocket connection.
@@ -645,7 +645,7 @@ func (w *messageWriter) close() error {
645645
if w.closed() {
646646
return fmt.Errorf("cannot use closed writer")
647647
}
648-
w.c.activeWriter = nil
648+
w.c.activeWriter.Store((*messageWriter)(nil))
649649

650650
_, err := w.c.writeFrame(w.c.writeMsgCtx, true, w.c.writeMsgOpcode, nil)
651651
if err != nil {
@@ -925,7 +925,3 @@ func (c *Conn) extractBufioWriterBuf(w io.Writer) {
925925

926926
c.bw.Reset(w)
927927
}
928-
929-
func (c *netConn) netConnReader(ctx context.Context) (MessageType, io.Reader, error) {
930-
return c.c.Reader(c.readContext)
931-
}

‎doc.go

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,13 +26,20 @@
2626
// See https://developer.mozilla.org/en-US/docs/Web/API/WebSocket
2727
//
2828
// Thus the unsupported features (not compiled in) for WASM are:
29+
//
2930
// - Accept and AcceptOptions
30-
// - Conn's Reader, Writer, SetReadLimit and Ping methods
31+
// - Conn.Ping
3132
// - HTTPClient and HTTPHeader fields in DialOptions
3233
//
3334
// The *http.Response returned by Dial will always either be nil or &http.Response{} as
3435
// we do not have access to the handshake response in the browser.
3536
//
37+
// The Writer method on the Conn buffers everything in memory and then sends it as a message
38+
// when the writer is closed.
39+
//
40+
// The Reader method also reads the entire response and then returns a reader that
41+
// reads from the byte slice.
42+
//
3643
// Writes are also always async so the passed context is no-op.
3744
//
3845
// Everything else is fully supported. This includes the wsjson and wspb helper packages.

‎netconn.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ func (c *netConn) Read(p []byte) (int, error) {
9999
}
100100

101101
if c.reader == nil {
102-
typ, r, err := c.netConnReader(c.readContext)
102+
typ, r, err := c.c.Reader(c.readContext)
103103
if err != nil {
104104
var ce CloseError
105105
if errors.As(err, &ce) && (ce.Code == StatusNormalClosure) || (ce.Code == StatusGoingAway) {

‎websocket_js.go

Lines changed: 74 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,16 @@ import (
1313
"sync/atomic"
1414
"syscall/js"
1515

16+
"nhooyr.io/websocket/internal/bpool"
1617
"nhooyr.io/websocket/internal/wsjs"
1718
)
1819

1920
// Conn provides a wrapper around the browser WebSocket API.
2021
type Conn struct {
2122
ws wsjs.WebSocket
2223

24+
msgReadLimit int64
25+
2326
readClosed int64
2427
closeOnce sync.Once
2528
closed chan struct{}
@@ -43,6 +46,7 @@ func (c *Conn) close(err error) {
4346
func (c *Conn) init() {
4447
c.closed = make(chan struct{})
4548
c.readch = make(chan wsjs.MessageEvent, 1)
49+
c.msgReadLimit = 32768
4650

4751
c.releaseOnClose = c.ws.OnClose(func(e wsjs.CloseEvent) {
4852
cerr := CloseError{
@@ -77,6 +81,10 @@ func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
7781
if err != nil {
7882
return 0, nil, fmt.Errorf("failed to read: %w", err)
7983
}
84+
if len(p) > c.msgReadLimit {
85+
r.c.Close(StatusMessageTooBig, fmt.Sprintf("read limited at %v bytes", r.c.msgReadLimit))
86+
return 0, nil, c.closeErr
87+
}
8088
return typ, p, nil
8189
}
8290

@@ -106,6 +114,11 @@ func (c *Conn) read(ctx context.Context) (MessageType, []byte, error) {
106114
func (c *Conn) Write(ctx context.Context, typ MessageType, p []byte) error {
107115
err := c.write(ctx, typ, p)
108116
if err != nil {
117+
// Have to ensure the WebSocket is closed after a write error
118+
// to match the Go API. It can only error if the message type
119+
// is unexpected or the passed bytes contain invalid UTF-8 for
120+
// MessageText.
121+
c.Close(StatusInternalError, "something went wrong")
109122
return fmt.Errorf("failed to write: %w", err)
110123
}
111124
return nil
@@ -216,8 +229,10 @@ func dial(ctx context.Context, url string, opts *DialOptions) (*Conn, *http.Resp
216229
return c, &http.Response{}, nil
217230
}
218231

219-
func (c *netConn) netConnReader(ctx context.Context) (MessageType, io.Reader, error) {
220-
typ, p, err := c.c.Read(ctx)
232+
// Reader attempts to read a message from the connection.
233+
// The maximum time spent waiting is bounded by the context.
234+
func (c *Conn) Reader(ctx context.Context) (MessageType, io.Reader, error) {
235+
typ, p, err := c.Read(ctx)
221236
if err != nil {
222237
return 0, nil, err
223238
}
@@ -228,3 +243,60 @@ func (c *netConn) netConnReader(ctx context.Context) (MessageType, io.Reader, er
228243
func (c *Conn) reader(ctx context.Context) {
229244
c.read(ctx)
230245
}
246+
247+
// Writer returns a writer to write a WebSocket data message to the connection.
248+
// It buffers the entire message in memory and then sends it when the writer
249+
// is closed.
250+
func (c *Conn) Writer(ctx context.Context, typ MessageType) (io.WriteCloser, error) {
251+
return writer{
252+
c: c,
253+
ctx: ctx,
254+
typ: typ,
255+
b: bpool.Get(),
256+
}, nil
257+
}
258+
259+
type writer struct {
260+
closed bool
261+
262+
c *Conn
263+
ctx context.Context
264+
typ MessageType
265+
266+
b *bytes.Buffer
267+
}
268+
269+
func (w writer) Write(p []byte) (int, error) {
270+
if w.closed {
271+
return 0, errors.New("cannot write to closed writer")
272+
}
273+
n, err := w.b.Write(p)
274+
if err != nil {
275+
return n, fmt.Errorf("failed to write message: %w", err)
276+
}
277+
return n, nil
278+
}
279+
280+
func (w writer) Close() error {
281+
if w.closed {
282+
return errors.New("cannot close closed writer")
283+
}
284+
w.closed = true
285+
defer bpool.Put(w.b)
286+
287+
err := w.c.Write(w.ctx, w.typ, w.b.Bytes())
288+
if err != nil {
289+
return fmt.Errorf("failed to close writer: %w", err)
290+
}
291+
return nil
292+
}
293+
294+
// SetReadLimit sets the max number of bytes to read for a single message.
295+
// It applies to the Reader and Read methods.
296+
//
297+
// By default, the connection has a message read limit of 32768 bytes.
298+
//
299+
// When the limit is hit, the connection will be closed with StatusMessageTooBig.
300+
func (c *Conn) SetReadLimit(n int64) {
301+
c.readLimit = n
302+
}

‎wsjson/wsjson.go

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
// +build !js
2-
31
// Package wsjson provides websocket helpers for JSON messages.
42
package wsjson // import "nhooyr.io/websocket/wsjson"
53

@@ -34,9 +32,7 @@ func read(ctx context.Context, c *websocket.Conn, v interface{}) error {
3432
}
3533

3634
b := bpool.Get()
37-
defer func() {
38-
bpool.Put(b)
39-
}()
35+
defer bpool.Put(b)
4036

4137
_, err = b.ReadFrom(r)
4238
if err != nil {

‎wsjson/wsjson_js.go

Lines changed: 0 additions & 58 deletions
This file was deleted.

‎wspb/wspb.go

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
// +build !js
2-
31
// Package wspb provides websocket helpers for protobuf messages.
42
package wspb // import "nhooyr.io/websocket/wspb"
53

@@ -36,9 +34,7 @@ func read(ctx context.Context, c *websocket.Conn, v proto.Message) error {
3634
}
3735

3836
b := bpool.Get()
39-
defer func() {
40-
bpool.Put(b)
41-
}()
37+
defer bpool.Put(b)
4238

4339
_, err = b.ReadFrom(r)
4440
if err != nil {

‎wspb/wspb_js.go

Lines changed: 0 additions & 67 deletions
This file was deleted.

0 commit comments

Comments
 (0)