🐛 Refresh Kinesis client credentials when using task role

lib/sequin/runtime/kinesis_pipeline.ex

@@ -9,9 +9,15 @@ defmodule Sequin.Runtime.KinesisPipeline do
   alias Sequin.Runtime.Routing
   alias Sequin.Runtime.SinkPipeline
 
+  require Logger
+
   # Kinesis PartitionKey has a maximum of 256 bytes
   @max_partition_key_bytes 256
 
+  # aws_credentials refreshes temporary credentials 5 minutes before they expire,
+  # so rebuilding the client at most once a minute always picks up a valid token.
+  @default_task_role_client_refresh_seconds 60
+
   @impl SinkPipeline
   def init(context, _opts) do
     %{consumer: consumer, test_pid: test_pid} = context
@@ -20,7 +26,9 @@ defmodule Sequin.Runtime.KinesisPipeline do
 
     case KinesisSink.aws_client(consumer.sink) do
       {:ok, client} ->
-        Map.put(context, :kinesis_client, client)
+        context
+        |> Map.put(:kinesis_client, client)
+        |> Map.put(:kinesis_client_built_at, now_seconds())
 
       {:error, reason} ->
         raise "Failed to initialize Kinesis client: #{inspect(reason)}"
@@ -54,16 +62,20 @@ defmodule Sequin.Runtime.KinesisPipeline do
   def handle_batch(:default, messages, %{batch_key: stream_arn}, context) do
     %{
       consumer: %SinkConsumer{} = consumer,
-      kinesis_client: kinesis_client,
       test_pid: test_pid
     } = context
 
     setup_allowances(test_pid)
 
+    # With task role (ECS task role, IRSA, instance profile) the credentials are
+    # temporary. aws_credentials refreshes them before expiry, but the client built
+    # in init/2 keeps the original token, so rebuild it periodically.
+    context = maybe_refresh_client(context)
+
     records =
       Enum.map(messages, fn message -> build_kinesis_record(consumer, message) end)
 
-    case Kinesis.put_records(kinesis_client, stream_arn, records) do
+    case Kinesis.put_records(context.kinesis_client, stream_arn, records) do
       :ok -> {:ok, messages, context}
       {:error, error} -> {:error, error}
     end
@@ -82,6 +94,38 @@ defmodule Sequin.Runtime.KinesisPipeline do
     }
   end
 
+  defp maybe_refresh_client(%{consumer: %SinkConsumer{sink: %KinesisSink{use_task_role: true} = sink}} = context) do
+    built_at = Map.get(context, :kinesis_client_built_at)
+
+    # monotonic time can be negative, so a missing timestamp is checked with is_nil
+    if is_nil(built_at) or now_seconds() - built_at >= task_role_client_refresh_seconds() do
+      case KinesisSink.aws_client(sink) do
+        {:ok, client} ->
+          context
+          |> Map.put(:kinesis_client, client)
+          |> Map.put(:kinesis_client_built_at, now_seconds())
+
+        {:error, reason} ->
+          Logger.warning("[KinesisPipeline] Failed to refresh AWS client for task role: #{inspect(reason)}")
+          context
+      end
+    else
+      context
+    end
+  end
+
+  defp maybe_refresh_client(context), do: context
+
+  defp now_seconds, do: System.monotonic_time(:second)
+
+  defp task_role_client_refresh_seconds do
+    Keyword.get(
+      Application.get_env(:sequin, __MODULE__, []),
+      :task_role_client_refresh_seconds,
+      @default_task_role_client_refresh_seconds
+    )
+  end
+
   defp setup_allowances(nil), do: :ok
 
   defp setup_allowances(test_pid) do

test/sequin/kinesis_pipeline_task_role_test.exs

