TUN-7125: Add management streaming logs WebSocket protocol

vendor/nhooyr.io/websocket/close_notjs.go

@@ -0,0 +1,211 @@
+// +build !js
+
+package websocket
+
+import (
+	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"log"
+	"time"
+
+	"nhooyr.io/websocket/internal/errd"
+)
+
+// Close performs the WebSocket close handshake with the given status code and reason.
+//
+// It will write a WebSocket close frame with a timeout of 5s and then wait 5s for
+// the peer to send a close frame.
+// All data messages received from the peer during the close handshake will be discarded.
+//
+// The connection can only be closed once. Additional calls to Close
+// are no-ops.
+//
+// The maximum length of reason must be 125 bytes. Avoid
+// sending a dynamic reason.
+//
+// Close will unblock all goroutines interacting with the connection once
+// complete.
+func (c *Conn) Close(code StatusCode, reason string) error {
+	return c.closeHandshake(code, reason)
+}
+
+func (c *Conn) closeHandshake(code StatusCode, reason string) (err error) {
+	defer errd.Wrap(&err, "failed to close WebSocket")
+
+	writeErr := c.writeClose(code, reason)
+	closeHandshakeErr := c.waitCloseHandshake()
+
+	if writeErr != nil {
+		return writeErr
+	}
+
+	if CloseStatus(closeHandshakeErr) == -1 {
+		return closeHandshakeErr
+	}
+
+	return nil
+}
+
+var errAlreadyWroteClose = errors.New("already wrote close")
+
+func (c *Conn) writeClose(code StatusCode, reason string) error {
+	c.closeMu.Lock()
+	wroteClose := c.wroteClose
+	c.wroteClose = true
+	c.closeMu.Unlock()
+	if wroteClose {
+		return errAlreadyWroteClose
+	}
+
+	ce := CloseError{
+		Code:   code,
+		Reason: reason,
+	}
+
+	var p []byte
+	var marshalErr error
+	if ce.Code != StatusNoStatusRcvd {
+		p, marshalErr = ce.bytes()
+		if marshalErr != nil {
+			log.Printf("websocket: %v", marshalErr)
+		}
+	}
+
+	writeErr := c.writeControl(context.Background(), opClose, p)
+	if CloseStatus(writeErr) != -1 {
+		// Not a real error if it's due to a close frame being received.
+		writeErr = nil
+	}
+
+	// We do this after in case there was an error writing the close frame.
+	c.setCloseErr(fmt.Errorf("sent close frame: %w", ce))
+
+	if marshalErr != nil {
+		return marshalErr
+	}
+	return writeErr
+}
+
+func (c *Conn) waitCloseHandshake() error {
+	defer c.close(nil)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	err := c.readMu.lock(ctx)
+	if err != nil {
+		return err
+	}
+	defer c.readMu.unlock()
+
+	if c.readCloseFrameErr != nil {
+		return c.readCloseFrameErr
+	}
+
+	for {
+		h, err := c.readLoop(ctx)
+		if err != nil {
+			return err
+		}
+
+		for i := int64(0); i < h.payloadLength; i++ {
+			_, err := c.br.ReadByte()
+			if err != nil {
+				return err
+			}
+		}
+	}
+}
+
+func parseClosePayload(p []byte) (CloseError, error) {
+	if len(p) == 0 {
+		return CloseError{
+			Code: StatusNoStatusRcvd,
+		}, nil
+	}
+
+	if len(p) < 2 {
+		return CloseError{}, fmt.Errorf("close payload %q too small, cannot even contain the 2 byte status code", p)
+	}
+
+	ce := CloseError{
+		Code:   StatusCode(binary.BigEndian.Uint16(p)),
+		Reason: string(p[2:]),
+	}
+
+	if !validWireCloseCode(ce.Code) {
+		return CloseError{}, fmt.Errorf("invalid status code %v", ce.Code)
+	}
+
+	return ce, nil
+}
+
+// See http://www.iana.org/assignments/websocket/websocket.xhtml#close-code-number
+// and https://tools.ietf.org/html/rfc6455#section-7.4.1
+func validWireCloseCode(code StatusCode) bool {
+	switch code {
+	case statusReserved, StatusNoStatusRcvd, StatusAbnormalClosure, StatusTLSHandshake:
+		return false
+	}
+
+	if code >= StatusNormalClosure && code <= StatusBadGateway {
+		return true
+	}
+	if code >= 3000 && code <= 4999 {
+		return true
+	}
+
+	return false
+}
+
+func (ce CloseError) bytes() ([]byte, error) {
+	p, err := ce.bytesErr()
+	if err != nil {
+		err = fmt.Errorf("failed to marshal close frame: %w", err)
+		ce = CloseError{
+			Code: StatusInternalError,
+		}
+		p, _ = ce.bytesErr()
+	}
+	return p, err
+}
+
+const maxCloseReason = maxControlPayload - 2
+
+func (ce CloseError) bytesErr() ([]byte, error) {
+	if len(ce.Reason) > maxCloseReason {
+		return nil, fmt.Errorf("reason string max is %v but got %q with length %v", maxCloseReason, ce.Reason, len(ce.Reason))
+	}
+
+	if !validWireCloseCode(ce.Code) {
+		return nil, fmt.Errorf("status code %v cannot be set", ce.Code)
+	}
+
+	buf := make([]byte, 2+len(ce.Reason))
+	binary.BigEndian.PutUint16(buf, uint16(ce.Code))
+	copy(buf[2:], ce.Reason)
+	return buf, nil
+}
+
+func (c *Conn) setCloseErr(err error) {
+	c.closeMu.Lock()
+	c.setCloseErrLocked(err)
+	c.closeMu.Unlock()
+}
+
+func (c *Conn) setCloseErrLocked(err error) {
+	if c.closeErr == nil {
+		c.closeErr = fmt.Errorf("WebSocket closed: %w", err)
+	}
+}
+
+func (c *Conn) isClosed() bool {
+	select {
+	case <-c.closed:
+		return true
+	default:
+		return false
+	}
+}