update

plugins/host_websocket.go

@@ -0,0 +1,400 @@
+package plugins
+
+import (
+	"context"
+	"encoding/binary"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	gorillaws "github.com/gorilla/websocket"
+	gonanoid "github.com/matoous/go-nanoid/v2"
+	"github.com/navidrome/navidrome/log"
+	"github.com/navidrome/navidrome/plugins/api"
+	"github.com/navidrome/navidrome/plugins/host/websocket"
+)
+
+// WebSocketConnection represents a WebSocket connection
+type WebSocketConnection struct {
+	Conn         *gorillaws.Conn
+	PluginName   string
+	ConnectionID string
+	Done         chan struct{}
+	mu           sync.Mutex
+}
+
+// WebSocketHostFunctions implements the websocket.WebSocketService interface
+type WebSocketHostFunctions struct {
+	ws          *websocketService
+	pluginID    string
+	permissions *webSocketPermissions
+}
+
+func (s WebSocketHostFunctions) Connect(ctx context.Context, req *websocket.ConnectRequest) (*websocket.ConnectResponse, error) {
+	return s.ws.connect(ctx, s.pluginID, req, s.permissions)
+}
+
+func (s WebSocketHostFunctions) SendText(ctx context.Context, req *websocket.SendTextRequest) (*websocket.SendTextResponse, error) {
+	return s.ws.sendText(ctx, s.pluginID, req)
+}
+
+func (s WebSocketHostFunctions) SendBinary(ctx context.Context, req *websocket.SendBinaryRequest) (*websocket.SendBinaryResponse, error) {
+	return s.ws.sendBinary(ctx, s.pluginID, req)
+}
+
+func (s WebSocketHostFunctions) Close(ctx context.Context, req *websocket.CloseRequest) (*websocket.CloseResponse, error) {
+	return s.ws.close(ctx, s.pluginID, req)
+}
+
+// websocketService implements the WebSocket service functionality
+type websocketService struct {
+	connections map[string]*WebSocketConnection
+	manager     *managerImpl
+	mu          sync.RWMutex
+}
+
+// newWebsocketService creates a new websocketService instance
+func newWebsocketService(manager *managerImpl) *websocketService {
+	return &websocketService{
+		connections: make(map[string]*WebSocketConnection),
+		manager:     manager,
+	}
+}
+
+// HostFunctions returns the WebSocketHostFunctions for the given plugin
+func (s *websocketService) HostFunctions(pluginID string, permissions *webSocketPermissions) WebSocketHostFunctions {
+	return WebSocketHostFunctions{
+		ws:          s,
+		pluginID:    pluginID,
+		permissions: permissions,
+	}
+}
+
+// Safe accessor methods
+
+// hasConnection safely checks if a connection exists
+func (s *websocketService) hasConnection(id string) bool {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	_, exists := s.connections[id]
+	return exists
+}
+
+// connectionCount safely returns the number of connections
+func (s *websocketService) connectionCount() int {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return len(s.connections)
+}
+
+// getConnection safely retrieves a connection by internal ID
+func (s *websocketService) getConnection(internalConnectionID string) (*WebSocketConnection, error) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	conn, exists := s.connections[internalConnectionID]
+
+	if !exists {
+		return nil, fmt.Errorf("connection not found")
+	}
+	return conn, nil
+}
+
+// internalConnectionID builds the internal connection ID from plugin and connection ID
+func internalConnectionID(pluginName, connectionID string) string {
+	return pluginName + ":" + connectionID
+}
+
+// extractConnectionID extracts the original connection ID from an internal ID
+func extractConnectionID(internalID string) (string, error) {
+	parts := strings.Split(internalID, ":")
+	if len(parts) != 2 {
+		return "", fmt.Errorf("invalid internal connection ID format: %s", internalID)
+	}
+	return parts[1], nil
+}
+
+// connect establishes a new WebSocket connection
+func (s *websocketService) connect(ctx context.Context, pluginID string, req *websocket.ConnectRequest, permissions *webSocketPermissions) (*websocket.ConnectResponse, error) {
+	if s.manager == nil {
+		return nil, fmt.Errorf("websocket service not properly initialized")
+	}
+
+	// Check permissions if they exist
+	if permissions != nil {
+		if err := permissions.IsConnectionAllowed(req.Url); err != nil {
+			log.Warn(ctx, "WebSocket connection blocked by permissions", "plugin", pluginID, "url", req.Url, err)
+			return &websocket.ConnectResponse{Error: "Connection blocked by plugin permissions: " + err.Error()}, nil
+		}
+	}
+
+	// Create websocket dialer with the headers
+	dialer := gorillaws.DefaultDialer
+	header := make(map[string][]string)
+	for k, v := range req.Headers {
+		header[k] = []string{v}
+	}
+
+	// Connect to the WebSocket server
+	conn, resp, err := dialer.DialContext(ctx, req.Url, header)
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to WebSocket server: %w", err)
+	}
+	defer resp.Body.Close()
+
+	// Generate a connection ID
+	if req.ConnectionId == "" {
+		req.ConnectionId, _ = gonanoid.New(10)
+	}
+	connectionID := req.ConnectionId
+	internal := internalConnectionID(pluginID, connectionID)
+
+	// Create the connection object
+	wsConn := &WebSocketConnection{
+		Conn:         conn,
+		PluginName:   pluginID,
+		ConnectionID: connectionID,
+		Done:         make(chan struct{}),
+	}
+
+	// Store the connection
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.connections[internal] = wsConn
+
+	log.Debug("WebSocket connection established", "plugin", pluginID, "connectionID", connectionID, "url", req.Url)
+
+	// Start the message handling goroutine
+	go s.handleMessages(internal, wsConn)
+
+	return &websocket.ConnectResponse{
+		ConnectionId: connectionID,
+	}, nil
+}
+
+// writeMessage is a helper to send messages to a websocket connection
+func (s *websocketService) writeMessage(pluginID string, connID string, messageType int, data []byte) error {
+	internal := internalConnectionID(pluginID, connID)
+
+	conn, err := s.getConnection(internal)
+	if err != nil {
+		return err
+	}
+
+	conn.mu.Lock()
+	defer conn.mu.Unlock()
+
+	if err := conn.Conn.WriteMessage(messageType, data); err != nil {
+		return fmt.Errorf("failed to send message: %w", err)
+	}
+
+	return nil
+}
+
+// sendText sends a text message over a WebSocket connection
+func (s *websocketService) sendText(ctx context.Context, pluginID string, req *websocket.SendTextRequest) (*websocket.SendTextResponse, error) {
+	if err := s.writeMessage(pluginID, req.ConnectionId, gorillaws.TextMessage, []byte(req.Message)); err != nil {
+		return &websocket.SendTextResponse{Error: err.Error()}, nil //nolint:nilerr
+	}
+	return &websocket.SendTextResponse{}, nil
+}
+
+// sendBinary sends binary data over a WebSocket connection
+func (s *websocketService) sendBinary(ctx context.Context, pluginID string, req *websocket.SendBinaryRequest) (*websocket.SendBinaryResponse, error) {
+	if err := s.writeMessage(pluginID, req.ConnectionId, gorillaws.BinaryMessage, req.Data); err != nil {
+		return &websocket.SendBinaryResponse{Error: err.Error()}, nil //nolint:nilerr
+	}
+	return &websocket.SendBinaryResponse{}, nil
+}
+
+// close closes a WebSocket connection
+func (s *websocketService) close(ctx context.Context, pluginID string, req *websocket.CloseRequest) (*websocket.CloseResponse, error) {
+	internal := internalConnectionID(pluginID, req.ConnectionId)
+
+	s.mu.Lock()
+	conn, exists := s.connections[internal]
+	if !exists {
+		s.mu.Unlock()
+		return &websocket.CloseResponse{Error: "connection not found"}, nil
+	}
+	delete(s.connections, internal)
+	s.mu.Unlock()
+
+	// Signal the message handling goroutine to stop
+	close(conn.Done)
+
+	// Close the connection with the specified code and reason
+	conn.mu.Lock()
+	defer conn.mu.Unlock()
+
+	err := conn.Conn.WriteControl(
+		gorillaws.CloseMessage,
+		gorillaws.FormatCloseMessage(int(req.Code), req.Reason),
+		time.Now().Add(time.Second),
+	)
+	if err != nil {
+		log.Error("Error sending close message", "plugin", pluginID, "error", err)
+	}
+
+	if err := conn.Conn.Close(); err != nil {
+		return nil, fmt.Errorf("error closing connection: %w", err)
+	}
+
+	log.Debug("WebSocket connection closed", "plugin", pluginID, "connectionID", req.ConnectionId)
+	return &websocket.CloseResponse{}, nil
+}
+
+// handleMessages processes incoming WebSocket messages
+func (s *websocketService) handleMessages(internalID string, conn *WebSocketConnection) {
+	// Get the original connection ID (without plugin prefix)
+	connectionID, err := extractConnectionID(internalID)
+	if err != nil {
+		log.Error("Invalid internal connection ID", "id", internalID, "error", err)
+		return
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	defer func() {
+		// Ensure the connection is removed from the map if not already removed
+		s.mu.Lock()
+		defer s.mu.Unlock()
+		delete(s.connections, internalID)
+
+		log.Debug("WebSocket message handler stopped", "plugin", conn.PluginName, "connectionID", connectionID)
+	}()
+
+	// Add connection info to context
+	ctx = log.NewContext(ctx,
+		"connectionID", connectionID,
+		"plugin", conn.PluginName,
+	)
+
+	for {
+		select {
+		case <-conn.Done:
+			// Connection was closed by a Close call
+			return
+		default:
+			// Set a read deadline
+			_ = conn.Conn.SetReadDeadline(time.Now().Add(time.Second * 60))
+
+			// Read the next message
+			messageType, message, err := conn.Conn.ReadMessage()
+			if err != nil {
+				s.notifyErrorCallback(ctx, connectionID, conn, err.Error())
+				return
+			}
+
+			// Reset the read deadline
+			_ = conn.Conn.SetReadDeadline(time.Time{})
+
+			// Process the message based on its type
+			switch messageType {
+			case gorillaws.TextMessage:
+				s.notifyTextCallback(ctx, connectionID, conn, string(message))
+			case gorillaws.BinaryMessage:
+				s.notifyBinaryCallback(ctx, connectionID, conn, message)
+			case gorillaws.CloseMessage:
+				code := gorillaws.CloseNormalClosure
+				reason := ""
+				if len(message) >= 2 {
+					code = int(binary.BigEndian.Uint16(message[:2]))
+					if len(message) > 2 {
+						reason = string(message[2:])
+					}
+				}
+				s.notifyCloseCallback(ctx, connectionID, conn, code, reason)
+				return
+			}
+		}
+	}
+}
+
+// executeCallback is a common function that handles the plugin loading and execution
+// for all types of callbacks
+func (s *websocketService) executeCallback(ctx context.Context, pluginID, methodName string, fn func(context.Context, api.WebSocketCallback) error) {
+	log.Debug(ctx, "WebSocket received")
+
+	start := time.Now()
+
+	// Get the plugin
+	p := s.manager.LoadPlugin(pluginID, CapabilityWebSocketCallback)
+	if p == nil {
+		log.Error(ctx, "Plugin not found for WebSocket callback")
+		return
+	}
+
+	_, _ = callMethod(ctx, p, methodName, func(inst api.WebSocketCallback) (struct{}, error) {
+		// Call the appropriate callback function
+		log.Trace(ctx, "Executing WebSocket callback")
+		if err := fn(ctx, inst); err != nil {
+			log.Error(ctx, "Error executing WebSocket callback", "elapsed", time.Since(start), err)
+			return struct{}{}, fmt.Errorf("error executing WebSocket callback: %w", err)
+		}
+		log.Debug(ctx, "WebSocket callback executed", "elapsed", time.Since(start))
+		return struct{}{}, nil
+	})
+}
+
+// notifyTextCallback notifies the plugin of a text message
+func (s *websocketService) notifyTextCallback(ctx context.Context, connectionID string, conn *WebSocketConnection, message string) {
+	req := &api.OnTextMessageRequest{
+		ConnectionId: connectionID,
+		Message:      message,
+	}
+
+	ctx = log.NewContext(ctx, "callback", "OnTextMessage", "size", len(message))
+
+	s.executeCallback(ctx, conn.PluginName, "OnTextMessage", func(ctx context.Context, plugin api.WebSocketCallback) error {
+		_, err := checkErr(plugin.OnTextMessage(ctx, req))
+		return err
+	})
+}
+
+// notifyBinaryCallback notifies the plugin of a binary message
+func (s *websocketService) notifyBinaryCallback(ctx context.Context, connectionID string, conn *WebSocketConnection, data []byte) {
+	req := &api.OnBinaryMessageRequest{
+		ConnectionId: connectionID,
+		Data:         data,
+	}
+
+	ctx = log.NewContext(ctx, "callback", "OnBinaryMessage", "size", len(data))
+
+	s.executeCallback(ctx, conn.PluginName, "OnBinaryMessage", func(ctx context.Context, plugin api.WebSocketCallback) error {
+		_, err := checkErr(plugin.OnBinaryMessage(ctx, req))
+		return err
+	})
+}
+
+// notifyErrorCallback notifies the plugin of an error
+func (s *websocketService) notifyErrorCallback(ctx context.Context, connectionID string, conn *WebSocketConnection, errorMsg string) {
+	req := &api.OnErrorRequest{
+		ConnectionId: connectionID,
+		Error:        errorMsg,
+	}
+
+	ctx = log.NewContext(ctx, "callback", "OnError", "error", errorMsg)
+
+	s.executeCallback(ctx, conn.PluginName, "OnError", func(ctx context.Context, plugin api.WebSocketCallback) error {
+		_, err := checkErr(plugin.OnError(ctx, req))
+		return err
+	})
+}
+
+// notifyCloseCallback notifies the plugin that the connection was closed
+func (s *websocketService) notifyCloseCallback(ctx context.Context, connectionID string, conn *WebSocketConnection, code int, reason string) {
+	req := &api.OnCloseRequest{
+		ConnectionId: connectionID,
+		Code:         int32(code),
+		Reason:       reason,
+	}
+
+	ctx = log.NewContext(ctx, "callback", "OnClose", "code", code, "reason", reason)
+
+	s.executeCallback(ctx, conn.PluginName, "OnClose", func(ctx context.Context, plugin api.WebSocketCallback) error {
+		_, err := checkErr(plugin.OnClose(ctx, req))
+		return err
+	})
+}