Add websocket reader

go.mod

@@ -1,23 +1,25 @@
 module github.com/cloudbase/garm-provider-common
 
-go 1.20
+go 1.22
 
 require (
-	github.com/google/uuid v1.3.0
-	github.com/gorilla/handlers v1.5.1
-	github.com/mattn/go-isatty v0.0.19
-	github.com/minio/sio v0.3.1
+	github.com/google/uuid v1.6.0
+	github.com/gorilla/handlers v1.5.2
+	github.com/gorilla/websocket v1.5.4-0.20240702125206-a62d9d2a8413
+	github.com/mattn/go-isatty v0.0.20
+	github.com/minio/sio v0.4.0
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.4
 	github.com/teris-io/shortid v0.0.0-20220617161101-71ec9f2aa569
-	golang.org/x/crypto v0.12.0
-	golang.org/x/sys v0.11.0
+	golang.org/x/crypto v0.25.0
+	golang.org/x/sys v0.22.0
 	gopkg.in/natefinch/lumberjack.v2 v2.2.1
 	gopkg.in/yaml.v3 v3.0.1
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/felixge/httpsnoop v1.0.1 // indirect
+	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	golang.org/x/net v0.27.0 // indirect
 )

go.sum

@@ -1,15 +1,17 @@
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
-github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
-github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
-github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/gorilla/handlers v1.5.1 h1:9lRY6j8DEeeBT10CvO9hGW0gmky0BprnvDI5vfhUHH4=
-github.com/gorilla/handlers v1.5.1/go.mod h1:t8XrUpc4KVXb7HGyJ4/cEnwQiaxrX/hz1Zv/4g96P1Q=
-github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA=
-github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
-github.com/minio/sio v0.3.1 h1:d59r5RTHb1OsQaSl1EaTWurzMMDRLA5fgNmjzD4eVu4=
-github.com/minio/sio v0.3.1/go.mod h1:S0ovgVgc+sTlQyhiXA1ppBLv7REM7TYi5yyq2qL/Y6o=
+github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
+github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/gorilla/handlers v1.5.2 h1:cLTUSsNkgcwhgRqvCNmdbRWG0A3N4F+M2nWKdScwyEE=
+github.com/gorilla/handlers v1.5.2/go.mod h1:dX+xVpaxdSw+q0Qek8SSsl3dfMk3jNddUkMzo0GtH0w=
+github.com/gorilla/websocket v1.5.4-0.20240702125206-a62d9d2a8413 h1:0Zn/h+BUQg6QHkybGvjFD7BnIbjjz3oWUObacn//1Go=
+github.com/gorilla/websocket v1.5.4-0.20240702125206-a62d9d2a8413/go.mod h1:r4w70xmWCQKmi1ONH4KIaBptdivuRPyosB9RmPlGEwA=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
+github.com/minio/sio v0.4.0 h1:u4SWVEm5lXSqU42ZWawV0D9I5AZ5YMmo2RXpEQ/kRhc=
+github.com/minio/sio v0.4.0/go.mod h1:oBSjJeGbBdRMZZwna07sX9EFzZy+ywu5aofRiV1g79I=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -18,11 +20,13 @@ github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcU
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 github.com/teris-io/shortid v0.0.0-20220617161101-71ec9f2aa569 h1:xzABM9let0HLLqFypcxvLmlvEciCHL7+Lv+4vwZqecI=
 github.com/teris-io/shortid v0.0.0-20220617161101-71ec9f2aa569/go.mod h1:2Ly+NIftZN4de9zRmENdYbvPQeaVIYKWpLFStLFEBgI=
-golang.org/x/crypto v0.12.0 h1:tFM/ta59kqch6LlvYnPa0yx5a83cL2nHflFhYKvv9Yk=
-golang.org/x/crypto v0.12.0/go.mod h1:NF0Gs7EO5K4qLn+Ylc+fih8BSTeIjAP05siRnAh98yw=
+golang.org/x/crypto v0.25.0 h1:ypSNr+bnYL2YhwoMt2zPxHFmbAN1KZs/njMG3hxUp30=
+golang.org/x/crypto v0.25.0/go.mod h1:T+wALwcMOSE0kXgUAnPAHqTLW+XHgcELELW8VaDgm/M=
+golang.org/x/net v0.27.0 h1:5K3Njcw06/l2y9vpGCSdcxWOYHOUk3dVNGDXN+FvAys=
+golang.org/x/net v0.27.0/go.mod h1:dDi0PyhWNoiUOrAS8uXv/vnScO4wnHQO4mj9fn/RytE=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.11.0 h1:eG7RXZHdqOJ1i+0lgLgCpSXAp6M3LYlAo6osgSi0xOM=
-golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI=
+golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc=

util/websocket/reader.go

