1import { and, count, eq, isNull, ne } from "drizzle-orm";
2import * as Context from "effect/Context";
3import * as Effect from "effect/Effect";
4import * as Layer from "effect/Layer";
5import * as Schema from "effect/Schema";
6
7import * as RelayDb from "../db.ts";
8import {
9 relayEnvironmentLinks,
10 relayManagedEndpointAllocations,
11 relayManagedTunnelLimits,
12} from "../persistence/schema.ts";
13
14/**
15 * Managed tunnels a user may hold at once unless a row in
16 * `relay_managed_tunnel_limits` overrides it for that user.
17 */
18export const DEFAULT_MANAGED_TUNNEL_LIMIT = 3;
19
20export class ManagedTunnelLimitPersistenceError extends Schema.TaggedError<ManagedTunnelLimitPersistenceError>()(
21 "ManagedTunnelLimitPersistenceError",
22 {
23 operation: Schema.Literals(["load-limit", "count-tunnels"]),
24 userId: Schema.String,
25 cause: Schema.Defect(),
26 },
27) {
28 override get message(): string {
29 return `Managed tunnel limit '${this.operation}' failed for user '${this.userId}'`;
30 }
31}
32
33export class ManagedTunnelLimitExceeded extends Schema.TaggedError<ManagedTunnelLimitExceeded>()(
34 "ManagedTunnelLimitExceeded",
35 {
36 userId: Schema.String,
37 environmentId: Schema.String,
38 maxTunnels: Schema.Number,
39 activeTunnels: Schema.Number,
40 },
41) {
42 override get message(): string {
43 return `Managed tunnel limit reached for user '${this.userId}': ${this.activeTunnels} of ${this.maxTunnels} tunnels in use`;
44 }
45}
46
47export class ManagedTunnelLimits extends Context.Service<
48 ManagedTunnelLimits,
49 {
50 readonly ensureCapacity: (input: {
51 readonly userId: string;
52 readonly environmentId: string;
53 }) => Effect.Effect<void, ManagedTunnelLimitExceeded | ManagedTunnelLimitPersistenceError>;
54 }
55>()("t3code-relay/environments/ManagedTunnelLimits") {}
56
57export const make = Effect.gen(function* () {
58 const db = yield* RelayDb.RelayDb;
59
60 return ManagedTunnelLimits.of({
61 ensureCapacity: Effect.fn("relay.managed_tunnel_limits.ensure_capacity")(function* (input) {
62 const overrides = yield* db
63 .select({ maxTunnels: relayManagedTunnelLimits.maxTunnels })
64 .from(relayManagedTunnelLimits)
65 .where(eq(relayManagedTunnelLimits.userId, input.userId))
66 .limit(1)
67 .pipe(
68 Effect.mapError(
69 (cause) =>
70 new ManagedTunnelLimitPersistenceError({
71 operation: "load-limit",
72 userId: input.userId,
73 cause,
74 }),
75 ),
76 );
77 const maxTunnels = overrides[0]?.maxTunnels ?? DEFAULT_MANAGED_TUNNEL_LIMIT;
78
79 // Allocations already held for this environment are excluded so that
80 // re-linking an environment the user already has a tunnel for stays
81 // idempotent even when the account is at its limit.
82 const counted = yield* db
83 .select({ activeTunnels: count() })
84 .from(relayManagedEndpointAllocations)
85 .innerJoin(
86 relayEnvironmentLinks,
87 and(
88 eq(relayEnvironmentLinks.userId, relayManagedEndpointAllocations.userId),
89 eq(relayEnvironmentLinks.environmentId, relayManagedEndpointAllocations.environmentId),
90 ),
91 )
92 .where(
93 and(
94 eq(relayManagedEndpointAllocations.userId, input.userId),
95 ne(relayManagedEndpointAllocations.environmentId, input.environmentId),
96 isNull(relayEnvironmentLinks.revokedAt),
97 eq(relayEnvironmentLinks.managedTunnelsEnabled, true),
98 ),
99 )
100 .pipe(
101 Effect.mapError(
102 (cause) =>
103 new ManagedTunnelLimitPersistenceError({
104 operation: "count-tunnels",
105 userId: input.userId,
106 cause,
107 }),
108 ),
109 );
110 const activeTunnels = counted[0]?.activeTunnels ?? 0;
111
112 if (activeTunnels >= maxTunnels) {
113 return yield* new ManagedTunnelLimitExceeded({
114 userId: input.userId,
115 environmentId: input.environmentId,
116 maxTunnels,
117 activeTunnels,
118 });
119 }
120 }),
121 });
122});
123
124export const layer = Layer.effect(ManagedTunnelLimits, make);