websocket: preserve user type

src/core/aws-net-websocket-registry-utils.adb

@@ -35,9 +35,9 @@ package body AWS.Net.WebSocket.Registry.Utils is
    -- Register --
    --------------
 
-   function Register (WebSocket : Object'Class) return Object_Class is
+   procedure Register (WebSocket : in out Object_Class) is
    begin
-      return Net.WebSocket.Registry.Register (WebSocket);
+      Net.WebSocket.Registry.Register (WebSocket);
    end Register;
 
    -----------

src/core/aws-net-websocket-registry-utils.ads

@@ -31,10 +31,11 @@
 
 package AWS.Net.WebSocket.Registry.Utils is
 
-   function Register (WebSocket : Object'Class) return Object_Class;
-   --  Register WebSocket, returns a pointer to the registered WebSocket or
-   --  null if it was not possible to register the WebSocket. This can happen
-   --  if the server has reached the limit of opened WebSocket for example.
+   procedure Register (WebSocket : in out Object_Class);
+   --  Register WebSocket.
+   --  Free it and set it to null if it was not possible to register the
+   --  WebSocket. This can happen if the server has reached the limit of opened
+   --  WebSocket for example.
 
    procedure Watch (WebSocket : in out Object_Class) with
      Pre => WebSocket /= null;

src/core/aws-net-websocket-registry.adb

@@ -90,6 +90,11 @@ package body AWS.Net.WebSocket.Registry is
      (Left.Id = Right.Id);
    --  Equality is based on the unique id
 
+   function Create_Default_Socket
+      (Request_Ignored : AWS.Status.Data) return Object_Class
+      is (new Object);
+   --  Default factory
+
    package WebSocket_Map is
      new Ada.Containers.Ordered_Maps (UID, Object_Class, "=" => Same_WS);
 
@@ -1187,7 +1192,7 @@ package body AWS.Net.WebSocket.Registry is
          end loop;
       end if;
 
-      return Create'Access;
+      return Create_Default_Socket'Access;
    end Constructor;
 
    ------------
@@ -1314,17 +1319,15 @@ package body AWS.Net.WebSocket.Registry is
       Factories.Insert (URI, Factory);
    end Register;
 
-   function Register (WebSocket : Object'Class) return Object_Class is
-      WS      : Object_Class := new Object'Class'(WebSocket);
+   procedure Register (WebSocket : in out Object_Class) is
       Success : Boolean;
    begin
-      DB.Register (WS, Success);
+      DB.Register (WebSocket, Success);
 
       if not Success then
-         Unchecked_Free (WS);
+         Free (WebSocket.all);
+         Unchecked_Free (WebSocket);
       end if;
-
-      return WS;
    end Register;
 
    ----------------------

src/core/aws-net-websocket-registry.ads

@@ -39,8 +39,13 @@ private with GNAT.Regexp;
 package AWS.Net.WebSocket.Registry is
 
    type Factory is not null access function
-     (Socket  : Socket_Access;
-      Request : AWS.Status.Data) return Object'Class;
+     (Request : AWS.Status.Data) return Object_Class;
+   --  Return a newly allocated object.
+   --  You can use AWS.Status.Parameters (Request) to check what additional
+   --  parameters were sent by the user.
+   --
+   --  This object will later be initialized automatically, via a call to
+   --  AWS.Net.WebSocket.Setup_Socket.
 
    --  Creating and Registering WebSockets
 
@@ -207,9 +212,9 @@ private
    procedure Shutdown;
    --  Stop the WebServer's servers
 
-   function Register (WebSocket : Object'Class) return Object_Class;
-   --  Register a new WebSocket, returns a reference to the registered
-   --  WebSocket or null if it was impossible to register it.
+   procedure Register (WebSocket : in out Object_Class);
+   --  Register a new WebSocket.
+   --  Sets it to null (and free memory) if it was impossible to register it.
 
    procedure Watch (WebSocket : in out Object_Class)
      with Pre => WebSocket /= null;

src/core/aws-net-websocket.adb

@@ -134,32 +134,6 @@ package body AWS.Net.WebSocket is
       Socket.On_Open ("WebSocket connected with " & URI);
    end Connect;
 
-   ------------
-   -- Create --
-   ------------
-
-   function Create
-     (Socket  : Socket_Access;
-      Request : AWS.Status.Data) return Object'Class
-   is
-      Result   : Object;
-      Protocol : Net.WebSocket.Protocol.State_Class;
-      Headers  : constant AWS.Headers.List :=
-                   AWS.Status.Header (Request);
-   begin
-      if Headers.Exist (Messages.Sec_WebSocket_Key1_Token)
-        and then Headers.Exist (Messages.Sec_WebSocket_Key2_Token)
-      then
-         Protocol := new Net.WebSocket.Protocol.Draft76.State;
-      else
-         Protocol := new Net.WebSocket.Protocol.RFC6455.State;
-      end if;
-
-      Initialize (Result, Socket, Protocol, Headers);
-      Result.Request := Request;
-      return Result;
-   end Create;
-
    --------------------
    -- End_Of_Message --
    --------------------
@@ -627,6 +601,30 @@ package body AWS.Net.WebSocket is
       Socket.P_State.State.Send (Socket, Message);
    end Send;
 
+   ------------------
+   -- Setup_Socket --
+   ------------------
+
+   procedure Setup_Socket
+     (WS      : not null Object_Class;
+      Socket  : not null Socket_Access;
+      Request : AWS.Status.Data)
+   is
+      Protocol : Net.WebSocket.Protocol.State_Class;
+      Headers  : constant AWS.Headers.List := AWS.Status.Header (Request);
+   begin
+      if Headers.Exist (Messages.Sec_WebSocket_Key1_Token)
+        and then Headers.Exist (Messages.Sec_WebSocket_Key2_Token)
+      then
+         Protocol := new Net.WebSocket.Protocol.Draft76.State;
+      else
+         Protocol := new Net.WebSocket.Protocol.RFC6455.State;
+      end if;
+
+      Initialize (WS.all, Socket, Protocol, Headers);
+      WS.Request := Request;
+   end Setup_Socket;
+
    --------------
    -- Shutdown --
    --------------

src/core/aws-net-websocket.ads

@@ -45,6 +45,10 @@ package AWS.Net.WebSocket is
 
    type Object is new Net.Socket_Type with private;
    type Object_Class is access all Object'Class;
+    --  To implement your own handling of messages, you need to extend this
+    --  type and override at least the On_Message primitive operation.
+    --  In addition, you need to register a factory (to create new objects
+    --  based on the URI) using AWS.Net.WebSocket.Registry.Register).
 
    No_Object : constant Object'Class;
 
@@ -77,18 +81,8 @@ package AWS.Net.WebSocket is
    --  the default Send implementation should be ok for most usages.
    --
 
-   function Create
-     (Socket  : Socket_Access;
-      Request : AWS.Status.Data) return Object'Class
-   with Pre => Socket /= null;
-   --  Create a new instance of the WebSocket, this is used by AWS internal
-   --  server to create a default WebSocket if no other constructor are
-   --  provided. It is also needed when deriving from WebSocket.
-   --
-   --  This function must be registered via AWS.Net.WebSocket.Registry.Register
-
    procedure On_Message (Socket : in out Object; Message : String) is null;
-   --  Default implementation does nothing, it needs to be overriden by the
+   --  Default implementation does nothing, it needs to be overridden by the
    --  end-user. This is the callback that will get activated for every server
    --  incoming data. It is also important to keep in mind that the thread
    --  handling this WebSocket won't be released until the procedure returns.
@@ -97,11 +91,11 @@ package AWS.Net.WebSocket is
 
    procedure On_Message (Socket : in out Object; Message : Unbounded_String);
    --  Same a above but takes an Unbounded_String. This is supposed to be
-   --  overriden when handling large messages otherwise a stack-overflow could
-   --  be raised. The default implementation of this procedure to to call the
+   --  overridden when handling large messages otherwise a stack-overflow could
+   --  be raised. The default implementation of this procedure to call the
    --  On_Message above with a string.
    --
-   --  So either this version is overriden to handle the incoming messages or
+   --  So either this version is overridden to handle the incoming messages or
    --  the one above if the messages are known to be small.
 
    procedure On_Open (Socket : in out Object; Message : String) is null;
@@ -119,7 +113,7 @@ package AWS.Net.WebSocket is
      (Socket    : in out Object;
       Message   : String;
       Is_Binary : Boolean := False);
-   --  This default implementation just send a message to the client. The
+   --  This default implementation just sends a message to the client. The
    --  message is sent in a single chunk (not fragmented).
 
    procedure Send
@@ -167,7 +161,7 @@ package AWS.Net.WebSocket is
    --
    --  These function waits until it either receives a close or an error, or
    --  the beginning of a message frame. In the latter case, the function
-   --  will then block until it has receives all chunks of that frame, which
+   --  will then block until it has received all chunks of that frame, which
    --  might take longer than Timeout.
    --
    --  The function will return early if it doesn't receive the beginning
@@ -257,8 +251,23 @@ package AWS.Net.WebSocket is
    --  Returns a unique id for the given socket. The uniqueness for this socket
    --  is guaranteed during the lifetime of the application.
 
+
    overriding function Is_Secure (Socket : Object) return Boolean;
 
+   -----------------------
+   -- Internal services --
+   -----------------------
+   --  These subprograms are used internally by AWS, and do not need to be
+   --  called explicitly in user code.
+
+   procedure Setup_Socket
+     (WS      : not null Object_Class;
+      Socket  : not null Socket_Access;
+      Request : AWS.Status.Data);
+   --  Setup WS.
+   --  It will be called automatically for any websocket returned by a factory,
+   --  so in general you do not need to call it explicitly.
+
 private
 
    type Internal_State is record
@@ -351,6 +360,8 @@ private
      (Socket : Object; Size : Natural) is null;
 
    overriding procedure Free (Socket : in out Object);
+   --  This is called automatically when the socket is no longer needed, do not
+   --  call directly from user code.
 
    No_UID    : constant UID := 0;
 

src/core/aws-server-http_utils.adb

@@ -37,6 +37,7 @@ with Ada.Strings.Fixed;
 with Ada.Strings.Maps;
 with Ada.Strings.Unbounded;
 with Ada.Text_IO;
+with Ada.Unchecked_Deallocation;
 
 with GNAT.MD5;
 with GNAT.OS_Lib;
@@ -1667,33 +1668,50 @@ package body AWS.Server.HTTP_Utils is
                      --  if the WebSocket is not to be accepted. In this case
                      --  a forbidden message is sent back.
 
-                     WS : constant Net.WebSocket.Object'Class :=
-                            Net.WebSocket.Registry.Constructor
-                              (Status.URI (C_Stat))
-                              (Socket  => Status.Socket (C_Stat),
-                               Request => C_Stat);
+                     procedure Unchecked_Free is
+                       new Ada.Unchecked_Deallocation
+                          (Net.WebSocket.Object'Class,
+                           Net.WebSocket.Object_Class);
+
+                     use type Net.WebSocket.Object_Class;
+                     WS : Net.WebSocket.Object_Class;
+                     Registered : Boolean := False;
                   begin
+                     WS := Net.WebSocket.Registry.Constructor
+                        (Status.URI (C_Stat)) (C_Stat);
+
+                     if WS /= null then
+                        Net.WebSocket.Setup_Socket
+                           (WS, Status.Socket (C_Stat), C_Stat);
+                     end if;
+
                      --  Register this new WebSocket
 
-                     if WS in Net.WebSocket.Handshake_Error.Object'Class then
+                     if WS = null then
+                        Send_WebSocket_Handshake_Error
+                          (Messages.S412, "no route defined");
+
+                     elsif WS.all
+                        in Net.WebSocket.Handshake_Error.Object'Class
+                     then
                         declare
                            E : constant Net.WebSocket.Handshake_Error.Object :=
-                                 Net.WebSocket.Handshake_Error.Object (WS);
+                                 Net.WebSocket.Handshake_Error.Object (WS.all);
                         begin
                            Send_WebSocket_Handshake_Error
                              (E.Status_Code, E.Reason_Phrase);
+                           WS.Free;
+                           Unchecked_Free (WS);
                         end;
 
                      else
                         --  First try to register the WebSocket object
 
-                        declare
-                           use type Net.WebSocket.Object_Class;
-                           W : Net.WebSocket.Object_Class;
                         begin
-                           W := Net.WebSocket.Registry.Utils.Register (WS);
+                           Net.WebSocket.Registry.Utils.Register (WS);
+                           Registered := True;
 
-                           if W = null then
+                           if WS = null then
                               Send_WebSocket_Handshake_Error
                                 (Messages.S412,
                                  "too many WebSocket registered");
@@ -1705,7 +1723,7 @@ package body AWS.Server.HTTP_Utils is
                               Socket_Taken := True;
                               Will_Close := False;
 
-                              Net.WebSocket.Registry.Utils.Watch (W);
+                              Net.WebSocket.Registry.Utils.Watch (WS);
                            end if;
                         end;
                      end if;
@@ -1715,7 +1733,16 @@ package body AWS.Server.HTTP_Utils is
                         Send_WebSocket_Handshake_Error
                           (Messages.S403,
                            Exception_Message (E));
-                        WS.Shutdown;
+
+                        if Registered then
+                           --  Close will automatically free the memory for WS
+                           --  itself, by looking up the pointer in the
+                           --  registry.
+                           Net.WebSocket.Registry.Close
+                              (WS.all, "closed on error");
+                        else
+                           Unchecked_Free (WS);
+                        end if;
                   end;
 
                exception