binding_transport.go

package main

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	_ "embed"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"github.com/lxzan/gws"
	"github.com/wailsapp/wails/v3/pkg/application"

	"voidraft/internal/services"
)

//go:embed binding_transport.js
var bindingTransportClientTemplate string

type bindingWebSocketTransport struct {
	gws.BuiltinEventHandler

	logger           *slog.Logger
	messageProcessor *application.MessageProcessor
	server           *http.Server
	jsClient         []byte
	upgrader         *gws.Upgrader
	authToken        string

	clients sync.Map
}

type bindingTransportClient struct {
	conn      *gws.Conn
	ctx       context.Context
	cancel    context.CancelFunc
	closeOnce sync.Once
}

type bindingTransportMessage struct {
	ID       string                      `json:"id,omitempty"`
	Type     string                      `json:"type"`
	Request  *application.RuntimeRequest `json:"request,omitempty"`
	Response *bindingTransportResponse   `json:"response,omitempty"`
	Event    *application.CustomEvent    `json:"event,omitempty"`
}

type bindingTransportResponse struct {
	StatusCode int    `json:"statusCode"`
	Data       any    `json:"data,omitempty"`
	Error      string `json:"error,omitempty"`
}

func newBindingTransport() *bindingWebSocketTransport {
	return &bindingWebSocketTransport{
		logger: slog.Default(),
	}
}

func (t *bindingWebSocketTransport) Start(ctx context.Context, processor *application.MessageProcessor) error {
	t.messageProcessor = processor
	authToken, err := generateBindingTransportToken()
	if err != nil {
		return fmt.Errorf("generate websocket transport auth token: %w", err)
	}
	t.authToken = authToken

	t.upgrader = gws.NewUpgrader(t, &gws.ServerOption{
		Recovery:     gws.Recovery,
		Authorize:    t.authorizeRequest,
		SubProtocols: []string{authToken},
	})

	mux := http.NewServeMux()
	mux.HandleFunc("/wails/ws", t.handleWebSocket)

	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return fmt.Errorf("listen websocket transport: %w", err)
	}

	wsURL := "ws://" + listener.Addr().String() + "/wails/ws"
	t.jsClient = []byte(strings.NewReplacer(
		"__WAILS_WS_URL__", strconv.Quote(wsURL),
		"__WAILS_WS_TOKEN__", strconv.Quote(authToken),
	).Replace(bindingTransportClientTemplate))
	t.server = &http.Server{Handler: mux}

	go func() {
		<-ctx.Done()
		if err := t.Stop(); err != nil {
			t.logger.Error("failed to stop binding websocket transport", "error", err)
		}
	}()

	go func() {
		if err := t.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
			t.logger.Error("binding websocket transport server stopped unexpectedly", "error", err)
		}
	}()

	t.logger.Info("binding websocket transport listening", "addr", listener.Addr().String())
	return nil
}

func (t *bindingWebSocketTransport) JSClient() []byte {
	return t.jsClient
}

func (t *bindingWebSocketTransport) Stop() error {
	t.clients.Range(func(key, value any) bool {
		client, ok := value.(*bindingTransportClient)
		if ok {
			client.close()
		}
		t.clients.Delete(key)
		return true
	})

	if t.server == nil {
		return nil
	}

	return t.server.Shutdown(context.Background())
}

func (t *bindingWebSocketTransport) DispatchWailsEvent(event *application.CustomEvent) {
	payload, err := json.Marshal(bindingTransportMessage{
		Type:  "event",
		Event: event,
	})
	if err != nil {
		t.logger.Warn("failed to encode websocket event", "event", event.Name, "error", err)
		return
	}

	broadcaster := gws.NewBroadcaster(gws.OpcodeText, payload)
	defer func() {
		_ = broadcaster.Close()
	}()

	t.clients.Range(func(_ any, value any) bool {
		client, ok := value.(*bindingTransportClient)
		if !ok {
			return true
		}

		if err := broadcaster.Broadcast(client.conn); err != nil {
			t.logger.Debug("failed to broadcast websocket event", "event", event.Name, "error", err)
		}

		return true
	})
}

func (t *bindingWebSocketTransport) handleWebSocket(rw http.ResponseWriter, req *http.Request) {
	conn, err := t.upgrader.Upgrade(rw, req)
	if err != nil {
		t.logger.Error("failed to upgrade websocket connection", "error", err)
		return
	}

	go conn.ReadLoop()
}

func (t *bindingWebSocketTransport) OnOpen(socket *gws.Conn) {
	t.clients.Store(socket, newBindingTransportClient(socket))
}

