diff --git a/lib/core/buckets/Bucket.ts b/lib/core/buckets/Bucket.ts index a7b59b8b..1851ae09 100644 --- a/lib/core/buckets/Bucket.ts +++ b/lib/core/buckets/Bucket.ts @@ -11,6 +11,7 @@ export interface Bucket { status: string; transfer: number; storage: number; + usedSpaceBytes?: number; created?: Date; maxFrameSize?: number; publicPermissions?: string[]; diff --git a/lib/core/buckets/MongoDBBucketsRepository.ts b/lib/core/buckets/MongoDBBucketsRepository.ts index 8ce69daf..157618d4 100644 --- a/lib/core/buckets/MongoDBBucketsRepository.ts +++ b/lib/core/buckets/MongoDBBucketsRepository.ts @@ -57,12 +57,34 @@ export class MongoDBBucketsRepository implements BucketsRepository { return formatFromMongoToBucket(rawModel); } + async sumUsedSpaceBytes(userId: Bucket['userId']): Promise { + const [result] = await this.model.aggregate([ + { $match: { userId } }, + { $group: { _id: null, total: { $sum: { $ifNull: ['$usedSpaceBytes', 0] } } } } + ]); + + return result ? result.total : 0; + } + async destroyByUser(userId: Bucket['userId']): Promise { await this.model.deleteMany({ userId, }); } + async setUsedSpaceBytes( + bucketId: Bucket['id'], + userId: Bucket['userId'], + usedSpaceBytes: number + ): Promise { + const result = await this.model.updateOne( + { _id: bucketId, userId }, + { $set: { usedSpaceBytes } } + ); + + return result.matchedCount > 0; + } + async removeByIdAndUser(bucketId: Bucket['id'], userId: Bucket['userId']): Promise { await this.model.deleteOne({ userId, diff --git a/lib/core/buckets/Repository.ts b/lib/core/buckets/Repository.ts index e24da97d..3dc400cc 100644 --- a/lib/core/buckets/Repository.ts +++ b/lib/core/buckets/Repository.ts @@ -7,7 +7,13 @@ export interface BucketsRepository { findByIds(ids: Bucket['id'][]): Promise; find(where: Partial): Promise; findUserBucketsFromDate(userId: Bucket['id'], date?: Date, limit?: number): Promise; + setUsedSpaceBytes( + bucketId: Bucket['id'], + userId: Bucket['userId'], + usedSpaceBytes: number + ): Promise; + sumUsedSpaceBytes(userId: Bucket['userId']): Promise; destroyByUser(userId: Bucket['userId']): Promise; removeAll(where: Partial): Promise; - removeByIdAndUser(bucketId: Bucket['id'], userId: Bucket['userId']): Promise + removeByIdAndUser(bucketId: Bucket['id'], userId: Bucket['userId']): Promise } diff --git a/lib/core/users/usecase.ts b/lib/core/users/usecase.ts index 973307e4..c27af73a 100644 --- a/lib/core/users/usecase.ts +++ b/lib/core/users/usecase.ts @@ -2,6 +2,7 @@ import { createHash, randomBytes } from 'crypto'; import { UsersRepository } from './Repository'; import { BucketsRepository } from '../buckets/Repository'; +import { BucketNotFoundError } from '../buckets/usecase'; import { MailUsecase } from '../mail/usecase'; import { EventBus, EventBusEvents } from '../../server/eventBus'; import { FramesRepository } from '../frames/Repository'; @@ -20,6 +21,11 @@ function isEmailValid(email: string) { export const RESET_PASSWORD_TOKEN_BYTES_LENGTH = 256; export const SHA256_HASH_BYTES_LENGTH = 32; +export interface UserSpaceSnapshot { + maxSpaceBytes: number; + totalUsedSpaceBytes: number; +} + export class UserAlreadyExistsError extends Error { constructor() { super('User already exists'); @@ -286,12 +292,43 @@ export class UsersUsecase { status: 'Active', transfer: 0, storage: 0, + usedSpaceBytes: 0, encryptionKey: '', }); return { id: created.id, name: created.name }; } + async setBucketUsage( + uuid: User['uuid'], + bucketId: string, + usedSpaceBytes: number + ): Promise { + const updated = await this.bucketsRepository.setUsedSpaceBytes( + bucketId, + uuid, + usedSpaceBytes + ); + + if (!updated) { + throw new BucketNotFoundError(bucketId); + } + + const [user, bucketsUsedSpaceBytes] = await Promise.all([ + this.usersRepository.findByUuid(uuid), + this.bucketsRepository.sumUsedSpaceBytes(uuid), + ]); + + if (!user) { + throw new UserNotFoundError(uuid); + } + + return { + maxSpaceBytes: user.maxSpaceBytes, + totalUsedSpaceBytes: user.totalUsedSpaceBytes + bucketsUsedSpaceBytes, + }; + } + async deleteBucket(uuid: User['uuid'], bucketId: string): Promise { const user = await this.usersRepository.findByUuid(uuid); diff --git a/lib/models/bucket.ts b/lib/models/bucket.ts index 8db590a2..1a2efa78 100644 --- a/lib/models/bucket.ts +++ b/lib/models/bucket.ts @@ -5,6 +5,7 @@ const errors = require("storj-service-error-types"); interface IBucket extends Document { storage: number; + usedSpaceBytes: number; transfer: number; status: "Active" | "Inactive"; pubkeys: string[]; @@ -20,6 +21,7 @@ interface IBucket extends Document { const BucketSchema = new Schema( { storage: { type: Number, default: 0 }, + usedSpaceBytes: { type: Number, default: 0 }, transfer: { type: Number, default: 0 }, status: { type: String, diff --git a/lib/server/http/gateway/controller.ts b/lib/server/http/gateway/controller.ts index fa86d21e..3f6e1cd2 100644 --- a/lib/server/http/gateway/controller.ts +++ b/lib/server/http/gateway/controller.ts @@ -1,7 +1,8 @@ import { Request, Response } from 'express'; import { Logger } from 'winston'; -import { EmailIsAlreadyInUseError, InvalidDataFormatError, UserAlreadyExistsError, UserNotFoundError, UsersUsecase } from '../../../core'; +import { EmailIsAlreadyInUseError, InvalidDataFormatError, UserAlreadyExistsError, UserNotFoundError, UserSpaceSnapshot, UsersUsecase } from '../../../core'; import { BucketEntriesUsecase } from '../../../core/bucketEntries/usecase'; +import { BucketNotFoundError } from '../../../core/buckets/usecase'; import { GatewayUsecase } from '../../../core/gateway/Usecase'; import { EventBus, EventBusEvents, UserStorageChangedPayload } from '../../eventBus'; @@ -16,6 +17,13 @@ type DeleteFilesInBulkResponse = { type CreateBucketBody = { name: string }; type CreateBucketResponse = { id: string; name: string }; +type SetBucketUsageBody = { usedSpaceBytes: number }; + +const OBJECT_ID_PATTERN = /^[a-f0-9]{24}$/i; + +const isValidUsedSpaceBytes = (value: unknown): value is number => + typeof value === 'number' && Number.isFinite(value) && value >= 0; + export class HTTPGatewayController { constructor( private gatewayUsecase: GatewayUsecase, @@ -165,6 +173,44 @@ export class HTTPGatewayController { } } + async setBucketUsage( + req: Request<{ uuid: string; id: string }, {}, Partial, {}>, + res: Response + ) { + const { uuid, id } = req.params; + const { usedSpaceBytes } = req.body; + + if (!uuid || !id || !OBJECT_ID_PATTERN.test(id)) { + return res.status(400).send({ message: 'Invalid params' }); + } + + if (!isValidUsedSpaceBytes(usedSpaceBytes)) { + return res + .status(400) + .send({ message: 'usedSpaceBytes must be a non-negative number' }); + } + + try { + const snapshot = await this.usersUsecase.setBucketUsage(uuid, id, usedSpaceBytes); + + return res.status(200).send(snapshot); + } catch (err) { + if (err instanceof UserNotFoundError || err instanceof BucketNotFoundError) { + return res.status(404).send({ message: err.message }); + } + + this.logger.error( + '[GATEWAY/BUCKET_USAGE] Error setting usage for bucket %s of user %s: %s. %s', + id, + uuid, + (err as Error).message, + (err as Error).stack || 'NO STACK' + ); + + return res.status(500).send({ message: 'Internal server error' }); + } + } + async deleteUserBucket( req: Request<{ uuid: string; id: string }>, res: Response<{ message: string }> diff --git a/lib/server/http/gateway/index.ts b/lib/server/http/gateway/index.ts index 3e237efd..90765bd9 100644 --- a/lib/server/http/gateway/index.ts +++ b/lib/server/http/gateway/index.ts @@ -11,6 +11,7 @@ export const createGatewayHTTPRouter = ( router.patch('/users/:uuid', jwtMiddleware, controller.updateUserEmail.bind(controller)); router.put('/storage/users/:uuid', jwtMiddleware, controller.changeStorage.bind(controller)); router.post('/users/:uuid/buckets', jwtMiddleware, controller.createUserBucket.bind(controller)); + router.put('/users/:uuid/buckets/:id/usage', jwtMiddleware, controller.setBucketUsage.bind(controller)); router.delete('/users/:uuid/buckets/:id', jwtMiddleware, controller.deleteUserBucket.bind(controller)); router.delete('/storage/files', jwtMiddleware, controller.deleteFilesInBulk.bind(controller)); diff --git a/tests/lib/core/users/usecase.test.ts b/tests/lib/core/users/usecase.test.ts index 7c10c772..8bfab207 100644 --- a/tests/lib/core/users/usecase.test.ts +++ b/tests/lib/core/users/usecase.test.ts @@ -18,6 +18,7 @@ import { import { MongoDBFramesRepository } from '../../../../lib/core/frames/MongoDBFramesRepository'; import { FramesRepository } from '../../../../lib/core/frames/Repository'; import { BucketsRepository } from '../../../../lib/core/buckets/Repository'; +import { BucketNotFoundError } from '../../../../lib/core/buckets/usecase'; import { Mailer, MailUsecase, SendGridMailUsecase } from '../../../../lib/core/mail/usecase'; import { EventBus, EventBusEvents } from '../../../../lib/server/eventBus'; import { User } from '../../../../lib/core/users/User'; @@ -451,6 +452,7 @@ describe('Users usecases', () => { status: 'Active', transfer: 0, storage: 0, + usedSpaceBytes: 0, encryptionKey: '', }]); expect(result).toStrictEqual({ id: createdBucket.id, name: createdBucket.name }); @@ -503,6 +505,51 @@ describe('Users usecases', () => { }); }); + describe('Setting bucket usage', () => { + it('When the bucket does not belong to the user or does not exist, then it throws BucketNotFoundError', async () => { + const setUsage = stub(bucketsRepository, 'setUsedSpaceBytes').resolves(false); + + try { + await usecase.setBucketUsage('user-uuid', 'bucket-id', 1024); + expect(true).toBeFalsy(); + } catch (err) { + expect(err).toBeInstanceOf(BucketNotFoundError); + } + + expect(setUsage.calledOnceWithExactly('bucket-id', 'user-uuid', 1024)).toBeTruthy(); + }); + + it('When the bucket belongs to the user, then it stores the reported usage without touching the user total', async () => { + const user = fixtures.getUser({ maxSpaceBytes: 10000, totalUsedSpaceBytes: 4000 }); + + const setUsage = stub(bucketsRepository, 'setUsedSpaceBytes').resolves(true); + const addUsage = stub(usersRepository, 'addTotalUsedSpaceBytes').resolves(); + stub(usersRepository, 'findByUuid').resolves(user); + stub(bucketsRepository, 'sumUsedSpaceBytes').resolves(2048); + + await usecase.setBucketUsage(user.uuid, 'bucket-id', 2048); + + expect(setUsage.calledOnceWithExactly('bucket-id', user.uuid, 2048)).toBeTruthy(); + expect(addUsage.called).toBeFalsy(); + }); + + it('When the usage is stored, then it returns the user space snapshot with the buckets usage summed in', async () => { + const user = fixtures.getUser({ maxSpaceBytes: 10000, totalUsedSpaceBytes: 4000 }); + + stub(bucketsRepository, 'setUsedSpaceBytes').resolves(true); + stub(usersRepository, 'findByUuid').resolves(user); + const sumUsage = stub(bucketsRepository, 'sumUsedSpaceBytes').resolves(1500); + + const snapshot = await usecase.setBucketUsage(user.uuid, 'bucket-id', 1500); + + expect(sumUsage.calledOnceWithExactly(user.uuid)).toBeTruthy(); + expect(snapshot).toStrictEqual({ + maxSpaceBytes: 10000, + totalUsedSpaceBytes: 5500, + }); + }); + }); + describe('Confirming user destruction', () => { it('When confirming a destruction of a user that exists, then it works', async () => { const user = fixtures.getUser(); diff --git a/tests/lib/e2e/gateway/gateway-v2.e2e-spec.ts b/tests/lib/e2e/gateway/gateway-v2.e2e-spec.ts index d8f73abe..49286951 100644 --- a/tests/lib/e2e/gateway/gateway-v2.e2e-spec.ts +++ b/tests/lib/e2e/gateway/gateway-v2.e2e-spec.ts @@ -115,6 +115,126 @@ describe('Gateway V2 e2e tests', () => { }) }) + describe('Setting bucket usage', () => { + const createBucketForUser = async (userUuid: string, jwt: string) => { + const { body } = await testServer + .post(`/v2/gateway/users/${userUuid}/buckets`) + .set('Authorization', `Bearer ${jwt}`) + .send({ name: `mail-account-${crypto.randomUUID()}` }) + + return body as { id: string; name: string } + } + + it('When reporting usage, then it is persisted on the bucket as an absolute value', async () => { + const testUser = await createTestUser() + const jwt = signRS256JWT('5m', engine._config.gateway.SIGN_JWT_SECRET) + const bucket = await createBucketForUser(testUser.uuid, jwt) + + const firstReport = await testServer + .put(`/v2/gateway/users/${testUser.uuid}/buckets/${bucket.id}/usage`) + .set('Authorization', `Bearer ${jwt}`) + .send({ usedSpaceBytes: 5000 }) + + expect(firstReport.status).toBe(200) + + let bucketInDatabase = await databaseConnection.models.Bucket.findOne({ _id: bucket.id }) + expect(bucketInDatabase.usedSpaceBytes).toBe(5000) + + const secondReport = await testServer + .put(`/v2/gateway/users/${testUser.uuid}/buckets/${bucket.id}/usage`) + .set('Authorization', `Bearer ${jwt}`) + .send({ usedSpaceBytes: 2000 }) + + expect(secondReport.status).toBe(200) + + bucketInDatabase = await databaseConnection.models.Bucket.findOne({ _id: bucket.id }) + expect(bucketInDatabase.usedSpaceBytes).toBe(2000) + }) + + it('When reporting usage, then it returns the user space snapshot with buckets usage summed in', async () => { + const testUser = await createTestUser() + const jwt = signRS256JWT('5m', engine._config.gateway.SIGN_JWT_SECRET) + const bucket = await createBucketForUser(testUser.uuid, jwt) + + const response = await testServer + .put(`/v2/gateway/users/${testUser.uuid}/buckets/${bucket.id}/usage`) + .set('Authorization', `Bearer ${jwt}`) + .send({ usedSpaceBytes: 5000 }) + + expect(response.status).toBe(200) + + const userInDatabase = await databaseConnection.models.User.findOne({ uuid: testUser.uuid }) + expect(response.body).toEqual({ + maxSpaceBytes: userInDatabase.maxSpaceBytes, + totalUsedSpaceBytes: userInDatabase.totalUsedSpaceBytes + 5000, + }) + }) + + it('When reporting usage, then the user drive usage is not affected', async () => { + const testUser = await createTestUser() + const jwt = signRS256JWT('5m', engine._config.gateway.SIGN_JWT_SECRET) + const bucket = await createBucketForUser(testUser.uuid, jwt) + + const userBefore = await databaseConnection.models.User.findOne({ uuid: testUser.uuid }) + + const response = await testServer + .put(`/v2/gateway/users/${testUser.uuid}/buckets/${bucket.id}/usage`) + .set('Authorization', `Bearer ${jwt}`) + .send({ usedSpaceBytes: 3000 }) + + expect(response.status).toBe(200) + + const userAfter = await databaseConnection.models.User.findOne({ uuid: testUser.uuid }) + expect(userAfter.totalUsedSpaceBytes).toBe(userBefore.totalUsedSpaceBytes) + }) + + it('When reporting usage for a bucket of another user, then it returns 404 and changes nothing', async () => { + const owner = await createTestUser() + const otherUser = await createTestUser() + const jwt = signRS256JWT('5m', engine._config.gateway.SIGN_JWT_SECRET) + const bucket = await createBucketForUser(owner.uuid, jwt) + + const response = await testServer + .put(`/v2/gateway/users/${otherUser.uuid}/buckets/${bucket.id}/usage`) + .set('Authorization', `Bearer ${jwt}`) + .send({ usedSpaceBytes: 5000 }) + + expect(response.status).toBe(404) + + const bucketInDatabase = await databaseConnection.models.Bucket.findOne({ _id: bucket.id }) + expect(bucketInDatabase.usedSpaceBytes).toBe(0) + }) + + it('When the reported usage is negative or not a number, then it returns 400', async () => { + const testUser = await createTestUser() + const jwt = signRS256JWT('5m', engine._config.gateway.SIGN_JWT_SECRET) + const bucket = await createBucketForUser(testUser.uuid, jwt) + + const negative = await testServer + .put(`/v2/gateway/users/${testUser.uuid}/buckets/${bucket.id}/usage`) + .set('Authorization', `Bearer ${jwt}`) + .send({ usedSpaceBytes: -1 }) + + const missing = await testServer + .put(`/v2/gateway/users/${testUser.uuid}/buckets/${bucket.id}/usage`) + .set('Authorization', `Bearer ${jwt}`) + .send({}) + + expect(negative.status).toBe(400) + expect(missing.status).toBe(400) + }) + + it('When no auth token is provided, then it returns 401', async () => { + const testUser = await createTestUser() + + const response = await testServer + .put(`/v2/gateway/users/${testUser.uuid}/buckets/${'a'.repeat(24)}/usage`) + .send({ usedSpaceBytes: 1000 }) + + expect(response.status).toBe(401) + }) + }) + describe('Deleting a user bucket', () => { it('When deleting an existing bucket, then it is removed from the database', async () => { const testUser = await createTestUser()