@@ -0,0 +1,138 @@
+defmodule Sequin.Runtime.KinesisPipelineTaskRoleTest do
+  # Not async: the refresh interval is set through the application env
+  use Sequin.DataCase, async: false
+
+  alias Sequin.Aws.HttpClient
+  alias Sequin.AwsMock
+  alias Sequin.Consumers
+  alias Sequin.Consumers.ConsumerEvent
+  alias Sequin.Databases.ConnectionCache
+  alias Sequin.Factory.AccountsFactory
+  alias Sequin.Factory.AwsFactory
+  alias Sequin.Factory.CharacterFactory
+  alias Sequin.Factory.ConsumersFactory
+  alias Sequin.Factory.DatabasesFactory
+  alias Sequin.Factory.ReplicationFactory
+  alias Sequin.Runtime.KinesisPipeline
+  alias Sequin.Runtime.SinkPipeline
+
+  setup do
+    account = AccountsFactory.insert_account!()
+
+    postgres_database =
+      DatabasesFactory.insert_configured_postgres_database!(account_id: account.id, tables: :character_tables)
+
+    ConnectionCache.cache_connection(postgres_database, Sequin.Repo)
+
+    replication =
+      ReplicationFactory.insert_postgres_replication!(
+        account_id: account.id,
+        postgres_database_id: postgres_database.id
+      )
+
+    consumer =
+      ConsumersFactory.insert_sink_consumer!(
+        account_id: account.id,
+        type: :kinesis,
+        batch_size: 10,
+        replication_slot_id: replication.id
+      )
+
+    {:ok, consumer} =
+      Consumers.update_sink_consumer(consumer, %{
+        sink: %{
+          type: :kinesis,
+          stream_arn: "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
+          region: "us-east-1",
+          use_task_role: true,
+          routing_mode: :static
+        }
+      })
+
+    # Each call to get_client returns a new set of temporary credentials,
+    # as aws_credentials does after refreshing an expiring token
+    {:ok, counter} = Agent.start_link(fn -> 0 end)
+
+    Mox.stub(AwsMock, :get_client, fn "us-east-1" ->
+      n = Agent.get_and_update(counter, &{&1 + 1, &1 + 1})
+
+      {:ok,
+       "task-role-key-#{n}"
+       |> AWS.Client.create("task-role-secret-#{n}", "us-east-1")
+       |> Map.put(:session_token, "task-role-token-#{n}")}
+    end)
+
+    test_pid = self()
+
+    Req.Test.stub(HttpClient, fn conn ->
+      {_key, auth} = Enum.find(conn.req_headers, fn {key, _value} -> key == "authorization" end)
+      [_, access_key] = Regex.run(~r/Credential=([^\/]+)\//, auth)
+      send(test_pid, {:kinesis_access_key, access_key})
+      Req.Test.json(conn, AwsFactory.kinesis_put_records_response_success())
+    end)
+
+    previous_env = Application.get_env(:sequin, KinesisPipeline)
+    on_exit(fn -> restore_env(previous_env) end)
+
+    {:ok, %{consumer: consumer}}
+  end
+
+  test "rebuilds the client with fresh task role credentials before delivering", %{consumer: consumer} do
+    Application.put_env(:sequin, KinesisPipeline, task_role_client_refresh_seconds: 0)
+
+    start_pipeline!(consumer)
+
+    ref1 = send_test_event(consumer)
+    assert_receive {:ack, ^ref1, [%{data: %ConsumerEvent{}}], []}, 1_000
+    assert_receive {:kinesis_access_key, key1}, 1_000
+
+    ref2 = send_test_event(consumer)
+    assert_receive {:ack, ^ref2, [%{data: %ConsumerEvent{}}], []}, 1_000
+    assert_receive {:kinesis_access_key, key2}, 1_000
+
+    # The credentials obtained on init are replaced on every batch
+    assert key1 != key2
+  end
+
+  test "delivers with task role credentials while the client is still fresh", %{consumer: consumer} do
+    Application.put_env(:sequin, KinesisPipeline, task_role_client_refresh_seconds: 3600)
+
+    start_pipeline!(consumer)
+
+    ref = send_test_batch(consumer, [deliverable_event(consumer), deliverable_event(consumer)])
+    assert_receive {:ack, ^ref, [_, _], []}, 1_000
+    assert_receive {:kinesis_access_key, "task-role-key-" <> _}, 1_000
+  end
+
+  defp restore_env(nil), do: Application.delete_env(:sequin, KinesisPipeline)
+  defp restore_env(env), do: Application.put_env(:sequin, KinesisPipeline, env)
+
+  defp deliverable_event(consumer) do
+    [consumer_id: consumer.id]
+    |> ConsumersFactory.insert_deliverable_consumer_event!()
+    |> Map.put(:data, ConsumersFactory.consumer_event_data(record: CharacterFactory.character_attrs()))
+  end
+
+  defp start_pipeline!(consumer) do
+    start_supervised!(
+      {SinkPipeline,
+       [
+         consumer_id: consumer.id,
+         producer: Broadway.DummyProducer,
+         test_pid: self()
+       ]}
+    )
+  end
+
+  defp send_test_event(consumer) do
+    Broadway.test_message(broadway(consumer), deliverable_event(consumer), metadata: %{topic: "test", headers: []})
+  end
+
+  defp send_test_batch(consumer, events) do
+    Broadway.test_batch(broadway(consumer), events, metadata: %{topic: "test", headers: []})
+  end
+
+  defp broadway(consumer) do
+    SinkPipeline.via_tuple(consumer.id)
+  end
+end