Merge pull request #189 from waifuvault/premium-buckets

add premium buckets
waifuvault · Sep 18, 2024 · d32fd13 · d32fd13
2 parents 3464a84 + c1687e0
commit d32fd13
Show file tree

Hide file tree

Showing 5 changed files with 94 additions and 12 deletions.
diff --git a/src/migrations/1726652384908-add_bucket_type.ts b/src/migrations/1726652384908-add_bucket_type.ts
@@ -0,0 +1,24 @@
+import { MigrationInterface, QueryRunner } from "typeorm";
+
+export class AddBucketType1726652384908 implements MigrationInterface {
+    name = 'AddBucketType1726652384908'
+
+    public async up(queryRunner: QueryRunner): Promise<void> {
+        await queryRunner.query(`DROP INDEX "IDX_74d7bc76abfc080b96e42ddac7"`);
+        await queryRunner.query(`CREATE TABLE "temporary_bucket_model" ("id" integer PRIMARY KEY AUTOINCREMENT NOT NULL, "createdAt" datetime NOT NULL DEFAULT (datetime('now')), "updatedAt" datetime NOT NULL DEFAULT (datetime('now')), "bucketToken" text NOT NULL, "ip" text NOT NULL, "type" text NOT NULL DEFAULT ('NORMAL'), CONSTRAINT "UQ_092bd6d7824bd16ceee201ac5c6" UNIQUE ("ip"))`);
+        await queryRunner.query(`INSERT INTO "temporary_bucket_model"("id", "createdAt", "updatedAt", "bucketToken", "ip") SELECT "id", "createdAt", "updatedAt", "bucketToken", "ip" FROM "bucket_model"`);
+        await queryRunner.query(`DROP TABLE "bucket_model"`);
+        await queryRunner.query(`ALTER TABLE "temporary_bucket_model" RENAME TO "bucket_model"`);
+        await queryRunner.query(`CREATE UNIQUE INDEX "IDX_74d7bc76abfc080b96e42ddac7" ON "bucket_model" ("bucketToken") `);
+    }
+
+    public async down(queryRunner: QueryRunner): Promise<void> {
+        await queryRunner.query(`DROP INDEX "IDX_74d7bc76abfc080b96e42ddac7"`);
+        await queryRunner.query(`ALTER TABLE "bucket_model" RENAME TO "temporary_bucket_model"`);
+        await queryRunner.query(`CREATE TABLE "bucket_model" ("id" integer PRIMARY KEY AUTOINCREMENT NOT NULL, "createdAt" datetime NOT NULL DEFAULT (datetime('now')), "updatedAt" datetime NOT NULL DEFAULT (datetime('now')), "bucketToken" text NOT NULL, "ip" text NOT NULL, CONSTRAINT "UQ_092bd6d7824bd16ceee201ac5c6" UNIQUE ("ip"))`);
+        await queryRunner.query(`INSERT INTO "bucket_model"("id", "createdAt", "updatedAt", "bucketToken", "ip") SELECT "id", "createdAt", "updatedAt", "bucketToken", "ip" FROM "temporary_bucket_model"`);
+        await queryRunner.query(`DROP TABLE "temporary_bucket_model"`);
+        await queryRunner.query(`CREATE UNIQUE INDEX "IDX_74d7bc76abfc080b96e42ddac7" ON "bucket_model" ("bucketToken") `);
+    }
+
+}
diff --git a/src/model/constants/BucketType.ts b/src/model/constants/BucketType.ts
@@ -0,0 +1,6 @@
+enum BucketType {
+    NORMAL = "NORMAL",
+    PREMIUM = "PREMIUM",
+}
+
+export default BucketType;
diff --git a/src/model/db/Bucket.model.ts b/src/model/db/Bucket.model.ts
@@ -1,6 +1,7 @@
 import { AbstractModel } from "./AbstractModel.js";
 import { Column, Entity, Index, OneToMany } from "typeorm";
 import type { FileUploadModel } from "./FileUpload.model.js";
