TUN-7125: Add management streaming logs WebSocket protocol

vendor/nhooyr.io/websocket/ws_js.go

@@ -0,0 +1,379 @@
+package websocket // import "nhooyr.io/websocket"
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"reflect"
+	"runtime"
+	"strings"
+	"sync"
+	"syscall/js"
+
+	"nhooyr.io/websocket/internal/bpool"
+	"nhooyr.io/websocket/internal/wsjs"
+	"nhooyr.io/websocket/internal/xsync"
+)
+
+// Conn provides a wrapper around the browser WebSocket API.
+type Conn struct {
+	ws wsjs.WebSocket
+
+	// read limit for a message in bytes.
+	msgReadLimit xsync.Int64
+
+	closingMu     sync.Mutex
+	isReadClosed  xsync.Int64
+	closeOnce     sync.Once
+	closed        chan struct{}
+	closeErrOnce  sync.Once
+	closeErr      error
+	closeWasClean bool
+
+	releaseOnClose   func()
+	releaseOnMessage func()
+
+	readSignal chan struct{}
+	readBufMu  sync.Mutex
+	readBuf    []wsjs.MessageEvent
+}
+
+func (c *Conn) close(err error, wasClean bool) {
+	c.closeOnce.Do(func() {
+		runtime.SetFinalizer(c, nil)
+
+		if !wasClean {
+			err = fmt.Errorf("unclean connection close: %w", err)
+		}
+		c.setCloseErr(err)
+		c.closeWasClean = wasClean
+		close(c.closed)
+	})
+}
+
+func (c *Conn) init() {
+	c.closed = make(chan struct{})
+	c.readSignal = make(chan struct{}, 1)
+
+	c.msgReadLimit.Store(32768)
+
+	c.releaseOnClose = c.ws.OnClose(func(e wsjs.CloseEvent) {
+		err := CloseError{
+			Code:   StatusCode(e.Code),
+			Reason: e.Reason,
+		}
+		// We do not know if we sent or received this close as
+		// its possible the browser triggered it without us
+		// explicitly sending it.
+		c.close(err, e.WasClean)
+
+		c.releaseOnClose()
+		c.releaseOnMessage()
+	})
+
+	c.releaseOnMessage = c.ws.OnMessage(func(e wsjs.MessageEvent) {
+		c.readBufMu.Lock()
+		defer c.readBufMu.Unlock()
+
+		c.readBuf = append(c.readBuf, e)
+
+		// Lets the read goroutine know there is definitely something in readBuf.
+		select {
+		case c.readSignal <- struct{}{}:
+		default:
+		}
+	})
+
+	runtime.SetFinalizer(c, func(c *Conn) {
+		c.setCloseErr(errors.New("connection garbage collected"))
+		c.closeWithInternal()
+	})
+}
+
+func (c *Conn) closeWithInternal() {
+	c.Close(StatusInternalError, "something went wrong")
+}
+
+// Read attempts to read a message from the connection.
+// The maximum time spent waiting is bounded by the context.
+func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
+	if c.isReadClosed.Load() == 1 {
+		return 0, nil, errors.New("WebSocket connection read closed")
+	}
+
+	typ, p, err := c.read(ctx)
+	if err != nil {
+		return 0, nil, fmt.Errorf("failed to read: %w", err)
+	}
+	if int64(len(p)) > c.msgReadLimit.Load() {
+		err := fmt.Errorf("read limited at %v bytes", c.msgReadLimit.Load())
+		c.Close(StatusMessageTooBig, err.Error())
+		return 0, nil, err
+	}
+	return typ, p, nil
+}
+
+func (c *Conn) read(ctx context.Context) (MessageType, []byte, error) {
+	select {
+	case <-ctx.Done():
+		c.Close(StatusPolicyViolation, "read timed out")
+		return 0, nil, ctx.Err()
+	case <-c.readSignal:
+	case <-c.closed:
+		return 0, nil, c.closeErr
+	}
+
+	c.readBufMu.Lock()
+	defer c.readBufMu.Unlock()
+
+	me := c.readBuf[0]
+	// We copy the messages forward and decrease the size
+	// of the slice to avoid reallocating.
+	copy(c.readBuf, c.readBuf[1:])
+	c.readBuf = c.readBuf[:len(c.readBuf)-1]
+
+	if len(c.readBuf) > 0 {
+		// Next time we read, we'll grab the message.
+		select {
+		case c.readSignal <- struct{}{}:
+		default:
+		}
+	}
+
+	switch p := me.Data.(type) {
+	case string:
+		return MessageText, []byte(p), nil
+	case []byte:
+		return MessageBinary, p, nil
+	default:
+		panic("websocket: unexpected data type from wsjs OnMessage: " + reflect.TypeOf(me.Data).String())
+	}
+}
+
+// Ping is mocked out for Wasm.
+func (c *Conn) Ping(ctx context.Context) error {
+	return nil
+}
+
+// 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)
+	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
+		// is unexpected or the passed bytes contain invalid UTF-8 for
+		// MessageText.
+		err := fmt.Errorf("failed to write: %w", err)
+		c.setCloseErr(err)
+		c.closeWithInternal()
+		return err
+	}
+	return nil
+}
+
+func (c *Conn) write(ctx context.Context, typ MessageType, p []byte) error {
+	if c.isClosed() {
+		return c.closeErr
+	}
+	switch typ {
+	case MessageBinary:
+		return c.ws.SendBytes(p)
+	case MessageText:
+		return c.ws.SendText(string(p))
+	default:
+		return fmt.Errorf("unexpected message type: %v", typ)
+	}
+}
+
+// Close closes the WebSocket with the given code and reason.
+// It will wait until the peer responds with a close frame
+// or the connection is closed.
+// It thus performs the full WebSocket close handshake.
+func (c *Conn) Close(code StatusCode, reason string) error {
+	err := c.exportedClose(code, reason)
+	if err != nil {
+		return fmt.Errorf("failed to close WebSocket: %w", err)
+	}
+	return nil
+}
+
+func (c *Conn) exportedClose(code StatusCode, reason string) error {
+	c.closingMu.Lock()
+	defer c.closingMu.Unlock()
+
+	ce := fmt.Errorf("sent close: %w", CloseError{
+		Code:   code,
+		Reason: reason,
+	})
+
+	if c.isClosed() {
+		return fmt.Errorf("tried to close with %q but connection already closed: %w", ce, c.closeErr)
+	}
+
+	c.setCloseErr(ce)
+	err := c.ws.Close(int(code), reason)
+	if err != nil {
+		return err
+	}
+
+	<-c.closed
+	if !c.closeWasClean {
+		return c.closeErr
+	}
+	return nil
+}
+
+// Subprotocol returns the negotiated subprotocol.
+// An empty string means the default protocol.
+func (c *Conn) Subprotocol() string {
+	return c.ws.Subprotocol()
+}
+
+// DialOptions represents the options available to pass to Dial.
+type DialOptions struct {
+	// Subprotocols lists the subprotocols to negotiate with the server.
+	Subprotocols []string
+}
+
+// Dial creates a new WebSocket connection to the given url with the given options.
+// The passed context bounds the maximum time spent waiting for the connection to open.
+// The returned *http.Response is always nil or a mock. It's only in the signature
+// to match the core API.
+func Dial(ctx context.Context, url string, opts *DialOptions) (*Conn, *http.Response, error) {
+	c, resp, err := dial(ctx, url, opts)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to WebSocket dial %q: %w", url, err)
+	}
+	return c, resp, nil
+}
+
+func dial(ctx context.Context, url string, opts *DialOptions) (*Conn, *http.Response, error) {
+	if opts == nil {
+		opts = &DialOptions{}
+	}
+
+	url = strings.Replace(url, "http://", "ws://", 1)
+	url = strings.Replace(url, "https://", "wss://", 1)
+
+	ws, err := wsjs.New(url, opts.Subprotocols)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	c := &Conn{
+		ws: ws,
+	}
+	c.init()
+
+	opench := make(chan struct{})
+	releaseOpen := ws.OnOpen(func(e js.Value) {
+		close(opench)
+	})
+	defer releaseOpen()
+
+	select {
+	case <-ctx.Done():
+		c.Close(StatusPolicyViolation, "dial timed out")
+		return nil, nil, ctx.Err()
+	case <-opench:
+		return c, &http.Response{
+			StatusCode: http.StatusSwitchingProtocols,
+		}, nil
+	case <-c.closed:
+		return nil, nil, c.closeErr
+	}
+}
+
+// Reader attempts to read a message from the connection.
+// The maximum time spent waiting is bounded by the context.
+func (c *Conn) Reader(ctx context.Context) (MessageType, io.Reader, error) {
+	typ, p, err := c.Read(ctx)
+	if err != nil {
+		return 0, nil, err
+	}
+	return typ, bytes.NewReader(p), nil
+}
+
+// Writer returns a writer to write a WebSocket data message to the connection.
+// It buffers the entire message in memory and then sends it when the writer
+// is closed.
+func (c *Conn) Writer(ctx context.Context, typ MessageType) (io.WriteCloser, error) {
+	return writer{
+		c:   c,
+		ctx: ctx,
+		typ: typ,
+		b:   bpool.Get(),
+	}, nil
+}
+
+type writer struct {
+	closed bool
+
+	c   *Conn
+	ctx context.Context
+	typ MessageType
+
+	b *bytes.Buffer
+}
+
+func (w writer) Write(p []byte) (int, error) {
+	if w.closed {
+		return 0, errors.New("cannot write to closed writer")
+	}
+	n, err := w.b.Write(p)
+	if err != nil {
+		return n, fmt.Errorf("failed to write message: %w", err)
+	}
+	return n, nil
+}
+
+func (w writer) Close() error {
+	if w.closed {
+		return errors.New("cannot close closed writer")
+	}
+	w.closed = true
+	defer bpool.Put(w.b)
+
+	err := w.c.Write(w.ctx, w.typ, w.b.Bytes())
+	if err != nil {
+		return fmt.Errorf("failed to close writer: %w", err)
+	}
+	return nil
+}
+
+// CloseRead implements *Conn.CloseRead for wasm.
+func (c *Conn) CloseRead(ctx context.Context) context.Context {
+	c.isReadClosed.Store(1)
+
+	ctx, cancel := context.WithCancel(ctx)
+	go func() {
+		defer cancel()
+		c.read(ctx)
+		c.Close(StatusPolicyViolation, "unexpected data message")
+	}()
+	return ctx
+}
+
+// SetReadLimit implements *Conn.SetReadLimit for wasm.
+func (c *Conn) SetReadLimit(n int64) {
+	c.msgReadLimit.Store(n)
+}
+
+func (c *Conn) setCloseErr(err error) {
+	c.closeErrOnce.Do(func() {
+		c.closeErr = fmt.Errorf("WebSocket closed: %w", err)
+	})
+}
+
+func (c *Conn) isClosed() bool {
+	select {
+	case <-c.closed:
+		return true
+	default:
+		return false
+	}
+}