1import {
2 AwsCredentialProviderError,
3 Credentials,
4 type ResolvedCredentials,
5} from "@distilled.cloud/aws/Credentials";
6import { Region } from "@distilled.cloud/aws/Region";
7import * as sts from "@distilled.cloud/aws/sts";
8import * as Effect from "effect/Effect";
9import * as Layer from "effect/Layer";
10import * as Redacted from "effect/Redacted";
11import * as Ref from "effect/Ref";
12import * as Schedule from "effect/Schedule";
13import * as Semaphore from "effect/Semaphore";
14import * as FetchHttpClient from "effect/http/FetchHttpClient";
15import { AWSEnvironment } from "./Environment.ts";
16
17/**
18 * Refresh assumed-role credentials this long before they actually expire, so a
19 * request never signs with credentials on the verge of expiring.
20 */
21const CREDENTIAL_REFRESH_WINDOW_MS = 5 * 60 * 1000;
22
23// NB: do not re-export `fromEnv` here — it is unavailable under the
24// browser/workerd bundle condition (`credentials.browser`), which would break
25// any Cloudflare Worker that transitively imports this module.
26export { Credentials, fromCredentials } from "@distilled.cloud/aws/Credentials";
27
28declare module "@distilled.cloud/aws/Credentials" {
29 interface Credentials {
30 readonly kind: "Credentials";
31 }
32}
33
34/**
35 * Lazy `Credentials` layer derived from the surrounding {@link AWSEnvironment}.
36 * Credentials are resolved on first access (not during layer construction),
37 * matching the existing @distilled.cloud/aws semantics.
38 */
39export const fromEnvironment = Layer.effect(
40 Credentials,
41 Effect.gen(function* () {
42 return Effect.flatMap(yield* AWSEnvironment, (env) => env.credentials);
43 }),
44);
45
46/**
47 * Build a single-flight, expiry-aware assumed-role credentials **resolver**,
48 * constructing the cache (a `Ref` + a 1-permit refresh semaphore) exactly once.
49 *
50 * The OUTER effect builds the cache and returns the INNER resolver; run the
51 * outer ONCE (e.g. when a binding initializes) and reuse the returned resolver
52 * for every request. That is what makes the cache actually cache: each request
53 * runs the shared resolver, which only calls `AssumeRole` when the cached
54 * credentials are missing or within the refresh window — never per request.
55 *
56 * `roleArn` and `base` are accepted as effects/layers so the long-lived signing
57 * credentials and role ARN can be read lazily at refresh time (e.g. from a
58 * deployed Worker's environment), rather than captured eagerly.
59 */
60export const makeAssumeRoleResolver = (options: {
61 /** ARN of the IAM Role to assume (resolved lazily on each refresh). */
62 readonly roleArn: Effect.Effect<string>;
63 /** Layer supplying the long-lived credentials used to sign `AssumeRole`. */
64 readonly base: Layer.Layer<Credentials>;
65 /** STS role session name. @default "alchemy" */
66 readonly roleSessionName?: string;
67 /**
68 * Region for the STS endpoint. STS `AssumeRole` is global, so this only
69 * selects the regional STS endpoint. @default "us-east-1"
70 */
71 readonly region?: string;
72}): Effect.Effect<
73 Effect.Effect<ResolvedCredentials, AwsCredentialProviderError>
74> =>
75 Effect.gen(function* () {
76 const resolve = Effect.gen(function* () {
77 const roleArn = yield* options.roleArn;
78 const response = yield* sts
79 .assumeRole({
80 RoleArn: roleArn,
81 RoleSessionName: options.roleSessionName ?? "alchemy",
82 })
83 .pipe(
84 // A freshly-created IAM user/role/access-key is eventually
85 // consistent: the first `AssumeRole` calls can fail with
86 // `AccessDenied` until the trust policy + permissions propagate.
87 // Retry those for up to ~30s.
88 Effect.retry({
89 while: (error) =>
90 (error as { name?: string }).name === "AccessDeniedException" ||
91 (error as { _tag?: string })._tag === "AccessDeniedException",
92 schedule: Schedule.exponential("1 second"),
93 times: 8,
94 }),
95 );
96 const creds = response.Credentials;
97 if (!creds) {
98 return yield* new AwsCredentialProviderError({
99 message: "STS AssumeRole returned no credentials",
100 provider: "assume-role",
101 });
102 }
103 return {
104 accessKeyId: Redacted.make(creds.AccessKeyId),
105 secretAccessKey: Redacted.isRedacted(creds.SecretAccessKey)
106 ? creds.SecretAccessKey
107 : Redacted.make(creds.SecretAccessKey),
108 sessionToken: Redacted.make(creds.SessionToken),
109 region: options.region ?? "us-east-1",
110 expiration: creds.Expiration?.getTime(),
111 } satisfies ResolvedCredentials;
112 }).pipe(
113 // Sign AssumeRole with the static base credentials, and provide a
114 // (regional) STS endpoint + HttpClient so the resolve effect is
115 // self-contained (`R = never`).
116 Effect.provide(
117 options.base.pipe(
118 Layer.provideMerge(
119 Layer.succeed(
120 Region,
121 Effect.succeed(options.region ?? "us-east-1"),
122 ),
123 ),
124 Layer.provideMerge(FetchHttpClient.layer),
125 ),
126 ),
127 Effect.mapError((cause) =>
128 cause instanceof AwsCredentialProviderError
129 ? cause
130 : new AwsCredentialProviderError({
131 message: "STS AssumeRole failed",
132 provider: "assume-role",
133 cause,
134 }),
135 ),
136 );
137
138 // Single-flight, expiry-aware cache. A `Ref` holds the last-resolved
139 // credentials; a 1-permit semaphore serializes refreshes so that N
140 // concurrent callers observing stale/expiring credentials trigger
141 // exactly ONE `AssumeRole` call (the rest await the permit and then see
142 // the freshly-cached value via the re-check inside the critical section).
143 const cache = yield* Ref.make<ResolvedCredentials | undefined>(undefined);
144 const refreshLock = yield* Semaphore.make(1);
145
146 const isFresh = (
147 creds: ResolvedCredentials | undefined,
148 ): creds is ResolvedCredentials =>
149 creds !== undefined &&
150 (creds.expiration === undefined ||
151 creds.expiration - CREDENTIAL_REFRESH_WINDOW_MS > Date.now());
152
153 return Effect.gen(function* () {
154 const current = yield* Ref.get(cache);
155 // Fast path: cached and not within the refresh window.
156 if (isFresh(current)) return current;
157 // Slow path: acquire the refresh permit, re-check (another fiber may
158 // have refreshed while we waited), then resolve once and cache.
159 return yield* refreshLock.withPermits(1)(
160 Effect.gen(function* () {
161 const latest = yield* Ref.get(cache);
162 if (isFresh(latest)) return latest;
163 const resolved = yield* resolve;
164 yield* Ref.set(cache, resolved);
165 return resolved;
166 }),
167 );
168 });
169 });
170
171/**
172 * A `Credentials` layer that assumes an IAM Role via STS and serves the
173 * returned temporary credentials (cached until shortly before expiry).
174 *
175 * `base` supplies the long-lived credentials used to *sign* the `AssumeRole`
176 * call — typically static IAM-user access keys (see
177 * {@link fromCredentials}). The `AssumeRole` request itself is signed against
178 * `base`, while the temporary credentials it returns become the resolved
179 * `Credentials` for every downstream AWS call.
180 *
181 * The cache is built once when the layer is constructed; build/provide this
182 * layer ONCE and reuse it. To share the same cache across many `Effect.provide`
183 * sites without re-providing the layer, build {@link makeAssumeRoleResolver}
184 * directly and supply the resolver via `Layer.succeed(Credentials, resolver)`.
185 */
186export const fromAssumeRole = (options: {
187 /** ARN of the IAM Role to assume. */
188 readonly roleArn: string;
189 /** Static credentials used to sign the `AssumeRole` call. */
190 readonly base: Layer.Layer<Credentials>;
191 /** STS role session name. @default "alchemy" */
192 readonly roleSessionName?: string;
193 /**
194 * Region for the STS endpoint. STS `AssumeRole` is global, so this only
195 * selects the regional STS endpoint. @default "us-east-1"
196 */
197 readonly region?: string;
198}): Layer.Layer<Credentials> =>
199 Layer.effect(
200 Credentials,
201 makeAssumeRoleResolver({
202 roleArn: Effect.succeed(options.roleArn),
203 base: options.base,
204 roleSessionName: options.roleSessionName,
205 region: options.region,
206 }),
207 );