diff --git a/packages/server/modules/core/authz/loaders/index.ts b/packages/server/modules/core/authz/loaders/index.ts index b584f3987..35987ce0d 100644 --- a/packages/server/modules/core/authz/loaders/index.ts +++ b/packages/server/modules/core/authz/loaders/index.ts @@ -28,6 +28,10 @@ export default defineModuleLoaders(async () => { const counts = await dataLoaders.streams.getCollaboratorCounts.load(projectId) return counts?.[role] || 0 }, + getProjectModelCount: async ({ projectId }, { dataLoaders }) => { + const db = await getProjectDbClient({ projectId }) + return await dataLoaders.forRegion({ db }).streams.getBranchCount.load(projectId) + }, getModel: async ({ projectId, modelId }, { dataLoaders }) => { const db = await getProjectDbClient({ projectId }) const model = await dataLoaders.forRegion({ db }).branches.getById.load(modelId) diff --git a/packages/shared/src/authz/domain/loaders.ts b/packages/shared/src/authz/domain/loaders.ts index 468da1230..abf3522ae 100644 --- a/packages/shared/src/authz/domain/loaders.ts +++ b/packages/shared/src/authz/domain/loaders.ts @@ -3,6 +3,7 @@ import { MaybeAsync } from '../../core/index.js' import type { GetServerRole } from './core/operations.js' import type { GetProject, + GetProjectModelCount, GetProjectRole, GetProjectRoleCounts } from './projects/operations.js' @@ -56,6 +57,7 @@ export const AuthCheckContextLoaderKeys = { getProject: 'getProject', getProjectRoleCounts: 'getProjectRoleCounts', getProjectRole: 'getProjectRole', + getProjectModelCount: 'getProjectModelCount', getServerRole: 'getServerRole', getWorkspace: 'getWorkspace', getWorkspaceRole: 'getWorkspaceRole', @@ -83,6 +85,7 @@ export type AllAuthCheckContextLoaders = AuthContextLoaderMappingDefinition<{ getProject: GetProject getProjectRole: GetProjectRole getProjectRoleCounts: GetProjectRoleCounts + getProjectModelCount: GetProjectModelCount getServerRole: GetServerRole getWorkspace: GetWorkspace getWorkspaceRole: GetWorkspaceRole diff --git a/packages/shared/src/authz/domain/projects/operations.ts b/packages/shared/src/authz/domain/projects/operations.ts index aac8d44aa..44c6adc6a 100644 --- a/packages/shared/src/authz/domain/projects/operations.ts +++ b/packages/shared/src/authz/domain/projects/operations.ts @@ -12,3 +12,5 @@ export type GetProjectRoleCounts = (args: { projectId: string role: StreamRoles }) => Promise + +export type GetProjectModelCount = (args: { projectId: string }) => Promise diff --git a/packages/shared/src/authz/fragments/workspaces.ts b/packages/shared/src/authz/fragments/workspaces.ts index 7afcc187d..95bc5820e 100644 --- a/packages/shared/src/authz/fragments/workspaces.ts +++ b/packages/shared/src/authz/fragments/workspaces.ts @@ -18,6 +18,7 @@ import { Loaders } from '../domain/loaders.js' import { Roles, WorkspaceRoles } from '../../core/constants.js' import { MaybeUserContext, + MaybeWorkspaceContext, ProjectContext, WorkspaceContext } from '../domain/context.js' @@ -223,7 +224,14 @@ export const ensureModelCanBeCreatedFragment: AuthPolicyEnsureFragment< | typeof Loaders.getWorkspaceLimits | typeof Loaders.getProject | typeof Loaders.getWorkspaceModelCount, - ProjectContext & MaybeUserContext, + ProjectContext & + MaybeWorkspaceContext & + MaybeUserContext & { + /** + * How many models we're testing being added. Defaults to 1 + */ + addedModelCount?: number + }, InstanceType< | typeof WorkspaceNoAccessError | typeof WorkspaceReadOnlyError @@ -232,11 +240,13 @@ export const ensureModelCanBeCreatedFragment: AuthPolicyEnsureFragment< > > = (loaders) => - async ({ projectId, userId }) => { + async ({ projectId, userId, addedModelCount, workspaceId }) => { + addedModelCount = addedModelCount ?? 1 const project = await loaders.getProject({ projectId }) if (!project) return err(new ProjectNotFoundError()) - const { workspaceId } = project + // Project may not be attached to a workspace yet, then we use the specified workspaceId + workspaceId = workspaceId || project.workspaceId || undefined if (!workspaceId) return ok() if (userId) { @@ -267,7 +277,7 @@ export const ensureModelCanBeCreatedFragment: AuthPolicyEnsureFragment< if (currentModelCount === null) return err(new WorkspaceNoAccessError()) - return currentModelCount < workspaceLimits.modelCount + return currentModelCount + addedModelCount <= workspaceLimits.modelCount ? ok() : err( new WorkspaceLimitsReachedError({ diff --git a/packages/shared/src/authz/policies/project/canMoveToWorkspace.spec.ts b/packages/shared/src/authz/policies/project/canMoveToWorkspace.spec.ts index 5a305f9d9..cad4e2c08 100644 --- a/packages/shared/src/authz/policies/project/canMoveToWorkspace.spec.ts +++ b/packages/shared/src/authz/policies/project/canMoveToWorkspace.spec.ts @@ -4,8 +4,6 @@ import { canMoveToWorkspacePolicy } from './canMoveToWorkspace.js' import { parseFeatureFlags } from '../../../environment/index.js' import { Project } from '../../domain/projects/types.js' import { Roles, SeatTypes } from '../../../core/constants.js' -import { Workspace } from '../../domain/workspaces/types.js' -import { WorkspacePlan } from '../../../workspaces/index.js' import { ProjectNotEnoughPermissionsError, ServerNotEnoughPermissionsError, @@ -14,24 +12,26 @@ import { WorkspaceProjectMoveInvalidError, WorkspacesNotEnabledError } from '../../domain/authErrors.js' +import { getProjectFake, getWorkspaceFake } from '../../../tests/fakes.js' const buildCanMoveToWorkspace = ( overrides?: Partial[0]> ) => canMoveToWorkspacePolicy({ getEnv: async () => parseFeatureFlags({}), - getProject: async () => { - return {} as Project - }, + getProject: getProjectFake({ + id: 'project-id', + workspaceId: null + }), getProjectRole: async () => { return Roles.Stream.Owner }, getServerRole: async () => { return Roles.Server.User }, - getWorkspace: async () => { - return {} as Workspace - }, + getWorkspace: getWorkspaceFake({ + id: 'workspace-id' + }), getWorkspaceRole: async () => { return Roles.Workspace.Admin }, @@ -44,8 +44,11 @@ const buildCanMoveToWorkspace = ( }, getWorkspacePlan: async () => { return { - status: 'valid' - } as WorkspacePlan + status: 'valid', + workspaceId: 'workspace-id', + createdAt: new Date(), + name: 'team' + } }, getWorkspaceLimits: async () => { return { @@ -58,6 +61,12 @@ const buildCanMoveToWorkspace = ( getWorkspaceProjectCount: async () => { return 0 }, + getWorkspaceModelCount: async () => { + return 0 + }, + getProjectModelCount: async () => { + return 0 + }, ...overrides }) @@ -126,7 +135,8 @@ describe('canMoveToWorkspacePolicy returns a function, that', () => { code: WorkspaceNotEnoughPermissionsError.code }) }) - it('forbids move if target workspace will exceed plan limits', async () => { + + it('forbids move if target workspace will exceed project limits', async () => { const result = await buildCanMoveToWorkspace({ getWorkspaceLimits: async () => { return { @@ -146,6 +156,34 @@ describe('canMoveToWorkspacePolicy returns a function, that', () => { payload: { limit: 'projectCount' } }) }) + + it('forbids move if target workspace will exceed model limits', async () => { + const result = await buildCanMoveToWorkspace({ + getWorkspaceLimits: async () => { + return { + projectCount: 10, + modelCount: 5, + versionsHistory: null, + commentHistory: null + } + }, + getWorkspaceProjectCount: async () => { + return 1 + }, + getProjectModelCount: async () => { + return 5 + }, + getWorkspaceModelCount: async () => { + return 1 + } + })(canMoveToWorkspaceArgs()) + + expect(result).toBeAuthErrorResult({ + code: WorkspaceLimitsReachedError.code, + payload: { limit: 'modelCount' } + }) + }) + it('allows move project if target workspace will be within limits', async () => { const result = await buildCanMoveToWorkspace({})(canMoveToWorkspaceArgs()) expect(result).toBeAuthOKResult() diff --git a/packages/shared/src/authz/policies/project/canMoveToWorkspace.ts b/packages/shared/src/authz/policies/project/canMoveToWorkspace.ts index c85f65be9..19a74f7db 100644 --- a/packages/shared/src/authz/policies/project/canMoveToWorkspace.ts +++ b/packages/shared/src/authz/policies/project/canMoveToWorkspace.ts @@ -24,6 +24,7 @@ import { AuthCheckContextLoaderKeys } from '../../domain/loaders.js' import { AuthPolicy } from '../../domain/policies.js' import { Roles } from '../../../core/constants.js' import { + ensureModelCanBeCreatedFragment, ensureWorkspaceProjectCanBeCreatedFragment, ensureWorkspaceRoleAndSessionFragment, ensureWorkspacesEnabledFragment @@ -43,6 +44,8 @@ type PolicyLoaderKeys = | typeof AuthCheckContextLoaderKeys.getWorkspacePlan | typeof AuthCheckContextLoaderKeys.getWorkspaceLimits | typeof AuthCheckContextLoaderKeys.getWorkspaceProjectCount + | typeof AuthCheckContextLoaderKeys.getProjectModelCount + | typeof AuthCheckContextLoaderKeys.getWorkspaceModelCount | typeof AuthCheckContextLoaderKeys.getWorkspaceSeat type PolicyArgs = MaybeUserContext & MaybeProjectContext & MaybeWorkspaceContext @@ -120,5 +123,22 @@ export const canMoveToWorkspacePolicy: AuthPolicy< } } + if (workspaceId && projectId) { + // Check whether this specific project can be moved to the workspace + // Does it maybe have too many models? + const projectModelCount = await loaders.getProjectModelCount({ + projectId + }) + const ensuredModelsAccepted = await ensureModelCanBeCreatedFragment(loaders)({ + projectId, + userId, + addedModelCount: projectModelCount, + workspaceId + }) + if (ensuredModelsAccepted.isErr) { + return err(ensuredModelsAccepted.error) + } + } + return ok() }