mirror of
https://github.com/fosrl/pangolin.git
synced 2026-03-01 16:26:39 +00:00
Handle new usage tracking with multi org
This commit is contained in:
@@ -85,10 +85,14 @@ export async function getOrgUsage(
|
||||
orgId,
|
||||
FeatureId.REMOTE_EXIT_NODES
|
||||
);
|
||||
const egressData = await usageService.getUsage(
|
||||
const organizations = await usageService.getUsage(
|
||||
orgId,
|
||||
FeatureId.EGRESS_DATA_MB
|
||||
FeatureId.ORGINIZATIONS
|
||||
);
|
||||
// const egressData = await usageService.getUsage(
|
||||
// orgId,
|
||||
// FeatureId.EGRESS_DATA_MB
|
||||
// );
|
||||
|
||||
if (sites) {
|
||||
usageData.push(sites);
|
||||
@@ -96,15 +100,18 @@ export async function getOrgUsage(
|
||||
if (users) {
|
||||
usageData.push(users);
|
||||
}
|
||||
if (egressData) {
|
||||
usageData.push(egressData);
|
||||
}
|
||||
// if (egressData) {
|
||||
// usageData.push(egressData);
|
||||
// }
|
||||
if (domains) {
|
||||
usageData.push(domains);
|
||||
}
|
||||
if (remoteExitNodes) {
|
||||
usageData.push(remoteExitNodes);
|
||||
}
|
||||
if (organizations) {
|
||||
usageData.push(organizations);
|
||||
}
|
||||
|
||||
const orgLimits = await db
|
||||
.select()
|
||||
|
||||
@@ -12,7 +12,14 @@
|
||||
*/
|
||||
|
||||
import { NextFunction, Request, Response } from "express";
|
||||
import { db, exitNodes, exitNodeOrgs, ExitNode, ExitNodeOrg } from "@server/db";
|
||||
import {
|
||||
db,
|
||||
exitNodes,
|
||||
exitNodeOrgs,
|
||||
ExitNode,
|
||||
ExitNodeOrg,
|
||||
orgs
|
||||
} from "@server/db";
|
||||
import HttpCode from "@server/types/HttpCode";
|
||||
import { z } from "zod";
|
||||
import { remoteExitNodes } from "@server/db";
|
||||
@@ -25,7 +32,7 @@ import { createRemoteExitNodeSession } from "#private/auth/sessions/remoteExitNo
|
||||
import { fromError } from "zod-validation-error";
|
||||
import { hashPassword, verifyPassword } from "@server/auth/password";
|
||||
import logger from "@server/logger";
|
||||
import { and, eq } from "drizzle-orm";
|
||||
import { and, eq, inArray, ne } from "drizzle-orm";
|
||||
import { getNextAvailableSubnet } from "@server/lib/exitNodes";
|
||||
import { usageService } from "@server/lib/billing/usageService";
|
||||
import { FeatureId } from "@server/lib/billing";
|
||||
@@ -169,7 +176,17 @@ export async function createRemoteExitNode(
|
||||
);
|
||||
}
|
||||
|
||||
let numExitNodeOrgs: ExitNodeOrg[] | undefined;
|
||||
const [org] = await db
|
||||
.select()
|
||||
.from(orgs)
|
||||
.where(eq(orgs.orgId, orgId))
|
||||
.limit(1);
|
||||
|
||||
if (!org) {
|
||||
return next(
|
||||
createHttpError(HttpCode.NOT_FOUND, "Organization not found")
|
||||
);
|
||||
}
|
||||
|
||||
await db.transaction(async (trx) => {
|
||||
if (!existingExitNode) {
|
||||
@@ -217,19 +234,43 @@ export async function createRemoteExitNode(
|
||||
});
|
||||
}
|
||||
|
||||
numExitNodeOrgs = await trx
|
||||
.select()
|
||||
.from(exitNodeOrgs)
|
||||
.where(eq(exitNodeOrgs.orgId, orgId));
|
||||
});
|
||||
// calculate if the node is in any other of the orgs before we count it as an add to the billing org
|
||||
if (org.billingOrgId) {
|
||||
const otherBillingOrgs = await trx
|
||||
.select()
|
||||
.from(orgs)
|
||||
.where(
|
||||
and(
|
||||
eq(orgs.billingOrgId, org.billingOrgId),
|
||||
ne(orgs.orgId, orgId)
|
||||
)
|
||||
);
|
||||
|
||||
if (numExitNodeOrgs) {
|
||||
await usageService.updateCount(
|
||||
orgId,
|
||||
FeatureId.REMOTE_EXIT_NODES,
|
||||
numExitNodeOrgs.length
|
||||
);
|
||||
}
|
||||
const billingOrgIds = otherBillingOrgs.map((o) => o.orgId);
|
||||
|
||||
const orgsInBillingDomainThatTheNodeIsStillIn = await trx
|
||||
.select()
|
||||
.from(exitNodeOrgs)
|
||||
.where(
|
||||
and(
|
||||
eq(
|
||||
exitNodeOrgs.exitNodeId,
|
||||
existingExitNode.exitNodeId
|
||||
),
|
||||
inArray(exitNodeOrgs.orgId, billingOrgIds)
|
||||
)
|
||||
);
|
||||
|
||||
if (orgsInBillingDomainThatTheNodeIsStillIn.length === 0) {
|
||||
await usageService.add(
|
||||
orgId,
|
||||
FeatureId.REMOTE_EXIT_NODES,
|
||||
1,
|
||||
trx
|
||||
);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
const token = generateSessionToken();
|
||||
await createRemoteExitNodeSession(token, remoteExitNodeId);
|
||||
|
||||
@@ -13,9 +13,9 @@
|
||||
|
||||
import { NextFunction, Request, Response } from "express";
|
||||
import { z } from "zod";
|
||||
import { db, ExitNodeOrg, exitNodeOrgs, exitNodes } from "@server/db";
|
||||
import { db, ExitNodeOrg, exitNodeOrgs, exitNodes, orgs } from "@server/db";
|
||||
import { remoteExitNodes } from "@server/db";
|
||||
import { and, count, eq } from "drizzle-orm";
|
||||
import { and, count, eq, inArray } from "drizzle-orm";
|
||||
import response from "@server/lib/response";
|
||||
import HttpCode from "@server/types/HttpCode";
|
||||
import createHttpError from "http-errors";
|
||||
@@ -50,7 +50,8 @@ export async function deleteRemoteExitNode(
|
||||
const [remoteExitNode] = await db
|
||||
.select()
|
||||
.from(remoteExitNodes)
|
||||
.where(eq(remoteExitNodes.remoteExitNodeId, remoteExitNodeId));
|
||||
.where(eq(remoteExitNodes.remoteExitNodeId, remoteExitNodeId))
|
||||
.limit(1);
|
||||
|
||||
if (!remoteExitNode) {
|
||||
return next(
|
||||
@@ -70,7 +71,17 @@ export async function deleteRemoteExitNode(
|
||||
);
|
||||
}
|
||||
|
||||
let numExitNodeOrgs: ExitNodeOrg[] | undefined;
|
||||
const [org] = await db.select().from(orgs).where(eq(orgs.orgId, orgId));
|
||||
|
||||
if (!org) {
|
||||
return next(
|
||||
createHttpError(
|
||||
HttpCode.NOT_FOUND,
|
||||
`Org with ID ${orgId} not found`
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
await db.transaction(async (trx) => {
|
||||
await trx
|
||||
.delete(exitNodeOrgs)
|
||||
@@ -81,38 +92,39 @@ export async function deleteRemoteExitNode(
|
||||
)
|
||||
);
|
||||
|
||||
const [remainingExitNodeOrgs] = await trx
|
||||
.select({ count: count() })
|
||||
.from(exitNodeOrgs)
|
||||
.where(eq(exitNodeOrgs.exitNodeId, remoteExitNode.exitNodeId!));
|
||||
// calculate if the user is in any other of the orgs before we count it as an remove to the billing org
|
||||
if (org.billingOrgId) {
|
||||
const otherBillingOrgs = await trx
|
||||
.select()
|
||||
.from(orgs)
|
||||
.where(eq(orgs.billingOrgId, org.billingOrgId));
|
||||
|
||||
if (remainingExitNodeOrgs.count === 0) {
|
||||
await trx
|
||||
.delete(remoteExitNodes)
|
||||
const billingOrgIds = otherBillingOrgs.map((o) => o.orgId);
|
||||
|
||||
const orgsInBillingDomainThatTheNodeIsStillIn = await trx
|
||||
.select()
|
||||
.from(exitNodeOrgs)
|
||||
.where(
|
||||
eq(remoteExitNodes.remoteExitNodeId, remoteExitNodeId)
|
||||
and(
|
||||
eq(
|
||||
exitNodeOrgs.exitNodeId,
|
||||
remoteExitNode.exitNodeId!
|
||||
),
|
||||
inArray(exitNodeOrgs.orgId, billingOrgIds)
|
||||
)
|
||||
);
|
||||
await trx
|
||||
.delete(exitNodes)
|
||||
.where(
|
||||
eq(exitNodes.exitNodeId, remoteExitNode.exitNodeId!)
|
||||
|
||||
if (orgsInBillingDomainThatTheNodeIsStillIn.length === 0) {
|
||||
await usageService.add(
|
||||
orgId,
|
||||
FeatureId.REMOTE_EXIT_NODES,
|
||||
-1,
|
||||
trx
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
numExitNodeOrgs = await trx
|
||||
.select()
|
||||
.from(exitNodeOrgs)
|
||||
.where(eq(exitNodeOrgs.orgId, orgId));
|
||||
});
|
||||
|
||||
if (numExitNodeOrgs) {
|
||||
await usageService.updateCount(
|
||||
orgId,
|
||||
FeatureId.REMOTE_EXIT_NODES,
|
||||
numExitNodeOrgs.length
|
||||
);
|
||||
}
|
||||
|
||||
return response(res, {
|
||||
data: null,
|
||||
success: true,
|
||||
|
||||
Reference in New Issue
Block a user