From 5e3e7b9c4f109231eb456c6b2edf384fedbf76e9 Mon Sep 17 00:00:00 2001
From: Tako <53405089+Takoooooo@users.noreply.github.com>
Date: Fri, 1 Sep 2023 17:14:18 +0000
Subject: [PATCH] Merge pull request #12757 from
 AvaloniaUI/fix-previewer-security-issue

Fix previewer security issue
---
 .../Remote/HtmlTransport/HtmlTransport.cs     | 43 +++++++++++++++----
 .../webapp/src/PreviewerServerConnection.ts   |  6 ++-
 .../HtmlTransport/webapp/src/index.html       |  5 ++-
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/src/Avalonia.DesignerSupport/Remote/HtmlTransport/HtmlTransport.cs b/src/Avalonia.DesignerSupport/Remote/HtmlTransport/HtmlTransport.cs
index 143980dfa5..47e89fc872 100644
--- a/src/Avalonia.DesignerSupport/Remote/HtmlTransport/HtmlTransport.cs
+++ b/src/Avalonia.DesignerSupport/Remote/HtmlTransport/HtmlTransport.cs
@@ -23,6 +23,8 @@ namespace Avalonia.DesignerSupport.Remote.HtmlTransport
         private SimpleWebSocket _pendingSocket;
         private bool _disposed;
         private object _lock = new object();
+        private Uri _listenUri;
+        private Guid _secretCookie;
         private AutoResetEvent _wakeup = new AutoResetEvent(false);
         private FrameMessage _lastFrameMessage = null;
         private FrameMessage _lastSentFrameMessage = null;
@@ -42,6 +44,7 @@ namespace Avalonia.DesignerSupport.Remote.HtmlTransport
             if (listenUri.Scheme != "http")
                 throw new ArgumentException("URI scheme is not HTTP.", nameof(listenUri));
 
+            _listenUri = listenUri;
             var resourcePrefix = "Avalonia.DesignerSupport.Remote.HtmlTransport.webapp.build.";
             _resources = typeof(HtmlWebSocketTransport).Assembly.GetManifestResourceNames()
                 .Where(r => r.StartsWith(resourcePrefix, StringComparison.OrdinalIgnoreCase)
@@ -57,10 +60,21 @@ namespace Avalonia.DesignerSupport.Remote.HtmlTransport
                         {
                             var ms = new MemoryStream();
                             s.CopyTo(ms);
-                            return ms.ToArray();
+                            var currentIndexBytes = ms.ToArray();
+                            if (r == resourcePrefix + "index.html.gz" && ms.CanWrite)
+                            { 
+                                _secretCookie = Guid.NewGuid();
+                                var resultIndexString = Encoding.Default.GetString(currentIndexBytes)
+                                                                        .Replace("PREVIEWER_SECURITY_COOKIE", _secretCookie.ToString());
+                                var resultIndexBytes = Encoding.UTF8.GetBytes(resultIndexString);
+                                
+                                ms.Write(resultIndexBytes, 0, resultIndexBytes.Length);
+                                return resultIndexBytes;
+                            }
+                            return currentIndexBytes;
                         }
                     });
-            
+
             _signalTransport = signalTransport;
             var address = IPAddress.Parse(listenUri.Host);
             
@@ -98,12 +112,19 @@ namespace Avalonia.DesignerSupport.Remote.HtmlTransport
                     }
                     else
                     {
-                        var socket = await req.AcceptWebSocket();
-                        SocketReceiveWorker(socket);
-                        lock (_lock)
+                        if (req.Headers.TryGetValue("Origin", out var origin) && origin == _listenUri.OriginalString)
+                        {
+                            var socket = await req.AcceptWebSocket();
+                            SocketReceiveWorker(socket);
+                            lock (_lock)
+                            {
+                                _pendingSocket?.Dispose();
+                                _pendingSocket = socket;
+                            }
+                        }
+                        else
                         {
-                            _pendingSocket?.Dispose();
-                            _pendingSocket = socket;
+                            throw new Exception("Origin doesen't match Url");
                         }
                     }
                 }
@@ -114,13 +135,19 @@ namespace Avalonia.DesignerSupport.Remote.HtmlTransport
         {
             try
             {
+                if((await socket.ReceiveMessage().ConfigureAwait(false)).AsString() != _secretCookie.ToString())
+                {
+                    socket.Dispose();
+                    return;
+                }
+
                 while (true)
                 {
                     var msg = await socket.ReceiveMessage().ConfigureAwait(false);
                     if(msg != null && msg.IsText)
                     {
                         var message = ParseMessage(msg.AsString());
-                        if (message != null)
+                        if (message != null) 
                             _onMessage?.Invoke(this, message);
                     }
                 }
diff --git a/src/Avalonia.DesignerSupport/Remote/HtmlTransport/webapp/src/PreviewerServerConnection.ts b/src/Avalonia.DesignerSupport/Remote/HtmlTransport/webapp/src/PreviewerServerConnection.ts
index 7f1ab84f99..d9ace9857f 100644
--- a/src/Avalonia.DesignerSupport/Remote/HtmlTransport/webapp/src/PreviewerServerConnection.ts
+++ b/src/Avalonia.DesignerSupport/Remote/HtmlTransport/webapp/src/PreviewerServerConnection.ts
@@ -41,7 +41,7 @@ export class PreviewerServerConnection {
 
         const onMessage = this.onMessage;
         conn.onmessage = msg => onMessage(msg);
-
+        conn.onopen = open => this.onOpen(open);
         const onClose = () => this.setFrame(null);
         conn.onclose = () => onClose();
         conn.onerror = (err: Event) => {
@@ -50,6 +50,10 @@ export class PreviewerServerConnection {
         }
     }
 
+    private onOpen(open: Event) {
+        this.conn.send(window["avaloniaPreviewerSecurityCookie"]);
+    }
+
     private onMessage = (msg: MessageEvent) => {
         if (typeof msg.data == 'string' || msg.data instanceof String) {
             const parts = msg.data.split(':');
diff --git a/src/Avalonia.DesignerSupport/Remote/HtmlTransport/webapp/src/index.html b/src/Avalonia.DesignerSupport/Remote/HtmlTransport/webapp/src/index.html
index 9b1be377ed..35c17185e8 100644
--- a/src/Avalonia.DesignerSupport/Remote/HtmlTransport/webapp/src/index.html
+++ b/src/Avalonia.DesignerSupport/Remote/HtmlTransport/webapp/src/index.html
@@ -10,6 +10,9 @@
         <center>Loading...</center>
     </div>
     <noscript>JavaScript is required</noscript>
+    <script>
+        window["avaloniaPreviewerSecurityCookie"] = "PREVIEWER_SECURITY_COOKIE";
+    </script>
     <script src="index.tsx"></script>
 </body>
-</html>
\ No newline at end of file
+</html>