server: ensure message socket is still available on websocket handler

Fixes #306

4 files changed, 15 insertions, 14 deletions

diff --git a/libsoup/server/soup-server-connection.c b/libsoup/server/soup-server-connection.c
index f14058c2..ac0f1aaa 100644
--- a/libsoup/server/soup-server-connection.c
+++ b/libsoup/server/soup-server-connection.c
@@ -538,25 +538,31 @@ soup_server_connection_get_socket (SoupServerConnection *conn)
         return priv->socket;
 }
 
-GSocket *
-soup_server_connection_steal_socket (SoupServerConnection *conn)
+GIOStream *
+soup_server_connection_steal (SoupServerConnection *conn)
 {
         SoupServerConnectionPrivate *priv;
-        GSocket *socket;
+        GIOStream *stream;
 
         g_return_val_if_fail (SOUP_IS_SERVER_CONNECTION (conn), NULL);
 
         priv = soup_server_connection_get_instance_private (conn);
 
+        stream = priv->io_data ? soup_server_message_io_steal (priv->io_data) : NULL;
+        if (stream) {
+                g_object_set_data_full (G_OBJECT (stream), "GSocket",
+                                        priv->socket, g_object_unref);
+        }
+
         /* Cache local and remote address */
         soup_server_connection_get_local_address (conn);
         soup_server_connection_get_remote_address (conn);
 
-        socket = g_steal_pointer (&priv->socket);
+        g_clear_pointer (&priv->io_data, soup_server_message_io_destroy);
         g_clear_object (&priv->conn);
         g_clear_object (&priv->iostream);
 
-        return socket;
+        return stream;
 }
 
 GIOStream *
diff --git a/libsoup/server/soup-server-connection.h b/libsoup/server/soup-server-connection.h
index 50127cf8..31c1f3bd 100644
--- a/libsoup/server/soup-server-connection.h
+++ b/libsoup/server/soup-server-connection.h
@@ -30,7 +30,7 @@ gboolean              soup_server_connection_is_ssl                          (So
 void                  soup_server_connection_disconnect                      (SoupServerConnection  *conn);
 gboolean              soup_server_connection_is_connected                    (SoupServerConnection  *conn);
 GSocket              *soup_server_connection_get_socket                      (SoupServerConnection  *conn);
-GSocket              *soup_server_connection_steal_socket                    (SoupServerConnection  *conn);
+GIOStream            *soup_server_connection_steal                           (SoupServerConnection  *conn);
 GIOStream            *soup_server_connection_get_iostream                    (SoupServerConnection  *conn);
 GSocketAddress       *soup_server_connection_get_local_address               (SoupServerConnection  *conn);
 GSocketAddress       *soup_server_connection_get_remote_address              (SoupServerConnection  *conn);
diff --git a/libsoup/server/soup-server-message.c b/libsoup/server/soup-server-message.c
index 61d62a53..8f40d2a0 100644
--- a/libsoup/server/soup-server-message.c
+++ b/libsoup/server/soup-server-message.c
@@ -1103,15 +1103,8 @@ soup_server_message_steal_connection (SoupServerMessage *msg)
         GIOStream *stream;
 
         g_object_ref (msg);
-        stream = msg->io_data ? soup_server_message_io_steal (g_steal_pointer (&msg->io_data)) : NULL;
-        if (stream) {
-                g_object_set_data_full (G_OBJECT (stream), "GSocket",
-                                        soup_server_connection_steal_socket (msg->conn),
-                                        g_object_unref);
-        }
-
+        stream = soup_server_connection_steal (msg->conn);
         g_signal_handlers_disconnect_by_data (msg, msg->conn);
-
         connection_disconnected (msg);
         g_object_unref (msg);
 
diff --git a/tests/websocket-test.c b/tests/websocket-test.c
index e351443a..c9c7cffb 100644
--- a/tests/websocket-test.c
+++ b/tests/websocket-test.c
@@ -1571,6 +1571,8 @@ test_client_context_got_server_connection (SoupServer              *server,
 	g_assert_cmpstr (remote_ip, ==, str);
 	g_free (str);
 
+        g_assert_nonnull (soup_server_message_get_socket (msg));
+
 	test->server = g_object_ref (connection);
 }