From 8bd38ea7edaa26ed996d27b0bbbddc4613124da4 Mon Sep 17 00:00:00 2001
From: Dan Schultzer <1254724+danschultzer@users.noreply.github.com>
Date: Fri, 27 Dec 2024 07:30:27 -0800
Subject: [PATCH] Fix spec errors

---
 lib/assent/strategies/oauth2/base.ex |  1 +
 lib/assent/strategies/oidc.ex        | 21 ++++++++++++++++-----
 test/assent/strategies/oidc_test.exs |  5 +++++
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/lib/assent/strategies/oauth2/base.ex b/lib/assent/strategies/oauth2/base.ex
index baeb013..f9bbb3e 100644
--- a/lib/assent/strategies/oauth2/base.ex
+++ b/lib/assent/strategies/oauth2/base.ex
@@ -61,6 +61,7 @@ defmodule Assent.Strategy.OAuth2.Base do
 
   @spec authorize_url(Keyword.t(), module()) ::
           {:ok, %{session_params: %{state: binary()}, url: binary()}}
+          | {:error, term()}
   def authorize_url(config, strategy) do
     config
     |> set_config(strategy)
diff --git a/lib/assent/strategies/oidc.ex b/lib/assent/strategies/oidc.ex
index 1167118..cd875c5 100644
--- a/lib/assent/strategies/oidc.ex
+++ b/lib/assent/strategies/oidc.ex
@@ -332,12 +332,23 @@ defmodule Assent.Strategy.OIDC do
   end
 
   defp peek_header(encoded, config) do
-    with [header, _, _] <- String.split(encoded, "."),
-         {:ok, json} <- Base.url_decode64(header, padding: false) do
+    with {:ok, header} <- split_header(encoded),
+         {:ok, json} <- decode_base64_url(header) do
       Config.json_library(config).decode(json)
-    else
-      {:error, error} -> {:error, error}
-      _any -> {:error, "The ID Token is not a valid JWT"}
+    end
+  end
+
+  defp split_header(encoded) do
+    case String.split(encoded, ".") do
+      [header, _, _] -> {:ok, header}
+      _ -> {:error, "The ID Token is not a valid JWT"}
+    end
+  end
+
+  defp decode_base64_url(encoded) do
+    case Base.url_decode64(encoded, padding: false) do
+      {:ok, decoded} -> {:ok, decoded}
+      :error -> {:error, "Invalid Base64URL"}
     end
   end
 
diff --git a/test/assent/strategies/oidc_test.exs b/test/assent/strategies/oidc_test.exs
index 1144262..6aa3254 100644
--- a/test/assent/strategies/oidc_test.exs
+++ b/test/assent/strategies/oidc_test.exs
@@ -370,6 +370,11 @@ defmodule Assent.Strategy.OIDCTest do
                {:error, "The ID Token is not a valid JWT"}
     end
 
+    test "with invalid base64 header in id_token", %{config: config} do
+      assert OIDC.validate_id_token(config, "@invalid.payload.signature") ==
+               {:error, "Invalid Base64URL"}
+    end
+
     test "with no `:client_secret`", %{config: config, id_token: id_token} do
       config = Keyword.delete(config, :client_secret)