fixes #5

Author: zookzook

lib/mongo/auth.ex

@@ -1,27 +1,30 @@
 defmodule Mongo.Auth do
   @moduledoc false
 
-  def run(opts, s) do
-    auth = setup(opts)
-    auther = mechanism(s)
-
-    auth_source = opts[:auth_source]
-    wire_version = s[:wire_version]
-
-    s = if auth_source != nil && wire_version > 0,
-          do: Map.put(s, :database, auth_source),
-        else: s
-
-    Enum.find_value(auth, fn opts ->
-      case auther.auth(opts, s) do
-        :ok ->
-          nil
+  def run(opts, state) do
+
+    db           = opts[:database]
+    auth         = setup(opts)
+    auther       = mechanism(state)
+    auth_source  = opts[:auth_source]
+    wire_version = state[:wire_version]
+
+    # change database for auth
+    state = case auth_source != nil && wire_version > 0 do
+      true  -> Map.put(state, :database, auth_source)
+      false -> state
+    end
+
+    # do auth
+    Enum.find_value(auth, fn credentials ->
+      case auther.auth(credentials, db, state) do
+        :ok -> nil    # everything is okay, then return nil
         error ->
-          {mod, socket} = s.connection
+          {mod, socket} = state.connection
           mod.close(socket)
           error
       end
-    end) || {:ok, Map.put(s,:database, opts[:database])}
+    end) || {:ok, Map.put(state, :database, opts[:database])} # restore old database
   end
 
   defp setup(opts) do
@@ -39,10 +42,7 @@ defmodule Mongo.Auth do
     if username && password, do: auth ++ [{username, password}], else: auth
   end
 
-  defp mechanism(%{wire_version: version, auth_mechanism: :x509}) when version >= 3,
-    do: Mongo.Auth.X509
-  defp mechanism(%{wire_version: version}) when version >= 3,
-    do: Mongo.Auth.SCRAM
-  defp mechanism(_),
-    do: Mongo.Auth.CR
+  defp mechanism(%{wire_version: version, auth_mechanism: :x509}) when version >= 3, do: Mongo.Auth.X509
+  defp mechanism(%{wire_version: version}) when version >= 3,  do: Mongo.Auth.SCRAM
+  defp mechanism(_), do: Mongo.Auth.CR
 end

lib/mongo/auth/cr.ex

@@ -2,7 +2,7 @@ defmodule Mongo.Auth.CR do
   @moduledoc false
   alias Mongo.MongoDBConnection.Utils
 
-  def auth({username, password}, s) do
+  def auth({username, password}, _db, s) do
     with {:ok, message} <- Utils.command(-2, [getnonce: 1], s),
          do: nonce(message, username, password, s)
   end

lib/mongo/auth/scram.ex

@@ -5,20 +5,20 @@ defmodule Mongo.Auth.SCRAM do
 
   alias Mongo.MongoDBConnection.Utils
 
-  def auth({username, password}, s) do
-    # TODO: Wrap and log error
+  def auth({username, password}, db, s) do
 
-    nonce      = nonce()
-    first_bare = first_bare(username, nonce)
-    payload    = first_message(first_bare)
-    message    = [saslStart: 1, mechanism: "SCRAM-SHA-1", payload: payload]
+    {mechanism, digest} = select_digest(db, username, s)
+    nonce               = nonce()
+    first_bare          = first_bare(username, nonce)
+    payload             = first_message(first_bare)
+    message             = [saslStart: 1, mechanism: mechanism, payload: payload]
 
     result =
-      with {:ok, %{"ok" => ok} = reply} when ok == 1 <- Utils.command(-2, message, s),
-           {message, signature} = first(reply, first_bare, username, password, nonce),
-           {:ok, %{"ok" => ok} = reply} when ok == 1 <- Utils.command(-3, message, s),
-           message = second(reply, signature),
+      with {:ok, %{"ok" => ok} = reply} when ok == 1 <- Utils.command(-3, message, s),
+           {message, signature} = first(reply, first_bare, username, password, nonce, digest),
            {:ok, %{"ok" => ok} = reply} when ok == 1 <- Utils.command(-4, message, s),
+           message = second(reply, signature),
+           {:ok, %{"ok" => ok} = reply} when ok == 1 <- Utils.command(-5, message, s),
            do: final(reply)
 
     case result do
@@ -31,21 +31,21 @@ defmodule Mongo.Auth.SCRAM do
     end
   end
 
