edge/pkg/http/sockjs.go

package http

import (
	"context"
	"encoding/json"
	"net/http"

	"forge.cadoles.com/arcad/edge/pkg/module"
	"github.com/igm/sockjs-go/v3/sockjs"
	"github.com/pkg/errors"
	"gitlab.com/wpetit/goweb/logger"
)

const (
	statusChannelClosed = iota
)

const (
	ContextKeySessionID     module.ContextKey = "sessionId"
	ContextKeyOriginRequest module.ContextKey = "originRequest"
)

func (h *Handler) handleSockJS(w http.ResponseWriter, r *http.Request) {
	h.mutex.RLock()
	defer h.mutex.RUnlock()

	h.sockjs.ServeHTTP(w, r)
}

func (h *Handler) handleSockJSSession(sess sockjs.Session) {
	ctx := logger.With(sess.Request().Context(),
		logger.F("sessionID", sess.ID()),
	)

	logger.Debug(ctx, "new sockjs session")

	defer func() {
		if sess.GetSessionState() == sockjs.SessionActive {
			if err := sess.Close(statusChannelClosed, "channel closed"); err != nil {
				logger.Error(ctx, "could not close sockjs session", logger.CapturedE(errors.WithStack(err)))
			}
		}
	}()

	go h.handleOutgoingMessages(ctx, sess)
	h.handleIncomingMessages(ctx, sess)
}

func (h *Handler) handleOutgoingMessages(ctx context.Context, sess sockjs.Session) {
	envelopes, err := h.bus.Subscribe(ctx, AddressOutgoingMessage)
	if err != nil {
		panic(errors.WithStack(err))
	}

	defer func() {
		h.bus.Unsubscribe(AddressOutgoingMessage, envelopes)

		logger.Debug(ctx, "unsubscribed")

		if sess.GetSessionState() != sockjs.SessionActive {
			return
		}

		if err := sess.Close(statusChannelClosed, "channel closed"); err != nil {
			logger.Error(ctx, "could not close sockjs session", logger.CapturedE(errors.WithStack(err)))
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return

		case env := <-envelopes:
			outgoingMessage, ok := env.Message().(*OutgoingMessage)
			if !ok {
				logger.Error(
					ctx,
					"unexpected outgoing message",
					logger.F("message", env.Message()),
				)
			}

			isDest := outgoingMessage.SessionID == "" || outgoingMessage.SessionID == sess.ID()
			if !isDest {
				continue
			}

			payload, err := json.Marshal(outgoingMessage.Data)
			if err != nil {
				logger.Error(
					ctx,
					"could not encode message",
					logger.CapturedE(errors.WithStack(err)),
				)

				continue
			}

			message := NewWebsocketMessage(
				WebsocketMessageTypeMessage,
				json.RawMessage(payload),
			)

			data, err := json.Marshal(message)
			if err != nil {
				logger.Error(
					ctx,
					"could not encode message",
					logger.CapturedE(errors.WithStack(err)),
				)

				continue
			}

			logger.Debug(ctx, "sending message")

			// Send message
			if err := sess.Send(string(data)); err != nil {
				logger.Error(
					ctx,
					"could not send message",
					logger.CapturedE(errors.WithStack(err)),
				)
			}
		}
	}
}

func (h *Handler) handleIncomingMessages(ctx context.Context, sess sockjs.Session) {
	for {
		select {
		case <-ctx.Done():
			logger.Debug(ctx, "context done")

			return

		default:
			logger.Debug(ctx, "waiting for websocket data")

			data, err := sess.RecvCtx(ctx)
			if err != nil {
				if errors.Is(err, sockjs.ErrSessionNotOpen) || errors.Is(err, context.Canceled) {
					break
				}

				logger.Error(
					ctx,
					"could not read message",
					logger.CapturedE(errors.WithStack(err)),
				)

				break
			}

			logger.Debug(ctx, "websocket data received", logger.F("data", data))

			message := &WebsocketMessage{}
			if err := json.Unmarshal([]byte(data), message); err != nil {
				logger.Error(
					ctx,
					"could not decode message",
					logger.CapturedE(errors.WithStack(err)),
				)

				break
			}

			switch {

			case message.Type == WebsocketMessageTypeMessage:
				var payload map[string]any
				if err := json.Unmarshal(message.Payload, &payload); err != nil {
					logger.Error(
						ctx,
						"could not decode payload",
						logger.CapturedE(errors.WithStack(err)),
					)

					return
				}

				ctx := logger.With(ctx, logger.F("payload", payload))
				ctx = module.WithContext(ctx, map[module.ContextKey]any{
					ContextKeySessionID:     sess.ID(),
					ContextKeyOriginRequest: sess.Request(),
				})

				incomingMessage := NewIncomingMessageEnvelope(ctx, payload)

				logger.Debug(ctx, "publishing new incoming message", logger.F("message", incomingMessage))

				if err := h.bus.Publish(incomingMessage); err != nil {
					logger.Error(ctx, "could not publish message",
						logger.CapturedE(errors.WithStack(err)),
						logger.F("message", incomingMessage),
					)

					return
				}

			default:
				logger.Error(
					ctx,
					"unsupported message type",
					logger.F("messageType", message.Type),
				)
			}
		}
	}
}

const (
	WebsocketMessageTypeMessage = "message"
)

type WebsocketMessage struct {
	Type    string          `json:"t"`
	Payload json.RawMessage `json:"p"`
}

type WebsocketMessagePayload struct {
	Data map[string]interface{} `json:"d"`
}

func NewWebsocketMessage(dataType string, payload json.RawMessage) *WebsocketMessage {
	return &WebsocketMessage{
		Type:    dataType,
		Payload: payload,
	}
}