Add prefix to generated keys (#2118)

lib/livebook/hubs/personal.ex

@@ -81,7 +81,7 @@ defmodule Livebook.Hubs.Personal do
   """
   @spec generate_secret_key() :: String.t()
   def generate_secret_key() do
-    :crypto.strong_rand_bytes(@secret_key_size) |> Base.url_encode64(padding: false)
+    Base.url_encode64(:crypto.strong_rand_bytes(@secret_key_size), padding: false)
   end
 end
 

lib/livebook/hubs/team.ex

@@ -22,12 +22,12 @@ defmodule Livebook.Hubs.Team do
 
   @type t :: %__MODULE__{
           id: String.t() | nil,
-          org_id: non_neg_integer(),
-          user_id: non_neg_integer(),
-          org_key_id: non_neg_integer(),
-          teams_key: String.t(),
-          org_public_key: String.t(),
-          session_token: String.t(),
+          org_id: non_neg_integer() | nil,
+          user_id: non_neg_integer() | nil,
+          org_key_id: non_neg_integer() | nil,
+          teams_key: String.t() | nil,
+          org_public_key: String.t() | nil,
+          session_token: String.t() | nil,
           hub_name: String.t() | nil,
           hub_emoji: String.t() | nil,
           offline: Offline.t() | nil
@@ -62,6 +62,7 @@ defmodule Livebook.Hubs.Team do
   @doc """
   Initializes a new Team hub.
   """
+  @spec new() :: t()
   def new() do
     %__MODULE__{
       user_id: nil,
@@ -95,12 +96,21 @@ defmodule Livebook.Hubs.Team do
       changeset
     end
   end
+
+  @doc """
+  Returns the public key prefix
+  """
+  @spec public_key_prefix() :: String.t()
+  def public_key_prefix(), do: "lb_opk_"
 end
 
 defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   alias Livebook.Hubs.TeamClient
   alias Livebook.Teams
 
+  @teams_key_prefix Teams.Org.teams_key_prefix()
+  @public_key_prefix Livebook.Hubs.Team.public_key_prefix()
+
   def load(team, fields) do
     struct(team, fields)
   end
@@ -146,7 +156,8 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
     # stamp requires access to the shared local key and an authenticated
     # request to the Teams server (which ensures team membership).
 
-    token = Livebook.Stamping.aead_encrypt(metadata, notebook_source, team.teams_key)
+    @teams_key_prefix <> teams_key = team.teams_key
+    token = Livebook.Stamping.aead_encrypt(metadata, notebook_source, teams_key)
 
     case Livebook.Teams.org_sign(team, token) do
       {:ok, token_signature} ->
@@ -161,8 +172,11 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   def verify_notebook_stamp(team, notebook_source, stamp) do
     %{"version" => 1, "token" => token, "token_signature" => token_signature} = stamp
 
-    if Livebook.Stamping.rsa_verify?(token_signature, token, team.org_public_key) do
-      Livebook.Stamping.aead_decrypt(token, notebook_source, team.teams_key)
+    @teams_key_prefix <> teams_key = team.teams_key
+    @public_key_prefix <> org_public_key = team.org_public_key
+
+    if Livebook.Stamping.rsa_verify?(token_signature, token, org_public_key) do
+      Livebook.Stamping.aead_decrypt(token, notebook_source, teams_key)
     else
       :error
     end

lib/livebook/teams.ex

@@ -9,6 +9,8 @@ defmodule Livebook.Teams do
   import Ecto.Changeset,
     only: [add_error: 3, apply_action: 2, apply_action!: 2, get_field: 2, change: 1]
 
+  @prefix Org.teams_key_prefix()
+
   @doc """
   Creates an Org.
 
@@ -204,7 +206,7 @@ defmodule Livebook.Teams do
   Derives the secret and sign secret from given `teams_key`.
   """
   @spec derive_keys(String.t()) :: {bitstring(), bitstring()}
-  def derive_keys(teams_key) do
+  def derive_keys(@prefix <> teams_key) do
     binary_key = Base.url_decode64!(teams_key, padding: false)
 
     <<secret::16-bytes, sign_secret::16-bytes>> =

lib/livebook/teams/org.ex

@@ -4,6 +4,8 @@ defmodule Livebook.Teams.Org do
   use Ecto.Schema
   import Ecto.Changeset
 
+  @prefix "lb_tk_"
+
   @type t :: %__MODULE__{
           id: pos_integer() | nil,
           emoji: String.t() | nil,
@@ -12,6 +14,8 @@ defmodule Livebook.Teams.Org do
           user_code: String.t() | nil
         }
 
+  @secret_key_size 32
+
   @primary_key {:id, :id, autogenerate: false}
   embedded_schema do
     field :emoji, :string
@@ -27,7 +31,10 @@ defmodule Livebook.Teams.Org do
   Generates a new teams key.
   """
   @spec teams_key() :: String.t()
-  def teams_key, do: Base.url_encode64(:crypto.strong_rand_bytes(32), padding: false)
+  def teams_key() do
+    key = :crypto.strong_rand_bytes(@secret_key_size)
+    @prefix <> Base.url_encode64(key, padding: false)
+  end
 
   @doc """
   Generates a hash key.
@@ -43,4 +50,10 @@ defmodule Livebook.Teams.Org do
     |> validate_required(@required_fields)
     |> validate_format(:name, ~r/^[a-z0-9][a-z0-9\-]*$/)
   end
+
+  @doc """
+  Returns the teams key prefix
+  """
+  @spec teams_key_prefix() :: String.t()
+  def teams_key_prefix(), do: @prefix
 end

test/support/factory.ex

@@ -23,7 +23,7 @@ defmodule Livebook.Factory do
       org_id: 1,
       user_id: 1,
       org_key_id: 1,
-      org_public_key: Livebook.Utils.random_id(),
+      org_public_key: Livebook.Hubs.Team.public_key_prefix() <> Livebook.Utils.random_id(),
       teams_key: org.teams_key,
       session_token: Livebook.Utils.random_short_id(),
       offline: nil

test/support/hub_helpers.ex

@@ -7,9 +7,11 @@ defmodule Livebook.HubHelpers do
 
   @offline_hub %Livebook.Hubs.Team{
     id: "team-org-number-3079",
-    teams_key: "A9TarFeAzmX3sDwSPm5JP5qbLPnNpLpzmjVZUCHXwmI",
+    teams_key:
+      Livebook.Teams.Org.teams_key_prefix() <> "A9TarFeAzmX3sDwSPm5JP5qbLPnNpLpzmjVZUCHXwmI",
     org_public_key:
-      "MIIBCgKCAQEA5v_qciaRGOZd5kgCQbhQDgFCnTnIKI5xzN4m4rVtLXMPH7RTA-K6C-e4wy2gn8zulXgSYX4vXDACSjFAG4PlFhXTPgb-v3rFLwbBrUHdaTMTyxRdK52NyNoDpYklQ7FaEU9vr3Z_-cpAQjdADOV1k45GmFe3bo4gImIfUSDYp1rRiEsYcIBt0Wc0S-vQHKSlmfcCexe254_UkvWjLW7KO790bem-PSWcBI_713oRr2mQoxXeeGKd5dSyFsIr5SZXVRWcRK3soQimCXB0ddBSXZ7d2Md3P9Ylo7TcYdBGHlwVIsrmB-P70KPHPYuAVgS9QsIiiMGXPwYVW77xNRTlcwIDAQAB",
+      Livebook.Hubs.Team.public_key_prefix() <>
+        "MIIBCgKCAQEA5v_qciaRGOZd5kgCQbhQDgFCnTnIKI5xzN4m4rVtLXMPH7RTA-K6C-e4wy2gn8zulXgSYX4vXDACSjFAG4PlFhXTPgb-v3rFLwbBrUHdaTMTyxRdK52NyNoDpYklQ7FaEU9vr3Z_-cpAQjdADOV1k45GmFe3bo4gImIfUSDYp1rRiEsYcIBt0Wc0S-vQHKSlmfcCexe254_UkvWjLW7KO790bem-PSWcBI_713oRr2mQoxXeeGKd5dSyFsIr5SZXVRWcRK3soQimCXB0ddBSXZ7d2Md3P9Ylo7TcYdBGHlwVIsrmB-P70KPHPYuAVgS9QsIiiMGXPwYVW77xNRTlcwIDAQAB",
     hub_name: "org-number-3079",
     user_id: 0,
     org_id: 0,