Add Teleport ZTA method (#2296)

docs/deployment/teleport.md

@@ -0,0 +1,24 @@
+# Authentication with Teleport
+
+Setting up Teleport authentication will protect all routes of your notebook. It is particularly useful for adding authentication to deployed notebooks. Teleport authentication is provided in addition to [Livebook's authentication](../authentication.md) for authoring notebooks.
+
+## How to
+
+To integrate Teleport authentication with Livebook,
+set the `LIVEBOOK_IDENTITY_PROVIDER` environment variable to `LIVEBOOK_IDENTITY_PROVIDER=teleport:https://[cluster-name]:3080`.
+
+```bash
+LIVEBOOK_IDENTITY_PROVIDER=teleport:https://[cluster-name]:3080 \
+livebook server
+```
+
+See https://goteleport.com/docs/application-access/jwt/introduction/ for more information
+on how Teleport authentication works.
+
+## Livebook Teams
+
+[Livebook Teams](https://livebook.dev/teams/) users have access to airgapped notebook deployment via Docker, with pre-configured Zero Trust Authentication, shared team secrets, and file storages.
+
+Furthermore, if you are deploying multi-session apps via [Livebook Teams](https://livebook.dev/teams/), you can programmatically access data from the authenticated user by calling [`Kino.Hub.app_info/0`](https://hexdocs.pm/kino/Kino.Hub.html#app_info/0).
+
+To get started, open up Livebook, click "Add Organization" on the sidebar, and visit the "Airgapped Deployment" section of your organization.

lib/livebook/config.ex

@@ -25,6 +25,12 @@ defmodule Livebook.Config do
       name: "Tailscale",
       value: "Tailscale CLI socket path",
       module: Livebook.ZTA.Tailscale
+    },
+    %{
+      type: :teleport,
+      name: "Teleport",
+      value: "Teleport cluster address (https://[cluster-name]:3080)",
+      module: Livebook.ZTA.Teleport
     }
   ]
 

lib/livebook/zta/teleport.ex

@@ -0,0 +1,91 @@
+defmodule Livebook.ZTA.Teleport do
+  use GenServer
+  require Logger
+
+  defstruct [:req_options, :jwks]
+
+  @renew_afer 24 * 60 * 60 * 1000
+  @fields %{"sub" => :id, "username" => :username}
+  @assertion "teleport-jwt-assertion"
+  @well_known_jwks_path "/.well-known/jwks.json"
+
+  def start_link(opts) do
+    url =
+      opts[:identity_key]
+      |> URI.parse()
+      |> URI.append_path(@well_known_jwks_path)
+      |> URI.to_string()
+
+    options = [req_options: [url: url]]
+
+    GenServer.start_link(__MODULE__, options, Keyword.take(opts, [:name]))
+  end
+
+  def authenticate(name, conn, _opts) do
+    token = Plug.Conn.get_req_header(conn, @assertion)
+
+    jwks = GenServer.call(name, :get_jwks, :infinity)
+
+    {conn, authenticate_user(token, jwks)}
+  end
+
+  @impl true
+  def init(options) do
+    state = struct!(__MODULE__, options)
+
+    {:ok, %{state | jwks: renew_jwks(state.req_options)}}
+  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}
+  end
+
+  defp authenticate_user(token, jwks) do
+    with [encoded_token] <- token,
+         {:ok, %{fields: %{"exp" => exp, "nbf" => nbf}} = token} <-
+           verify_token(encoded_token, jwks),
+         :ok <- verify_timestamps(exp, nbf) do
+      for({k, v} <- token.fields, new_k = @fields[k], do: {new_k, v}, into: %{})
+    else
+      _ ->
+        nil
+    end
+  end
+
+  defp verify_token(token, keys) do
+    Enum.find_value(keys, :error, fn key ->
+      case JOSE.JWT.verify(key, token) do
+        {true, token, _s} -> {:ok, token}
+        _ -> nil
+      end
+    end)
+  end
+
+  defp verify_timestamps(exp, nbf) do
+    now = DateTime.utc_now()
+
+    with {:ok, exp} <- DateTime.from_unix(exp),
+         {:ok, nbf} <- DateTime.from_unix(nbf),
+         true <- DateTime.after?(exp, now),
+         true <- DateTime.after?(now, nbf) do
+      :ok
+    else
+      _ -> :error
+    end
+  end
+
+  defp renew_jwks(req_options) do
+    keys = Req.request!(req_options).body["keys"]
+
+    jwks = JOSE.JWK.from_map(keys)
+
+    Process.send_after(self(), :renew_jwks, @renew_afer)
+    jwks
+  end
+end

test/livebook/zta/teleport_test.exs

@@ -0,0 +1,126 @@
+defmodule Livebook.ZTA.TeleportTest do
+  use ExUnit.Case, async: true
+  use Plug.Test
+
+  alias Livebook.ZTA.Teleport
+
+  @fields [:id, :name, :email]
+  @name Context.Test.Teleport
+
+  @public_key """
+  -----BEGIN PUBLIC KEY-----
+  MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEEVs/o5+uQbTjL3chynL4wXgUg2R9
+  q9UU8I5mEovUf86QZ7kOBIjJwqnzD1omageEHWwHdBO6B+dFabmdT9POxg==
+  -----END PUBLIC KEY-----
+  """
+
+  @private_key """
+  -----BEGIN PRIVATE KEY-----
+  MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgevZzL1gdAFr88hb2
+  OF/2NxApJCzGCEDdfSp6VQO30hyhRANCAAQRWz+jn65BtOMvdyHKcvjBeBSDZH2r
+  1RTwjmYSi9R/zpBnuQ4EiMnCqfMPWiZqB4QdbAd0E7oH50VpuZ1P087G
+  -----END PRIVATE KEY-----
+  """
+
+  setup do
+    bypass = Bypass.open()
+
+    options = [
+      name: @name,
+      identity_key: "http://localhost:#{bypass.port}"
+    ]
+
+    Bypass.expect(bypass, "GET", "/.well-known/jwks.json", fn conn ->
+      conn
+      |> put_resp_content_type("application/json")
+      |> send_resp(200, Jason.encode!(%{keys: get_well_known_jwks()}))
+    end)
+
+    token = create_token()
+
+    conn = conn(:get, "/") |> put_req_header("teleport-jwt-assertion", token)
+
+    {:ok, bypass: bypass, options: options, conn: conn, token: token}
+  end
+
+  test "returns the user when it's valid", %{options: options, conn: conn} do
+    start_supervised!({Teleport, options})
+    {_conn, user} = Teleport.authenticate(@name, conn, fields: @fields)
+    assert %{id: "my-user-id", username: "myusername"} = user
+  end
+
+  test "returns nil when the exp is in the past", %{options: options, conn: conn} do
+    iat = DateTime.utc_now() |> DateTime.add(-10000)
+    exp = DateTime.utc_now() |> DateTime.add(-1000)
+    conn = put_req_header(conn, "teleport-jwt-assertion", create_token(iat, exp))
+    start_supervised!({Teleport, options})
+
+    assert {_conn, nil} = Teleport.authenticate(@name, conn, fields: @fields)
+  end
+
+  test "returns nil when the nbf is not reached yet", %{options: options, conn: conn} do
+    iat = DateTime.utc_now() |> DateTime.add(1000)
+    exp = DateTime.utc_now() |> DateTime.add(10000)
+    conn = put_req_header(conn, "teleport-jwt-assertion", create_token(iat, exp))
+    start_supervised!({Teleport, options})
+
+    assert {_conn, nil} = Teleport.authenticate(@name, conn, fields: @fields)
+  end
+
+  test "returns nil when the token is invalid", %{options: options} do
+    conn = conn(:get, "/") |> put_req_header("teleport-jwt-assertion", "invalid_token")
+    start_supervised!({Teleport, options})
+
+    assert {_conn, nil} = Teleport.authenticate(@name, conn, fields: @fields)
+  end
+
+  test "returns nil when the assertion is invalid", %{options: options} do
+    conn = conn(:get, "/") |> put_req_header("invalid_assertion", create_token())
+    start_supervised!({Teleport, options})
+
+    assert {_conn, nil} = Teleport.authenticate(@name, conn, fields: @fields)
+  end
+
+  test "fails to start the process when the key is invalid", %{bypass: bypass, options: options} do
+    Bypass.expect(bypass, "GET", "/.well-known/jwks.json", fn conn ->
+      conn
+      |> put_resp_content_type("application/json")
+      |> send_resp(200, Jason.encode!(%{keys: ["invalid_key"]}))
+    end)
+
+    assert_raise RuntimeError, fn ->
+      start_supervised!({Teleport, options})
+    end
+  end
+
+  defp get_well_known_jwks() do
+    jwk = @public_key |> JOSE.JWK.from_pem() |> JOSE.JWK.to_map() |> elem(1) |> Map.put("kid", "")
+    [jwk]
+  end
+
+  defp create_token(
+         iat \\ DateTime.utc_now(),
+         exp \\ DateTime.add(DateTime.utc_now(), 1000)
+       ) do
+    iat = DateTime.to_unix(iat)
+    exp = DateTime.to_unix(exp)
+
+    payload = %{
+      "aud" => ["http://localhost:4000"],
+      "exp" => exp,
+      "iat" => iat,
+      "iss" => "my-teleport-custer",
+      "nbf" => iat,
+      "roles" => ["access", "editor", "member"],
+      "sub" => "my-user-id",
+      "traits" => %{"host_user_gid" => [""], "host_user_uid" => [""]},
+      "username" => "myusername"
+    }
+
+    @private_key
+    |> JOSE.JWK.from_pem()
+    |> JOSE.JWT.sign(payload)
+    |> JOSE.JWS.compact()
+    |> elem(1)
+  end
+end