feat: support auth_query with md5 (#175)

* feat: support auth_query with md5
* support unicode in a username, separate auth_request for scram and md5

lib/supavisor.ex

@@ -202,6 +202,7 @@ defmodule Supavisor do
           upstream_verify: tenant_record.upstream_verify,
           upstream_tls_ca: H.upstream_cert(tenant_record.upstream_tls_ca),
           require_user: tenant_record.require_user,
+          method: method,
           secrets: secrets
         }
 

lib/supavisor/client_handler.ex

@@ -106,7 +106,7 @@ defmodule Supavisor.ClientHandler do
     case decode_startup_packet(bin) do
       {:ok, hello} ->
         Logger.debug("Client startup message: #{inspect(hello)}")
-        {user, external_id} = parse_user_info(hello.payload["user"])
+        {user, external_id} = parse_user_info(hello.payload)
         Logger.metadata(project: external_id, user: user, mode: data.mode)
         {:keep_state, data, {:next_event, :internal, {:hello, {user, external_id}}}}
 
@@ -393,15 +393,19 @@ defmodule Supavisor.ClientHandler do
 
   ## Internal functions
 
-  @spec parse_user_info(String.t()) :: {String.t() | nil, String.t()}
-  def parse_user_info(username) do
-    case :binary.matches(username, ".") do
+  @spec parse_user_info(map) :: {String.t() | nil, String.t()}
+  def parse_user_info(%{"user" => user, "options" => %{"reference" => ref}}) do
+    {user, ref}
+  end
+
+  def parse_user_info(%{"user" => user}) do
+    case :binary.matches(user, ".") do
       [] ->
-        {username, nil}
+        {user, nil}
 
       matches ->
-        {pos, _} = List.last(matches)
-        {name, "." <> external_id} = String.split_at(username, pos)
+        {pos, 1} = List.last(matches)
+        <<name::size(pos)-binary, ?., external_id::binary>> = user
         {name, external_id}
     end
   end
@@ -433,7 +437,10 @@ defmodule Supavisor.ClientHandler do
       map =
         fields
         |> Enum.chunk_every(2)
-        |> Enum.map(fn [k, v] -> {k, v} end)
+        |> Enum.map(fn
+          ["options" = k, v] -> {k, URI.decode_query(v)}
+          [k, v] -> {k, v}
+        end)
         |> Map.new()
 
       # We only do light validation on the fields in the payload. The only field we use at the
@@ -447,8 +454,25 @@ defmodule Supavisor.ClientHandler do
   end
 
   @spec handle_exchange(S.sock(), {atom(), fun()}) :: {:ok, binary() | nil} | {:error, String.t()}
+  def handle_exchange({_, socket} = sock, {:auth_query_md5 = method, secrets}) do
+    salt = :crypto.strong_rand_bytes(4)
+    :ok = sock_send(sock, Server.md5_request(salt))
+
+    with {:ok,
+          %{
+            tag: :password_message,
+            payload: {:md5, client_md5}
+          }, _} <- receive_next(socket, "Timeout while waiting for the md5 exchange"),
+         {:ok, key} <- authenticate_exchange(method, client_md5, secrets.().secret, salt) do
+      {:ok, key}
+    else
+      {:error, message} -> {:error, message}
+      other -> {:error, "Unexpected message #{inspect(other)}"}
+    end
+  end
+
   def handle_exchange({_, socket} = sock, {method, secrets}) do
