rudder-unity-sdk/Packages/rudder.sdk/Runtime/Sources/Realtime/JsWebSocketAdapter.cs

#if UNITY_WEBGL && !UNITY_EDITOR
using System.Runtime.InteropServices;
#endif
using System;
using System.Collections.Concurrent;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;

namespace RudderSdk.Unity
{
    /// <summary>
    /// WebGL WebSocket adapter using JavaScript interop via DllImport("__Internal__").
    /// Uses SendMessage callbacks from JS to route WebSocket events to Unity's main thread.
    /// Only functional in WebGL builds; throws PlatformNotSupportedException elsewhere.
    /// </summary>
    public sealed class JsWebSocketAdapter : IWebSocketAdapter, IDisposable
    {
#if UNITY_WEBGL && !UNITY_EDITOR
        [DllImport("__Internal__")]
        private static extern void RudderWebSocketCreate(string gameObjectName, string url);

        [DllImport("__Internal__")]
        private static extern void RudderWebSocketSend(string gameObjectName, byte[] data, int length);

        [DllImport("__Internal__")]
        private static extern void RudderWebSocketClose(string gameObjectName);
#endif

        private CancellationTokenSource _cts;
        private readonly ConcurrentQueue<Action> _eventQueue = new();
        private JsWebSocketBridge _bridge;
        private TaskCompletionSource<bool> _connectionTcs;
        private volatile bool _isConnected;
        private volatile bool _isConnecting;
        private bool _disposed;

        /// <summary>
        /// Fired when the WebSocket connection is closed.
        /// </summary>
        public event Action Closed;

        /// <summary>
        /// Fired when binary data is received from the server.
        /// </summary>
        public event Action<ArraySegment<byte>> Received;

        /// <summary>
        /// Fired when an error occurs on the WebSocket.
        /// </summary>
        public event Action<Exception> ReceivedError;

        /// <summary>
        /// True when the WebSocket connection is open.
        /// </summary>
        public bool IsConnected => _isConnected;

        /// <summary>
        /// True while the WebSocket connection is being established.
        /// </summary>
        public bool IsConnecting => _isConnecting;

        public JsWebSocketAdapter()
        {
            var go = new GameObject("JsWebSocketBridge");
            UnityEngine.Object.DontDestroyOnLoad(go);
            _bridge = go.AddComponent<JsWebSocketBridge>();
            _bridge.Init(this);
        }

        /// <summary>
        /// Connects to the specified WebSocket URI via the JavaScript WebSocket API.
        /// </summary>
        /// <param name="uri">The WebSocket server URI (ws:// or wss://).</param>
        /// <param name="timeoutSeconds">Connection timeout in seconds.</param>
        /// <param name="cancellationToken">Optional cancellation token.</param>
        /// <exception cref="ArgumentNullException">Thrown when uri is null.</exception>
        /// <exception cref="ArgumentOutOfRangeException">Thrown when timeoutSeconds is not positive.</exception>
        /// <exception cref="TimeoutException">Thrown when the connection attempt times out.</exception>
        /// <exception cref="PlatformNotSupportedException">Thrown when not running in WebGL.</exception>
        public Task ConnectAsync(Uri uri, int timeoutSeconds, CancellationToken cancellationToken = default)
        {
            if (uri == null)
                throw new ArgumentNullException(nameof(uri));
            if (timeoutSeconds <= 0)
                throw new ArgumentOutOfRangeException(nameof(timeoutSeconds), "Timeout must be greater than zero.");

#if UNITY_WEBGL && !UNITY_EDITOR
            _isConnecting = true;
            _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
            _connectionTcs = new TaskCompletionSource<bool>();

            // Register cancellation (including timeout)
            _cts.Token.Register(() =>
            {
                if (_connectionTcs?.TrySetException(
                    new TimeoutException($"WebSocket connection timed out after {timeoutSeconds} seconds.")) == true)
                {
                    _isConnecting = false;
                }
            });
            _cts.CancelAfter(TimeSpan.FromSeconds(timeoutSeconds));

            RudderWebSocketCreate(_bridge.gameObject.name, uri.ToString());

            return _connectionTcs.Task;
#else
            // Not available outside WebGL builds; editor test uses NativeWebSocketAdapter.
            return Task.FromException(new PlatformNotSupportedException(
                "JsWebSocketAdapter is only supported on WebGL. Use NativeWebSocketAdapter for other platforms."));
#endif
        }

        /// <summary>
        /// Gracefully closes the WebSocket connection via the JavaScript WebSocket API.
        /// </summary>
        public Task CloseAsync()
        {
#if UNITY_WEBGL && !UNITY_EDITOR
            RudderWebSocketClose(_bridge.gameObject.name);
#endif
            _cts?.Cancel();
            _isConnected = false;
            _isConnecting = false;
            EnqueueEvent(() => Closed?.Invoke());
            return Task.CompletedTask;
        }

