New step: checksum (#265)

lib/req.ex

@@ -327,19 +327,22 @@ defmodule Req do
     %Req.Request{
       registered_options:
         MapSet.new([
+          # request steps
           :user_agent,
           :compressed,
           :range,
-          :http_errors,
           :base_url,
           :params,
           :path_params,
           :auth,
           :form,
           :json,
           :compress_body,
-          :compressed,
+          :checksum,
+
+          # response steps
           :raw,
+          :http_errors,
           :decode_body,
           :decode_json,
           :redirect,
@@ -381,11 +384,13 @@ defmodule Req do
       put_range: &Req.Steps.put_range/1,
       cache: &Req.Steps.cache/1,
       put_plug: &Req.Steps.put_plug/1,
-      compress_body: &Req.Steps.compress_body/1
+      compress_body: &Req.Steps.compress_body/1,
+      checksum: &Req.Steps.checksum/1
     )
     |> Req.Request.prepend_response_steps(
       retry: &Req.Steps.retry/1,
       redirect: &Req.Steps.redirect/1,
+      verify_checksum: &Req.Steps.verify_checksum/1,
       decompress_body: &Req.Steps.decompress_body/1,
       decode_body: &Req.Steps.decode_body/1,
       handle_http_errors: &Req.Steps.handle_http_errors/1,

lib/req/steps.ex

@@ -1092,8 +1092,169 @@ defmodule Req.Steps do
     plug.(conn)
   end
 
+  defmodule CollectableWithChecksum do
+    @moduledoc false
+
+    defstruct [:collectable, :hash]
+
+    defimpl Collectable do
+      def into(%{collectable: collectable, hash: hash}) do
+        {acc, collector} = Collectable.into(collectable)
+
+        new_collector = fn
+          {acc, hash}, {:cont, element} ->
+            hash = :crypto.hash_update(hash, element)
+            {collector.(acc, {:cont, element}), hash}
+
+          {acc, hash}, :done ->
+            # verification happens in verify_checksum step (so it happens after retries,
+            # redirects, etc) and there's no other way to put the data there.
+            Process.put(:req_checksum_hash, hash)
+            collector.(acc, :done)
+
+          {acc, _hash}, :halt ->
+            collector.(acc, :halt)
+        end
+
+        {{acc, hash}, new_collector}
+      end
+    end
+  end
+
+  @doc """
+  Sets expected response body checksum.
+
+  ## Request Options
+
+    * `:checksum` - if set, this is the expected response body checksum.
+
+      Can be one of:
+
+        * `"sha1:(...)"`
+        * `"sha256:(...)"`
+
+  ## Examples
+
+      iex> resp = Req.get!("https://httpbin.org/json", checksum: "sha1:9274ffd9cf273d4a008750f44540c4c5d4c8227c")
+      iex> resp.status
+      200
+
+      iex> Req.get!("https://httpbin.org/json", checksum: "sha1:bad")
+      ** (RuntimeError) checksum mismatch
+      expected: sha1:bad
+      actual:   sha1:9274ffd9cf273d4a008750f44540c4c5d4c8227c
+  """
+  @doc step: :request
+  def checksum(request) do
+    case Req.Request.get_option(request, :checksum) do
+      nil ->
+        request
+
+      checksum when is_binary(checksum) ->
+        type = checksum_type(checksum)
+
+        case request.into do
+          nil ->
+            hash = hash_init(type)
+
+            into =
+              fn {:data, chunk}, {req, resp} ->
+                req = update_in(req.private.req_checksum_hash, &:crypto.hash_update(&1, chunk))
+                resp = update_in(resp.body, &(&1 <> chunk))
+                {:cont, {req, resp}}
+              end
+
+            request
+            |> Req.Request.put_private(:req_checksum_type, type)
+            |> Req.Request.put_private(:req_checksum_expected, checksum)
+            |> Req.Request.put_private(:req_checksum_hash, hash)
+            |> Map.replace!(:into, into)
+
+          fun when is_function(fun, 2) ->
+            hash = hash_init(type)
+
+            into =
+              fn {:data, chunk}, {req, resp} ->
+                req = update_in(req.private.req_checksum_hash, &:crypto.hash_update(&1, chunk))
+                fun.({:data, chunk}, {req, resp})
+              end
+
+            request
+            |> Req.Request.put_private(:req_checksum_type, type)
+            |> Req.Request.put_private(:req_checksum_expected, checksum)
+            |> Req.Request.put_private(:req_checksum_hash, hash)
+            |> Map.replace!(:into, into)
+
+          collectable ->
+            hash = hash_init(type)
+
+            into =
+              %CollectableWithChecksum{
+                collectable: collectable,
+                hash: hash
+              }
+
+            request
+            |> Req.Request.put_private(:req_checksum_type, type)
+            |> Req.Request.put_private(:req_checksum_expected, checksum)
+            |> Req.Request.put_private(:req_checksum_hash, :pdict)
+            |> Map.replace!(:into, into)
+        end
+    end
+  end
+
+  defp checksum_type("sha1:" <> _), do: :sha1
+  defp checksum_type("sha256:" <> _), do: :sha256
+
+  defp hash_init(:sha1), do: hash_init(:sha)
+  defp hash_init(type), do: :crypto.hash_init(type)
+
   ## Response steps
 
+  @doc """
+  Verifies the response body checksum.
+
+  See `checksum/1` for more information.
+  """
+  @doc step: :response
+  def verify_checksum({request, response}) do
+    if hash = request.private[:req_checksum_hash] do
+      hash =
+        if hash == :pdict do
+          Process.delete(:req_checksum_hash)
+        else
+          hash
+        end
+
+      type = request.private.req_checksum_type
+      expected = request.private.req_checksum_expected
+
+      actual =
+        "#{type}:" <>
+          (hash
+           |> :crypto.hash_final()
+           |> Base.encode16(case: :lower, padding: false))
+
+      if expected != actual do
+        raise """
+        checksum mismatch
+        expected: #{expected}
+        actual:   #{actual}\
+        """
+      end
+
+      request =
+        update_in(
+          request.private,
+          &Map.drop(&1, [:req_checksum_hash, :req_checksum_expected, :req_checksum_type])
+        )
+
+      {request, response}
+    else
+      {request, response}
+    end
+  end
+
   @doc """
   Decompresses the response body based on the `content-encoding` header.
 

test/req/integration_test.exs

@@ -26,6 +26,7 @@ defmodule Req.IntegrationTest do
   doctest Req.Steps,
     only: [
       auth: 1,
+      checksum: 1,
       put_user_agent: 1,
       compressed: 1,
       put_base_url: 1,

test/req/steps_test.exs

@@ -272,6 +272,93 @@ defmodule Req.StepsTest do
     end
   end
 
+  describe "checksum" do
+    @foo_sha1 "sha1:0beec7b5ea3f0fdbc95d0dd47f3c5bc275da8a33"
+    @foo_sha256 "sha256:2c26b46b68ffc68ff99b453c1d30413413422d706483bfa0f98a5e886266e7ae"
+
+    test "into: binary", c do
+      Bypass.stub(c.bypass, "GET", "/", fn conn ->
+        Plug.Conn.send_resp(conn, 200, "foo")
+      end)
+
+      req = Req.new(url: c.url)
+
+      resp = Req.get!(req, checksum: @foo_sha1)
+      assert resp.body == "foo"
+
+      resp = Req.get!(req, checksum: @foo_sha256)
+      assert resp.body == "foo"
+
+      assert_raise RuntimeError,
+                   """
+                   checksum mismatch
+                   expected: sha1:bad
+                   actual:   #{@foo_sha1}\
+                   """,
+                   fn ->
+                     Req.get!(req, checksum: "sha1:bad")
+                   end
+    end
+
+    test "into: fun", c do
+      Bypass.stub(c.bypass, "GET", "/", fn conn ->
+        Plug.Conn.send_resp(conn, 200, "foo")
+      end)
+
+      req =
+        Req.new(
+          url: c.url,
+          into: fn {:data, chunk}, {req, resp} ->
+            {:cont, {req, update_in(resp.body, &(&1 <> chunk))}}
+          end
+        )
+
+      resp = Req.get!(req, checksum: @foo_sha1)
+      assert resp.body == "foo"
+
+      resp = Req.get!(req, checksum: @foo_sha256)
+      assert resp.body == "foo"
+
+      assert_raise RuntimeError,
+                   """
+                   checksum mismatch
+                   expected: sha1:bad
+                   actual:   #{@foo_sha1}\
+                   """,
+                   fn ->
+                     Req.get!(req, checksum: "sha1:bad")
+                   end
+    end
+
+    test "into: collectable", c do
+      Bypass.stub(c.bypass, "GET", "/", fn conn ->
+        Plug.Conn.send_resp(conn, 200, "foo")
+      end)
+
+      req =
+        Req.new(
+          url: c.url,
+          into: []
+        )
+
+      resp = Req.get!(req, checksum: @foo_sha1)
+      assert resp.body == ["foo"]
+
+      resp = Req.get!(req, checksum: @foo_sha256)
+      assert resp.body == ["foo"]
+
+      assert_raise RuntimeError,
+                   """
+                   checksum mismatch
+                   expected: sha1:bad
+                   actual:   #{@foo_sha1}\
+                   """,
+                   fn ->
+                     Req.get!(req, checksum: "sha1:bad")
+                   end
+    end
+  end
+
   ## Response steps
 
   describe "decompress_body" do