Add ZTA cache

lib/livebook/application.ex

@@ -2,6 +2,7 @@ defmodule Livebook.Application do
   use Application
 
   def start(_type, _args) do
+    Livebook.ZTA.init()
     setup_optional_dependencies()
     ensure_directories!()
     set_local_file_system!()

lib/livebook/zta.ex

@@ -0,0 +1,14 @@
+defmodule Livebook.ZTA do
+  @doc false
+  def init do
+    :ets.new(__MODULE__, [:named_table, :public, :set, read_concurrency: true])
+  end
+
+  def get(name) do
+    :ets.lookup_element(__MODULE__, name, 2)
+  end
+
+  def put(name, value) do
+    :ets.insert(__MODULE__, [{name, value}])
+  end
+end

lib/livebook/zta/cloudflare.ex

@@ -7,41 +7,38 @@ defmodule Livebook.ZTA.Cloudflare do
   @renew_afer 24 * 60 * 60 * 1000
   @fields %{"user_uuid" => :id, "name" => :name, "email" => :email}
 
-  defstruct [:req_options, :identity, :keys]
+  defstruct [:req_options, :identity, :name]
 
   def start_link(opts) do
     identity = opts[:custom_identity] || identity(opts[:identity_key])
-    options = [req_options: [url: identity.certs], identity: identity, keys: nil]
-    GenServer.start_link(__MODULE__, options, Keyword.take(opts, [:name]))
+    name = Keyword.fetch!(opts, :name)
+    options = [req_options: [url: identity.certs], identity: identity, name: name]
+    GenServer.start_link(__MODULE__, options, name: name)
   end
 
   def authenticate(name, conn, _opts) do
     token = get_req_header(conn, @assertion)
-    {identity, keys} = GenServer.call(name, :info, :infinity)
+    {identity, keys} = Livebook.ZTA.get(name)
     {conn, authenticate_user(token, identity, keys)}
   end
 
   @impl true
   def init(options) do
     state = struct!(__MODULE__, options)
-    {:ok, %{state | keys: keys(state)}}
-  end
-
-  @impl true
-  def handle_call(:info, _from, state) do
-    {:reply, {state.identity, state.keys}, state}
+    {:ok, renew(state)}
   end
 
   @impl true
   def handle_info(:renew, state) do
-    {:noreply, %{state | keys: keys(state)}}
+    {:noreply, renew(state)}
   end
 
-  defp keys(state) do
+  defp renew(state) do
     Logger.debug("[#{inspect(__MODULE__)}] requesting #{inspect(state.req_options)}")
     keys = Req.request!(state.req_options).body["keys"]
     Process.send_after(self(), :renew, @renew_afer)
-    keys
+    Livebook.ZTA.put(state.name, {state.identity, keys})
+    state
   end
 
   defp authenticate_user(token, identity, keys) do

lib/livebook/zta/google_iap.ex

@@ -7,41 +7,38 @@ defmodule Livebook.ZTA.GoogleIAP do
   @renew_afer 24 * 60 * 60 * 1000
   @fields %{"sub" => :id, "name" => :name, "email" => :email}
 
-  defstruct [:req_options, :identity, :keys]
+  defstruct [:req_options, :identity, :name]
 
   def start_link(opts) do
     identity = opts[:custom_identity] || identity(opts[:identity_key])
-    options = [req_options: [url: identity.certs], identity: identity, keys: nil]
-    GenServer.start_link(__MODULE__, options, Keyword.take(opts, [:name]))
+    name = Keyword.fetch!(opts, :name)
+    options = [req_options: [url: identity.certs], identity: identity, name: name]
+    GenServer.start_link(__MODULE__, options, name: name)
   end
 
   def authenticate(name, conn, _opts) do
     token = get_req_header(conn, @assertion)
-    {identity, keys} = GenServer.call(name, :info, :infinity)
+    {identity, keys} = Livebook.ZTA.get(name)
     {conn, authenticate_user(token, identity, keys)}
   end
 
   @impl true
   def init(options) do
     state = struct!(__MODULE__, options)