        /// <summary>
        /// Sends binary data over the WebSocket via the JavaScript WebSocket API.
        /// </summary>
        public Task SendAsync(ArraySegment<byte> data, CancellationToken cancellationToken = default)
        {
            if (!_isConnected) return Task.CompletedTask;

#if UNITY_WEBGL && !UNITY_EDITOR
            byte[] arr;
            int offset;
            int count;
            if (data.Offset == 0 && data.Count == data.Array.Length)
            {
                arr = data.Array;
                offset = 0;
                count = data.Count;
            }
            else
            {
                // Copy to a contiguous buffer since DllImport needs a pinned array from offset 0
                arr = new byte[data.Count];
                System.Buffer.BlockCopy(data.Array, data.Offset, arr, 0, data.Count);
                offset = 0;
                count = data.Count;
            }

            RudderWebSocketSend(_bridge.gameObject.name, arr, count);
#endif
            return Task.CompletedTask;
        }

        /// <summary>
        /// Disposes the adapter, cleaning up the bridge GameObject.
        /// </summary>
        public void Dispose()
        {
            if (_disposed) return;
            _disposed = true;

#if UNITY_WEBGL && !UNITY_EDITOR
            RudderWebSocketClose(_bridge.gameObject.name);
#endif
            _cts?.Cancel();
            _cts?.Dispose();

            if (_bridge != null && _bridge.gameObject != null)
            {
                UnityEngine.Object.Destroy(_bridge.gameObject);
            }
        }

        /// <summary>
        /// Called by JsWebSocketBridge when the JS WebSocket onopen fires.
        /// </summary>
        internal void HandleOpen()
        {
            _isConnected = true;
            _isConnecting = false;
            _connectionTcs?.TrySetResult(true);
        }

        /// <summary>
        /// Called by JsWebSocketBridge when the JS WebSocket onclose fires.
        /// </summary>
        internal void HandleClose()
        {
            _isConnected = false;
            _isConnecting = false;
            _connectionTcs?.TrySetCanceled();
            Closed?.Invoke();
        }

        /// Called by JsWebSocketBridge when the JS WebSocket onerror fires.
        /// </summary>
        internal void HandleError(string errorMessage)
        {
            _isConnecting = false;
            _connectionTcs?.TrySetException(new Exception(errorMessage));
            ReceivedError?.Invoke(new Exception(errorMessage));
        }

        /// Called by JsWebSocketBridge when binary message data is received.
        /// </summary>
        internal void HandleMessage(ArraySegment<byte> data)
        {
            Received?.Invoke(data);
        }

        internal void HandleMessageError(Exception ex)
        {
            ReceivedError?.Invoke(ex);
        }

        internal void EnqueueEvent(Action action)
        {
            _eventQueue.Enqueue(action);
        }

        internal void ProcessEvents()
        {
            while (_eventQueue.TryDequeue(out var action))
            {
                try
                {
                    action?.Invoke();
                }
                catch (Exception ex)
                {
                    Debug.LogException(ex);
                }
            }
        }
    }

    /// <summary>
    /// MonoBehaviour that receives SendMessage callbacks from the JavaScript WebSocket bridge
    /// and dispatches events on the Unity main thread.
    /// </summary>
    internal class JsWebSocketBridge : MonoBehaviour
    {
        private JsWebSocketAdapter _adapter;

        public void Init(JsWebSocketAdapter adapter)
        {
            _adapter = adapter;
        }

        /// <summary>
        /// Called from JavaScript via SendMessage when the WebSocket connection opens.
        /// </summary>
        public void OnOpen(string _)
        {
            _adapter?.EnqueueEvent(() => _adapter?.HandleOpen());
        }

        /// <summary>
        /// Called from JavaScript via SendMessage when binary data is received.
        /// Data is base64-encoded by the JS bridge and decoded here.
        /// </summary>
        public void OnMessage(string base64Data)
        {
            if (string.IsNullOrEmpty(base64Data) || _adapter == null) return;

            try
            {
                var data = Convert.FromBase64String(base64Data);
                var segment = new ArraySegment<byte>(data);
                _adapter.EnqueueEvent(() => _adapter.HandleMessage(segment));
            }
            catch (FormatException ex)
            {
                Debug.LogError($"JsWebSocketBridge: Failed to decode base64 message: {ex.Message}");
                _adapter.EnqueueEvent(() => _adapter.HandleMessageError(ex));
            }
        }

        /// <summary>
        /// Called from JavaScript via SendMessage when the WebSocket connection closes.
        /// </summary>
        public void OnClose(string _)
        {
            _adapter?.EnqueueEvent(() => _adapter?.HandleClose());
        }

        /// <summary>
        /// Called from JavaScript via SendMessage when a WebSocket error occurs.
        /// </summary>
        public void OnError(string error)
        {
            _adapter?.EnqueueEvent(() => _adapter?.HandleError(error ?? "Unknown WebSocket error"));
        }

        private void Update()
        {
            _adapter?.ProcessEvents();
        }

        private void OnDestroy()
        {
            _adapter?.Dispose();
            _adapter = null;
        }
    }
}