func (t *bindingWebSocketTransport) OnClose(socket *gws.Conn, err error) {
	if client, ok := t.loadClient(socket); ok {
		client.close()
	}
	t.clients.Delete(socket)

	if err != nil && !errors.Is(err, net.ErrClosed) {
		t.logger.Debug("binding websocket closed", "error", err)
	}
}

func (t *bindingWebSocketTransport) OnMessage(socket *gws.Conn, message *gws.Message) {
	defer message.Close()

	body := append([]byte(nil), message.Bytes()...)

	var payload bindingTransportMessage
	if err := json.Unmarshal(body, &payload); err != nil {
		t.logger.Warn("failed to decode websocket request", "error", err)
		return
	}

	if payload.Type != "request" || payload.Request == nil {
		return
	}

	if payload.Request.Args == nil {
		payload.Request.Args = &application.Args{}
	}

	client, ok := t.loadClient(socket)
	if !ok {
		return
	}

	go t.handleRequest(client, payload.ID, payload.Request)
}

func (t *bindingWebSocketTransport) handleRequest(client *bindingTransportClient, id string, req *application.RuntimeRequest) {
	if t.messageProcessor == nil {
		t.sendResponse(client, id, nil, http.StatusServiceUnavailable, errors.New("binding transport not ready"))
		return
	}

	resp, err := t.messageProcessor.HandleRuntimeCallWithIDs(client.ctx, req)
	statusCode := http.StatusOK
	if err != nil {
		statusCode = http.StatusUnprocessableEntity
		if errors.Is(err, services.ErrDocumentRevisionConflict) {
			statusCode = http.StatusConflict
		}
	}

	t.sendResponse(client, id, resp, statusCode, err)
}

func (t *bindingWebSocketTransport) sendResponse(client *bindingTransportClient, id string, data any, statusCode int, err error) {
	response := &bindingTransportResponse{
		StatusCode: statusCode,
		Data:       data,
	}
	if err != nil {
		response.Error = err.Error()
	}

	if err := client.writeJSON(bindingTransportMessage{
		ID:       id,
		Type:     "response",
		Response: response,
	}); err != nil {
		t.logger.Debug("failed to encode websocket response", "id", id, "error", err)
	}
}

func (t *bindingWebSocketTransport) loadClient(socket *gws.Conn) (*bindingTransportClient, bool) {
	value, ok := t.clients.Load(socket)
	if !ok {
		return nil, false
	}

	client, ok := value.(*bindingTransportClient)
	return client, ok
}

func (t *bindingWebSocketTransport) authorizeRequest(req *http.Request, _ gws.SessionStorage) bool {
	if !isLoopbackRequest(req) {
		t.logger.Warn("rejected non-loopback websocket connection", "remoteAddr", req.RemoteAddr)
		return false
	}

	protocols := parseWebSocketProtocols(req.Header.Get("Sec-WebSocket-Protocol"))
	if len(protocols) == 1 && secureCompareString(protocols[0], t.authToken) {
		return true
	}

	t.logger.Debug("rejected websocket connection with invalid auth token", "remoteAddr", req.RemoteAddr)
	return false
}

func generateBindingTransportToken() (string, error) {
	token := make([]byte, 32)
	if _, err := rand.Read(token); err != nil {
		return "", err
	}

	return hex.EncodeToString(token), nil
}

func isLoopbackRequest(req *http.Request) bool {
	host, _, err := net.SplitHostPort(req.RemoteAddr)
	if err != nil {
		return false
	}

	ip := net.ParseIP(host)
	return ip != nil && ip.IsLoopback()
}

func secureCompareString(actual string, expected string) bool {
	if len(actual) == 0 || len(actual) != len(expected) {
		return false
	}

	return subtle.ConstantTimeCompare([]byte(actual), []byte(expected)) == 1
}

func parseWebSocketProtocols(header string) []string {
	if header == "" {
		return nil
	}

	rawProtocols := strings.Split(header, ",")
	protocols := make([]string, 0, len(rawProtocols))
	for _, protocol := range rawProtocols {
		protocol = strings.TrimSpace(protocol)
		if protocol == "" {
			continue
		}

		protocols = append(protocols, protocol)
	}

	return protocols
}

func newBindingTransportClient(conn *gws.Conn) *bindingTransportClient {
	ctx, cancel := context.WithCancel(context.Background())
	return &bindingTransportClient{
		conn:   conn,
		ctx:    ctx,
		cancel: cancel,
	}
}

func (c *bindingTransportClient) close() {
	c.closeOnce.Do(func() {
		c.cancel()
		_ = c.conn.WriteClose(1001, nil)
	})
}

func (c *bindingTransportClient) writeJSON(message bindingTransportMessage) error {
	payload, err := json.Marshal(message)
	if err != nil {
		return err
	}

	return c.conn.WriteMessage(gws.OpcodeText, payload)
}