-    :ok = sock_send(sock, Server.auth_request())
+    :ok = sock_send(sock, Server.scram_request())
 
     with {:ok,
           %{
@@ -473,8 +497,7 @@ defmodule Supavisor.ClientHandler do
 
   defp receive_next(socket, timeout_message) do
     receive do
-      {_proto, ^socket, bin} ->
-        Server.decode_pkt(bin)
+      {_proto, ^socket, bin} -> Server.decode_pkt(bin)
     after
       15_000 -> {:error, timeout_message}
     end
@@ -504,6 +527,14 @@ defmodule Supavisor.ClientHandler do
     end
   end
 
+  defp authenticate_exchange(:auth_query_md5, client_hash, server_hash, salt) do
+    if "md5" <> H.md5([server_hash, salt]) == client_hash do
+      {:ok, nil}
+    else
+      {:error, "Wrong password"}
+    end
+  end
+
   @spec db_checkout(:on_connect | :on_query, map()) :: pid() | nil
   defp db_checkout(_, %{mode: :session, db_pid: db_pid}) when is_pid(db_pid) do
     db_pid
@@ -579,11 +610,8 @@ defmodule Supavisor.ClientHandler do
     case Cachex.fetch(Supavisor.Cache, cache_key, fn _key ->
            {:commit, {:cached, get_secrets(info, db_user)}, ttl: 15_000}
          end) do
-      {_, {:cached, value}} ->
-        value
-
-      {_, {:cached, value}, _} ->
-        value
+      {_, {:cached, value}} -> value
+      {_, {:cached, value}, _} -> value
     end
   end
 
@@ -615,7 +643,8 @@ defmodule Supavisor.ClientHandler do
     resp =
       case H.get_user_secret(conn, tenant.auth_query, db_user) do
         {:ok, secret} ->
-          {:ok, {:auth_query, fn -> Map.put(secret, :alias, user.db_user_alias) end}}
+          t = if secret.digest == :md5, do: :auth_query_md5, else: :auth_query
+          {:ok, {t, fn -> Map.put(secret, :alias, user.db_user_alias) end}}
 
         {:error, reason} ->
           {:error, reason}

lib/supavisor/db_handler.ex

@@ -179,17 +179,18 @@ defmodule Supavisor.DbHandler do
           acc
 
         %{payload: {:authentication_md5_password, salt}}, {ps, _} ->
-          password = data.auth.password
-          user = data.auth.user
+          Logger.debug("dec_pkt, #{inspect(dec_pkt, pretty: true)}")
 
-          digest = [password.(), user] |> :erlang.md5() |> Base.encode16(case: :lower)
-          digest = [digest, salt] |> :erlang.md5() |> Base.encode16(case: :lower)
-          payload = ["md5", digest, 0]
+          digest =
+            if data.auth.method == :password do
+              H.md5([data.auth.password.(), data.auth.user])
+            else
+              data.auth.secrets.().secret
+            end
 
+          payload = ["md5", H.md5([digest, salt]), 0]
           bin = [?p, <<IO.iodata_length(payload) + 4::signed-32>>, payload]
-
           :ok = sock_send(data.sock, bin)
-
           {ps, :authentication_md5}
 
         %{tag: :error_response, payload: error}, _ ->

lib/supavisor/helpers.ex

@@ -124,6 +124,10 @@ defmodule Supavisor.Helpers do
     end
   end
 
+  def parse_secret("md5" <> secret, user) do
+    {:ok, %{digest: :md5, secret: secret, user: user}}
+  end
+
   def parse_postgres_secret(_), do: {:error, "Digest not supported"}
 
   ## Internal functions
@@ -303,4 +307,11 @@ defmodule Supavisor.Helpers do
   def parse_server_first(message, nonce) do
     :pgo_scram.parse_server_first(message, nonce) |> Map.new()
   end
+
+  @spec md5([String.t()]) :: String.t()
+  def md5(strings) do
+    strings
+    |> :erlang.md5()
+    |> Base.encode16(case: :lower)
+  end
 end

lib/supavisor/protocol/server.ex

@@ -12,7 +12,7 @@ defmodule Supavisor.Protocol.Server do
   @authentication_ok <<?R, 8::32, 0::32>>
   @ready_for_query <<?Z, 5::32, ?I>>
   @ssl_request <<8::32, 1234::16, 5679::16>>
-  @auth_request <<?R, 23::32, 10::32, "SCRAM-SHA-256", 0, 0>>
+  @scram_request <<?R, 23::32, 10::32, "SCRAM-SHA-256", 0, 0>>
 
   defmodule Pkt do
     @moduledoc "Representing a packet structure with tag, length, and payload fields."
@@ -185,6 +185,13 @@ defmodule Supavisor.Protocol.Server do
     end
   end
 
+  def decode_payload(:password_message, "md5" <> _ = bin) do
+    case String.split(bin, <<0>>) do
+      [digest, ""] -> {:md5, digest}
+      _ -> :undefined
+    end
+  end
+
   def decode_payload(:password_message, bin) do
     case kv_to_map(bin) do
       {:ok, map} -> {:first_msg_response, map}
@@ -258,12 +265,17 @@ defmodule Supavisor.Protocol.Server do
     end
   end
 
-  @spec auth_request() :: iodata()
-  def auth_request() do
-    @auth_request
+  @spec scram_request() :: iodata
+  def scram_request() do
+    @scram_request
+  end
+
+  @spec md5_request(<<_::32>>) :: iodata
+  def md5_request(salt) do
+    [<<?R, 12::32, 5::32>>, salt]
   end
 
-  @spec exchange_first_message(binary(), binary() | boolean(), pos_integer()) :: binary()
+  @spec exchange_first_message(binary, binary | boolean, pos_integer) :: binary
   def exchange_first_message(nonce, salt \\ false, iterations \\ 4096) do
     secret =
       if salt do

test/supavisor/client_handler_test.exs

@@ -5,16 +5,33 @@ defmodule Supavisor.ClientHandlerTest do
 
   describe "parse_user_info/1" do
     test "extracts the external_id from the username" do
-      username = "test.user.external_id"
-      {name, external_id} = ClientHandler.parse_user_info(username)
+      payload = %{"user" => "test.user.external_id"}
+      {name, external_id} = ClientHandler.parse_user_info(payload)
       assert name == "test.user"
       assert external_id == "external_id"
     end
 
     test "username consists only of username" do
       username = "username"
-      {user, nil} = ClientHandler.parse_user_info(username)
+      payload = %{"user" => username}
+      {user, nil} = ClientHandler.parse_user_info(payload)
       assert username == user
     end
+
+    test "external_id in options" do
+      user = "test.user"
+      external_id = "external_id"
+      payload = %{"options" => %{"reference" => external_id}, "user" => user}
+      {user1, external_id1} = ClientHandler.parse_user_info(payload)
+      assert user1 == user
+      assert external_id1 == external_id
+    end
+
+    test "unicode in username" do
+      payload = %{"user" => "тестовe.імʼя.external_id"}
+      {name, external_id} = ClientHandler.parse_user_info(payload)
+      assert name == "тестовe.імʼя"
+      assert external_id == "external_id"
+    end
   end
 end

test/supavisor/db_handler_test.exs

@@ -139,7 +139,8 @@ defmodule Supavisor.DbHandlerTest do
       data = %{
         auth: %{
           password: fn -> "some_password" end,
-          user: "some_user"
+          user: "some_user",
+          method: :password
         },
         sock: {:gen_tcp, sock_port}
       }