Fix WebSocket memory free to be thread safe.

Also skip WebSocket that have a deferred free status.

For U504-028.

Author: TurboGit

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;
 
       -----------------
@@ -442,7 +446,7 @@ package body AWS.Net.WebSocket.Registry is
                end;
 
                WebSocket.Shutdown;
-               Unchecked_Free (WebSocket);
+               DB.Free (WebSocket);
             end if;
          end Close_To;
 
@@ -499,7 +503,7 @@ package body AWS.Net.WebSocket.Registry is
          Socket.On_Close (Message);
          Socket.Shutdown;
 
-         Unchecked_Free (W);
+         DB.Free (W);
       end Close;
 
       ----------------
@@ -516,9 +520,11 @@ package body AWS.Net.WebSocket.Registry is
             --  Add watched sockets
 
             for Id of Watched loop
-               FD_Set.Add
-                 (Result,
-                  Registered (Id).all, Registered (Id), FD_Set.Input);
+               if not Registered (Id).To_Free then
+                  FD_Set.Add
+                    (Result,
+                     Registered (Id).all, Registered (Id), FD_Set.Input);
+               end if;
             end loop;
          end return;
       end Create_Set;
@@ -555,7 +561,7 @@ package body AWS.Net.WebSocket.Registry is
             end;
 
             WebSocket.Shutdown;
-            Unchecked_Free (WebSocket);
+            DB.Free (WebSocket);
          end On_Close;
 
       begin
@@ -568,6 +574,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 +762,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;
 
       ------------
@@ -869,15 +902,15 @@ package body AWS.Net.WebSocket.Registry is
                      begin
                         if Error = null then
                            DB.Unregister (W);
-                           Unchecked_Free (W);
+                           DB.Free (W);
 
                         else
                            Error (W.all, A);
 
                            case A is
                               when Close =>
                                  DB.Unregister (W);
-                                 Unchecked_Free (W);
+                                 DB.Free (W);
                               when None =>
                                  null;
                            end case;
@@ -931,9 +964,7 @@ package body AWS.Net.WebSocket.Registry is
                            WS.Close (Exception_Message (E), Going_Away);
                            WS.On_Close (Exception_Message (E));
 
-                           --  ??? if we free it now, there might be a reader
-                           --  in parallel that is using this socket...
-                           Unchecked_Free (WS);
+                           DB.Free (WS);
 
                            --  No more data to send from this socket
                            Pending := 0;
@@ -996,7 +1027,7 @@ package body AWS.Net.WebSocket.Registry is
                         WebSocket.Close (Exception_Message (E), Going_Away);
 
                         --  Do not free, it might be used by another
-                        Unchecked_Free (WebSocket);
+                        Free (WebSocket);
                   end;
 
                else
@@ -1267,7 +1298,7 @@ package body AWS.Net.WebSocket.Registry is
                   DB.Unregister (WS);
                   WebSocket_Exception
                     (WS, Exception_Message (E), Protocol_Error);
-                  Unchecked_Free (WS);
+                  DB.Free (WS);
                   --  No more data to send from this socket
                   Pending := 0;
             end Read_Send;
@@ -1539,7 +1570,7 @@ package body AWS.Net.WebSocket.Registry is
       DB.Watch (WebSocket);
    exception
       when others =>
-         Unchecked_Free (WebSocket);
+         DB.Free (WebSocket);
          raise;
    end Watch;
 

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