From e8d6b75623295a98714203269df6269ac9363ebe Mon Sep 17 00:00:00 2001
From: Michael Davis <mcarsondavis@gmail.com>
Date: Mon, 28 Jun 2021 17:07:23 -0500
Subject: [PATCH 1/2] rewrite masking function to be tail-recursive and use
 binary matches

---
 lib/mint/web_socket/frame.ex | 33 +++++++++++++++++++++------------
 1 file changed, 21 insertions(+), 12 deletions(-)

diff --git a/lib/mint/web_socket/frame.ex b/lib/mint/web_socket/frame.ex
index 30560e1d..e33060b5 100644
--- a/lib/mint/web_socket/frame.ex
+++ b/lib/mint/web_socket/frame.ex
@@ -134,20 +134,29 @@ defmodule Mint.WebSocket.Frame do
   # bytes (where the mask bytes repeat).
   # This is an "involution" function: applying the mask will mask
   # the data and applying the mask again will unmask it.
-  def apply_mask(payload, nil), do: payload
-
-  def apply_mask(payload, _mask = <<a, b, c, d>>) do
-    [a, b, c, d]
-    |> Stream.cycle()
-    |> Enum.reduce_while({payload, _acc = <<>>}, fn
-      _mask_key, {<<>>, acc} ->
-        {:halt, acc}
-
-      mask_key, {<<part_key::integer, payload_rest::binary>>, acc} ->
-        {:cont, {payload_rest, <<acc::binary, Bitwise.bxor(mask_key, part_key)::integer>>}}
-    end)
+  def apply_mask(payload, mask, acc \\ <<>>)
+
+  def apply_mask(payload, nil, _acc), do: payload
+
+  # n=4 is the happy path
+  # n=3..1 catches cases where the remaining byte_size/1 of the payload is shorter
+  # than the mask
+  for n <- 4..1 do
+    def apply_mask(
+          <<part_key::integer-size(8)-unit(unquote(n)), payload_rest::binary>>,
+          <<mask_key::integer-size(8)-unit(unquote(n)), _::binary>> = mask,
+          acc
+        ) do
+      apply_mask(
+        payload_rest,
+        mask,
+        <<acc::binary, Bitwise.bxor(mask_key, part_key)::integer-size(8)-unit(unquote(n))>>
+      )
+    end
   end
 
+  def apply_mask(<<>>, _mask, acc), do: acc
+
   @spec decode(Mint.WebSocket.t(), binary()) ::
           {:ok, Mint.WebSocket.t(), [Mint.WebSocket.frame() | {:error, term()}]}
           | {:error, Mint.WebSocket.t(), any()}

From 16a8e83a0a35b0edba2e26f7de238a3362b9f59e Mon Sep 17 00:00:00 2001
From: Michael Davis <mcarsondavis@gmail.com>
Date: Mon, 28 Jun 2021 20:55:34 -0500
Subject: [PATCH 2/2] compile apply_mask/2 and apply_mask/3 inline

---
 lib/mint/web_socket/frame.ex | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/lib/mint/web_socket/frame.ex b/lib/mint/web_socket/frame.ex
index e33060b5..6079d8e0 100644
--- a/lib/mint/web_socket/frame.ex
+++ b/lib/mint/web_socket/frame.ex
@@ -8,6 +8,8 @@ defmodule Mint.WebSocket.Frame do
   alias Mint.WebSocket.{Utils, Extension}
   alias Mint.WebSocketError
 
+  @compile {:inline, apply_mask: 2, apply_mask: 3}
+
   shared = [{:reserved, <<0::size(3)>>}, :mask, :data, :fin?]
 
   defrecord :continuation, shared