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.handleServerMessages(ctx, sess)
	h.handleClientMessages(ctx, sess)
}

func (h *Handler) handleServerMessages(ctx context.Context, sess sockjs.Session) {
	messages, err := h.bus.Subscribe(ctx, module.MessageNamespaceServer)
	if err != nil {
		panic(errors.WithStack(err))
	}

	defer func() {
		// Close messages subscriber
		h.bus.Unsubscribe(ctx, module.MessageNamespaceServer, messages)

		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 msg := <-messages:
			serverMessage, ok := msg.(*module.ServerMessage)
			if !ok {
				logger.Error(
					ctx,
					"unexpected server message",
					logger.F("message", msg),
				)

				continue
			}

			sessionID := module.ContextValue[string](serverMessage.Context, ContextKeySessionID)

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

			payload, err := json.Marshal(serverMessage.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) handleClientMessages(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) {
					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]interface{}
				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(),
				})

				clientMessage := module.NewClientMessage(ctx, payload)

				logger.Debug(ctx, "publishing new client message", logger.F("message", clientMessage))

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

					return
				}

				logger.Debug(ctx, "new client message published", logger.F("message", clientMessage))

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