+import BucketType from "../constants/BucketType.js";
 
 @Entity()
 @Index(["bucketToken"], {
@@ -20,6 +21,13 @@ export class BucketModel extends AbstractModel {
     })
     public ip: string;
 
+    @Column({
+        nullable: false,
+        type: "text",
+        default: BucketType.NORMAL,
+    })
+    public type: BucketType;
+
     @OneToMany("FileUploadModel", "bucket", {
         cascade: true,
         eager: true,

diff --git a/src/public/secure/files.ejs b/src/public/secure/files.ejs
@@ -37,9 +37,15 @@
     <div class="row">
         <div class="col-md-12">
             <div class="card border-primary mb-3">
-                <div class="card-header">Current Entries
-                    <% if(loginType === "bucket") { %>
-                        for bucket: <%- bucket.bucketToken; -%>
+                <div class="card-header">
+                    <div>
+                        Current Entries
+                        <% if(loginType === "bucket") { %>
+                            for bucket: <%- bucket.bucketToken; -%>
+                        <% } %>
+                    </div>
+                    <% if(loginType === "bucket" && bucket.type === "PREMIUM") { %>
+                        <div class="alert alert-info mt-3"><i class="bi bi-question-circle"></i> This is a <%-bucket.type-%> bucket, this means that you get unlimited expiry.</div>
                     <% } %>
                 </div>
                 <div class="card-body">

diff --git a/src/services/FileUploadService.ts b/src/services/FileUploadService.ts
@@ -23,6 +23,8 @@ import { EntryModificationDto } from "../model/dto/EntryModificationDto.js";
 import { FileUploadQueryParameters } from "../model/rest/FileUploadQueryParameters.js";
 import { ProcessUploadException } from "../model/exceptions/ProcessUploadException.js";
 import { FileService } from "./FileService.js";
+import { BucketService } from "./BucketService.js";
+import BucketType from "../model/constants/BucketType.js";
 
 @Service()
 export class FileUploadService {
@@ -41,6 +43,7 @@ export class FileUploadService {
         @Inject() private encryptionService: EncryptionService,
         @Inject() private recordInfoSocket: RecordInfoSocket,
         @Inject() private fileService: FileService,
+        @Inject() private bucketService: BucketService,
     ) {}
 
     public async processUpload({
@@ -91,11 +94,8 @@ export class FileUploadService {
             }
             uploadEntry.originalFileName(originalFileName);
             uploadEntry.checksum(checksum);
-            if (expires) {
-                this.calculateCustomExpires(uploadEntry, expires, secretToken);
-            } else if (secretToken !== this.secret) {
-                uploadEntry.expires(FileUtils.getExpiresBySize(fileSize));
-            }
+
+            await this.setExpires(uploadEntry, fileSize, secretToken, bucketToken, expires);
 
             if (password) {
                 try {
@@ -120,6 +120,22 @@ export class FileUploadService {
         }
     }
 
+    private async setExpires(
+        entryBuilder: IBuilder<FileUploadModel>,
+        fileSize: number,
+        secretToken?: string,
+        bucketToken?: string,
+        customExpires?: string,
+    ): Promise<void> {
+        if (customExpires) {
+            await this.calculateCustomExpires(entryBuilder, customExpires, secretToken, bucketToken);
+        } else if (await this.hasUnlimitedExpire(secretToken, bucketToken)) {
+            entryBuilder.expires(null);
+        } else {
+            entryBuilder.expires(FileUtils.getExpiresBySize(fileSize));
+        }
+    }
+
     private hashPassword(password: string): Promise<string> {
         return argon2.hash(password);
     }
@@ -220,7 +236,7 @@ export class FileUploadService {
             builder.settings(newSettings);
         }
         if (dto.customExpiry) {
-            this.calculateCustomExpires(builder, dto.customExpiry);
+            await this.calculateCustomExpires(builder, dto.customExpiry);
         } else if (dto.customExpiry === "") {
             const fileSize = await FileUtils.getFileSize(entryToModify);
             builder.expires(FileUtils.getExpiresBySize(fileSize, entryToModify.createdAt.getTime()));
@@ -229,7 +245,12 @@ export class FileUploadService {
         return FileUploadResponseDto.fromModel(updatedEntry, this.baseUrl);
     }
 
-    private calculateCustomExpires(entry: IBuilder<FileUploadModel>, expires: string, secretToken?: string): void {
+    private async calculateCustomExpires(
+        entry: IBuilder<FileUploadModel>,
+        expires: string,
+        secretToken?: string,
+        bucketToken?: string,
+    ): Promise<void> {
         let value: number = ObjectUtils.getNumber(expires);
         let timeFactor: TimeUnit = TimeUnit.minutes;
 
@@ -242,15 +263,32 @@ export class FileUploadService {
             timeFactor = TimeUnit.hours;
         }
         value = ObjectUtils.convertToMilli(value, timeFactor);
-        const maxExp: number | null =
-            this.secret && this.secret === secretToken ? null : FileUtils.getTimeLeftBySize(entry.fileSize());
+
+        const unlimitedExpire = await this.hasUnlimitedExpire(secretToken, bucketToken);
+
+        const maxExp: number | null = unlimitedExpire ? null : FileUtils.getTimeLeftBySize(entry.fileSize());
 
         if (maxExp !== null && value > maxExp) {
             throw new BadRequest(`Cannot extend time remaining beyond ${ObjectUtils.timeToHuman(maxExp)}`);
         }
         entry.expires(Date.now() + value);
     }
 
+    private async hasUnlimitedExpire(secretToken?: string, bucketToken?: string): Promise<boolean> {
+        if (secretToken === this.secret) {
+            return true;
+        }
+
+        if (bucketToken) {
+            const bucket = await this.bucketService.getBucket(bucketToken);
+            if (bucket && bucket.type === BucketType.PREMIUM) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
     private async getFileHash(resourcePath: string): Promise<string> {
         const fileBuffer = await fs.readFile(resourcePath);
         const hashSum = crypto.createHash("md5");