wsjson.go

Documentation: nhooyr.io/websocket/wsjson

     1  // Package wsjson provides helpers for reading and writing JSON messages.
     2  package wsjson // import "nhooyr.io/websocket/wsjson"
     3  
     4  import (
     5  	"context"
     6  	"encoding/json"
     7  	"fmt"
     8  
     9  	"nhooyr.io/websocket"
    10  	"nhooyr.io/websocket/internal/bpool"
    11  	"nhooyr.io/websocket/internal/errd"
    12  	"nhooyr.io/websocket/internal/util"
    13  )
    14  
    15  // Read reads a JSON message from c into v.
    16  // It will reuse buffers in between calls to avoid allocations.
    17  func Read(ctx context.Context, c *websocket.Conn, v interface{}) error {
    18  	return read(ctx, c, v)
    19  }
    20  
    21  func read(ctx context.Context, c *websocket.Conn, v interface{}) (err error) {
    22  	defer errd.Wrap(&err, "failed to read JSON message")
    23  
    24  	_, r, err := c.Reader(ctx)
    25  	if err != nil {
    26  		return err
    27  	}
    28  
    29  	b := bpool.Get()
    30  	defer bpool.Put(b)
    31  
    32  	_, err = b.ReadFrom(r)
    33  	if err != nil {
    34  		return err
    35  	}
    36  
    37  	err = json.Unmarshal(b.Bytes(), v)
    38  	if err != nil {
    39  		c.Close(websocket.StatusInvalidFramePayloadData, "failed to unmarshal JSON")
    40  		return fmt.Errorf("failed to unmarshal JSON: %w", err)
    41  	}
    42  
    43  	return nil
    44  }
    45  
    46  // Write writes the JSON message v to c.
    47  // It will reuse buffers in between calls to avoid allocations.
    48  func Write(ctx context.Context, c *websocket.Conn, v interface{}) error {
    49  	return write(ctx, c, v)
    50  }
    51  
    52  func write(ctx context.Context, c *websocket.Conn, v interface{}) (err error) {
    53  	defer errd.Wrap(&err, "failed to write JSON message")
    54  
    55  	// json.Marshal cannot reuse buffers between calls as it has to return
    56  	// a copy of the byte slice but Encoder does as it directly writes to w.
    57  	err = json.NewEncoder(util.WriterFunc(func(p []byte) (int, error) {
    58  		err := c.Write(ctx, websocket.MessageText, p)
    59  		if err != nil {
    60  			return 0, err
    61  		}
    62  		return len(p), nil
    63  	})).Encode(v)
    64  	if err != nil {
    65  		return fmt.Errorf("failed to marshal JSON: %w", err)
    66  	}
    67  	return nil
    68  }
    69
View as plain text