feat: storage/file quota + urls quota

prisma/schema.prisma

@@ -30,16 +30,37 @@ model User {
   totpSecret String?
   passkeys   UserPasskey[]
 
+  quota UserQuota?
+
   files          File[]
   urls           Url[]
   folders        Folder[]
-  limits         UserLimit[]
   invites        Invite[]
   tags           Tag[]
   oauthProviders OAuthProvider[]
   IncompleteFile IncompleteFile[]
 }
 
+model UserQuota {
+  id        String   @id @default(cuid())
+  createdAt DateTime @default(now())
+  updatedAt DateTime @updatedAt
+
+  filesQuota UserFilesQuota
+  maxBytes   String?
+  maxFiles   Int?
+
+  maxUrls Int?
+
+  User   User?   @relation(fields: [userId], references: [id], onDelete: Cascade)
+  userId String? @unique
+}
+
+enum UserFilesQuota {
+  BY_BYTES
+  BY_FILES
+}
+
 model UserPasskey {
   id        String   @id @default(cuid())
   createdAt DateTime @default(now())
@@ -85,35 +106,6 @@ enum OAuthProviderType {
   AUTHENTIK
 }
 
-model UserLimit {
-  id        String   @id @default(cuid())
-  createdAt DateTime @default(now())
-  updatedAt DateTime @updatedAt
-
-  type      LimitType      @unique
-  value     Int
-  timeframe LimitTimeframe
-
-  User   User   @relation(fields: [userId], references: [id], onDelete: Cascade, onUpdate: Cascade)
-  userId String
-}
-
-enum LimitType {
-  UPLOAD_COUNT
-  UPLOAD_SIZE
-  SHORTEN_COUNT
-}
-
-enum LimitTimeframe {
-  SECONDLY
-  MINUTELY
-  HOURLY
-  DAILY
-  WEEKLY
-  MONTHLY
-  YEARLY
-}
-
 model File {
   id        String    @id @default(cuid())
   createdAt DateTime  @default(now())

src/components/pages/dashboard/index.tsx

@@ -20,6 +20,27 @@ export default function DashboardHome() {
       <Text size='sm' c='dimmed'>
         You have <b>{statsLoading ? '...' : stats?.filesUploaded}</b> files uploaded.
       </Text>
+      {user?.quota && (user.quota.maxBytes || user.quota.maxFiles) ? (
+        <Text size='sm' c='dimmed'>
+          {user.quota.filesQuota === 'BY_BYTES' ? (
+            <>
+              You have used <b>{statsLoading ? '...' : bytes(stats!.storageUsed)}</b> out of{' '}
+              <b>{user.quota.maxBytes}</b> of storage
+            </>
+          ) : (
+            <>
+              You have uploaded <b>{statsLoading ? '...' : stats?.filesUploaded}</b> files out of{' '}
+              <b>{user.quota.maxFiles}</b> files allowed.
+            </>
+          )}
+        </Text>
+      ) : null}
+      {user?.quota && user.quota.maxUrls ? (
+        <Text size='sm' c='dimmed'>
+          You have created <b>{statsLoading ? '...' : stats?.urlsCreated}</b> links out of{' '}
+          <b>{user.quota.maxUrls}</b> links allowed.
+        </Text>
+      ) : null}
 
       <Title order={2} mt='md' mb='xs'>
         Recent files

src/components/pages/upload/uploadPartialFiles.tsx

@@ -227,6 +227,7 @@ export async function uploadPartialFiles(
       req.setRequestHeader('x-zipline-p-filename', file.name);
       req.setRequestHeader('x-zipline-p-lastchunk', j === chunks.length - 1 ? 'true' : 'false');
       req.setRequestHeader('x-zipline-p-content-type', file.type);
+      req.setRequestHeader('x-zipline-p-content-length', file.size.toString());
       req.setRequestHeader('content-range', `bytes ${chunks[j].start}-${chunks[j].end}/${file.size}`);
 
       req.send(body);

src/components/pages/users/EditUserModal.tsx

@@ -1,14 +1,17 @@
 import { Response } from '@/lib/api/response';
+import { readToDataURL } from '@/lib/base64';
+import { bytes } from '@/lib/bytes';
 import { User } from '@/lib/db/models/user';
 import { fetchApi } from '@/lib/fetchApi';
-import { readToDataURL } from '@/lib/base64';
 import { canInteract } from '@/lib/role';
 import { useUserStore } from '@/lib/store/user';
 import {
   ActionIcon,
   Button,
+  Divider,
   FileInput,
   Modal,
+  NumberInput,
   PasswordInput,
   Select,
   Stack,
@@ -20,6 +23,7 @@ import {
 import { useForm } from '@mantine/form';
 import { notifications } from '@mantine/notifications';
 import { IconPhotoMinus, IconUserCancel, IconUserEdit } from '@tabler/icons-react';
+import { useEffect } from 'react';
 import { mutate } from 'swr';
 
 export default function EditUserModal({
@@ -38,18 +42,40 @@ export default function EditUserModal({
     password: string;
     role: 'USER' | 'ADMIN' | 'SUPERADMIN';
     avatar: File | null;
+    fileType: 'BY_BYTES' | 'BY_FILES' | 'NONE';
+    maxFiles: number;
+    maxBytes: string;
+    maxUrls: number;
   }>({
     initialValues: {
       username: user?.username || '',
       password: '',
       role: user?.role || 'USER',
       avatar: null,
+      fileType: user?.quota?.filesQuota || 'NONE',
+      maxFiles: user?.quota?.maxFiles || 0,
+      maxBytes: user?.quota?.maxBytes || '',
+      maxUrls: user?.quota?.maxUrls || 0,
+    },
+    validate: {
+      maxBytes(value, values) {
+        if (values.fileType !== 'BY_BYTES') return;
+        if (typeof value !== 'string') return 'Invalid value';
+        const byte = bytes(value);
+        if (!bytes || byte < 0) return 'Invalid byte format';
+      },
+      maxFiles(value, values) {
+        if (values.fileType !== 'BY_FILES') return;
+        if (typeof value !== 'number' || value < 0) return 'Invalid value';
+      },
     },
   });
 
   const onSubmit = async (values: typeof form.values) => {
     if (!user) return;
 
+    console.log(values);
+
     let avatar64: string | null = null;
     if (values.avatar) {
       if (!values.avatar.type.startsWith('image/')) return form.setFieldError('avatar', 'Invalid file type');
@@ -64,11 +90,35 @@ export default function EditUserModal({
       }
     }
 
+    const finalQuota: {
+      filesType?: 'BY_BYTES' | 'BY_FILES' | 'NONE';
+      maxFiles?: number | null;
+      maxBytes?: string | null;
+
+      maxUrls?: number | null;
+    } = {};
+
+    if (values.fileType === 'NONE') {
+      finalQuota.filesType = 'NONE';
+      finalQuota.maxFiles = null;
+      finalQuota.maxBytes = null;
+      finalQuota.maxUrls = null;
+    } else if (values.fileType === 'BY_BYTES') {
+      finalQuota.filesType = 'BY_BYTES';
+      finalQuota.maxBytes = values.maxBytes;
+    } else {
+      finalQuota.filesType = 'BY_FILES';
+      finalQuota.maxFiles = values.maxFiles;
+    }
+
+    if (values.maxUrls) finalQuota.maxUrls = values.maxUrls > 0 ? values.maxUrls : null;
+
     const { data, error } = await fetchApi<Response['/api/users/[id]']>(`/api/users/${user.id}`, 'PATCH', {
       ...(values.username !== user.username && { username: values.username }),
       ...(values.password && { password: values.password }),
       ...(values.role !== user.role && { role: values.role }),
       ...(avatar64 && { avatar: avatar64 }),
+      quota: finalQuota,
     });
 
     if (error) {
@@ -92,6 +142,19 @@ export default function EditUserModal({
     }
   };
 
+  useEffect(() => {
+    form.setValues({
+      username: user?.username || '',
+      password: '',
+      role: user?.role || 'USER',
+      avatar: null,
+      fileType: user?.quota?.filesQuota || 'NONE',
+      maxFiles: user?.quota?.maxFiles || 0,
+      maxBytes: user?.quota?.maxBytes || '',
+      maxUrls: user?.quota?.maxUrls || 0,
+    });
+  }, [user]);
+
   return (
     <Modal centered title={<Title>Edit {user?.username ?? ''}</Title>} onClose={onClose} opened={opened}>
       <Text size='sm' c='dimmed'>
@@ -142,6 +205,44 @@ export default function EditUserModal({
               {...form.getInputProps('role')}
             />
 
+            <Divider />
+            <Title order={4}>Quota</Title>
+
+            <Select
+              label='File Quota Type'
+              description='Whether to set a quota on files by total bytes or the total number of files.'
+              data={[
+                { value: 'BY_BYTES', label: 'By Bytes' },
+                { value: 'BY_FILES', label: 'By File Count' },
+                { value: 'NONE', label: 'No File Quota' },
+              ]}
+              {...form.getInputProps('fileType')}
+            />
+            {form.values.fileType === 'BY_FILES' ? (
+              <NumberInput
+                label='Max Files'
+                description='The maximum number of files the user can upload.'
+                placeholder='Enter a number...'
+                mx='lg'
+                min={0}
+                {...form.getInputProps('maxFiles')}
+              />
+            ) : form.values.fileType === 'BY_BYTES' ? (
+              <TextInput
+                label='Max Bytes'
+                description='The maximum number of bytes the user can upload.'
+                placeholder='Enter a human readable byte-format...'
+                mx='lg'
+                {...form.getInputProps('maxBytes')}
+              />
+            ) : null}
+
+            <NumberInput
+              label='Max URLs'
+              placeholder='Enter a number...'
+              {...form.getInputProps('maxUrls')}
+            />
+
             <Button
               type='submit'
               variant='outline'

src/lib/bytes.ts

@@ -4,5 +4,5 @@ export function bytes(value: string): number;
 export function bytes(value: number, options?: BytesOptions): string;
 export function bytes(value: string | number, options?: BytesOptions): string | number {
   if (typeof value === 'string') return bytesFn(value);
-  return bytesFn(value, { ...options, unitSeparator: ' ' });
+  return bytesFn(Number(value), { ...options, unitSeparator: ' ' });
 }

src/lib/db/models/user.ts

@@ -1,4 +1,4 @@
-import { OAuthProvider, UserPasskey } from '@prisma/client';
+import { OAuthProvider, UserPasskey, UserQuota } from '@prisma/client';
 import { z } from 'zod';
 
 export type User = {
@@ -14,6 +14,8 @@ export type User = {
   totpSecret?: string | null;
   passkeys?: UserPasskey[];
 
+  quota?: UserQuota | null;
+
   avatar?: string | null;
   password?: string | null;
   token?: string | null;
@@ -29,6 +31,7 @@ export const userSelect = {
   oauthProviders: true,
   totpSecret: true,
   passkeys: true,
+  quota: true,
 };
 
 export type UserViewSettings = z.infer<typeof userViewSchema>;

src/lib/middleware/functions.ts

@@ -46,6 +46,14 @@ export function functions() {
         });
       };
 
+      res.tooLarge = (message: string = 'Payload Too Large', data: ErrorBody = {}) => {
+        return res.status(413).json({
+          code: 413,
+          message,
+          ...data,
+        });
+      };
+
       res.ratelimited = (retryAfter: number, message: string = 'Ratelimited', data: ErrorBody = {}) => {
         res.setHeader('Retry-After', retryAfter);
         return res.status(429).json({

src/lib/response.ts

@@ -35,6 +35,7 @@ export interface NextApiRes<Data = any> extends NextApiResponse {
   unauthorized: (message?: string, data?: ErrorBody) => void;
   forbidden: (message?: string, data?: ErrorBody) => void;
   notFound: (message?: string, data?: ErrorBody) => void;
+  tooLarge: (message?: string, data?: ErrorBody) => void;
   ratelimited: (retryAfter: number, message?: string, data?: ErrorBody) => void;
   serverError: (message?: string, data?: ErrorBody) => void;
   methodNotAllowed: (message?: string, data?: ErrorBody) => void;

src/lib/uploader/parseHeaders.ts

@@ -59,6 +59,7 @@ export type UploadHeaders = {
   'x-zipline-p-content-type'?: string;
   'x-zipline-p-identifier'?: string;
   'x-zipline-p-lastchunk'?: StringBoolean;
+  'x-zipline-p-content-length'?: string;
 };
 
 export type UploadOptions = {
@@ -88,6 +89,7 @@ export type UploadOptions = {
     identifier: string;
     lastchunk: boolean;
     range: [number, number, number]; // start, end, total
+    contentLength: number;
   };
 };
 
@@ -222,6 +224,7 @@ export function parseHeaders(headers: UploadHeaders, fileConfig: Config['files']
       identifier: headers['x-zipline-p-identifier']!,
       lastchunk: headers['x-zipline-p-lastchunk'] === 'true',
       range: [start, end, total],
+      contentLength: Number(headers['x-zipline-p-content-length']!),
     };
   }
 

src/pages/api/upload.ts

@@ -1,6 +1,8 @@
 import { handlePartialUpload } from '@/lib/api/partialUpload';
 import { handleFile } from '@/lib/api/upload';
+import { bytes } from '@/lib/bytes';
 import { config as zconfig } from '@/lib/config';
+import { prisma } from '@/lib/db';
 import { log } from '@/lib/logger';
 import { combine } from '@/lib/middleware/combine';
 import { file } from '@/lib/middleware/file';
@@ -32,6 +34,36 @@ export async function handler(req: NextApiReq<any, any, UploadHeaders>, res: Nex
 
   if (options.header) return res.badRequest('', options);
 
+  if (req.user.quota) {
+    const totalFileSize = options.partial
+      ? options.partial.contentLength
+      : req.files.reduce((acc, x) => acc + x.size, 0);
+
+    const userAggregateStats = await prisma.file.aggregate({
+      where: {
+        userId: req.user.id,
+      },
+      _sum: {
+        size: true,
+      },
+      _count: {
+        _all: true,
+      },
+    });
+    const aggSize = userAggregateStats!._sum?.size === null ? 0 : userAggregateStats!._sum?.size;
+
+    if (req.user.quota.filesQuota === 'BY_BYTES' && aggSize + totalFileSize > bytes(req.user.quota.maxBytes!))
+      return res.tooLarge(
+        `uploading will exceed your storage quota of ${bytes(req.user.quota.maxBytes!)} bytes`,
+      );
+
+    if (
+      req.user.quota.filesQuota === 'BY_FILES' &&
+      userAggregateStats!._count?._all + req.files.length > req.user.quota.maxFiles!
+    )
+      return res.tooLarge(`uploading will exceed your file count quota of ${req.user.quota.maxFiles} files`);
+  }
+
   const response: ApiUploadResponse = {
     files: [],
     ...(options.deletesAt && { deletesAt: options.deletesAt.toISOString() }),

src/pages/api/user/urls/index.ts

@@ -44,6 +44,14 @@ export async function handler(req: NextApiReq<Body, Query, Headers>, res: NextAp
     const { vanity, destination } = req.body;
     const noJson = !!req.headers['x-zipline-no-json'];
 
+    const countUrls = await prisma.url.count({
+      where: {
+        userId: req.user.id,
+      },
+    });
+    if (req.user.quota && req.user.quota.maxUrls && countUrls + 1 > req.user.quota.maxUrls)
+      return res.forbidden(`shortenning this url would exceed your quota of ${req.user.quota.maxUrls} urls`);
+
     let maxViews: number | undefined;
     const returnDomain = req.headers['x-zipline-domain'];
 

src/pages/api/users/[id].ts

@@ -1,3 +1,4 @@
+import { bytes } from '@/lib/bytes';
 import { hashPassword } from '@/lib/crypto';
 import { datasource } from '@/lib/datasource';
 import { prisma } from '@/lib/db';
@@ -8,6 +9,7 @@ import { method } from '@/lib/middleware/method';
 import { ziplineAuth } from '@/lib/middleware/ziplineAuth';
 import { NextApiReq, NextApiRes } from '@/lib/response';
 import { canInteract } from '@/lib/role';
+import { UserFilesQuota } from '@prisma/client';
 import { z } from 'zod';
 
 export type ApiUsersIdResponse = User;
@@ -17,6 +19,13 @@ type Body = {
   password?: string;
   avatar?: string;
   role?: 'USER' | 'ADMIN' | 'SUPERADMIN';
+  quota?: {
+    filesType?: UserFilesQuota & 'NONE';
+    maxFiles?: number;
+    maxBytes?: string;
+
+    maxUrls?: number;
+  };
 
   delete?: boolean;
 };
@@ -26,6 +35,7 @@ type Query = {
 };
 
 const logger = log('api').c('users').c('[id]');
+const zNumber = z.number();
 
 export async function handler(req: NextApiReq<Body, Query>, res: NextApiRes<ApiUsersIdResponse>) {
   const user = await prisma.user.findUnique({
@@ -37,13 +47,55 @@ export async function handler(req: NextApiReq<Body, Query>, res: NextApiRes<ApiU
   if (!user) return res.notFound('User not found');
 
   if (req.method === 'PATCH') {
-    const { username, password, avatar, role } = req.body;
+    const { username, password, avatar, role, quota } = req.body;
 
     if (role && !z.enum(['USER', 'ADMIN']).safeParse(role).success)
       return res.badRequest('Invalid role (USER, ADMIN)');
 
     if (role && !canInteract(req.user.role, role)) return res.forbidden('You cannot create this role');
 
+    let finalQuota:
+      | {
+          filesQuota?: UserFilesQuota;
+          maxFiles?: number | null;
+          maxBytes?: string | null;
+          maxUrls?: number | null;
+        }
+      | undefined = undefined;
+    if (quota) {
+      if (quota.filesType && !z.enum(['BY_BYTES', 'BY_FILES', 'NONE']).safeParse(quota.filesType).success)
+        return res.badRequest('Invalid filesType (BY_BYTES, BY_FILES, NONE)');
+
+      if (quota.maxFiles && !zNumber.safeParse(quota.maxFiles).success)
+        return res.badRequest('Invalid maxFiles');
+      if (quota.maxUrls && !zNumber.safeParse(quota.maxUrls).success)
+        return res.badRequest('Invalid maxUrls');
+
+      if (quota.filesType === 'BY_BYTES' && quota.maxBytes === undefined)
+        return res.badRequest('maxBytes is required');
+      if (quota.filesType === 'BY_FILES' && quota.maxFiles === undefined)
+        return res.badRequest('maxFiles is required');
+
+      finalQuota = {
+        ...(quota.filesType === 'BY_BYTES' && {
+          filesQuota: 'BY_BYTES',
+          maxBytes: bytes(quota.maxBytes || '0') > 0 ? quota.maxBytes : null,
+          maxFiles: null,
+        }),
+        ...(quota.filesType === 'BY_FILES' && {
+          filesQuota: 'BY_FILES',
+          maxFiles: quota.maxFiles,
+          maxBytes: null,
+        }),
+        ...(quota.filesType === 'NONE' && {
+          filesQuota: 'BY_BYTES',
+          maxFiles: null,
+          maxBytes: null,
+        }),
+        maxUrls: (quota.maxUrls || 0) > 0 ? quota.maxUrls : null,
+      };
+    }
+
     const updatedUser = await prisma.user.update({
       where: {
         id: user.id,
@@ -53,6 +105,22 @@ export async function handler(req: NextApiReq<Body, Query>, res: NextApiRes<ApiU
         ...(password && { password: await hashPassword(password) }),
         ...(role !== undefined && { role: 'USER' }),
         ...(avatar && { avatar }),
+        ...(finalQuota && {
+          quota: {
+            upsert: {
+              where: {
+                userId: user.id,
+              },
+              create: {
+                filesQuota: finalQuota.filesQuota || 'BY_BYTES',
+                maxFiles: finalQuota.maxFiles ?? null,
+                maxBytes: finalQuota.maxBytes ?? null,
+                maxUrls: finalQuota.maxUrls ?? null,
+              },
+              update: finalQuota,
+            },
+          },
+        }),
       },
       select: {
         ...userSelect,