From 6aadb1cb409a80632d17bba487cfabfdb0b13d34 Mon Sep 17 00:00:00 2001
From: Alexander Strizhakov <alex.strizhakov@gmail.com>
Date: Thu, 26 Nov 2020 11:12:44 +0300
Subject: [PATCH] digest algorithm is taken from header

---
 lib/pleroma/web/plugs/digest_plug.ex        | 18 +++++++-
 test/pleroma/web/plugs/digest_plug_test.exs | 48 +++++++++++++++++++++
 2 files changed, 64 insertions(+), 2 deletions(-)
 create mode 100644 test/pleroma/web/plugs/digest_plug_test.exs

diff --git a/lib/pleroma/web/plugs/digest_plug.ex b/lib/pleroma/web/plugs/digest_plug.ex
index b521b3073..fb2723b97 100644
--- a/lib/pleroma/web/plugs/digest_plug.ex
+++ b/lib/pleroma/web/plugs/digest_plug.ex
@@ -7,8 +7,22 @@ defmodule Pleroma.Web.Plugs.DigestPlug do
   require Logger
 
   def read_body(conn, opts) do
+    digest_algorithm =
+      with [digest_header] <- Conn.get_req_header(conn, "digest") do
+        digest_header
+        |> String.split("=", parts: 2)
+        |> List.first()
+      else
+        _ -> "SHA-256"
+      end
+
+    unless String.downcase(digest_algorithm) == "sha-256" do
+      raise ArgumentError,
+        message: "invalid value for digest algorithm, got: #{digest_algorithm}"
+    end
+
     {:ok, body, conn} = Conn.read_body(conn, opts)
-    digest = "SHA-256=" <> (:crypto.hash(:sha256, body) |> Base.encode64())
-    {:ok, body, Conn.assign(conn, :digest, digest)}
+    encoded_digest = :crypto.hash(:sha256, body) |> Base.encode64()
+    {:ok, body, Conn.assign(conn, :digest, "#{digest_algorithm}=#{encoded_digest}")}
   end
 end
diff --git a/test/pleroma/web/plugs/digest_plug_test.exs b/test/pleroma/web/plugs/digest_plug_test.exs
new file mode 100644
index 000000000..629c28c93
--- /dev/null
+++ b/test/pleroma/web/plugs/digest_plug_test.exs
@@ -0,0 +1,48 @@
+defmodule Pleroma.Web.Plugs.DigestPlugTest do
+  use ExUnit.Case, async: true
+  use Plug.Test
+
+  test "digest algorithm is taken from digest header" do
+    body = "{\"hello\": \"world\"}"
+    digest = "X48E9qOokqqrvdts8nOJRJN3OWDUoyWxBf7kbu9DBPE="
+
+    {:ok, ^body, conn} =
+      :get
+      |> conn("/", body)
+      |> put_req_header("content-type", "application/json")
+      |> put_req_header("digest", "sha-256=" <> digest)
+      |> Pleroma.Web.Plugs.DigestPlug.read_body([])
+
+    assert conn.assigns[:digest] == "sha-256=" <> digest
+
+    {:ok, ^body, conn} =
+      :get
+      |> conn("/", body)
+      |> put_req_header("content-type", "application/json")
+      |> put_req_header("digest", "SHA-256=" <> digest)
+      |> Pleroma.Web.Plugs.DigestPlug.read_body([])
+
+    assert conn.assigns[:digest] == "SHA-256=" <> digest
+  end
+
+  test "error if digest algorithm is invalid" do
+    body = "{\"hello\": \"world\"}"
+    digest = "X48E9qOokqqrvdts8nOJRJN3OWDUoyWxBf7kbu9DBPE="
+
+    assert_raise ArgumentError, "invalid value for digest algorithm, got: MD5", fn ->
+      :get
+      |> conn("/", body)
+      |> put_req_header("content-type", "application/json")
+      |> put_req_header("digest", "MD5=" <> digest)
+      |> Pleroma.Web.Plugs.DigestPlug.read_body([])
+    end
+
+    assert_raise ArgumentError, "invalid value for digest algorithm, got: md5", fn ->
+      :get
+      |> conn("/", body)
+      |> put_req_header("content-type", "application/json")
+      |> put_req_header("digest", "md5=" <> digest)
+      |> Pleroma.Web.Plugs.DigestPlug.read_body([])
+    end
+  end
+end