-  defp first(%{"conversationId" => 1, "payload" => server_payload, "done" => false},
-             first_bare, username, password, client_nonce) do
+  defp first(%{"conversationId" => 1, "payload" => server_payload, "done" => false}, first_bare, username, password, client_nonce, digest) do
+
     params          = parse_payload(server_payload)
     server_nonce    = params["r"]
     salt            = params["s"] |> Base.decode64!
     iter            = params["i"] |> String.to_integer
-    pass            = Utils.digest_password(username, password)
-    salted_password = hi(pass, salt, iter)
+    pass            = Utils.digest_password(username, password, digest)
+    salted_password = hi(pass, salt, iter, digest)
 
     <<^client_nonce::binary(24), _::binary>> = server_nonce
 
     client_message       = "c=biws,r=#{server_nonce}"
     auth_message         = "#{first_bare},#{server_payload.binary},#{client_message}"
-    server_signature     = generate_signature(salted_password, auth_message)
-    proof                = generate_proof(salted_password, auth_message)
+    server_signature     = generate_signature(salted_password, auth_message, digest)
+    proof                = generate_proof(salted_password, auth_message, digest)
     client_final_message = %BSON.Binary{binary: "#{client_message},#{proof}"}
     message              = [saslContinue: 1, conversationId: 1, payload: client_final_message]
 
@@ -70,32 +70,28 @@ defmodule Mongo.Auth.SCRAM do
     "n=#{encode_username(username)},r=#{nonce}"
   end
 
-  defp hi(password, salt, iterations) do
-    Mongo.PBKDF2Cache.pbkdf2(password, salt, iterations)
+  defp hi(password, salt, iterations, digest) do
+    Mongo.PBKDF2Cache.pbkdf2(password, salt, iterations, digest)
   end
 
-  defp generate_proof(salted_password, auth_message) do
-    client_key   = :crypto.hmac(:sha, salted_password, "Client Key")
-    stored_key   = :crypto.hash(:sha, client_key)
-    signature    = :crypto.hmac(:sha, stored_key, auth_message)
+  defp generate_proof(salted_password, auth_message, digest) do
+    client_key   = :crypto.hmac(digest, salted_password, "Client Key")
+    stored_key   = :crypto.hash(digest, client_key)
+    signature    = :crypto.hmac(digest, stored_key, auth_message)
     client_proof = xor_keys(client_key, signature, "")
     "p=#{Base.encode64(client_proof)}"
   end
 
-  defp generate_signature(salted_password, auth_message) do
-    server_key = :crypto.hmac(:sha, salted_password, "Server Key")
-    :crypto.hmac(:sha, server_key, auth_message)
+  defp generate_signature(salted_password, auth_message, digest) do
+    server_key = :crypto.hmac(digest, salted_password, "Server Key")
+    :crypto.hmac(digest, server_key, auth_message)
   end
 
-  defp xor_keys("", "", result),
-    do: result
-  defp xor_keys(<<fa, ra::binary>>, <<fb, rb::binary>>, result),
-    do: xor_keys(ra, rb, <<result::binary, fa ^^^ fb>>)
-
+  defp xor_keys("", "", result), do: result
+  defp xor_keys(<<fa, ra::binary>>, <<fb, rb::binary>>, result), do: xor_keys(ra, rb, <<result::binary, fa ^^^ fb>>)
 
   defp nonce do
-    :crypto.strong_rand_bytes(18)
-    |> Base.encode64
+    :crypto.strong_rand_bytes(18) |> Base.encode64
   end
 
   defp encode_username(username) do
@@ -109,4 +105,22 @@ defmodule Mongo.Auth.SCRAM do
     |> String.split(",")
     |> Enum.into(%{}, &List.to_tuple(String.split(&1, "=", parts: 2)))
   end
+
+  ##
+  # selects the supported sasl mechanism
+  # It calls isMaster with saslSupportedMechs option to ask for the selected user which mechanism is supported
+  #
+  defp select_digest(database, username, state) do
+    with {:ok, reply} <- Utils.command(-2, [isMaster: 1, saslSupportedMechs: database <> "." <> username], state ) do
+      select_digest(reply)
+    end
+  end
+  defp select_digest(%{"saslSupportedMechs" => mechs}) do
+    case Enum.member?(mechs, "SCRAM-SHA-256") do
+      true  -> {"SCRAM-SHA-256", :sha256}
+      false -> {"SCRAM-SHA-1", :sha}
+    end
+  end
+  defp select_digest(_), do: {"SCRAM-SHA-1", :sha}
+
 end

lib/mongo/auth/x509.ex

@@ -2,7 +2,7 @@ defmodule Mongo.Auth.X509 do
   @moduledoc false
   alias Mongo.MongoDBConnection.Utils
 
