TUN-7125: Add management streaming logs WebSocket protocol

management/events.go

@@ -1,5 +1,67 @@
 package management
 
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/rs/zerolog"
+	"nhooyr.io/websocket"
+)
+
+var (
+	errInvalidMessageType = fmt.Errorf("invalid message type was provided")
+)
+
+// ServerEventType represents the event types that can come from the server
+type ServerEventType string
+
+// ClientEventType represents the event types that can come from the client
+type ClientEventType string
+
+const (
+	UnknownClientEventType ClientEventType = ""
+	StartStreaming         ClientEventType = "start_streaming"
+	StopStreaming          ClientEventType = "stop_streaming"
+
+	UnknownServerEventType ServerEventType = ""
+	Logs                   ServerEventType = "logs"
+)
+
+// ServerEvent is the base struct that informs, based of the Type field, which Event type was provided from the server.
+type ServerEvent struct {
+	Type ServerEventType `json:"type,omitempty"`
+	// The raw json message is provided to allow better deserialization once the type is known
+	event jsoniter.RawMessage
+}
+
+// ClientEvent is the base struct that informs, based of the Type field, which Event type was provided from the client.
+type ClientEvent struct {
+	Type ClientEventType `json:"type,omitempty"`
+	// The raw json message is provided to allow better deserialization once the type is known
+	event jsoniter.RawMessage
+}
+
+// EventStartStreaming signifies that the client wishes to start receiving log events.
+// Additional filters can be provided to augment the log events requested.
+type EventStartStreaming struct {
+	ClientEvent
+	Filters []string `json:"filters"`
+}
+
+// EventStopStreaming signifies that the client wishes to halt receiving log events.
+type EventStopStreaming struct {
+	ClientEvent
+}
+
+// EventLog is the event that the server sends to the client with the log events.
+type EventLog struct {
+	ServerEvent
+	Logs []Log `json:"logs"`
+}
+
 // LogEventType is the way that logging messages are able to be filtered.
 // Example: assigning LogEventType.Cloudflared to a zerolog event will allow the client to filter for only
 // the Cloudflared-related events.
@@ -38,3 +100,113 @@ const (
 	Warn  LogLevel = "warn"
 	Error LogLevel = "error"
 )
+
+// Log is the basic structure of the events that are sent to the client.
+type Log struct {
+	Event     LogEventType `json:"event"`
+	Timestamp string       `json:"timestamp"`
+	Level     LogLevel     `json:"level"`
+	Message   string       `json:"message"`
+}
+
+// IntoClientEvent unmarshals the provided ClientEvent into the proper type.
+func IntoClientEvent[T EventStartStreaming | EventStopStreaming](e *ClientEvent, eventType ClientEventType) (*T, bool) {
+	if e.Type != eventType {
+		return nil, false
+	}
+	event := new(T)
+	err := json.Unmarshal(e.event, event)
+	if err != nil {
+		return nil, false
+	}
+	return event, true
+}
+
+// IntoServerEvent unmarshals the provided ServerEvent into the proper type.
+func IntoServerEvent[T EventLog](e *ServerEvent, eventType ServerEventType) (*T, bool) {
+	if e.Type != eventType {
+		return nil, false
+	}
+	event := new(T)
+	err := json.Unmarshal(e.event, event)
+	if err != nil {
+		return nil, false
+	}
+	return event, true
+}
+
+// ReadEvent will read a message from the websocket connection and parse it into a valid ServerEvent.
+func ReadServerEvent(c *websocket.Conn, ctx context.Context) (*ServerEvent, error) {
+	message, err := readMessage(c, ctx)
+	if err != nil {
+		return nil, err
+	}
+	event := ServerEvent{}
+	if err := json.Unmarshal(message, &event); err != nil {
+		return nil, err
+	}
+	switch event.Type {
+	case Logs:
+		event.event = message
+		return &event, nil
+	case UnknownServerEventType:
+		return nil, errInvalidMessageType
+	default:
+		return nil, fmt.Errorf("invalid server message type was provided: %s", event.Type)
+	}
+}
+
+// ReadEvent will read a message from the websocket connection and parse it into a valid ClientEvent.
+func ReadClientEvent(c *websocket.Conn, ctx context.Context) (*ClientEvent, error) {
+	message, err := readMessage(c, ctx)
+	if err != nil {
+		return nil, err
+	}
+	event := ClientEvent{}
+	if err := json.Unmarshal(message, &event); err != nil {
+		return nil, err
+	}
+	switch event.Type {
+	case StartStreaming, StopStreaming:
+		event.event = message
+		return &event, nil
+	case UnknownClientEventType:
+		return nil, errInvalidMessageType
+	default:
+		return nil, fmt.Errorf("invalid client message type was provided: %s", event.Type)
+	}
+}
+
+// readMessage will read a message from the websocket connection and return the payload.
+func readMessage(c *websocket.Conn, ctx context.Context) ([]byte, error) {
+	messageType, reader, err := c.Reader(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if messageType != websocket.MessageText {
+		return nil, errInvalidMessageType
+	}
+	return io.ReadAll(reader)
+}
+
+// WriteEvent will write a Event type message to the websocket connection.
+func WriteEvent(c *websocket.Conn, ctx context.Context, event any) error {
+	payload, err := json.Marshal(event)
+	if err != nil {
+		return err
+	}
+	return c.Write(ctx, websocket.MessageText, payload)
+}
+
+// IsClosed returns true if the websocket error is a websocket.CloseError; returns false if not a
+// websocket.CloseError
+func IsClosed(err error, log *zerolog.Logger) bool {
+	var closeErr websocket.CloseError
+	if errors.As(err, &closeErr) {
+		if closeErr.Code != websocket.StatusNormalClosure {
+			log.Debug().Msgf("connection is already closed: (%d) %s", closeErr.Code, closeErr.Reason)
+		}
+		return true
+	}
+	return false
+}