go-stomp websocket ReadWriteCloser

wsstomp.go

package wsstomp

import (
	"io"

	"github.com/gorilla/websocket"
)

type WebsocketSTOMP struct {
	connection *websocket.Conn

	readerBuffer []byte
	writeBuffer  []byte
}

func (w *WebsocketSTOMP) Read(p []byte) (int, error) {
	// if we have no more data, read the next message from the websocket
	if len(w.readerBuffer) == 0 {
		_, msg, err := w.connection.ReadMessage()
		if err != nil {
			return 0, err
		}
		w.readerBuffer = msg
	}

	n := copy(p, w.readerBuffer)
	w.readerBuffer = w.readerBuffer[n:]
	return n, nil
}

func (w *WebsocketSTOMP) Write(p []byte) (int, error) {
	var err error
	w.writeBuffer = append(w.writeBuffer, p...)
	// if we reach a null byte or the entire message is a newline (heartbeat), send the message
	if p[len(p)-1] == 0x00 || (len(w.writeBuffer) == 1 && len(p) == 1 && p[0] == 0x0a) {
		err = w.connection.WriteMessage(1, w.writeBuffer)
		w.writeBuffer = []byte{}
	}
	return len(p), err
}

func (w *WebsocketSTOMP) Close() error {
	return w.connection.Close()
}

func Connect(url string) (io.ReadWriteCloser, error) {
	con, _, err := websocket.DefaultDialer.Dial(url, nil)
	return &WebsocketSTOMP{connection: con}, err
}