-  def auth({username, _password}, s) do
+  def auth({username, _password}, _db, s) do
     cmd = [authenticate: 1, user: username, mechanism: "MONGODB-X509"]
     with {:ok, _message} <- Utils.command(-2, cmd, s) do
       :ok

lib/mongo/pbkdf2.ex

@@ -25,8 +25,8 @@ defmodule Mongo.PBKDF2 do
   """
   def generate(secret, salt, opts \\ []) do
     iterations = Keyword.get(opts, :iterations, 1000)
-    length = Keyword.get(opts, :length, 32)
-    digest = Keyword.get(opts, :digest, :sha256)
+    length     = Keyword.get(opts, :length, 32)
+    digest     = Keyword.get(opts, :digest, :sha256)
 
     if length > @max_length do
       raise ArgumentError, "length must be less than or equal to #{@max_length}"
@@ -35,8 +35,7 @@ defmodule Mongo.PBKDF2 do
     end
   end
 
-  defp generate(_fun, _salt, _iterations, max_length, _block_index, acc, length)
-      when length >= max_length do
+  defp generate(_fun, _salt, _iterations, max_length, _block_index, acc, length) when length >= max_length do
     key = acc |> Enum.reverse |> IO.iodata_to_binary
     <<bin::binary-size(max_length), _::binary>> = key
     bin
@@ -45,8 +44,7 @@ defmodule Mongo.PBKDF2 do
   defp generate(fun, salt, iterations, max_length, block_index, acc, length) do
     initial = fun.(<<salt::binary, block_index::integer-size(32)>>)
     block   = iterate(fun, iterations - 1, initial, initial)
-    generate(fun, salt, iterations, max_length, block_index + 1,
-             [block | acc], byte_size(block) + length)
+    generate(fun, salt, iterations, max_length, block_index + 1, [block | acc], byte_size(block) + length)
   end
 
   defp iterate(_fun, 0, _prev, acc), do: acc

lib/mongo/pbkdf2_cache.ex

@@ -7,8 +7,8 @@ defmodule Mongo.PBKDF2Cache do
     GenServer.start_link(__MODULE__, [], name: @name)
   end
 
-  def pbkdf2(password, salt, iterations) do
-    GenServer.call(@name, {password, salt, iterations})
+  def pbkdf2(password, salt, iterations, digest) do
+    GenServer.call(@name, {password, salt, iterations, digest})
   end
 
   def init([]) do
@@ -17,10 +17,8 @@ defmodule Mongo.PBKDF2Cache do
 
   def handle_call(key, from, s) do
     cond do
-      salted_password = s.cache[key] ->
-        {:reply, salted_password, s}
-      list = s.pending[key] ->
-        {:noreply, put_in(s.pending[key], [from|list])}
+      salted_password = s.cache[key] -> {:reply, salted_password, s}
+      list = s.pending[key] -> {:noreply, put_in(s.pending[key], [from|list])}
       true ->
         _ = run_task(key)
         {:noreply, put_in(s.pending[key], [from])}
@@ -41,12 +39,15 @@ defmodule Mongo.PBKDF2Cache do
     {:noreply, s}
   end
 
-  defp run_task({password, salt, iterations} = key) do
+  defp run_task({password, salt, iterations, :sha256} = key) do
     Task.async(fn ->
-      result = Mongo.PBKDF2.generate(password, salt,
-                                     iterations: iterations,
-                                     length: 20,
-                                     digest: :sha)
+      result = Mongo.PBKDF2.generate(password, salt, iterations: iterations, length: 32, digest: :sha256)
+      {key, result}
+    end)
+  end
+  defp run_task({password, salt, iterations, :sha} = key) do
+    Task.async(fn ->
+      result = Mongo.PBKDF2.generate(password, salt, iterations: iterations, length: 20, digest: :sha)
       {key, result}
     end)
   end

lib/mongo_db_connection/utils.ex

@@ -90,12 +90,15 @@ defmodule Mongo.MongoDBConnection.Utils do
   def namespace(coll, _, database), do: [database, ?. | coll]
 
   def digest(nonce, username, password) do
-    :crypto.hash(:md5, [nonce, username, digest_password(username, password)])
+    :crypto.hash(:md5, [nonce, username, digest_password(username, password, :sha)])
     |> Base.encode16(case: :lower)
   end
 
-  def digest_password(username, password) do
+  def digest_password(username, password, :sha) do
     :crypto.hash(:md5, [username, ":mongo:", password])
     |> Base.encode16(case: :lower)
   end
+  def digest_password(username, password, :sha256) do
+    password
+  end
 end