-    {:ok, %{state | keys: keys(state)}}
-  end
-
-  @impl true
-  def handle_call(:info, _from, state) do
-    {:reply, {state.identity, state.keys}, state}
+    {:ok, renew(state)}
   end
 
   @impl true
   def handle_info(:renew, state) do
-    {:noreply, %{state | keys: keys(state)}}
+    {:noreply, renew(state)}
   end
 
-  defp keys(state) do
+  defp renew(state) do
     Logger.debug("[#{inspect(__MODULE__)}] requesting #{inspect(state.req_options)}")
     keys = Req.request!(state.req_options).body["keys"]
     Process.send_after(self(), :renew, @renew_afer)
-    keys
+    Livebook.ZTA.put(state.name, {state.identity, keys})
+    state
   end
 
   defp authenticate_user(token, identity, keys) do

lib/livebook/zta/tailscale.ex

@@ -15,13 +15,13 @@ defmodule Livebook.ZTA.Tailscale do
       raise "invalid Tailscale ZTA configuration"
     end
 
-    :persistent_term.put({__MODULE__, name}, address)
+    Livebook.ZTA.put(name, address)
     :ignore
   end
 
   def authenticate(name, conn, _opts) do
     remote_ip = to_string(:inet_parse.ntoa(conn.remote_ip))
-    tailscale_address = :persistent_term.get({__MODULE__, name})
+    tailscale_address = Livebook.ZTA.get(name)
     user = authenticate_ip(remote_ip, tailscale_address)
     {conn, user}
   end

lib/livebook/zta/teleport.ex

@@ -2,7 +2,7 @@ defmodule Livebook.ZTA.Teleport do
   use GenServer
   require Logger
 
-  defstruct [:req_options, :jwks]
+  defstruct [:req_options, :name]
 
   @renew_afer 24 * 60 * 60 * 1000
   @fields %{"sub" => :id, "username" => :username}
@@ -16,34 +16,26 @@ defmodule Livebook.ZTA.Teleport do
       |> URI.append_path(@well_known_jwks_path)
       |> URI.to_string()
 
-    options = [req_options: [url: url]]
-
-    GenServer.start_link(__MODULE__, options, Keyword.take(opts, [:name]))
+    name = Keyword.fetch!(opts, :name)
+    options = [req_options: [url: url], name: name]
+    GenServer.start_link(__MODULE__, options, name: name)
   end
 
   def authenticate(name, conn, _opts) do
     token = Plug.Conn.get_req_header(conn, @assertion)
-
-    jwks = GenServer.call(name, :get_jwks, :infinity)
-
+    jwks = Livebook.ZTA.get(name)
     {conn, authenticate_user(token, jwks)}
   end
 
   @impl true
   def init(options) do
     state = struct!(__MODULE__, options)
-
-    {:ok, %{state | jwks: renew_jwks(state.req_options)}}
+    {:ok, renew(state)}
   end
 
   @impl true
-  def handle_info(:renew_jwks, state) do
-    {:noreply, %{state | jwks: renew_jwks(state.req_options)}}
-  end
-
-  @impl true
-  def handle_call(:get_jwks, _, state) do
-    {:reply, state.jwks, state}
+  def handle_info(:renew, state) do
+    {:noreply, renew(state)}
   end
 
   defp authenticate_user(token, jwks) do
@@ -85,12 +77,11 @@ defmodule Livebook.ZTA.Teleport do
     end
   end
 
-  defp renew_jwks(req_options) do
-    keys = Req.request!(req_options).body["keys"]
-
+  defp renew(state) do
+    keys = Req.request!(state.req_options).body["keys"]
     jwks = JOSE.JWK.from_map(keys)
-
-    Process.send_after(self(), :renew_jwks, @renew_afer)
-    jwks
+    Process.send_after(self(), :renew, @renew_afer)
+    Livebook.ZTA.put(state.name, jwks)
+    state
   end
 end