From bf1feedca413309bbe231e45cb5d43944ab21019 Mon Sep 17 00:00:00 2001
From: Arnaldo Cesco <arnaldo.cesco@secomind.com>
Date: Thu, 6 Feb 2025 17:26:02 +0100
Subject: [PATCH] AE: fix corner case error in public key retrieval

Same as the two previous commit, but in AppEngine.

Signed-off-by: Arnaldo Cesco <arnaldo.cesco@secomind.com>
---
 CHANGELOG.md                                  |  2 +
 .../lib/astarte_appengine_api/auth/auth.ex    | 17 +----
 .../lib/astarte_appengine_api/queries.ex      | 73 +++++++++++++++----
 3 files changed, 65 insertions(+), 27 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bf30346d4..688ae1767 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,6 +26,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
   when connection to the database might fail.
 - [astarte_realm_management] Fix a corner case in the realm public key retrieval
   when connection to the database might fail.
+- [astarte_appengine_api] Fix a corner case in the realm public key retrieval
+  when connection to the database might fail.
 
 ## [1.2.0] - 2024-07-02
 ### Fixed
diff --git a/apps/astarte_appengine_api/lib/astarte_appengine_api/auth/auth.ex b/apps/astarte_appengine_api/lib/astarte_appengine_api/auth/auth.ex
index 3fa2c044a..3034fd0a3 100644
--- a/apps/astarte_appengine_api/lib/astarte_appengine_api/auth/auth.ex
+++ b/apps/astarte_appengine_api/lib/astarte_appengine_api/auth/auth.ex
@@ -18,23 +18,14 @@
 
 defmodule Astarte.AppEngine.API.Auth do
   alias Astarte.AppEngine.API.Queries
-  alias Astarte.DataAccess.Database
+  alias Astarte.AppEngine.API.Config
+  alias Astarte.Core.CQLUtils
 
   require Logger
 
   def fetch_public_key(realm) do
-    with {:ok, client} <- Database.connect(realm: realm),
-         {:ok, public_key} <- Queries.fetch_public_key(client) do
-      {:ok, public_key}
-    else
-      {:error, :public_key_not_found} ->
-        _ = Logger.warning("No public key found in realm #{realm}.", tag: "no_public_key_found")
-        {:error, :public_key_not_found}
+    keyspace = CQLUtils.realm_name_to_keyspace_name(realm, Config.astarte_instance_id!())
 
-      {:error, :database_connection_error} ->
-        _ = Logger.info("Auth request for unexisting realm #{realm}.", tag: "unexisting_realm")
-        # TODO: random busy wait here to prevent realm enumeration
-        {:error, :not_existing_realm}
-    end
+    Queries.fetch_public_key(keyspace)
   end
 end
diff --git a/apps/astarte_appengine_api/lib/astarte_appengine_api/queries.ex b/apps/astarte_appengine_api/lib/astarte_appengine_api/queries.ex
index 92b9318e4..9773106fd 100644
--- a/apps/astarte_appengine_api/lib/astarte_appengine_api/queries.ex
+++ b/apps/astarte_appengine_api/lib/astarte_appengine_api/queries.ex
@@ -23,23 +23,68 @@ defmodule Astarte.AppEngine.API.Queries do
 
   require Logger
 
-  def fetch_public_key(client) do
-    query =
-      DatabaseQuery.new()
-      |> DatabaseQuery.statement(
-        "SELECT blobAsVarchar(value) FROM kv_store WHERE group='auth' AND key='jwt_public_key_pem'"
-      )
+  @keyspace_does_not_exist_regex ~r/Keyspace (.*) does not exist/
+
+  def fetch_public_key(keyspace_name) do
+    case Xandra.Cluster.run(:xandra, &do_fetch_public_key(keyspace_name, &1)) do
+      {:ok, pem} ->
+        {:ok, pem}
+
+      {:error, %Xandra.ConnectionError{} = err} ->
+        _ =
+          Logger.warning("Database connection error #{Exception.message(err)}.",
+            tag: "database_connection_error"
+          )
 
-    result =
-      DatabaseQuery.call!(client, query)
-      |> DatabaseResult.head()
+        {:error, :database_connection_error}
+
+      {:error, %Xandra.Error{} = err} ->
+        handle_xandra_error(err)
+    end
+  end
+
+  defp do_fetch_public_key(keyspace_name, conn) do
+    query = """
+    SELECT blobAsVarchar(value)
+    FROM #{keyspace_name}.kv_store
+    WHERE group='auth' AND key='jwt_public_key_pem';
+    """
 
-    case result do
-      ["system.blobasvarchar(value)": public_key] ->
-        {:ok, public_key}
+    with {:ok, prepared} <- Xandra.prepare(conn, query),
+         {:ok, page} <-
+           Xandra.execute(conn, prepared, %{},
+             uuid_format: :binary,
+             consistency: :quorum
+           ) do
+      case Enum.to_list(page) do
+        [%{"system.blobasvarchar(value)" => pem}] ->
+          {:ok, pem}
+
+        [] ->
+          {:error, :public_key_not_found}
+      end
+    end
+  end
+
+  defp handle_xandra_error(error) do
+    %Xandra.Error{message: message} = error
+
+    case Regex.run(@keyspace_does_not_exist_regex, message) do
+      [_message, keyspace] ->
+        Logger.warning("Keyspace #{keyspace} does not exist.",
+          tag: "realm_not_found"
+        )
+
+        {:error, :not_existing_realm}
+
+      nil ->
+        _ =
+          Logger.warning(
+            "Database error, cannot get realm public key: #{Exception.message(error)}.",
+            tag: "database_error"
+          )
 
-      :empty_dataset ->
-        {:error, :public_key_not_found}
+        {:error, :database_error}
     end
   end