From aac1931569314df5b879b8b3272b1266ee4e07e3 Mon Sep 17 00:00:00 2001
From: sleepycatcoding <131554884+sleepycatcoding@users.noreply.github.com>
Date: Fri, 4 Aug 2023 12:12:49 +0300
Subject: [PATCH] avm1: Prevent panic when close is called in event handler

Also adds a test case for this.
---
 core/src/socket.rs                            |  13 ++++++++--
 .../xml_socket_close_in_handler/output.txt    |   6 +++++
 .../xml_socket_close_in_handler/socket.json   |  13 ++++++++++
 .../avm1/xml_socket_close_in_handler/test.as  |  23 ++++++++++++++++++
 .../avm1/xml_socket_close_in_handler/test.swf | Bin 0 -> 426 bytes
 .../xml_socket_close_in_handler/test.toml     |   1 +
 6 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 tests/tests/swfs/avm1/xml_socket_close_in_handler/output.txt
 create mode 100644 tests/tests/swfs/avm1/xml_socket_close_in_handler/socket.json
 create mode 100644 tests/tests/swfs/avm1/xml_socket_close_in_handler/test.as
 create mode 100644 tests/tests/swfs/avm1/xml_socket_close_in_handler/test.swf
 create mode 100644 tests/tests/swfs/avm1/xml_socket_close_in_handler/test.toml

diff --git a/core/src/socket.rs b/core/src/socket.rs
index 8cc82eb71..7520e5270 100644
--- a/core/src/socket.rs
+++ b/core/src/socket.rs
@@ -321,10 +321,13 @@ impl<'gc> Sockets<'gc> {
                             let xml_socket =
                                 XmlSocket::cast(target.into()).expect("target should be XmlSocket");
 
+                            // Get XmlSocket read buffer.
                             let mut buffer = xml_socket.read_buffer();
                             buffer.extend(data);
 
-                            // Check for a message.
+                            let mut messages = vec![];
+
+                            // Check for messages.
                             while let Some((index, _)) =
                                 buffer.iter().enumerate().find(|(_, &b)| b == 0)
                             {
@@ -332,8 +335,14 @@ impl<'gc> Sockets<'gc> {
                                 // Remove null byte.
                                 let _ = buffer.drain(..1);
 
-                                let message = AvmString::new_utf8_bytes(activation.gc(), &message);
+                                // Store the message in cache.
+                                messages.push(AvmString::new_utf8_bytes(activation.gc(), &message));
+                            }
 
+                            // Drop the reference to the buffer, to make sure the SWF can call with close.
+                            drop(buffer);
+
+                            for message in messages {
                                 let _ = target.call_method(
                                     "onData".into(),
                                     &[message.into()],
diff --git a/tests/tests/swfs/avm1/xml_socket_close_in_handler/output.txt b/tests/tests/swfs/avm1/xml_socket_close_in_handler/output.txt
new file mode 100644
index 000000000..a4811a6de
--- /dev/null
+++ b/tests/tests/swfs/avm1/xml_socket_close_in_handler/output.txt
@@ -0,0 +1,6 @@
+connected
+true
+data:
+Hello!
+Closing socket in onData handler
+Successfully closed socket
diff --git a/tests/tests/swfs/avm1/xml_socket_close_in_handler/socket.json b/tests/tests/swfs/avm1/xml_socket_close_in_handler/socket.json
new file mode 100644
index 000000000..c493046e4
--- /dev/null
+++ b/tests/tests/swfs/avm1/xml_socket_close_in_handler/socket.json
@@ -0,0 +1,13 @@
+[
+    {
+        "type": "Send",
+        "payload": [ 72, 101, 108, 108, 111, 33, 0 ]
+    },
+    {
+        "type": "Send",
+        "payload": [ 60, 114, 111, 111, 116, 62, 60, 105, 116, 101, 109, 32, 47, 62, 60, 47, 114, 111, 111, 116, 62, 0 ]
+    },
+    {
+        "type": "WaitForDisconnect"
+    }
+]
diff --git a/tests/tests/swfs/avm1/xml_socket_close_in_handler/test.as b/tests/tests/swfs/avm1/xml_socket_close_in_handler/test.as
new file mode 100644
index 000000000..d20a645af
--- /dev/null
+++ b/tests/tests/swfs/avm1/xml_socket_close_in_handler/test.as
@@ -0,0 +1,23 @@
+class Test {
+    static function main() {
+        var socket = new XMLSocket();
+
+        socket.onConnect = function (status:Boolean) {
+            trace("connected");
+            trace(status);
+        };
+
+        socket.onData = function (data:String) {
+            trace("data:")
+            trace(data);
+            trace("Closing socket in onData handler");
+            socket.close();
+            trace("Successfully closed socket");
+        };
+        socket.onClose = function () {
+            trace("closed");
+        };
+
+        socket.connect("localhost", 8001);
+    }
+}
diff --git a/tests/tests/swfs/avm1/xml_socket_close_in_handler/test.swf b/tests/tests/swfs/avm1/xml_socket_close_in_handler/test.swf
new file mode 100644
index 0000000000000000000000000000000000000000..3799ad47709aa4cee05fe6b973bb1c1da541da88
GIT binary patch
literal 426
zcmV;b0agA(S5pUU0ssJb+DuVRZWA#S{>C%OOrT8L2x){A!K@;c1qVPyZAgWLM4C~i
z3wDv2m?WwhN3kQ(6VO$bv+Ff_fF7gb80^cl9a_A2em~#)p8dU_2>t-A-@xktL-s(Y
z(*~U)AajzwBuic~off>h<a43lT?lRPkltHBaeap_TxL3#>Km6>t(8&cY9;W4FKcL3
zT?&J*qc5fRR6SI6Eh>Y`-(p6&Gj2F$WZq*am0s5K+;`^7I``##!Rwh6Yg7afC^uCl
zwEn)4@+x<!*{;Qoa1ft`R27#Cr5`6^P?o|JYqk0$`COx6NYLwcZx~vE4T7FaZ4;|R
z_=?Og`t7)h5Fugb6a<+MLXTlbxJA=J62&%5TBpaBA#yzJ?SU0BJOa~(n~lbv8yWY>
zMyx-6f+pMaAJmG0w+`6Z0j>Y6-^B1ke9QAmX0N^fX<j_Os_h6;DyE1BZ>VcOW>3ur
zlWg4VBeuI_Q4oi=boe2|N%)c6>mL1abTKH0)6wAbIm}|I{kGk1ok&)hY$samnG)v6
U!%2oy3jYEC0RR630Bi|^@|_>fuK)l5

literal 0
HcmV?d00001

diff --git a/tests/tests/swfs/avm1/xml_socket_close_in_handler/test.toml b/tests/tests/swfs/avm1/xml_socket_close_in_handler/test.toml
new file mode 100644
index 000000000..0d5c2c67f
--- /dev/null
+++ b/tests/tests/swfs/avm1/xml_socket_close_in_handler/test.toml
@@ -0,0 +1 @@
+num_ticks = 10