picoclaw/pkg/channels/webui.go

package channels

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"

	"github.com/sipeed/picoclaw/pkg/bus"
	"github.com/sipeed/picoclaw/pkg/config"
	"github.com/sipeed/picoclaw/pkg/logger"
)

type WebUIChannel struct {
	*BaseChannel
	cfg        config.GatewayConfig
	httpServer *http.Server
	mu         sync.RWMutex
	clients    map[*webUIClient]struct{}
}

type webUIClient struct {
	conn    *websocket.Conn
	chatID  string
	sender  string
	writeMu sync.Mutex
}

type webUIInboundMessage struct {
	ChatID   string `json:"chat_id"`
	SenderID string `json:"sender_id"`
	Content  string `json:"content"`
}

type webUIOutboundMessage struct {
	Type    string `json:"type"`
	ChatID  string `json:"chat_id"`
	Content string `json:"content"`
}

func NewWebUIChannel(cfg config.GatewayConfig, messageBus *bus.MessageBus) (*WebUIChannel, error) {
	base := NewBaseChannel("webui", cfg, messageBus, nil)
	return &WebUIChannel{
		BaseChannel: base,
		cfg:         cfg,
		clients:     make(map[*webUIClient]struct{}),
	}, nil
}

func (c *WebUIChannel) Start(ctx context.Context) error {
	addr, err := c.cfg.ResolvedAddr()
	if err != nil {
		return err
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/ws", c.handleWS)
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})
	mux.Handle("/", c.staticHandler())

	c.httpServer = &http.Server{
		Addr:              addr,
		Handler:           mux,
		ReadHeaderTimeout: 10 * time.Second,
	}

	go func() {
		logger.InfoCF("webui", "WebUI server listening", map[string]interface{}{
			"addr": addr,
			"bind": c.cfg.Bind,
		})
		if err := c.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logger.ErrorCF("webui", "WebUI server error", map[string]interface{}{"error": err.Error()})
		}
	}()

	c.setRunning(true)
	return nil
}

func (c *WebUIChannel) Stop(ctx context.Context) error {
	c.setRunning(false)

	c.mu.Lock()
	for cl := range c.clients {
		cl.conn.Close()
		delete(c.clients, cl)
	}
	c.mu.Unlock()

	if c.httpServer != nil {
		shutdownCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
		defer cancel()
		_ = c.httpServer.Shutdown(shutdownCtx)
	}
	return nil
}

func (c *WebUIChannel) Send(ctx context.Context, msg bus.OutboundMessage) error {
	if !c.IsRunning() {
		return fmt.Errorf("webui channel not running")
	}

	out := webUIOutboundMessage{Type: "message", ChatID: msg.ChatID, Content: msg.Content}
	payload, err := json.Marshal(out)
	if err != nil {
		return err
	}

	c.mu.RLock()
	clients := make([]*webUIClient, 0, len(c.clients))
	for cl := range c.clients {
		clients = append(clients, cl)
	}
	c.mu.RUnlock()

	var sendErr error
	for _, cl := range clients {
		if msg.ChatID != "" && cl.chatID != "" && cl.chatID != msg.ChatID {
			continue
		}
		cl.writeMu.Lock()
		err := cl.conn.WriteMessage(websocket.TextMessage, payload)
		cl.writeMu.Unlock()
		if err != nil {
			sendErr = err
			c.removeClient(cl)
		}
	}

	return sendErr
}

func (c *WebUIChannel) removeClient(cl *webUIClient) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if _, ok := c.clients[cl]; ok {
		_ = cl.conn.Close()
		delete(c.clients, cl)
	}
}

func (c *WebUIChannel) isAuthorized(u *url.URL) bool {
	expected := strings.TrimSpace(c.cfg.Token)
	if expected == "" {
		return true
	}
	provided := strings.TrimSpace(u.Query().Get("token"))
	return provided != "" && provided == expected
}

func (c *WebUIChannel) handleWS(w http.ResponseWriter, r *http.Request) {
	if !c.isAuthorized(r.URL) {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}

	upgrader := websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			// If token auth is enabled, only allow same-origin WS to reduce cross-site hijacking.
			// If token is not set, keep permissive behavior for local/LAN usage.
			if strings.TrimSpace(c.cfg.Token) == "" {
				return true
			}
			origin := r.Header.Get("Origin")
			if origin == "" {
				return true
			}
			u, err := url.Parse(origin)
			if err != nil {
				return false
			}
			return strings.EqualFold(u.Host, r.Host)
		},
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}

	client := &webUIClient{
		conn:   conn,
		chatID: r.URL.Query().Get("chat_id"),
		sender: r.RemoteAddr,
	}
	if client.chatID == "" {
		client.chatID = "browser"
	}

	c.mu.Lock()
	c.clients[client] = struct{}{}
	c.mu.Unlock()

	defer c.removeClient(client)

	for {
		_, data, err := conn.ReadMessage()
		if err != nil {
			return
		}

		var in webUIInboundMessage
		if err := json.Unmarshal(data, &in); err != nil {
			continue
		}
		content := strings.TrimSpace(in.Content)
		if content == "" {
			continue
		}

		chatID := strings.TrimSpace(in.ChatID)
		if chatID == "" {
			chatID = client.chatID
		}
		senderID := strings.TrimSpace(in.SenderID)
		if senderID == "" {
			senderID = client.sender
		}

		client.chatID = chatID
		client.sender = senderID

		c.HandleMessage(senderID, chatID, content, nil, map[string]string{"source": "webui"})
	}
}

func (c *WebUIChannel) staticHandler() http.Handler {
	root := c.findUIRoot()
	fs := http.FileServer(http.Dir(root))

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/" {
			fs.ServeHTTP(w, r)
			return
		}

		reqPath := filepath.Clean(strings.TrimPrefix(r.URL.Path, "/"))
		full := filepath.Join(root, reqPath)
		if st, err := os.Stat(full); err == nil && !st.IsDir() {
			fs.ServeHTTP(w, r)
			return
		}

		index := filepath.Join(root, "index.html")
		if _, err := os.Stat(index); err == nil {
			http.ServeFile(w, r, index)
			return
		}

		http.NotFound(w, r)
	})
}

func (c *WebUIChannel) findUIRoot() string {
	candidates := []string{
		filepath.Join(string(os.PathSeparator), "usr", "local", "share", "picoclaw", "ui", "dist"),
		filepath.Join(string(os.PathSeparator), "usr", "share", "picoclaw", "ui", "dist"),
		filepath.Join("ui", "dist"),
		"ui",
	}
	for _, p := range candidates {
		if st, err := os.Stat(p); err == nil && st.IsDir() {
			return p
		}
	}
	return "ui"
}