portainer/api/ws/stream.go

package ws

import (
	"io"
	"time"
	"unicode/utf8"

	"github.com/gorilla/websocket"
	"github.com/rs/zerolog/log"
)

const ReaderBufferSize = 2048

// MessageHandler processes a WebSocket message before it is forwarded to the writer.
// It returns true if the message was handled and should not be forwarded to the writer.
type MessageHandler func(messageType int, data []byte) bool

func StreamFromWebsocketToWriter(websocketConn *websocket.Conn, writer io.Writer, errorChan chan error, handlers ...MessageHandler) {
	for {
		messageType, in, err := websocketConn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway) {
				log.Debug().Err(err).Msg("unexpected close error")
			}
			errorChan <- err

			return
		}

		if messageType != websocket.TextMessage && messageType != websocket.BinaryMessage {
			continue
		}

		handled := false
		for _, h := range handlers {
			if h(messageType, in) {
				handled = true
				break
			}
		}

		if handled {
			continue
		}

		if _, err := writer.Write(in); err != nil {
			log.Debug().Err(err).Msg("writing error")
			errorChan <- err

			return
		}
	}
}

func StreamFromReaderToWebsocket(websocketConn *websocket.Conn, reader io.Reader, errorChan chan error) {
	out := make([]byte, ReaderBufferSize)

	for {
		n, err := reader.Read(out)
		if err != nil {
			errorChan <- err

			break
		}

		processedOutput := ValidString(string(out[:n]))

		if err := websocketConn.SetWriteDeadline(time.Now().Add(WriteWait)); err != nil {
			errorChan <- err

			break
		}

		if err := websocketConn.WriteMessage(websocket.TextMessage, []byte(processedOutput)); err != nil {
			errorChan <- err

			break
		}
	}
}

func ValidString(s string) string {
	if utf8.ValidString(s) {
		return s
	}

	v := make([]rune, 0, len(s))

	for i, r := range s {
		if r == utf8.RuneError {
			_, size := utf8.DecodeRuneInString(s[i:])
			if size == 1 {
				continue
			}
		}

		v = append(v, r)
	}

	return string(v)
}