pangolin_mirror/pangolin
1
0
Fork 0
mirror of https://github.com/fosrl/pangolin.git synced 2026-03-01 16:26:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Handle new usage tracking with multi org

Browse Source
This commit is contained in:
Owen
2026-02-17 17:09:48 -08:00
parent 79cf7c84dc
commit 4d6240c987
17 changed files with 432 additions and 584 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
server/lib/billing/features.ts
View File

@@ -4,6 +4,7 @@ export enum FeatureId {
EGRESS_DATA_MB = "egressDataMb",
DOMAINS = "domains",
REMOTE_EXIT_NODES = "remoteExitNodes",
ORGINIZATIONS = "organizations",
TIER1 = "tier1"
}
@@ -19,6 +20,8 @@ export async function getFeatureDisplayName(featureId: FeatureId): Promise<strin
return "Domains";
case FeatureId.REMOTE_EXIT_NODES:
return "Remote Exit Nodes";
case FeatureId.ORGINIZATIONS:
return "Organizations";
case FeatureId.TIER1:
return "Home Lab";
default:

17
server/lib/billing/limitSet.ts
View File

@@ -7,18 +7,12 @@ export type LimitSet = Partial<{
};
}>;
export const sandboxLimitSet: LimitSet = {
[FeatureId.USERS]: { value: 1, description: "Sandbox limit" },
[FeatureId.SITES]: { value: 1, description: "Sandbox limit" },
[FeatureId.DOMAINS]: { value: 0, description: "Sandbox limit" },
[FeatureId.REMOTE_EXIT_NODES]: { value: 0, description: "Sandbox limit" },
};
export const freeLimitSet: LimitSet = {
[FeatureId.SITES]: { value: 5, description: "Basic limit" },
[FeatureId.USERS]: { value: 5, description: "Basic limit" },
[FeatureId.DOMAINS]: { value: 5, description: "Basic limit" },
[FeatureId.REMOTE_EXIT_NODES]: { value: 1, description: "Basic limit" },
[FeatureId.ORGINIZATIONS]: { value: 1, description: "Basic limit" },
};
export const tier1LimitSet: LimitSet = {
@@ -26,6 +20,7 @@ export const tier1LimitSet: LimitSet = {
[FeatureId.SITES]: { value: 10, description: "Home limit" },
[FeatureId.DOMAINS]: { value: 10, description: "Home limit" },
[FeatureId.REMOTE_EXIT_NODES]: { value: 1, description: "Home limit" },
[FeatureId.ORGINIZATIONS]: { value: 1, description: "Home limit" },
};
export const tier2LimitSet: LimitSet = {
@@ -45,6 +40,10 @@ export const tier2LimitSet: LimitSet = {
value: 3,
description: "Team limit"
},
[FeatureId.ORGINIZATIONS]: {
value: 1,
description: "Team limit"
}
};
export const tier3LimitSet: LimitSet = {
@@ -64,4 +63,8 @@ export const tier3LimitSet: LimitSet = {
value: 20,
description: "Business limit"
},
[FeatureId.ORGINIZATIONS]: {
value: 20,
description: "Business limit"
},
};

183
server/lib/billing/usageService.ts
View File

@@ -1,12 +1,8 @@
import { eq, sql, and } from "drizzle-orm";
import { v4 as uuidv4 } from "uuid";
import { PutObjectCommand } from "@aws-sdk/client-s3";
import {
db,
usage,
customers,
sites,
newts,
limits,
Usage,
Limit,
@@ -15,21 +11,9 @@ import {
} from "@server/db";
import { FeatureId, getFeatureMeterId } from "./features";
import logger from "@server/logger";
import { sendToClient } from "#dynamic/routers/ws";
import { build } from "@server/build";
import { s3Client } from "@server/lib/s3";
import cache from "@server/lib/cache";
interface StripeEvent {
identifier?: string;
timestamp: number;
event_name: string;
payload: {
value: number;
stripe_customer_id: string;
};
}
export function noop() {
if (build !== "saas") {
return true;
@@ -38,41 +22,11 @@ export function noop() {
}
export class UsageService {
// private bucketName: string | undefined;
// private events: StripeEvent[] = [];
// private lastUploadTime: number = Date.now();
// private isUploading: boolean = false;
constructor() {
if (noop()) {
return;
}
// this.bucketName = process.env.S3_BUCKET || undefined;
// // Periodically check and upload events
// setInterval(() => {
// this.checkAndUploadEvents().catch((err) => {
// logger.error("Error in periodic event upload:", err);
// });
// }, 30000); // every 30 seconds
// // Handle graceful shutdown on SIGTERM
// process.on("SIGTERM", async () => {
// logger.info(
// "SIGTERM received, uploading events before shutdown..."
// );
// await this.forceUpload();
// logger.info("Events uploaded, proceeding with shutdown");
// });
// // Handle SIGINT as well (Ctrl+C)
// process.on("SIGINT", async () => {
// logger.info("SIGINT received, uploading events before shutdown...");
// await this.forceUpload();
// logger.info("Events uploaded, proceeding with shutdown");
// process.exit(0);
// });
}
/**
@@ -103,16 +57,6 @@ export class UsageService {
while (attempt <= maxRetries) {
try {
// // Get subscription data for this org (with caching)
// const customerId = await this.getCustomerId(orgIdToUse, featureId);
// if (!customerId) {
// logger.warn(
// `No subscription data found for org ${orgIdToUse} and feature ${featureId}`
// );
// return null;
// }
let usage;
if (transaction) {
usage = await this.internalAddUsage(
@@ -132,11 +76,6 @@ export class UsageService {
});
}
// Log event for Stripe
// if (privateConfig.getRawPrivateConfig().flags.usage_reporting) {
// await this.logStripeEvent(featureId, value, customerId);
// }
return usage || null;
} catch (error: any) {
// Check if this is a deadlock error
@@ -191,13 +130,14 @@ export class UsageService {
featureId,
orgId,
meterId,
instantaneousValue: value,
latestValue: value,
updatedAt: Math.floor(Date.now() / 1000)
})
.onConflictDoUpdate({
target: usage.usageId,
set: {
latestValue: sql`${usage.latestValue} + ${value}`
instantaneousValue: sql`${usage.instantaneousValue} + ${value}`
}
})
.returning();
@@ -228,17 +168,6 @@ export class UsageService {
let orgIdToUse = await this.getBillingOrg(orgId);
try {
// if (!customerId) {
// customerId =
// (await this.getCustomerId(orgIdToUse, featureId)) || undefined;
// if (!customerId) {
// logger.warn(
// `No subscription data found for org ${orgIdToUse} and feature ${featureId}`
// );
// return;
// }
// }
// Truncate value to 11 decimal places if provided
if (value !== undefined && value !== null) {
value = this.truncateValue(value);
@@ -523,114 +452,6 @@ export class UsageService {
return hasExceededLimits;
}
// private async logStripeEvent(
// featureId: FeatureId,
// value: number,
// customerId: string
// ): Promise<void> {
// // Truncate value to 11 decimal places before sending to Stripe
// const truncatedValue = this.truncateValue(value);
// const event: StripeEvent = {
// identifier: uuidv4(),
// timestamp: Math.floor(new Date().getTime() / 1000),
// event_name: featureId,
// payload: {
// value: truncatedValue,
// stripe_customer_id: customerId
// }
// };
// this.addEventToMemory(event);
// await this.checkAndUploadEvents();
// }
// private addEventToMemory(event: StripeEvent): void {
// if (!this.bucketName) {
// logger.warn(
// "S3 bucket name is not configured, skipping event storage."
// );
// return;
// }
// this.events.push(event);
// }
// private async checkAndUploadEvents(): Promise<void> {
// const now = Date.now();
// const timeSinceLastUpload = now - this.lastUploadTime;
// // Check if at least 1 minute has passed since last upload
// if (timeSinceLastUpload >= 60000 && this.events.length > 0) {
// await this.uploadEventsToS3();
// }
// }
// private async uploadEventsToS3(): Promise<void> {
// if (!this.bucketName) {
// logger.warn(
// "S3 bucket name is not configured, skipping S3 upload."
// );
// return;
// }
// if (this.events.length === 0) {
// return;
// }
// // Check if already uploading
// if (this.isUploading) {
// logger.debug("Already uploading events, skipping");
// return;
// }
// this.isUploading = true;
// try {
// // Take a snapshot of current events and clear the array
// const eventsToUpload = [...this.events];
// this.events = [];
// this.lastUploadTime = Date.now();
// const fileName = this.generateEventFileName();
// const fileContent = JSON.stringify(eventsToUpload, null, 2);
// // Upload to S3
// const uploadCommand = new PutObjectCommand({
// Bucket: this.bucketName,
// Key: fileName,
// Body: fileContent,
// ContentType: "application/json"
// });
// await s3Client.send(uploadCommand);
// logger.info(
// `Uploaded ${fileName} to S3 with ${eventsToUpload.length} events`
// );
// } catch (error) {
// logger.error("Failed to upload events to S3:", error);
// // Note: Events are lost if upload fails. In a production system,
// // you might want to add the events back to the array or implement retry logic
// } finally {
// this.isUploading = false;
// }
// }
// private generateEventFileName(): string {
// const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
// const uuid = uuidv4().substring(0, 8);
// return `events-${timestamp}-${uuid}.json`;
// }
// public async forceUpload(): Promise<void> {
// if (this.events.length > 0) {
// // Force upload regardless of time
// this.lastUploadTime = 0; // Reset to force upload
// await this.uploadEventsToS3();
// }
// }
}
export const usageService = new UsageService();

206
server/lib/createUserAccountOrg.ts
View File

@@ -1,206 +0,0 @@
import { isValidCIDR } from "@server/lib/validators";
import { getNextAvailableOrgSubnet } from "@server/lib/ip";
import {
actions,
apiKeyOrg,
apiKeys,
db,
domains,
Org,
orgDomains,
orgs,
roleActions,
roles,
userOrgs
} from "@server/db";
import { eq } from "drizzle-orm";
import { defaultRoleAllowedActions } from "@server/routers/role";
import { FeatureId, limitsService, sandboxLimitSet } from "@server/lib/billing";
import { createCustomer } from "#dynamic/lib/billing";
import { usageService } from "@server/lib/billing/usageService";
import config from "@server/lib/config";
import { generateCA } from "@server/private/lib/sshCA";
import { encrypt } from "@server/lib/crypto";
export async function createUserAccountOrg(
userId: string,
userEmail: string
): Promise<{
success: boolean;
org?: {
orgId: string;
name: string;
subnet: string;
};
error?: string;
}> {
// const subnet = await getNextAvailableOrgSubnet();
const orgId = "org_" + userId;
const name = `${userEmail}'s Organization`;
// if (!isValidCIDR(subnet)) {
// return {
// success: false,
// error: "Invalid subnet format. Please provide a valid CIDR notation."
// };
// }
// // make sure the subnet is unique
// const subnetExists = await db
// .select()
// .from(orgs)
// .where(eq(orgs.subnet, subnet))
// .limit(1);
// if (subnetExists.length > 0) {
// return { success: false, error: `Subnet ${subnet} already exists` };
// }
// make sure the orgId is unique
const orgExists = await db
.select()
.from(orgs)
.where(eq(orgs.orgId, orgId))
.limit(1);
if (orgExists.length > 0) {
return {
success: false,
error: `Organization with ID ${orgId} already exists`
};
}
let error = "";
let org: Org | null = null;
await db.transaction(async (trx) => {
const allDomains = await trx
.select()
.from(domains)
.where(eq(domains.configManaged, true));
const utilitySubnet = config.getRawConfig().orgs.utility_subnet_group;
// Generate SSH CA keys for the org
const ca = generateCA(`${orgId}-ca`);
const encryptionKey = config.getRawConfig().server.secret!;
const encryptedCaPrivateKey = encrypt(ca.privateKeyPem, encryptionKey);
const newOrg = await trx
.insert(orgs)
.values({
orgId,
name,
// subnet
subnet: "100.90.128.0/24", // TODO: this should not be hardcoded - or can it be the same in all orgs?
utilitySubnet: utilitySubnet,
createdAt: new Date().toISOString(),
sshCaPrivateKey: encryptedCaPrivateKey,
sshCaPublicKey: ca.publicKeyOpenSSH
})
.returning();
if (newOrg.length === 0) {
error = "Failed to create organization";
trx.rollback();
return;
}
org = newOrg[0];
// Create admin role within the same transaction
const [insertedRole] = await trx
.insert(roles)
.values({
orgId: newOrg[0].orgId,
isAdmin: true,
name: "Admin",
description: "Admin role with the most permissions"
})
.returning({ roleId: roles.roleId });
if (!insertedRole || !insertedRole.roleId) {
error = "Failed to create Admin role";
trx.rollback();
return;
}
const roleId = insertedRole.roleId;
// Get all actions and create role actions
const actionIds = await trx.select().from(actions).execute();
if (actionIds.length > 0) {
await trx.insert(roleActions).values(
actionIds.map((action) => ({
roleId,
actionId: action.actionId,
orgId: newOrg[0].orgId
}))
);
}
if (allDomains.length) {
await trx.insert(orgDomains).values(
allDomains.map((domain) => ({
orgId: newOrg[0].orgId,
domainId: domain.domainId
}))
);
}
await trx.insert(userOrgs).values({
userId,
orgId: newOrg[0].orgId,
roleId: roleId,
isOwner: true
});
const memberRole = await trx
.insert(roles)
.values({
name: "Member",
description: "Members can only view resources",
orgId
})
.returning();
await trx.insert(roleActions).values(
defaultRoleAllowedActions.map((action) => ({
roleId: memberRole[0].roleId,
actionId: action,
orgId
}))
);
});
await limitsService.applyLimitSetToOrg(orgId, sandboxLimitSet);
if (!org) {
return { success: false, error: "Failed to create org" };
}
if (error) {
return {
success: false,
error: `Failed to create org: ${error}`
};
}
// make sure we have the stripe customer
const customerId = await createCustomer(orgId, userEmail);
if (customerId) {
await usageService.updateCount(orgId, FeatureId.USERS, 1, customerId); // Only 1 because we are crating the org
}
return {
org: {
orgId,
name,
// subnet
subnet: "100.90.128.0/24"
},
success: true
};
}

27
server/lib/deleteOrg.ts
View File

@@ -19,6 +19,8 @@ import { sendToClient } from "#dynamic/routers/ws";
import { deletePeer } from "@server/routers/gerbil/peers";
import { OlmErrorCodes } from "@server/routers/olm/error";
import { sendTerminateClient } from "@server/routers/client/terminate";
import { usageService } from "./billing/usageService";
import { FeatureId } from "./billing";
export type DeleteOrgByIdResult = {
deletedNewtIds: string[];
@@ -74,9 +76,7 @@ export async function deleteOrgById(
deletedNewtIds.push(deletedNewt.newtId);
await trx
.delete(newtSessions)
.where(
eq(newtSessions.newtId, deletedNewt.newtId)
);
.where(eq(newtSessions.newtId, deletedNewt.newtId));
}
}
}
@@ -137,6 +137,9 @@ export async function deleteOrgById(
.where(inArray(domains.domainId, domainIdsToDelete));
}
await trx.delete(resources).where(eq(resources.orgId, orgId));
await usageService.add(orgId, FeatureId.ORGINIZATIONS, -1, trx); // here we are decreasing the org count BEFORE deleting the org because we need to still be able to get the org to get the billing org inside of here
await trx.delete(orgs).where(eq(orgs.orgId, orgId));
});
@@ -155,15 +158,13 @@ export function sendTerminationMessages(result: DeleteOrgByIdResult): void {
);
}
for (const olmId of result.olmsToTerminate) {
sendTerminateClient(
0,
OlmErrorCodes.TERMINATED_REKEYED,
olmId
).catch((error) => {
logger.error(
"Failed to send termination message to olm:",
error
);
});
sendTerminateClient(0, OlmErrorCodes.TERMINATED_REKEYED, olmId).catch(
(error) => {
logger.error(
"Failed to send termination message to olm:",
error
);
}
);
}
}