cloudflared-mirror/management/events.go

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.
type LogEventType int

const (
	// Cloudflared events are signficant to cloudflared operations like connection state changes.
	// Cloudflared is also the default event type for any events that haven't been separated into a proper event type.
	Cloudflared LogEventType = iota
	HTTP
	TCP
	UDP
)

func (l LogEventType) String() string {
	switch l {
	case Cloudflared:
		return "cloudflared"
	case HTTP:
		return "http"
	case TCP:
		return "tcp"
	case UDP:
		return "udp"
	default:
		return ""
	}
}

// LogLevel corresponds to the zerolog logging levels
// "panic", "fatal", and "trace" are exempt from this list as they are rarely used and, at least
// the the first two are limited to failure conditions that lead to cloudflared shutting down.
type LogLevel string

const (
	Debug LogLevel = "debug"
	Info  LogLevel = "info"
	Warn  LogLevel = "warn"
	Error LogLevel = "error"
)

const (
	// TimeKey aligns with the zerolog.TimeFieldName
	TimeKey = "time"
	// LevelKey aligns with the zerolog.LevelFieldName
	LevelKey = "level"
	// LevelKey aligns with the zerolog.MessageFieldName
	MessageKey = "message"
	// EventTypeKey is the custom JSON key of the LogEventType in ZeroLogEvent
	EventTypeKey = "event"
	// FieldsKey is a custom JSON key to match and store every other key for a zerolog event
	FieldsKey = "fields"
)

// Log is the basic structure of the events that are sent to the client.
type Log struct {
	Time    string                 `json:"time,omitempty"`
	Level   LogLevel               `json:"level,omitempty"`
	Message string                 `json:"message,omitempty"`
	Event   LogEventType           `json:"event,omitempty"`
	Fields  map[string]interface{} `json:"fields,omitempty"`
}

// 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
}

func AsClosed(err error) *websocket.CloseError {
	var closeErr websocket.CloseError
	if errors.As(err, &closeErr) {
		return &closeErr
	}
	return nil
}