Skip to content

Commit

Permalink
Fix WebSocket memory free to be thread safe.
Browse files Browse the repository at this point in the history
Also skip WebSocket that have a deferred free status.

For U504-028.
  • Loading branch information
TurboGit committed May 5, 2021
1 parent 9e5bad5 commit 4ea6d48
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 21 deletions.
69 changes: 50 additions & 19 deletions src/core/aws-net-websocket-registry.adb
Original file line number Diff line number Diff line change
@@ -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 --
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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;

-----------------
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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;

----------------
Expand All @@ -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;
Expand Down Expand Up @@ -555,7 +561,7 @@ package body AWS.Net.WebSocket.Registry is
end;

WebSocket.Shutdown;
Unchecked_Free (WebSocket);
DB.Free (WebSocket);
end On_Close;

begin
Expand All @@ -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 --
----------------
Expand Down Expand Up @@ -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;

------------
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;

Expand Down
6 changes: 4 additions & 2 deletions src/core/aws-net-websocket.ads
Original file line number Diff line number Diff line change
@@ -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 --
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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

Expand Down

0 comments on commit 4ea6d48

Please sign in to comment.