@@ -0,0 +1,184 @@
+package websocket
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"net/url"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+const (
+	// Time allowed to write a message to the peer.
+	writeWait = 10 * time.Second
+
+	// Time allowed to read the next pong message from the peer.
+	pongWait = 60 * time.Second
+
+	// Send pings to peer with this period. Must be less than pongWait.
+	pingPeriod = (pongWait * 9) / 10
+
+	// Maximum message size allowed from peer.
+	maxMessageSize = 16384 // 16 KB
+)
+
+// MessageHandler is a function that processes a message received from a websocket connection.
+type MessageHandler func(msgType int, msg []byte) error
+
+type APIErrorResponse struct {
+	Error   string `json:"error"`
+	Details string `json:"details"`
+}
+
+// NewReader creates a new websocket reader. The reader will pass on any message it receives to the
+// handler function. The handler function should return an error if it fails to process the message.
+func NewReader(ctx context.Context, baseURL, pth, token string, handler MessageHandler) (*Reader, error) {
+	parsedURL, err := url.Parse(baseURL)
+	if err != nil {
+		return nil, err
+	}
+
+	wsScheme := "ws"
+	if parsedURL.Scheme == "https" {
+		wsScheme = "wss"
+	}
+	u := url.URL{Scheme: wsScheme, Host: parsedURL.Host, Path: pth}
+	header := http.Header{}
+	header.Add("Authorization", fmt.Sprintf("Bearer %s", token))
+
+	return &Reader{
+		ctx:     ctx,
+		url:     u,
+		header:  header,
+		handler: handler,
+		done:    make(chan struct{}),
+	}, nil
+}
+
+type Reader struct {
+	ctx    context.Context
+	url    url.URL
+	header http.Header
+
+	done    chan struct{}
+	running bool
+
+	handler MessageHandler
+
+	conn     *websocket.Conn
+	mux      sync.Mutex
+	writeMux sync.Mutex
+}
+
+func (w *Reader) Stop() {
+	w.mux.Lock()
+	defer w.mux.Unlock()
+	if !w.running {
+		return
+	}
+	w.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+	w.conn.Close()
+	close(w.done)
+	w.running = false
+}
+
+func (w *Reader) Done() <-chan struct{} {
+	return w.done
+}
+
+func (w *Reader) WriteMessage(messageType int, data []byte) error {
+	// The websocket package does not support concurrent writes and panics if it
+	// detects that one has occurred, so we need to lock the writeMux to prevent
+	// concurrent writes to the same connection.
+	w.writeMux.Lock()
+	defer w.writeMux.Unlock()
+	if !w.running {
+		return fmt.Errorf("websocket is not running")
+	}
+	if err := w.conn.SetWriteDeadline(time.Now().Add(writeWait)); err != nil {
+		return err
+	}
+	return w.conn.WriteMessage(messageType, data)
+}
+
+func (w *Reader) Start() error {
+	w.mux.Lock()
+	defer w.mux.Unlock()
+	if w.running {
+		return nil
+	}
+
+	c, response, err := websocket.DefaultDialer.Dial(w.url.String(), w.header)
+	if err != nil {
+		var resp APIErrorResponse
+		var msg string
+		var status string
+		if response != nil {
+			if response.Body != nil {
+				if err := json.NewDecoder(response.Body).Decode(&resp); err == nil {
+					msg = resp.Details
+				}
+			}
+			status = response.Status
+		}
+		return fmt.Errorf("failed to stream logs: %q %s (%s)", err, msg, status)
+	}
+	w.conn = c
+	w.running = true
+	go w.loop()
+	go w.handlerReader()
+	return nil
+}
+
+func (w *Reader) handlerReader() {
+	defer w.Stop()
+	w.writeMux.Lock()
+	w.conn.SetReadLimit(maxMessageSize)
+	w.conn.SetReadDeadline(time.Now().Add(pongWait))
+	w.conn.SetPongHandler(func(string) error { w.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
+	w.writeMux.Unlock()
+	for {
+		msgType, message, err := w.conn.ReadMessage()
+		if err != nil {
+			if IsErrorOfInterest(err) {
+				// TODO(gabriel-samfira): we should allow for an error channel that can be used to signal
+				// the caller that the connection has been closed.
+				slog.With(slog.Any("error", err)).Error("reading log message")
+			}
+			return
+		}
+		if w.handler != nil {
+			if err := w.handler(msgType, message); err != nil {
+				slog.With(slog.Any("error", err)).Error("handling log message")
+			}
+		}
+	}
+}
+
+func (w *Reader) loop() {
+	defer w.Stop()
+	ticker := time.NewTicker(pingPeriod)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-w.ctx.Done():
+			return
+		case <-w.Done():
+			return
+		case <-ticker.C:
+			w.writeMux.Lock()
+			w.conn.SetWriteDeadline(time.Now().Add(writeWait))
+			err := w.conn.WriteMessage(websocket.PingMessage, nil)
+			if err != nil {
+				w.writeMux.Unlock()
+				return
+			}
+			w.writeMux.Unlock()
+		}
+	}
+}

util/websocket/util.go

@@ -0,0 +1,37 @@
+package websocket
+
+import (
+	"errors"
+	"net"
+
+	"github.com/gorilla/websocket"
+)
+
+func IsErrorOfInterest(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	if errors.Is(err, websocket.ErrCloseSent) {
+		return false
+	}
+
+	if errors.Is(err, websocket.ErrBadHandshake) {
+		return false
+	}
+
+	if errors.Is(err, net.ErrClosed) {
+		return false
+	}
+
+	asCloseErr, ok := err.(*websocket.CloseError)
+	if ok {
+		switch asCloseErr.Code {
+		case websocket.CloseNormalClosure, websocket.CloseGoingAway,
+			websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure:
+			return false
+		}
+	}
+
+	return true
+}