ZTA - user identity (#2015)

* Get user identity

* Moves current_user_id to identity_data

* Renames cookies to session_identity

* Keep zta keys in the state

* User id from GoogleIap

* Update cloudflare.ex

* Update googleiap.ex

* Email only for ZTA

* Get the current_user_id from the identity_data

* Applying suggestions

* Applying suggestions

* Fix verify_token

---------

Co-authored-by: José Valim <jose.valim@gmail.com>

lib/livebook.ex

@@ -198,7 +198,7 @@ defmodule Livebook do
     config :livebook,
            :identity_provider,
            Livebook.Config.identity_provider!("LIVEBOOK_IDENTITY_PROVIDER") ||
-             {LivebookWeb.Cookies, :unused}
+             {LivebookWeb.SessionIdentity, :unused}
   end
 
   @doc """

lib/livebook/config.ex

@@ -181,7 +181,7 @@ defmodule Livebook.Config do
   """
   @spec identity_readonly?() :: boolean()
   def identity_readonly?() do
-    not match?({LivebookWeb.Cookies, _}, Livebook.Config.identity_provider())
+    not match?({LivebookWeb.SessionIdentity, _}, Livebook.Config.identity_provider())
   end
 
   @doc """

lib/livebook/users/user.ex

@@ -17,6 +17,7 @@ defmodule Livebook.Users.User do
   @type t :: %__MODULE__{
           id: id(),
           name: String.t() | nil,
+          email: String.t() | nil,
           hex_color: hex_color()
         }
 
@@ -25,6 +26,7 @@ defmodule Livebook.Users.User do
 
   embedded_schema do
     field :name, :string
+    field :email, :string
     field :hex_color, Livebook.EctoTypes.HexColor
   end
 
@@ -36,13 +38,14 @@ defmodule Livebook.Users.User do
     %__MODULE__{
       id: Utils.random_id(),
       name: nil,
+      email: nil,
       hex_color: Livebook.EctoTypes.HexColor.random()
     }
   end
 
   def changeset(user, attrs \\ %{}) do
     user
-    |> cast(attrs, [:id, :name, :hex_color])
+    |> cast(attrs, [:id, :name, :email, :hex_color])
     |> validate_required([:id, :name, :hex_color])
   end
 end

lib/livebook/zta/cloudflare.ex

@@ -8,70 +8,57 @@ defmodule Livebook.ZTA.Cloudflare do
   @assertion "cf-access-jwt-assertion"
   @renew_afer 24 * 60 * 60 * 1000
 
-  defstruct [:name, :req_options, :identity]
+  defstruct [:name, :req_options, :identity, :keys]
 
   def start_link(opts) do
     identity = identity(opts[:identity][:key])
-    options = [req_options: [url: identity.certs], identity: identity]
+    options = [req_options: [url: identity.certs], identity: identity, keys: nil]
     GenServer.start_link(__MODULE__, options, name: opts[:name])
   end
 
-  def authenticate(name, conn) do
+  def authenticate(name, conn, fields: fields) do
     token = get_req_header(conn, @assertion)
-    GenServer.call(name, {:authenticate, token})
+    user = GenServer.call(name, {:authenticate, token, fields})
+    {conn, user}
   end
 
   @impl true
   def init(options) do
-    :ets.new(options[:name], [:public, :named_table])
-    {:ok, struct!(__MODULE__, options)}
+    state = struct!(__MODULE__, options)
+    {:ok, %{state | keys: keys(state)}}
   end
 
   @impl true
-  def handle_call(:get_keys, _from, state) do
-    keys = get_from_ets(state.name) || request_and_store_in_ets(state)
-    {:reply, keys, state}
-  end
-
-  def handle_call({:authenticate, token}, _from, state) do
-    keys = get_from_ets(state.name) || request_and_store_in_ets(state)
-    user = authenticate(token, state.identity, keys)
+  def handle_call({:authenticate, token, fields}, _from, state) do
+    user = authenticated_user(token, fields, state.identity, state.keys)
     {:reply, user, state}
   end
 
   @impl true
-  def handle_info(:request, state) do
-    request_and_store_in_ets(state)
-    {:noreply, state}
+  def handle_info(:renew, state) do
+    {:noreply, %{state | keys: keys(state)}}
   end
 
-  defp request_and_store_in_ets(state) do
+  defp keys(state) do
     Logger.debug("[#{inspect(__MODULE__)}] requesting #{inspect(state.req_options)}")
     keys = Req.request!(state.req_options).body["keys"]
-    :ets.insert(state.name, keys: keys)
-    Process.send_after(self(), :request, @renew_afer)
+    Process.send_after(self(), :renew, @renew_afer)
     keys
   end
 
-  defp get_from_ets(name) do
-    case :ets.lookup(name, :keys) do
-      [keys: keys] -> keys
-      [] -> nil
-    end
-  end
-
-  defp authenticate(token, identity, keys) do
-    with [token] <- token,
-         {:ok, token} <- verify_token(token, keys),
-         :ok <- verify_iss(token, identity.iss) do
-      %{name: token.fields["email"]}
+  defp authenticated_user(token, fields, identity, keys) do
+    with [encoded_token] <- token,
+         {:ok, token} <- verify_token(encoded_token, keys),
+         :ok <- verify_iss(token, identity.iss),
+         {:ok, user} <- get_user_identity(encoded_token, fields, identity.user_identity) do
+      Map.new(user, fn {k, v} -> {String.to_atom(k), to_string(v)} end)
     else
       _ -> nil
     end
   end
 
   defp verify_token(token, keys) do
-    Enum.find_value(keys, fn key ->
+    Enum.find_value(keys, :error, fn key ->
       case JOSE.JWT.verify(key, token) do
         {true, token, _s} -> {:ok, token}
         {_, _t, _s} -> nil
@@ -80,7 +67,14 @@ defmodule Livebook.ZTA.Cloudflare do
   end
 
   defp verify_iss(%{fields: %{"iss" => iss}}, iss), do: :ok
-  defp verify_iss(_, _), do: nil
+  defp verify_iss(_, _), do: :error
+
+  defp get_user_identity(token, fields, url) do
+    token = "CF_Authorization=#{token}"
+    fields = Enum.map(fields, &Atom.to_string/1)
+    resp = Req.request!(url: url, headers: [{"cookie", token}])
+    if resp.status == 200, do: {:ok, Map.take(resp.body, fields)}, else: :error
+  end
 
   defp identity(key) do
     %{
@@ -89,7 +83,8 @@ defmodule Livebook.ZTA.Cloudflare do
       iss: "https://#{key}.cloudflareaccess.com",
       certs: "https://#{key}.cloudflareaccess.com/cdn-cgi/access/certs",
       assertion: "cf-access-jwt-assertion",
-      email: "cf-access-authenticated-user-email"
+      email: "cf-access-authenticated-user-email",
+      user_identity: "https://#{key}.cloudflareaccess.com/cdn-cgi/access/get-identity"
     }
   end
 end

lib/livebook/zta/googleiap.ex

@@ -5,73 +5,60 @@ defmodule Livebook.ZTA.GoogleIAP do
   require Logger
   import Plug.Conn
 
-  @assertion "cf-access-jwt-assertion"
+  @assertion "x-goog-iap-jwt-assertion"
   @renew_afer 24 * 60 * 60 * 1000
 
-  defstruct [:name, :req_options, :identity]
+  defstruct [:name, :req_options, :identity, :keys]
 
   def start_link(opts) do
     identity = identity(opts[:identity][:key])
-    options = [req_options: [url: identity.certs], identity: identity]
+    options = [req_options: [url: identity.certs], identity: identity, keys: nil]
     GenServer.start_link(__MODULE__, options, name: opts[:name])
   end
 
-  def authenticate(name, conn) do
+  def authenticate(name, conn, fields: fields) do
     token = get_req_header(conn, @assertion)
-    GenServer.call(name, {:authenticate, token})
+    user = GenServer.call(name, {:authenticate, token, fields})
+    {conn, user}
   end
 
   @impl true
   def init(options) do
-    :ets.new(options[:name], [:public, :named_table])
-    {:ok, struct!(__MODULE__, options)}
+    state = struct!(__MODULE__, options)
+    {:ok, %{state | keys: keys(state)}}
   end
 
   @impl true
-  def handle_call(:get_keys, _from, state) do
-    keys = get_from_ets(state.name) || request_and_store_in_ets(state)
-    {:reply, keys, state}
-  end
-
-  def handle_call({:authenticate, token}, _from, state) do
-    keys = get_from_ets(state.name) || request_and_store_in_ets(state)
-    user = authenticate(token, state.identity, keys)
+  def handle_call({:authenticate, token, fields}, _from, state) do
+    user = authenticated_user(token, fields, state.identity, state.keys)
     {:reply, user, state}
   end
 
   @impl true
-  def handle_info(:request, state) do
-    request_and_store_in_ets(state)
-    {:noreply, state}
+  def handle_info(:renew, state) do
+    {:noreply, %{state | keys: keys(state)}}
   end
 
-  defp request_and_store_in_ets(state) do
+  defp keys(state) do
     Logger.debug("[#{inspect(__MODULE__)}] requesting #{inspect(state.req_options)}")
     keys = Req.request!(state.req_options).body["keys"]
-    :ets.insert(state.name, keys: keys)
-    Process.send_after(self(), :request, @renew_afer)
+    Process.send_after(self(), :renew, @renew_afer)
     keys
   end
 
-  defp get_from_ets(name) do
-    case :ets.lookup(name, :keys) do
-      [keys: keys] -> keys
-      [] -> nil
-    end
-  end
-
-  defp authenticate(token, identity, keys) do
-    with [token] <- token,
-         {:ok, token} <- verify_token(token, keys),
-         :ok <- verify_iss(token, identity.iss) do
-      %{name: token.fields["email"]}
+  defp authenticated_user(token, fields, identity, keys) do
+    with [encoded_token] <- token,
+         {:ok, token} <- verify_token(encoded_token, keys),
+         :ok <- verify_iss(token, identity.iss),
+         {:ok, user} <- get_user_identity(token, fields, identity.user_identity) do
+      user
     else
       _ -> nil
     end
   end
 
   defp verify_token(token, keys) do
-    Enum.find_value(keys, fn key ->
+    Enum.find_value(keys, :error, fn key ->
       case JOSE.JWT.verify(key, token) do
         {true, token, _s} -> {:ok, token}
         {_, _t, _s} -> nil
@@ -80,7 +67,19 @@ defmodule Livebook.ZTA.GoogleIAP do
   end
 
   defp verify_iss(%{fields: %{"iss" => iss}}, iss), do: :ok
-  defp verify_iss(_, _), do: nil
+  defp verify_iss(_, _), do: :error
+
+  defp get_user_identity(%{fields: %{"gcip" => gcip}}, _, _) do
+    user = %{name: gcip["name"], email: gcip["email"], id: gcip["sub"]}
+    {:ok, user}
+  end
+
+  defp get_user_identity(%{fields: fields}, _, _url) do
+    user = %{name: fields["email"], email: fields["email"], id: fields["sub"]}
+    {:ok, user}
+  end
+
+  defp get_user_identity(_, _, _), do: :error
 
   defp identity(key) do
     %{
@@ -89,7 +88,8 @@ defmodule Livebook.ZTA.GoogleIAP do
       iss: "https://cloud.google.com/iap",
       certs: "https://www.gstatic.com/iap/verify/public_key",
       assertion: "x-goog-iap-jwt-assertion",
-      email: "x-goog-authenticated-user-email"
+      email: "x-goog-authenticated-user-email",
+      user_identity: "https://www.googleapis.com/plus/v1/people/me"
     }
   end
 end

lib/livebook_web/cookies.ex

@@ -1,12 +0,0 @@
-defmodule LivebookWeb.Cookies do
-  # This module implements the ZTA contract specific to Livebook cookies
-  @moduledoc false
-
-  def authenticate(_, _conn) do
-    %{}
-  end
-
-  def child_spec(_opts) do
-    %{id: __MODULE__, start: {Function, :identity, [:ignore]}}
-  end
-end

lib/livebook_web/live/hooks/user_hook.ex

@@ -4,9 +4,9 @@ defmodule LivebookWeb.UserHook do
 
   alias Livebook.Users.User
 
-  def on_mount(:default, _params, %{"current_user_id" => current_user_id} = session, socket) do
+  def on_mount(:default, _params, %{"identity_data" => identity_data} = session, socket) do
     if connected?(socket) do
-      Livebook.Users.subscribe(current_user_id)
+      Livebook.Users.subscribe(identity_data.id)
     end
 
     socket =
@@ -33,8 +33,7 @@ defmodule LivebookWeb.UserHook do
   # attributes if the socket is connected. Otherwise uses
   # `user_data` from session.
   defp build_current_user(session, socket) do
-    %{"current_user_id" => current_user_id} = session
-    user = %{User.new() | id: current_user_id}
+    user = User.new()
     identity_data = Map.new(session["identity_data"], fn {k, v} -> {Atom.to_string(k), v} end)
 
     connect_params = get_connect_params(socket) || %{}

lib/livebook_web/live/user_component.ex

@@ -41,6 +41,9 @@ defmodule LivebookWeb.UserComponent do
             spellcheck="false"
             disabled={Livebook.Config.identity_readonly?()}
           />
+          <%= if @user.email do %>
+            <.text_field field={f[:email]} label="email" spellcheck="false" disabled="true" />
+          <% end %>
           <.hex_color_field
             field={f[:hex_color]}
             label="Cursor color"

lib/livebook_web/plugs/user_plug.ex

@@ -4,8 +4,8 @@ defmodule LivebookWeb.UserPlug do
   # Initializes the session and cookies with user-related info.
   #
   # The first time someone visits Livebook
-  # this plug stores a new random user id
-  # in the session under `:current_user_id`.
+  # this plug stores a new random user id or the ZTA user
+  # in the session under `:identity_data`.
   #
   # Additionally the cookies are checked for the presence
   # of `"user_data"` and if there is none, a new user
@@ -26,24 +26,16 @@ defmodule LivebookWeb.UserPlug do
   @impl true
   def call(conn, _opts) do
     conn
-    |> ensure_current_user_id()
     |> ensure_user_identity()
     |> ensure_user_data()
     |> mirror_user_data_in_session()
   end
 
-  defp ensure_current_user_id(conn) do
-    if get_session(conn, :current_user_id) do
-      conn
-    else
-      user_id = Livebook.Utils.random_id()
-      put_session(conn, :current_user_id, user_id)
-    end
-  end
-
   defp ensure_user_identity(conn) do
     {module, _} = Livebook.Config.identity_provider()
-    identity_data = module.authenticate(LivebookWeb.ZTA, conn)
+
+    {conn, identity_data} =
+      module.authenticate(LivebookWeb.ZTA, conn, fields: [:id, :name, :email])
 
     if identity_data do
       put_session(conn, :identity_data, identity_data)
@@ -51,7 +43,7 @@ defmodule LivebookWeb.UserPlug do
       conn
       |> put_status(:forbidden)
       |> put_view(LivebookWeb.ErrorHTML)
-      |> render("403.html")
+      |> render("403.html", %{status: 403})
       |> halt()
     end
   end

lib/livebook_web/session_identity.ex

@@ -0,0 +1,19 @@
+defmodule LivebookWeb.SessionIdentity do
+  # This module implements the ZTA contract specific to Livebook cookies
+  @moduledoc false
+
+  import Plug.Conn
+
+  def authenticate(_, conn, _) do
+    if id = get_session(conn, :current_user_id) do
+      {conn, %{id: id}}
+    else
+      user_id = Livebook.Utils.random_id()
+      {put_session(conn, :current_user_id, user_id), %{id: user_id}}
+    end
+  end
+
+  def child_spec(_opts) do
+    %{id: __MODULE__, start: {Function, :identity, [:ignore]}}
+  end
+end