porter/api/server/shared/websocket/response_writer.go

package websocket

import (
	"errors"
	"net/http"
	"sync"
	"syscall"

	"github.com/gorilla/websocket"
)

type WebsocketSafeReadWriter struct {
	conn *websocket.Conn
	mu   sync.Mutex
}

func (w *WebsocketSafeReadWriter) WriteJSON(v interface{}) error {
	w.mu.Lock()
	defer w.mu.Unlock()
	err := w.conn.WriteJSON(v)
	if err != nil {
		if errOr(err, websocket.ErrCloseSent, syscall.EPIPE, syscall.ECONNRESET) {
			// if close has been sent, or error is broken pipe error or connection reset, we want to
			// send a message to the error channel to ensure closure but we ignore the error
			return nil
		}

		return err
	}

	return nil
}

func (w *WebsocketSafeReadWriter) Write(data []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	err := w.conn.WriteMessage(websocket.TextMessage, data)
	if err != nil {
		if errOr(err, websocket.ErrCloseSent, syscall.EPIPE, syscall.ECONNRESET) {
			// if close has been sent, or error is broken pipe error or connection reset, we want to
			// send a message to the error channel to ensure closure but we ignore the error
			return 0, nil
		} else if err != nil {
			return 0, err
		}
	}

	return len(data), nil
}

func (w *WebsocketSafeReadWriter) ReadMessage() (messageType int, p []byte, err error) {
	return w.conn.ReadMessage()
}

func (w *WebsocketSafeReadWriter) Close() error {
	return w.conn.Close()
}

type WebsocketResponseWriter struct {
	conn       *websocket.Conn
	safeWriter *WebsocketSafeReadWriter
}

// no HTTP headers in websocket protocol
func (w *WebsocketResponseWriter) Header() http.Header {
	return nil
}

// Write attempts to write a message to the websocket connection
func (w *WebsocketResponseWriter) Write(data []byte) (int, error) {
	return w.safeWriter.Write(data)
}

// no-op; no HTTP headers in websocket protocol
func (w *WebsocketResponseWriter) WriteHeader(statusCode int) {
	return
}

// helper that returns true when `err` matches any of the candidates
func errOr(err error, candidates ...error) bool {
	res := false

	for _, cErr := range candidates {
		res = res || errors.Is(err, cErr)
	}

	return res
}