From cb3521678d6260778a1b63378dc8220982300897 Mon Sep 17 00:00:00 2001
From: Felipe Ripoll <ferigis@gmail.com>
Date: Tue, 3 Jul 2018 09:28:31 -0600
Subject: [PATCH] [#20] removing client when the WS connection finishes

---
 lib/poa_backend/receiver.ex            |  6 +++++-
 lib/poa_backend/receivers/dashboard.ex | 17 +++++++++++++++--
 mix.exs                                |  3 ++-
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/lib/poa_backend/receiver.ex b/lib/poa_backend/receiver.ex
index 601be01..8d8b9d7 100644
--- a/lib/poa_backend/receiver.ex
+++ b/lib/poa_backend/receiver.ex
@@ -52,7 +52,7 @@ defmodule POABackend.Receiver do
       config :poa_backend,
              :subscriptions,
              [
-               # {:my_receiver, [:ethereum_metrics]}
+               {:my_receiver, [:ethereum_metrics]}
              ]
 
   ## Implementing A Receiver Plugin
@@ -80,6 +80,10 @@ defmodule POABackend.Receiver do
           {:ok, state}
         end
 
+        def handle_message(_message, state) do
+          {:ok, state}
+        end
+
         def terminate(_state) do
           :ok
         end
diff --git a/lib/poa_backend/receivers/dashboard.ex b/lib/poa_backend/receivers/dashboard.ex
index e75a983..c130060 100644
--- a/lib/poa_backend/receivers/dashboard.ex
+++ b/lib/poa_backend/receivers/dashboard.ex
@@ -29,6 +29,8 @@ defmodule POABackend.Receivers.Dashboard do
 
     @moduledoc false
 
+    require Logger
+
     alias POABackend.Protocol.Message
 
     @behaviour :cowboy_websocket_handler
@@ -40,9 +42,13 @@ defmodule POABackend.Receivers.Dashboard do
     def websocket_init(_type, req, %{receiver_id: receiver_id, ws_secret: ws_secret}) do
       case :cowboy_req.parse_header("wssecret", req) do
         {_, ^ws_secret, req} ->
+          Logger.info("Websocket connection with dashboard client accepted")
+
           send(receiver_id, {:add_client, self()})
-          {:ok, req, %{}}
+          {:ok, req, %{receiver_id: receiver_id}}
         {_, _, req} ->
+          Logger.info("Websocket connection with dashboard client unauthorized")
+
           {:shutdown, req}
       end
     end
@@ -55,7 +61,10 @@ defmodule POABackend.Receivers.Dashboard do
       {:reply, {:text, build_message(message)}, req, state}
     end
 
-    def websocket_terminate(_reason, _req, _state) do
+    def websocket_terminate(_reason, _req, %{receiver_id: receiver_id}) do
+      Logger.warn("Websocket connection with dashboard client lost")
+
+      send(receiver_id, {:remove_client, self()})
       :ok
     end
 
@@ -98,6 +107,10 @@ defmodule POABackend.Receivers.Dashboard do
     {:ok, %{state | clients: [client | clients]}}
   end
 
+  def handle_message({:remove_client, client}, %{clients: clients} = state) do
+    {:ok, %{state | clients: List.delete(clients, client)}}
+  end
+
   def terminate(_) do
     :ok
   end
diff --git a/mix.exs b/mix.exs
index a50518a..6b958ae 100644
--- a/mix.exs
+++ b/mix.exs
@@ -66,7 +66,8 @@ defmodule POABackend.MixProject do
         ],
         "Receivers": [
           POABackend.Receiver,
-          POABackend.Receivers.DynamoDB
+          POABackend.Receivers.DynamoDB,
+          POABackend.Receivers.Dashboard
         ]
       ]
     ]