Implement basic shared key auth

README.md

@@ -2,7 +2,7 @@
 
 ## Running
 ### Starting servers
-Server ID (`SERVER_ID`) and remote servers (`SERVERS`) are provided as environment variables.
+Server ID (`SERVER_ID`) and remote servers (`SERVERS`) are provided as environment variables along with shared auth key (`AUTH_KEY`).
 
 #### Server 1
 ```bash
@@ -28,3 +28,9 @@ curl -X POST http://localhost:4000/api/messages -H "Content-Type: application/js
 ```curl
 curl -X POST http://localhost:4001/api/messages -H "Content-Type: application/json" -d '{"from": "2-nibbler", "to": "1-bender", "message": "I am Nibbler, agent of the Nibblonian fleet."}'
 ```
+
+## Security
+This application uses a shared auth key (`AUTH_KEY`) to authenticate requests between servers. The key is provided as an environment variable and must be the same on all servers.
+
+### Next steps
+HTTPS with client certificates should be implemented to ensure secure communication between servers and prevent unauthorized access and possible man-in-the-middle attacks.

lib/message_server/auth.ex

@@ -0,0 +1,42 @@
+defmodule MessageServer.Auth do
+  @spec generate_auth_token(String.t()) :: String.t()
+  def generate_auth_token(server_id) do
+    timestamp = System.system_time(:second)
+    payload = "#{server_id}:#{timestamp}"
+    signature = :crypto.mac(:hmac, :sha256, get_shared_secret(), payload) |> Base.encode64()
+    "#{payload}:#{signature}"
+  end
+
+  @spec verify_auth_token(String.t()) :: {:ok, String.t()} | {:error, String.t()}
+  def verify_auth_token(token) do
+    case String.split(token, ":") do
+      [server_id, timestamp_str, signature] ->
+        payload = "#{server_id}:#{timestamp_str}"
+
+        expected_signature =
+          :crypto.mac(:hmac, :sha256, get_shared_secret(), payload) |> Base.encode64()
+
+        timestamp = String.to_integer(timestamp_str)
+        current_time = System.system_time(:second)
+
+        cond do
+          signature != expected_signature ->
+            {:error, :auth_failure, "Invalid signature"}
+
+          # 5 minute window
+          current_time - timestamp > 300 ->
+            {:error, :auth_failure, "Token expired"}
+
+          true ->
+            {:ok, server_id}
+        end
+
+      _ ->
+        {:error, :auth_failure, "Invalid token format"}
+    end
+  end
+
+  def get_shared_secret do
+    System.get_env("AUTH_KEY", "default-secret")
+  end
+end

lib/message_server/remote_client.ex

@@ -1,7 +1,7 @@
 defmodule MessageServer.RemoteClient do
   require Logger
 
-  alias MessageServer.{MessageRequest, ServerRegistry}
+  alias MessageServer.{MessageRequest, ServerRegistry, Auth}
 
   @spec send_message_to_server(String.t(), MessageRequest.t()) :: :ok | {:error, String.t()}
   def send_message_to_server(target_server_id, payload) do
@@ -32,10 +32,15 @@ defmodule MessageServer.RemoteClient do
   @spec make_request(map(), binary()) :: {:ok, Req.Response.t()} | {:error, String.t()}
   defp make_request(%{host: host, port: port}, serialized_payload) do
     url = "http://#{host}:#{port}/api/remote/messages"
+    local_server_id = ServerRegistry.get_server_id()
+    auth_token = Auth.generate_auth_token(local_server_id)
 
     case Req.post(url,
            body: serialized_payload,
-           headers: [{"content-type", "application/octet-stream"}]
+           headers: [
+             {"content-type", "application/octet-stream"},
+             {"authorization", "Bearer #{auth_token}"}
+           ]
          ) do
       {:ok, %Req.Response{status: status} = response} when status in 200..299 ->
         {:ok, response}

lib/message_server/router.ex

@@ -5,7 +5,8 @@ defmodule MessageServer.Router do
   alias MessageServer.{
     MessageRequest,
     RemoteHandler,
-    MessageHandler
+    MessageHandler,
+    Auth
   }
 
   plug(:match)
@@ -41,8 +42,8 @@ defmodule MessageServer.Router do
   end
 
   post "/api/remote/messages" do
-    case handle_etf_body(conn) do
-      {:ok, payload} ->
+    with {:ok, _server_id} <- authenticate_request(conn),
+         {:ok, payload} <- handle_etf_body(conn) do
       payload
       |> RemoteHandler.handle_remote_message()
       |> case do
@@ -58,6 +59,11 @@ defmodule MessageServer.Router do
           |> put_resp_content_type("application/json")
           |> send_resp(400, Jason.encode!(%{error: reason}))
       end
+    else
+      {:error, :auth_failure, reason} ->
+        conn
+        |> put_resp_content_type("application/json")
+        |> send_resp(401, Jason.encode!(%{error: "Unauthorized", message: reason}))
 
       {:error, reason} ->
         conn
@@ -98,6 +104,16 @@ defmodule MessageServer.Router do
     end
   end
 
+  defp authenticate_request(conn) do
+    case get_req_header(conn, "authorization") do
+      ["Bearer " <> token] ->
+        Auth.verify_auth_token(token)
+
+      _ ->
+        {:error, :unauthorized}
+    end
+  end
+
   @spec validate_message_request(map()) :: {:ok, map()} | {:error, String.t()}
   defp validate_message_request(params) do
     required_fields = ["from", "to", "message"]

test/message_server/storage_test.exs

@@ -17,7 +17,6 @@ defmodule MessageServer.StorageTest do
   end
 
   test "creates storage directory on startup", %{storage_dir: storage_dir} do
-    IO.puts(storage_dir)
     assert File.exists?(storage_dir)
     assert File.dir?(storage_dir)
   end