Fix WebSocket memory free to be thread safe.

For U504-028.
AdaCore · May 5, 2021 · e7b05e2 · e7b05e2
1 parent 9e5bad5
commit e7b05e2
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 7 deletions.
diff --git a/src/core/aws-net-websocket-registry.adb b/src/core/aws-net-websocket-registry.adb
@@ -1,7 +1,7 @@
 ------------------------------------------------------------------------------
 --                              Ada Web Server                              --
 --                                                                          --
---                     Copyright (C) 2012-2019, AdaCore                     --
+--                     Copyright (C) 2012-2021, AdaCore                     --
 --                                                                          --
 --  This library is free software;  you can redistribute it and/or modify   --
 --  it under terms of the  GNU General Public License  as published by the  --
@@ -161,10 +161,14 @@ package body AWS.Net.WebSocket.Registry is
       entry Get_Socket (WebSocket : out Object_Class);
       --  Get a WebSocket having some data to be sent
 
-      procedure Release_Socket (WebSocket : Object_Class);
+      procedure Release_Socket (WebSocket : in out Object_Class);
       --  Release a socket retrieved with Get_Socket above, this socket will be
       --  then available again.
 
+      procedure Free (WebSocket : in out Object_Class);
+      --  Free WebSocket immediately if not taken by another task, otherwise
+      --  record it to be freed as soon as it is released.
+
       entry Not_Empty;
       --  Returns if the Set is not empty
 
@@ -329,7 +333,7 @@ package body AWS.Net.WebSocket.Registry is
 
       procedure Do_Free (WebSocket : in out Object_Class) is
       begin
-         Unchecked_Free (WebSocket);
+         DB.Free (WebSocket);
       end Do_Free;
 
       -----------------
@@ -568,6 +572,24 @@ package body AWS.Net.WebSocket.Registry is
          Registered.Clear;
       end Finalize;
 
+      ----------
+      -- Free --
+      ----------
+
+      procedure Free (WebSocket : in out Object_Class) is
+      begin
+         --  If WebSocket is in Sending it means that it has been
+         --  taken by the Get_Socket call. We cannot free it now, we
+         --  record this socket to be freed as soon as it is released
+         --  (Release_Socket) call.
+
+         if Sending.Contains (WebSocket.Id) then
+            WebSocket.To_Free := True;
+         else
+            Unchecked_Free (WebSocket);
+         end if;
+      end Free;
+
       ----------------
       -- Get_Socket --
       ----------------
@@ -738,10 +760,19 @@ package body AWS.Net.WebSocket.Registry is
       -- Release_Socket --
       --------------------
 
-      procedure Release_Socket (WebSocket : Object_Class) is
+      procedure Release_Socket (WebSocket : in out Object_Class) is
       begin
          Sending.Exclude (WebSocket.Id);
-         New_Pending := True;
+
+         --  The socket has been recorded to be freed. It is not anymore
+         --  in the registry, we just need to free it now that it has
+         --  been released.
+
+         if WebSocket.To_Free then
+            Unchecked_Free (WebSocket);
+         else
+            New_Pending := True;
+         end if;
       end Release_Socket;
 
       ------------

diff --git a/src/core/aws-net-websocket.ads b/src/core/aws-net-websocket.ads
@@ -1,7 +1,7 @@
 ------------------------------------------------------------------------------
 --                              Ada Web Server                              --
 --                                                                          --
---                     Copyright (C) 2012-2020, AdaCore                     --
+--                     Copyright (C) 2012-2021, AdaCore                     --
 --                                                                          --
 --  This library is free software;  you can redistribute it and/or modify   --
 --  it under terms of the  GNU General Public License  as published by the  --
@@ -289,6 +289,7 @@ private
       Messages : Message_List.List;
       Mem_Sock : Net.Socket_Access;
       In_Mem   : Boolean := False;
+      To_Free  : Boolean := False;
 
       Connection : AWS.Client.HTTP_Connection_Access;
       --  Only set when the web socket is initialized as a client.
@@ -366,7 +367,8 @@ private
                     Messages   => Message_List.Empty_List,
                     Mem_Sock   => null,
                     Connection => null,
-                    In_Mem     => False);
+                    In_Mem     => False,
+                    To_Free    => False);
 
    --  Error codes corresponding to all errors