Fix relay credential lookup for unlinked environments

infra/relay/src/environments/EnvironmentCredentials.test.ts

@@ -1,6 +1,6 @@
 import * as NodeCryptoLayer from "@effect/platform-node/NodeCrypto";
 import { describe, expect, it } from "@effect/vitest";
-import { PgDialect, QueryBuilder } from "drizzle-orm/pg-core";
+import { PgDialect } from "drizzle-orm/pg-core";
 import * as Effect from "effect/Effect";
 import * as Layer from "effect/Layer";
 
@@ -58,13 +58,17 @@ describe("EnvironmentCredentials", () => {
   it.effect("does not retain credential tokens when lookup persistence fails", () => {
     const cause = new Error("database unavailable");
     const token = "t3env_sensitive-credential-token";
+    const whereConditions: Array<unknown> = [];
     const fakeDb = {
       select: () => ({
         from: (table: unknown) => {
           expect(table).toBe(relayEnvironmentCredentials);
           return {
-            where: () => ({
-              limit: () => Effect.fail(cause),
+            where: (condition: unknown) => ({
+              limit: () => {
+                whereConditions.push(condition);
+                return Effect.fail(cause);
+              },
             }),
           };
         },
@@ -81,6 +85,13 @@ describe("EnvironmentCredentials", () => {
       });
       expect(error.cause).toBe(cause);
       expect(error).not.toHaveProperty("token");
+      expect(whereConditions).toHaveLength(1);
+
+      const query = new PgDialect().sqlToQuery(whereConditions[0] as never);
+      expect(query.sql).toContain("exists");
+      expect(query.sql).toContain('"relay_environment_links"."environment_id"');
+      expect(query.sql).toContain('"relay_environment_links"."environment_public_key"');
+      expect(query.sql).toContain('"relay_environment_links"."revoked_at" is null');
     }).pipe(
       Effect.provide(
         EnvironmentCredentials.layer.pipe(
@@ -180,7 +191,6 @@ describe("EnvironmentCredentials", () => {
     const updateValues: Array<Record<string, unknown>> = [];
     const whereConditions: Array<unknown> = [];
     const fakeDb = {
-      select: (fields: Parameters<QueryBuilder["select"]>[0]) => new QueryBuilder().select(fields),
       update: (table: unknown) => {
         expect(table).toBe(relayEnvironmentCredentials);
         return {

infra/relay/src/environments/EnvironmentCredentials.ts

@@ -6,7 +6,8 @@ import * as Encoding from "effect/Encoding";
 import * as Layer from "effect/Layer";
 import * as Option from "effect/Option";
 import * as Schema from "effect/Schema";
-import { and, eq, isNull, ne, notExists } from "drizzle-orm";
+import { and, eq, exists, isNull, ne, notExists } from "drizzle-orm";
+import { QueryBuilder } from "drizzle-orm/pg-core";
 
 import * as RelayDb from "../db.ts";
 import { relayEnvironmentCredentials, relayEnvironmentLinks } from "../persistence/schema.ts";
@@ -196,6 +197,24 @@ const make = Effect.gen(function* () {
           and(
             eq(relayEnvironmentCredentials.credentialHash, credentialHash),
             isNull(relayEnvironmentCredentials.revokedAt),
+            exists(
+              new QueryBuilder()
+                .select({ userId: relayEnvironmentLinks.userId })
+                .from(relayEnvironmentLinks)
+                .where(
+                  and(
+                    eq(
+                      relayEnvironmentLinks.environmentId,
+                      relayEnvironmentCredentials.environmentId,
+                    ),
+                    eq(
+                      relayEnvironmentLinks.environmentPublicKey,
+                      relayEnvironmentCredentials.environmentPublicKey,
+                    ),
+                    isNull(relayEnvironmentLinks.revokedAt),
+                  ),
+                ),
+            ),
           ),
         )
         .limit(1)
@@ -238,7 +257,7 @@ const make = Effect.gen(function* () {
             eq(relayEnvironmentCredentials.environmentPublicKey, input.environmentPublicKey),
             isNull(relayEnvironmentCredentials.revokedAt),
             notExists(
-              db
+              new QueryBuilder()
                 .select({ userId: relayEnvironmentLinks.userId })
                 .from(relayEnvironmentLinks)
                 .where(

infra/relay/src/environments/ManagedTunnelLimits.test.ts

@@ -4,6 +4,7 @@ import * as Layer from "effect/Layer";
 
 import * as RelayDb from "../db.ts";
 import {
+  relayEnvironmentLinks,
   relayManagedEndpointAllocations,
   relayManagedTunnelLimits,
 } from "../persistence/schema.ts";
@@ -28,7 +29,16 @@ function makeFakeDb(input: {
         }
         expect(table).toBe(relayManagedEndpointAllocations);
         return {
-          where: () => input.countRows ?? Effect.succeed([{ activeTunnels: 0 }]),
+          innerJoin: (table: unknown, condition: unknown) => {
+            expect(table).toBe(relayEnvironmentLinks);
+            expect(condition).toBeDefined();
+            return {
+              where: (where: unknown) => {
+                expect(where).toBeDefined();
+                return input.countRows ?? Effect.succeed([{ activeTunnels: 0 }]);
+              },
+            };
+          },
         };
       },
     }),

infra/relay/src/environments/ManagedTunnelLimits.ts

@@ -1,11 +1,12 @@
-import { and, count, eq, ne } from "drizzle-orm";
+import { and, count, eq, isNull, ne } from "drizzle-orm";
 import * as Context from "effect/Context";
 import * as Effect from "effect/Effect";
 import * as Layer from "effect/Layer";
 import * as Schema from "effect/Schema";
 
 import * as RelayDb from "../db.ts";
 import {
+  relayEnvironmentLinks,
   relayManagedEndpointAllocations,
   relayManagedTunnelLimits,
 } from "../persistence/schema.ts";
@@ -81,10 +82,19 @@ export const make = Effect.gen(function* () {
       const counted = yield* db
         .select({ activeTunnels: count() })
         .from(relayManagedEndpointAllocations)
+        .innerJoin(
+          relayEnvironmentLinks,
+          and(
+            eq(relayEnvironmentLinks.userId, relayManagedEndpointAllocations.userId),
+            eq(relayEnvironmentLinks.environmentId, relayManagedEndpointAllocations.environmentId),
+          ),
+        )
         .where(
           and(
             eq(relayManagedEndpointAllocations.userId, input.userId),
             ne(relayManagedEndpointAllocations.environmentId, input.environmentId),
+            isNull(relayEnvironmentLinks.revokedAt),
+            eq(relayEnvironmentLinks.managedTunnelsEnabled, true),
           ),
         )
         .pipe(

infra/relay/src/http/Api.test.ts

@@ -22,12 +22,17 @@ import {
   relayDocsRedirectRoute,
   relayEnvironmentAuthLayer,
   relayNotFoundRoute,
+  revokeEnvironmentLinkRecord,
   traceRelayHttpRequestWith,
+  unlinkEnvironmentRecord,
   verifyRelayClientBearerToken,
   withoutCapturedParentSpan,
 } from "./Api.ts";
 import * as RelayConfiguration from "../Config.ts";
+import * as RelayDb from "../db.ts";
 import * as EnvironmentCredentials from "../environments/EnvironmentCredentials.ts";
+import * as EnvironmentLinks from "../environments/EnvironmentLinks.ts";
+import * as ManagedEndpointProvider from "../environments/ManagedEndpointProvider.ts";
 
 vi.mock("@clerk/backend", () => ({
   createClerkClient: vi.fn(),
@@ -155,6 +160,185 @@ describe("relay environment authentication", () => {
   });
 });
 
+describe("relay environment unlink", () => {
+  it.effect("revokes the link and its credentials in one database transaction", () => {
+    const calls: Array<string> = [];
+    const db = {
+      $client: {
+        withTransaction: <A, E, R>(effect: Effect.Effect<A, E, R>) => {
+          calls.push("transaction");
+          return effect;
+        },
+      },
+    } as unknown as RelayDb.RelayDb["Service"];
+    const links = {
+      revokeForUser: () =>
+        Effect.sync(() => {
+          calls.push("link");
+          return true;
+        }),
+    } as unknown as EnvironmentLinks.EnvironmentLinks["Service"];
+    const credentials = {
+      revokeForEnvironmentPublicKey: () =>
+        Effect.sync(() => {
+          calls.push("credential");
+          return true;
+        }),
+    } as unknown as EnvironmentCredentials.EnvironmentCredentials["Service"];
+
+    return Effect.gen(function* () {
+      expect(
+        yield* revokeEnvironmentLinkRecord({
+          db,
+          links,
+          credentials,
+          userId: "user-1",
+          environmentId: "environment-1",
+          environmentPublicKey: "public-key",
+        }),
+      ).toBe(true);
+      expect(calls).toEqual(["transaction", "link", "credential"]);
+    });
+  });
+
+  it.effect("commits database revocation before deprovisioning the managed endpoint", () => {
+    const calls: Array<string> = [];
+    const db = {
+      $client: {
+        withTransaction: <A, E, R>(effect: Effect.Effect<A, E, R>) => {
+          calls.push("transaction");
+          return effect;
+        },
+      },
+    } as unknown as RelayDb.RelayDb["Service"];
+    const links = {
+      getForUser: () =>
+        Effect.sync(() => {
+          calls.push("lookup");
+          return {
+            environmentId: "environment-1",
+            environmentPublicKey: "public-key",
+          };
+        }),
+      revokeForUser: () =>
+        Effect.sync(() => {
+          calls.push("link");
+          return true;
+        }),
+    } as unknown as EnvironmentLinks.EnvironmentLinks["Service"];
+    const credentials = {
+      revokeForEnvironmentPublicKey: () =>
+        Effect.sync(() => {
+          calls.push("credential");
+          return true;
+        }),
+    } as unknown as EnvironmentCredentials.EnvironmentCredentials["Service"];
+    const managedEndpointProvider = {
+      deprovision: () =>
+        Effect.sync(() => {
+          calls.push("deprovision");
+        }),
+    } as unknown as ManagedEndpointProvider.ManagedEndpointProvider["Service"];
+
+    return Effect.gen(function* () {
+      expect(
+        yield* unlinkEnvironmentRecord({
+          db,
+          links,
+          credentials,
+          managedEndpointProvider,
+          userId: "user-1",
+          environmentId: "environment-1",
+        }),
+      ).toBe(true);
+      expect(calls).toEqual(["lookup", "transaction", "link", "credential", "deprovision"]);
+    });
+  });
+
+  it.effect("does not deprovision when database revocation fails", () => {
+    const calls: Array<string> = [];
+    const failure = new EnvironmentCredentials.EnvironmentCredentialRevokePersistenceError({
+      environmentId: "environment-1",
+      cause: "database unavailable",
+    });
+    const db = {
+      $client: {
+        withTransaction: <A, E, R>(effect: Effect.Effect<A, E, R>) => {
+          calls.push("transaction");
+          return effect;
+        },
+      },
+    } as unknown as RelayDb.RelayDb["Service"];
+    const links = {
+      getForUser: () =>
+        Effect.succeed({
+          environmentId: "environment-1",
+          environmentPublicKey: "public-key",
+        }),
+      revokeForUser: () =>
+        Effect.sync(() => {
+          calls.push("link");
+          return true;
+        }),
+    } as unknown as EnvironmentLinks.EnvironmentLinks["Service"];
+    const credentials = {
+      revokeForEnvironmentPublicKey: () =>
+        Effect.sync(() => {
+          calls.push("credential");
+        }).pipe(Effect.andThen(Effect.fail(failure))),
+    } as unknown as EnvironmentCredentials.EnvironmentCredentials["Service"];
+    const managedEndpointProvider = {
+      deprovision: () =>
+        Effect.sync(() => {
+          calls.push("deprovision");
+        }),
+    } as unknown as ManagedEndpointProvider.ManagedEndpointProvider["Service"];
+
+    return Effect.gen(function* () {
+      expect(
+        yield* Effect.flip(
+          unlinkEnvironmentRecord({
+            db,
+            links,
+            credentials,
+            managedEndpointProvider,
+            userId: "user-1",
+            environmentId: "environment-1",
+          }),
+        ),
+      ).toBe(failure);
+      expect(calls).toEqual(["transaction", "link", "credential"]);
+    });
+  });
+
+  it.effect("retries deprovisioning after the link is already revoked", () => {
+    const calls: Array<string> = [];
+    const links = {
+      getForUser: () => Effect.succeed(null),
+    } as unknown as EnvironmentLinks.EnvironmentLinks["Service"];
+    const managedEndpointProvider = {
+      deprovision: () =>
+        Effect.sync(() => {
+          calls.push("deprovision");
+        }),
+    } as unknown as ManagedEndpointProvider.ManagedEndpointProvider["Service"];
+
+    return Effect.gen(function* () {
+      expect(
+        yield* unlinkEnvironmentRecord({
+          db: {} as RelayDb.RelayDb["Service"],
+          links,
+          credentials: {} as EnvironmentCredentials.EnvironmentCredentials["Service"],
+          managedEndpointProvider,
+          userId: "user-1",
+          environmentId: "environment-1",
+        }),
+      ).toBe(false);
+      expect(calls).toEqual(["deprovision"]);
+    });
+  });
+});
+
 describe("relay request tracing", () => {
   it.effect(
     "does not parent endpoint spans to an ambient parent captured while building handlers",