websocket.go

package main

import (
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	runtime "github.com/wailsapp/wails/v2/pkg/runtime"
)

// WebSocket support: a per-App registry of live connections. Each connection
// streams incoming text messages to the frontend via Wails runtime events
// "ws:message:<id>", and signals close via "ws:close:<id>".

// WSConnect dials a ws:// or wss:// URL and starts reading. It returns a
// connection id used to send/close and to scope the emitted events.
func (a *App) WSConnect(url string, headers map[string]string) (string, error) {
	hdr := http.Header{}
	for k, v := range headers {
		hdr.Set(k, v)
	}

	dialer := websocket.Dialer{HandshakeTimeout: 15 * time.Second}
	conn, _, err := dialer.Dial(url, hdr)
	if err != nil {
		return "", err
	}

	id := newID()
	a.wsMu.Lock()
	if a.wsConns == nil {
		a.wsConns = make(map[string]*websocket.Conn)
	}
	a.wsConns[id] = conn
	a.wsMu.Unlock()

	go a.wsRead(id, conn)
	return id, nil
}

func (a *App) wsRead(id string, conn *websocket.Conn) {
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			runtime.EventsEmit(a.ctx, "ws:close:"+id, err.Error())
			a.wsMu.Lock()
			delete(a.wsConns, id)
			a.wsMu.Unlock()
			_ = conn.Close()
			return
		}
		runtime.EventsEmit(a.ctx, "ws:message:"+id, string(msg))
	}
}

// WSSend writes a text message to the connection with the given id.
func (a *App) WSSend(id, msg string) error {
	a.wsMu.Lock()
	conn := a.wsConns[id]
	a.wsMu.Unlock()
	if conn == nil {
		return websocket.ErrCloseSent
	}
	return conn.WriteMessage(websocket.TextMessage, []byte(msg))
}

// WSClose closes the connection with the given id.
func (a *App) WSClose(id string) error {
	a.wsMu.Lock()
	conn := a.wsConns[id]
	delete(a.wsConns, id)
	a.wsMu.Unlock()
	if conn == nil {
		return nil
	}
	return conn.Close()
}