From b911bd16dd55ba0e212b8ad6d7415582818c06ce Mon Sep 17 00:00:00 2001
From: Andrew Dryga <andrew@dryga.com>
Date: Tue, 12 Sep 2023 16:50:22 -0600
Subject: [PATCH] Return HTTP 401 status code for invalid tokens (#1988)

Closes https://github.com/firezone/product/issues/651
---
 elixir/apps/api/lib/api/gateway/socket.ex |  3 +++
 elixir/apps/api/lib/api/relay/socket.ex   |  3 +++
 elixir/apps/api/lib/api/sockets.ex        | 10 +++++-----
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/elixir/apps/api/lib/api/gateway/socket.ex b/elixir/apps/api/lib/api/gateway/socket.ex
index b66a9346f..9d56848ab 100644
--- a/elixir/apps/api/lib/api/gateway/socket.ex
+++ b/elixir/apps/api/lib/api/gateway/socket.ex
@@ -33,6 +33,9 @@ defmodule API.Gateway.Socket do
 
       {:ok, socket}
     else
+      {:error, :invalid_token} ->
+        {:error, :invalid_token}
+
       {:error, reason} ->
         Logger.debug("Error connecting gateway websocket: #{inspect(reason)}")
         {:error, reason}
diff --git a/elixir/apps/api/lib/api/relay/socket.ex b/elixir/apps/api/lib/api/relay/socket.ex
index c35f9428a..29ceee9dd 100644
--- a/elixir/apps/api/lib/api/relay/socket.ex
+++ b/elixir/apps/api/lib/api/relay/socket.ex
@@ -33,6 +33,9 @@ defmodule API.Relay.Socket do
 
       {:ok, socket}
     else
+      {:error, :invalid_token} ->
+        {:error, :invalid_token}
+
       {:error, reason} ->
         Logger.debug("Error connecting relay websocket: #{inspect(reason)}")
         {:error, reason}
diff --git a/elixir/apps/api/lib/api/sockets.ex b/elixir/apps/api/lib/api/sockets.ex
index b6f58b635..745e0369a 100644
--- a/elixir/apps/api/lib/api/sockets.ex
+++ b/elixir/apps/api/lib/api/sockets.ex
@@ -16,18 +16,18 @@ defmodule API.Sockets do
     ]
   end
 
+  def handle_error(conn, :invalid_token),
+    do: Plug.Conn.send_resp(conn, 401, "Invalid token")
+
   def handle_error(conn, :unauthenticated),
     do: Plug.Conn.send_resp(conn, 403, "Forbidden")
 
-  def handle_error(conn, :invalid_token),
-    do: Plug.Conn.send_resp(conn, 422, "Unprocessable Entity")
+  def handle_error(conn, %Ecto.Changeset{}),
+    do: Plug.Conn.send_resp(conn, 422, "Invalid or missing connection parameters")
 
   def handle_error(conn, :rate_limit),
     do: Plug.Conn.send_resp(conn, 429, "Too many requests")
 
-  def handle_error(conn, %Ecto.Changeset{}),
-    do: Plug.Conn.send_resp(conn, 422, "Invalid or missing connection parameters")
-
   def real_ip(x_headers, peer_data) do
     real_ip =
       if is_list(x_headers) and length(x_headers) > 0 do