From f0e1a7d0e5a5e11c30dd660969aa986658b00138 Mon Sep 17 00:00:00 2001
From: moon <moon@shipoclu.com>
Date: Tue, 3 Dec 2024 13:36:31 -0500
Subject: [PATCH] jwt stuff works now

---
 lib/balls_pds/jwt.ex            | 65 ++++++++++++++++++++++++---------
 lib/balls_pds/wac.ex            |  4 +-
 lib/mix/tasks/generate_token.ex | 14 +++++++
 3 files changed, 64 insertions(+), 19 deletions(-)
 create mode 100644 lib/mix/tasks/generate_token.ex

diff --git a/lib/balls_pds/jwt.ex b/lib/balls_pds/jwt.ex
index 1e1ab29..1c0cd49 100644
--- a/lib/balls_pds/jwt.ex
+++ b/lib/balls_pds/jwt.ex
@@ -1,16 +1,24 @@
 defmodule BallsPDS.JWT do
   alias BallsPDS.WAC
   require Logger
+  use Joken.Config
 
   def generate_my_jwk() do
-    raw_private_key = Base.decode16!(Application.get_env(:balls_pds, :owner_private_key))
+    raw_private_key =
+      Base.decode16!(Application.get_env(:balls_pds, :owner_private_key), case: :lower)
+
     generate_jwk(raw_private_key)
   end
 
   def generate_jwk(raw_private_key) when is_binary(raw_private_key) do
-    private_key = raw_private_key |> Base.encode64()
-    {:ok, public_key} = :crypto.generate_key(:eddsa, :ed25519, private_key)
-    JOSE.JWK.from_key({:okp, :Ed25519, public_key, private_key})
+    public_key = :crypto.generate_key(:eddsa, :ed25519, raw_private_key) |> elem(0)
+
+    %{
+      "kty" => "OKP",
+      "crv" => "Ed25519",
+      "d" => Base.url_encode64(raw_private_key, padding: false),
+      "x" => Base.url_encode64(public_key, padding: false)
+    }
   end
 
   def query_public_jwk(ap_id, key_id) do
@@ -30,9 +38,23 @@ defmodule BallsPDS.JWT do
     end
   end
 
+  defp get_kid(jwt) when is_binary(jwt) do
+    with {:kid, %{fields: %{"kid" => kid}}} <- {:kid, JOSE.JWT.peek_protected(jwt)} do
+      kid
+    else
+      {:kid, protected} ->
+        Logger.debug("protected: #{inspect(protected)}")
+        nil
+
+      error ->
+        Logger.error(inspect(error))
+        nil
+    end
+  end
+
   def extract_key_info(jwt) when is_binary(jwt) do
     with {:subject, {:ok, %{"sub" => subject}}} <- {:subject, JOSE.JWT.peek_payload(jwt)},
-         {:kid, {:ok, %{"kid" => kid}}} <- {:kid, JOSE.JWT.peek_protected(jwt)} do
+         {:kid, kid} <- {:kid, get_kid(jwt)} do
       {:ok, %{subject: subject, id: kid}}
     else
       {err, {:error, error}} ->
@@ -41,28 +63,37 @@ defmodule BallsPDS.JWT do
     end
   end
 
-  def generate_jwt(days \\ 30) when is_integer(days) and days > 0 do
-    jwk = generate_my_jwk()
-    signer = Joken.Signer.create("EdDSA", jwk)
+  def generate_jwt(jwk, days) when is_integer(days) and days > 0 do
+    Logger.debug("Creating signer.")
+    signer = Joken.Signer.create("EdDSA", jwk, %{"kid" => "key-1"})
 
     id = Application.get_env(:balls_pds, :owner_ap_id)
 
-    claims = %{
-      "iss" => id,
-      "sub" => id,
-      "aud" => Application.get_env(:balls_pds, :owner_ap_id),
-      "iat" => DateTime.utc_now() |> DateTime.to_unix(),
-      "exp" => DateTime.utc_now() |> DateTime.add(30, :day) |> DateTime.to_unix()
-    }
+    claims =
+      default_claims(%{
+        "aud" => Application.get_env(:balls_pds, :owner_ap_id),
+        "iat" => DateTime.utc_now() |> DateTime.to_unix(),
+        "exp" => DateTime.utc_now() |> DateTime.add(30, :day) |> DateTime.to_unix()
+      })
 
-    Joken.generate_and_sign!(claims, signer)
+    additional_claims = %{"sub" => id}
+
+    Logger.debug("Signing with claims: #{inspect(claims)}")
+    Joken.generate_and_sign!(claims, additional_claims, signer)
+  end
+
+  def generate_jwt(days) when is_integer(days) and days > 0 do
+    jwk = generate_my_jwk()
+    generate_jwt(jwk, days)
   end
 
   def verify_jwt(jwt, jwk) do
     public_jwk = Map.drop(jwk, ["d"])
     signer = Joken.Signer.create("EdDSA", public_jwk)
 
-    case Joken.verify_and_validate(signer, jwt) do
+    Logger.debug("KID: #{get_kid(jwt)}")
+
+    case Joken.verify_and_validate(public_jwk, jwt, signer) do
       {:ok, claims} -> {:ok, claims}
       {:error, reason} -> {:error, reason}
     end
diff --git a/lib/balls_pds/wac.ex b/lib/balls_pds/wac.ex
index 3eb38a3..d9918f9 100644
--- a/lib/balls_pds/wac.ex
+++ b/lib/balls_pds/wac.ex
@@ -24,7 +24,7 @@ defmodule BallsPDS.WAC do
   end
 
   def cached_query_public_key("did:web:" <> _ = did, id) do
-    cache_key = "KEY:#{did}"
+    cache_key = "KEY:#{did}:#{id}"
 
     case Cachex.get(:balls_cache, cache_key) do
       key when is_binary(key) ->
@@ -43,7 +43,7 @@ defmodule BallsPDS.WAC do
   end
 
   def cached_query_public_key(url, id) when is_binary(url) do
-    cache_key = "KEY:#{url}"
+    cache_key = "KEY:#{url}:#{id}"
 
     case Cachex.get(:balls_cache, cache_key) do
       key when is_binary(key) ->
diff --git a/lib/mix/tasks/generate_token.ex b/lib/mix/tasks/generate_token.ex
new file mode 100644
index 0000000..ca2df07
--- /dev/null
+++ b/lib/mix/tasks/generate_token.ex
@@ -0,0 +1,14 @@
+defmodule Mix.Tasks.GenerateToken do
+  use Mix.Task
+
+  def run(_) do
+    jwk = BallsPDS.JWT.generate_my_jwk()
+    IO.puts("JWK: #{inspect(jwk)}")
+
+    jwt = BallsPDS.JWT.generate_jwt(jwk, 30)
+    IO.puts("JWT: #{jwt}")
+
+    result = BallsPDS.JWT.verify_jwt(jwt, jwk)
+    IO.puts("Validated: #{inspect(result)}")
+  end
+end