From bf79caecfd81ca47f1c6771b101bdec9d66c90bc Mon Sep 17 00:00:00 2001
From: Ian Young <ian@iangreenleaf.com>
Date: Thu, 2 Jul 2020 22:33:41 -0500
Subject: [PATCH] Allow custom token fetching defined in config

---
 README.md                                     | 28 +++++++++++++++++++
 lib/waffle/storage/google/cloud_storage.ex    |  6 ++--
 .../storage/google/token/default_fetcher.ex   |  9 ++++++
 lib/waffle/storage/google/token/fetcher.ex    |  3 ++
 4 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 lib/waffle/storage/google/token/default_fetcher.ex
 create mode 100644 lib/waffle/storage/google/token/fetcher.ex

diff --git a/README.md b/README.md
index 4688c5c..3433d31 100644
--- a/README.md
+++ b/README.md
@@ -62,6 +62,34 @@ hard-coded string (e.g. `"gcs-bucket"`) or a system env tuple (e.g.
 `{:system, "WAFFLE_BUCKET"}`). You can also override this in your definition
 module (e.g. `def bucket(), do: "my-bucket"`).
 
+### Custom Token Generation ###
+
+By default, the credentials provided to Goth will be used to generate tokens.
+If you have multiple sets of credentials in Goth or otherwise need more control
+over token generation, you can define your own module:
+
+```elixir
+defmodule MyCredentials do
+  @behaviour Waffle.Storage.Google.TokenFetcher
+  @impl Waffle.Storage.Google.TokenFetcher
+  def get_token(scopes) when is_list(scopes), do: get_token(Enum.join(scopes, " "))
+  @impl Waffle.Storage.Google.TokenFetcher
+  def get_token(scope) when is_binary(scope) do
+    {:ok, token} = Goth.Token.for_scope({"my-user@my-gcs-account.com", scope})
+    token.token
+  end
+end
+```
+
+And configure it to use this new module instead of the default token generation:
+
+```elixir
+config :waffle,
+  storage: Waffle.Storage.Google,
+  bucket: "gcs-bucket-name",
+  token_fetcher: MyCredentials
+```
+
 ## URL Signing
 
 If your bucket/object permissions do not allow for public access, you will need
diff --git a/lib/waffle/storage/google/cloud_storage.ex b/lib/waffle/storage/google/cloud_storage.ex
index d85428f..6a3ef28 100644
--- a/lib/waffle/storage/google/cloud_storage.ex
+++ b/lib/waffle/storage/google/cloud_storage.ex
@@ -65,8 +65,10 @@ defmodule Waffle.Storage.Google.CloudStorage do
   """
   @spec conn(String.t) :: Tesla.Env.client
   def conn(scope \\ @full_control_scope) do
-    {:ok, token} = Goth.Token.for_scope(scope)
-    Connection.new(token.token)
+    token_store = Application.get_env(:waffle, :token_fetcher, Waffle.Storage.Google.Token.DefaultFetcher)
+
+    token_store.get_token(scope)
+    |> Connection.new()
   end
 
   @doc """
diff --git a/lib/waffle/storage/google/token/default_fetcher.ex b/lib/waffle/storage/google/token/default_fetcher.ex
new file mode 100644
index 0000000..49da6e1
--- /dev/null
+++ b/lib/waffle/storage/google/token/default_fetcher.ex
@@ -0,0 +1,9 @@
+defmodule Waffle.Storage.Google.Token.DefaultFetcher do
+  @behaviour Waffle.Storage.Google.Token.Fetcher
+
+  @impl Waffle.Storage.Google.Token.Fetcher
+  def get_token(scope) do
+    {:ok, token} = Goth.Token.for_scope(scope)
+    token.token
+  end
+end
diff --git a/lib/waffle/storage/google/token/fetcher.ex b/lib/waffle/storage/google/token/fetcher.ex
new file mode 100644
index 0000000..b57e135
--- /dev/null
+++ b/lib/waffle/storage/google/token/fetcher.ex
@@ -0,0 +1,3 @@
+defmodule Waffle.Storage.Google.Token.Fetcher do
+  @callback get_token(binary) :: binary
+end