diff --git a/lib/sequin/runtime/kinesis_pipeline.ex b/lib/sequin/runtime/kinesis_pipeline.ex index 3c435adc1..69c68c24a 100644 --- a/lib/sequin/runtime/kinesis_pipeline.ex +++ b/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 diff --git a/test/sequin/kinesis_pipeline_task_role_test.exs b/test/sequin/kinesis_pipeline_task_role_test.exs new file mode 100644 index 000000000..64cac097b --- /dev/null +++ b/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