feat(server): support selfhost licenses (#8947)

packages/backend/server/src/plugins/payment/service.ts

@@ -11,7 +11,9 @@ import {
   CustomerPortalCreateFailed,
   InternalServerError,
   InvalidCheckoutParameters,
+  InvalidLicenseSessionId,
   InvalidSubscriptionParameters,
+  LicenseRevealed,
   Mutex,
   OnEvent,
   SameSubscriptionRecurring,
@@ -38,6 +40,11 @@ import {
   WorkspaceSubscriptionIdentity,
   WorkspaceSubscriptionManager,
 } from './manager';
+import {
+  SelfhostTeamCheckoutArgs,
+  SelfhostTeamSubscriptionIdentity,
+  SelfhostTeamSubscriptionManager,
+} from './manager/selfhost';
 import { ScheduleManager } from './schedule';
 import {
   decodeLookupKey,
@@ -56,11 +63,13 @@ import {
 export const CheckoutExtraArgs = z.union([
   UserSubscriptionCheckoutArgs,
   WorkspaceSubscriptionCheckoutArgs,
+  SelfhostTeamCheckoutArgs,
 ]);
 
 export const SubscriptionIdentity = z.union([
   UserSubscriptionIdentity,
   WorkspaceSubscriptionIdentity,
+  SelfhostTeamSubscriptionIdentity,
 ]);
 
 export { CheckoutParams };
@@ -78,6 +87,7 @@ export class SubscriptionService implements OnApplicationBootstrap {
     private readonly models: Models,
     private readonly userManager: UserSubscriptionManager,
     private readonly workspaceManager: WorkspaceSubscriptionManager,
+    private readonly selfhostManager: SelfhostTeamSubscriptionManager,
     private readonly mutex: Mutex
   ) {}
 
@@ -92,6 +102,8 @@ export class SubscriptionService implements OnApplicationBootstrap {
       case SubscriptionPlan.Pro:
       case SubscriptionPlan.AI:
         return this.userManager;
+      case SubscriptionPlan.SelfHostedTeam:
+        return this.selfhostManager;
       default:
         throw new UnsupportedSubscriptionPlan({ plan });
     }
@@ -122,7 +134,7 @@ export class SubscriptionService implements OnApplicationBootstrap {
     if (
       this.config.deploy &&
       this.config.affine.canary &&
-      !this.feature.isStaff(args.user.email)
+      (!('user' in args) || !this.feature.isStaff(args.user.email))
     ) {
       throw new ActionForbidden();
     }
@@ -291,10 +303,133 @@ export class SubscriptionService implements OnApplicationBootstrap {
     return newSubscription;
   }
 
-  async createCustomerPortal(id: string) {
+  async updateSubscriptionQuantity(
+    identity: z.infer<typeof SubscriptionIdentity>,
+    count: number
+  ) {
+    this.assertSubscriptionIdentity(identity);
+
+    const subscription = await this.select(identity.plan).getSubscription(
+      identity
+    );
+
+    if (!subscription) {
+      throw new SubscriptionNotExists({ plan: identity.plan });
+    }
+
+    if (!subscription.stripeSubscriptionId) {
+      throw new CantUpdateOnetimePaymentSubscription();
+    }
+
+    const stripeSubscription = await this.stripe.subscriptions.retrieve(
+      subscription.stripeSubscriptionId
+    );
+
+    const lookupKey =
+      retriveLookupKeyFromStripeSubscription(stripeSubscription);
+
+    await this.stripe.subscriptions.update(stripeSubscription.id, {
+      items: [
+        {
+          id: stripeSubscription.items.data[0].id,
+          quantity: count,
+        },
+      ],
+      payment_behavior: 'pending_if_incomplete',
+      proration_behavior:
+        lookupKey?.recurring === SubscriptionRecurring.Yearly
+          ? 'always_invoice'
+          : 'none',
+    });
+
+    if (subscription.stripeScheduleId) {
+      const schedule = await this.scheduleManager.fromSchedule(
+        subscription.stripeScheduleId
+      );
+      await schedule.updateQuantity(count);
+    }
+  }
+
+  async generateLicenseKey(stripeCheckoutSessionId: string) {
+    if (!stripeCheckoutSessionId) {
+      throw new InvalidLicenseSessionId();
+    }
+
+    let session: Stripe.Checkout.Session;
+    try {
+      session = await this.stripe.checkout.sessions.retrieve(
+        stripeCheckoutSessionId
+      );
+    } catch {
+      throw new InvalidLicenseSessionId();
+    }
+
+    // session should be complete and have a subscription
+    if (session.status !== 'complete' || !session.subscription) {
+      throw new InvalidLicenseSessionId();
+    }
+
+    const subscription =
+      typeof session.subscription === 'string'
+        ? await this.stripe.subscriptions.retrieve(session.subscription)
+        : session.subscription;
+
+    const knownSubscription = await this.parseStripeSubscription(subscription);
+
+    // invalid subscription triple
+    if (
+      !knownSubscription ||
+      knownSubscription.lookupKey.plan !== SubscriptionPlan.SelfHostedTeam
+    ) {
+      throw new InvalidLicenseSessionId();
+    }
+
+    let subInDB = await this.db.subscription.findUnique({
+      where: {
+        stripeSubscriptionId: subscription.id,
+      },
+    });
+
+    // subscription not found in db
+    if (!subInDB) {
+      subInDB =
+        await this.selfhostManager.saveStripeSubscription(knownSubscription);
+    }
+
+    const license = await this.db.license.findUnique({
+      where: {
+        key: subInDB.targetId,
+      },
+    });
+
+    // subscription and license are created in a transaction
+    // there is no way a sub exist but the license is not created
+    if (!license) {
+      throw new Error(
+        'unaccessible path. if you see this error, there must be a bug in the codebase.'
+      );
+    }
+
+    if (!license.revealedAt) {
+      await this.db.license.update({
+        where: {
+          key: license.key,
+        },
+        data: {
+          revealedAt: new Date(),
+        },
+      });
+
+      return license.key;
+    }
+
+    throw new LicenseRevealed();
+  }
+
+  async createCustomerPortal(userId: string) {
     const user = await this.db.userStripeCustomer.findUnique({
       where: {
-        userId: id,
+        userId: userId,
       },
     });
 
@@ -416,15 +551,18 @@ export class SubscriptionService implements OnApplicationBootstrap {
 
   private async retrieveUserFromCustomer(
     customer: string | Stripe.Customer | Stripe.DeletedCustomer
-  ) {
+  ): Promise<{ id?: string; email: string } | null> {
     const userStripeCustomer = await this.db.userStripeCustomer.findUnique({
       where: {
         stripeCustomerId: typeof customer === 'string' ? customer : customer.id,
       },
+      select: {
+        user: true,
+      },
     });
 
     if (userStripeCustomer) {
-      return userStripeCustomer.userId;
+      return userStripeCustomer.user;
     }
 
     if (typeof customer === 'string') {
@@ -438,17 +576,13 @@ export class SubscriptionService implements OnApplicationBootstrap {
     const user = await this.models.user.getPublicUserByEmail(customer.email);
 
     if (!user) {
-      return null;
+      return {
+        id: undefined,
+        email: customer.email,
+      };
     }
 
-    await this.db.userStripeCustomer.create({
-      data: {
-        userId: user.id,
-        stripeCustomerId: customer.id,
-      },
-    });
-
-    return user.id;
+    return user;
   }
 
   private async listStripePrices(): Promise<KnownStripePrice[]> {
@@ -489,14 +623,9 @@ export class SubscriptionService implements OnApplicationBootstrap {
       invoice.customer_email
     );
 
-    // TODO(@forehalo): the email may actually not appear to be AFFiNE user
-    // There is coming feature that allow anonymous user with only email provided to buy selfhost licenses
-    if (!user) {
-      return null;
-    }
-
     return {
-      userId: user.id,
+      userId: user?.id,
+      userEmail: invoice.customer_email,
       stripeInvoice: invoice,
       lookupKey,
       metadata: invoice.subscription_details?.metadata ?? {},
@@ -512,14 +641,18 @@ export class SubscriptionService implements OnApplicationBootstrap {
       return null;
     }
 
-    const userId = await this.retrieveUserFromCustomer(subscription.customer);
+    const user = await this.retrieveUserFromCustomer(subscription.customer);
 
-    if (!userId) {
+    // stripe customer got deleted or customer email is null
+    // it's an invalid status
+    // maybe we need to check stripe dashboard
+    if (!user) {
       return null;
     }
 
     return {
-      userId,
+      userId: user.id,
+      userEmail: user.email,
       lookupKey,
       stripeSubscription: subscription,
       quantity: subscription.items.data[0]?.quantity ?? 1,