From 21c840d2b59c8185c4164fe7bd03712c2e23f46b Mon Sep 17 00:00:00 2001 From: Aliaksandr Stsiapanay Date: Fri, 17 Jul 2026 11:54:17 +0300 Subject: [PATCH] feat: add per-resource limits and skill validation hardening for folder resources #1637 Add configurable maxFiles/maxTotalBytes/maxFileSizeBytes settings for folder-as-resource complex resources, enforced on whole-resource and single-file PUT so a rejection leaves nothing observable and bounds lock-hold time. Consolidate per-file size/etag tracking into a single fileMetadata map on the .dial-resource marker. --- README.md | 16 ++ .../com/epam/aidial/core/server/AiDial.java | 4 +- .../data/folder/FolderResourceMarker.java | 20 +- .../resource/ComplexResourceService.java | 178 +++++++++++---- .../src/main/resources/aidial.settings.json | 5 + .../core/server/SkillResourceApiTest.java | 204 ++++++++++++++++++ .../resource/ComplexResourceServiceTest.java | 109 +++++++++- .../ComplexResourceSweepServiceTest.java | 2 +- .../src/test/resources/aidial.settings.json | 5 + 9 files changed, 493 insertions(+), 50 deletions(-) diff --git a/README.md b/README.md index d5ea2f242..14154da40 100644 --- a/README.md +++ b/README.md @@ -228,6 +228,22 @@ Priority order: +
+Complex Resource Limits Configurations + +Per-resource limits for folder-as-resource complex resources (e.g. skills), enforced on whole-resource +PUT and single-file PUT. Exceeding a limit rejects the request and leaves nothing observable written. +These limits also bound listing reads, archive (ZIP download) size and how long a mutation holds the +resource lock (a single-file PUT copies the whole current version under the lock). + +| Setting | Default | Required | Description | +|-------------------------------------|:--------:|:--------:|--------------------------------------------------------------------| +| complexResource.maxFiles | 100 | No | Max number of files a complex resource may contain. | +| complexResource.maxTotalBytes | 16777216 | No | Max aggregate size in bytes of all files in a complex resource. | +| complexResource.maxFileSizeBytes | 1048576 | No | Max size in bytes of a single file within a complex resource. | + +
+
Complex Resource Sweep Configurations diff --git a/server/src/main/java/com/epam/aidial/core/server/AiDial.java b/server/src/main/java/com/epam/aidial/core/server/AiDial.java index 3cebfa700..04652669a 100644 --- a/server/src/main/java/com/epam/aidial/core/server/AiDial.java +++ b/server/src/main/java/com/epam/aidial/core/server/AiDial.java @@ -311,8 +311,10 @@ vertx, settings("config"), null, ResponseMappingService responseMappingService = new ResponseMappingService(vertx, generator, resourceService); responseMappingService.init(taskExecutor); + ComplexResourceService.Settings complexResourceSettings = Json.decodeValue( + settings("complexResource").toBuffer(), ComplexResourceService.Settings.class); ComplexResourceService complexResourceService = new ComplexResourceService( - resourceService, lockService, shareService, invitationService, storage); + resourceService, lockService, shareService, invitationService, storage, complexResourceSettings); ComplexResourceSweepService.Settings complexResourceSweepSettings = Json.decodeValue( settings("complexResourceSweep").toBuffer(), ComplexResourceSweepService.Settings.class); complexResourceSweepService = new ComplexResourceSweepService(timerService, storage, redis, lockService, diff --git a/server/src/main/java/com/epam/aidial/core/server/data/folder/FolderResourceMarker.java b/server/src/main/java/com/epam/aidial/core/server/data/folder/FolderResourceMarker.java index 687cb7288..ff5fedc10 100644 --- a/server/src/main/java/com/epam/aidial/core/server/data/folder/FolderResourceMarker.java +++ b/server/src/main/java/com/epam/aidial/core/server/data/folder/FolderResourceMarker.java @@ -1,6 +1,7 @@ package com.epam.aidial.core.server.data.folder; import com.fasterxml.jackson.annotation.JsonInclude; +import lombok.AllArgsConstructor; import lombok.Data; import lombok.NoArgsConstructor; @@ -40,10 +41,12 @@ public class FolderResourceMarker { */ private String etag; /** - * Per-file etags of the current version (relative path -> etag). The aggregate {@link #etag} is derived - * from these, so single-file mutations can recompute it without reading file content. + * Per-file metadata of the current version (relative path -> size/etag). The aggregate {@link #etag} is + * derived from the per-file etags, so single-file mutations can recompute it without reading file + * content; the sizes let {@code maxTotalBytes}/{@code maxFiles} be enforced on a single-file mutation + * without re-reading every existing file's content or metadata. */ - private Map files; + private Map fileMetadata; private Long createdAt; private Long updatedAt; private Long deletedAt; @@ -52,4 +55,15 @@ public class FolderResourceMarker { * Type-specific metadata extracted by the per-type handler (e.g. skill name/description/version). */ private Map metadata; + + /** + * Per-file size and etag, keyed by relative path in {@link #fileMetadata}. + */ + @Data + @NoArgsConstructor + @AllArgsConstructor + public static class ResourceFileMetadata { + private long size; + private String etag; + } } diff --git a/server/src/main/java/com/epam/aidial/core/server/service/resource/ComplexResourceService.java b/server/src/main/java/com/epam/aidial/core/server/service/resource/ComplexResourceService.java index 16fba3833..5a580c39b 100644 --- a/server/src/main/java/com/epam/aidial/core/server/service/resource/ComplexResourceService.java +++ b/server/src/main/java/com/epam/aidial/core/server/service/resource/ComplexResourceService.java @@ -19,7 +19,9 @@ import com.epam.aidial.core.storage.service.ResourceService; import com.epam.aidial.core.storage.util.EtagBuilder; import com.epam.aidial.core.storage.util.EtagHeader; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import io.vertx.core.buffer.Buffer; +import lombok.Data; import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; @@ -81,15 +83,30 @@ public class ComplexResourceService { private final ShareService shareService; private final InvitationService invitationService; private final BlobStorage blobStorage; + private final Settings settings; public ComplexResourceService(ResourceService resourceService, LockService lockService, ShareService shareService, InvitationService invitationService, - BlobStorage blobStorage) { + BlobStorage blobStorage, Settings settings) { this.resourceService = resourceService; this.lockService = lockService; this.shareService = shareService; this.invitationService = invitationService; this.blobStorage = blobStorage; + this.settings = settings; + } + + /** + * Per-resource limits for a folder resource, bounding abuse (upload size/count), listing reads, archive + * size and lock-hold time (a single-file mutation copies the whole current version under the resource + * lock, so a smaller resource copies - and holds the lock - for less time). + */ + @Data + @JsonIgnoreProperties(ignoreUnknown = true) + public static class Settings { + private int maxFiles = 100; + private long maxTotalBytes = 16L * 1024 * 1024; + private long maxFileSizeBytes = 1024 * 1024; } /** @@ -118,6 +135,9 @@ public String put(ResourceDescriptor resource, ComplexResourceHandler handler, files.put(normalizePath(entry.getKey()), entry.getValue().getBytes()); } + // Enforced before any write (and before the handler reads file content), so a rejection never + // takes the bucket/resource lock and leaves nothing observable. + validateLimits(files); handler.validate(files); // Creating a resource is a structural change: the bucket lock serializes the ancestor walk @@ -134,14 +154,15 @@ public String put(ResourceDescriptor resource, ComplexResourceHandler handler, etag.validate(readAggregateEtag(marker)); try { - Map fileEtags = new TreeMap<>(); + Map fileMetadata = new TreeMap<>(); for (Map.Entry entry : files.entrySet()) { String relativePath = entry.getKey(); byte[] content = entry.getValue(); ResourceDescriptor file = versionFile(resource, versionId, relativePath); FileMetadata written = resourceService.putFile(file, content, EtagHeader.ANY, BlobStorageUtil.getContentType(relativePath), author); - fileEtags.put(relativePath, written.getEtag()); + fileMetadata.put(relativePath, + new FolderResourceMarker.ResourceFileMetadata(content.length, written.getEtag())); } FolderResourceMarker existing = readMarker(marker, false); if (existing == null) { @@ -149,8 +170,8 @@ public String put(ResourceDescriptor resource, ComplexResourceHandler handler, // its reference, which is keyed by URL, not by version. writeReference(resource); } - return commitMarker(marker, resource, versionId, aggregateEtag(fileEtags), - handler.buildMarkerMetadata(files), fileEtags, existing, author); + return commitMarker(marker, resource, versionId, aggregateEtag(fileMetadata), + handler.buildMarkerMetadata(files), fileMetadata, existing, author); } catch (Exception e) { // Nothing is observable until the marker is committed; drop the orphan version files. deleteVersion(versionFolder(resource, versionId)); @@ -160,6 +181,33 @@ public String put(ResourceDescriptor resource, ComplexResourceHandler handler, }); } + /** + * Enforces {@code maxFiles}/{@code maxFileSizeBytes}/{@code maxTotalBytes} against a whole file set. + * + * @throws HttpException BAD_REQUEST if the file count exceeds {@code maxFiles}, or + * REQUEST_ENTITY_TOO_LARGE if any single file or the aggregate size exceeds + * {@code maxFileSizeBytes}/{@code maxTotalBytes} + */ + private void validateLimits(Map files) { + if (files.size() > settings.getMaxFiles()) { + throw new HttpException(HttpStatus.BAD_REQUEST, + "Resource contains %d files which exceeds the limit of %d".formatted(files.size(), settings.getMaxFiles())); + } + long totalBytes = 0; + for (Map.Entry entry : files.entrySet()) { + long size = entry.getValue().length; + if (size > settings.getMaxFileSizeBytes()) { + throw new HttpException(HttpStatus.REQUEST_ENTITY_TOO_LARGE, + "File '%s' size %d exceeds max file size of %d".formatted(entry.getKey(), size, settings.getMaxFileSizeBytes())); + } + totalBytes += size; + } + if (totalBytes > settings.getMaxTotalBytes()) { + throw new HttpException(HttpStatus.REQUEST_ENTITY_TOO_LARGE, + "Resource size %d exceeds max total size of %d".formatted(totalBytes, settings.getMaxTotalBytes())); + } + } + /** * Streams a single file of the current version, or returns {@code null} if the resource is absent or * in the {@code deleting} state, or the file does not exist in the current version (both surface as 404). @@ -179,18 +227,43 @@ public ResourceService.ResourceStream getFileStream(ResourceDescriptor resource, * Adds or replaces a single file via copy-on-write of a new version: the current version is copied * server-side to a fresh {@code v/{versionId}/} prefix, the single file is written, then the marker is * committed under the folder-scoped lock. Returns the new aggregate etag. + * + * @throws HttpException REQUEST_ENTITY_TOO_LARGE if the file or the resulting aggregate size exceeds + * {@code maxFileSizeBytes}/{@code maxTotalBytes}, or BAD_REQUEST if adding a new file + * would exceed {@code maxFiles} */ public String putFile(ResourceDescriptor resource, ComplexResourceHandler handler, String relativePath, byte[] content, EtagHeader etag, String author) { String normalized = normalizePath(relativePath); + if (content.length > settings.getMaxFileSizeBytes()) { + throw new HttpException(HttpStatus.REQUEST_ENTITY_TOO_LARGE, + "File size %d exceeds max file size of %d".formatted(content.length, settings.getMaxFileSizeBytes())); + } handler.validateFileMutation(normalized, ComplexResourceHandler.FileMutation.PUT); - return copyOnWrite(resource, etag, author, (newVersionId, fileEtags) -> { - FileMetadata written = resourceService.putFile(versionFile(resource, newVersionId, normalized), - content, EtagHeader.ANY, BlobStorageUtil.getContentType(normalized), author); - fileEtags.put(normalized, written.getEtag()); - // Only a change to a metadata-bearing file (e.g. the manifest) refreshes the marker metadata. - return handler.refreshMetadataOnPut(normalized, content); - }); + return copyOnWrite(resource, etag, author, + fileMetadata -> { + boolean newFile = !fileMetadata.containsKey(normalized); + if (newFile && fileMetadata.size() + 1 > settings.getMaxFiles()) { + throw new HttpException(HttpStatus.BAD_REQUEST, + "Resource would contain more than %d files".formatted(settings.getMaxFiles())); + } + long existingTotal = fileMetadata.values().stream() + .mapToLong(FolderResourceMarker.ResourceFileMetadata::getSize).sum(); + long oldSize = newFile ? 0L : fileMetadata.get(normalized).getSize(); + long prospectiveTotal = existingTotal - oldSize + content.length; + if (prospectiveTotal > settings.getMaxTotalBytes()) { + throw new HttpException(HttpStatus.REQUEST_ENTITY_TOO_LARGE, + "Resource size %d exceeds max total size of %d".formatted(prospectiveTotal, settings.getMaxTotalBytes())); + } + }, + (newVersionId, fileMetadata) -> { + FileMetadata written = resourceService.putFile(versionFile(resource, newVersionId, normalized), + content, EtagHeader.ANY, BlobStorageUtil.getContentType(normalized), author); + fileMetadata.put(normalized, + new FolderResourceMarker.ResourceFileMetadata(content.length, written.getEtag())); + // Only a change to a metadata-bearing file (e.g. the manifest) refreshes the marker metadata. + return handler.refreshMetadataOnPut(normalized, content); + }); } /** @@ -202,26 +275,33 @@ public String deleteFile(ResourceDescriptor resource, ComplexResourceHandler han EtagHeader etag, String author) { String normalized = normalizePath(relativePath); handler.validateFileMutation(normalized, ComplexResourceHandler.FileMutation.DELETE); - return copyOnWrite(resource, etag, author, (newVersionId, fileEtags) -> { - if (fileEtags.remove(normalized) == null) { - throw new HttpException(HttpStatus.NOT_FOUND, "File not found: " + normalized); - } - resourceService.deleteResource(versionFile(resource, newVersionId, normalized), EtagHeader.ANY); - // A deletable file never carries marker metadata (the manifest cannot be deleted). - return null; - }); + return copyOnWrite(resource, etag, author, + fileMetadata -> { }, + (newVersionId, fileMetadata) -> { + if (fileMetadata.remove(normalized) == null) { + throw new HttpException(HttpStatus.NOT_FOUND, "File not found: " + normalized); + } + resourceService.deleteResource(versionFile(resource, newVersionId, normalized), EtagHeader.ANY); + // A deletable file never carries marker metadata (the manifest cannot be deleted). + return null; + }); } /** * Shared copy-on-write engine for single-file mutations. Under the folder-scoped lock it validates the - * {@code If-Match} precondition on the aggregate etag, copies the current version server-side to a fresh - * one, applies the single change to the new version and commits the marker. The aggregate etag is - * recomputed from the per-file etags carried in the marker, so no file content is read. + * {@code If-Match} precondition on the aggregate etag, checks the prospective limits (before copying + * anything, so a rejection never pays for the copy or extends the lock hold time), copies the current + * version server-side to a fresh one, applies the single change to the new version and commits the + * marker. The aggregate etag is recomputed from the per-file etags carried in the marker, so no file + * content is read. * - * @param mutation applies the change to the fresh version and to the per-file etag map, returning the - * refreshed marker metadata, or {@code null} to keep the existing metadata + * @param preflight validates the prospective change against the current per-file metadata map before + * any copy happens; throws {@link HttpException} to reject + * @param mutation applies the change to the fresh version and to the per-file metadata map, returning + * the refreshed marker metadata, or {@code null} to keep the existing metadata */ - private String copyOnWrite(ResourceDescriptor resource, EtagHeader etag, String author, VersionMutation mutation) { + private String copyOnWrite(ResourceDescriptor resource, EtagHeader etag, String author, + LimitsPreflight preflight, VersionMutation mutation) { ResourceDescriptor marker = markerDescriptor(resource); String newVersionId = UUID.randomUUID().toString().replace("-", ""); @@ -232,15 +312,16 @@ private String copyOnWrite(ResourceDescriptor resource, EtagHeader etag, String } etag.validate(current.getEtag()); - Map fileEtags = currentFileEtags(resource, current); + Map fileMetadata = currentFileMetadata(resource, current); + preflight.check(fileMetadata); try { resourceService.copyFolder(versionFolder(resource, current.getCurrentVersion()), versionFolder(resource, newVersionId), false); - Map refreshed = mutation.apply(newVersionId, fileEtags); + Map refreshed = mutation.apply(newVersionId, fileMetadata); Map metadata = refreshed != null ? refreshed : current.getMetadata(); - String aggregateEtag = commitMarker(marker, resource, newVersionId, aggregateEtag(fileEtags), - metadata, fileEtags, current, author); + String aggregateEtag = commitMarker(marker, resource, newVersionId, aggregateEtag(fileMetadata), + metadata, fileMetadata, current, author); deleteVersion(versionFolder(resource, current.getCurrentVersion())); return aggregateEtag; } catch (Exception e) { @@ -251,33 +332,42 @@ private String copyOnWrite(ResourceDescriptor resource, EtagHeader etag, String } } + @FunctionalInterface + private interface LimitsPreflight { + void check(Map fileMetadata); + } + @FunctionalInterface private interface VersionMutation { @Nullable - Map apply(String newVersionId, Map fileEtags); + Map apply(String newVersionId, Map fileMetadata); } /** - * Per-file etag map of the current version. Reconstructed from per-file metadata for legacy markers that - * predate the stored map (still no file content is read). + * Per-file metadata map of the current version. Reconstructed from per-file blob metadata for legacy + * markers that predate the stored map (still no file content is read). */ - private Map currentFileEtags(ResourceDescriptor resource, FolderResourceMarker current) { - if (current.getFiles() != null) { - return new TreeMap<>(current.getFiles()); + private Map currentFileMetadata( + ResourceDescriptor resource, FolderResourceMarker current) { + if (current.getFileMetadata() != null) { + return new TreeMap<>(current.getFileMetadata()); } - Map fileEtags = new TreeMap<>(); + Map fileMetadata = new TreeMap<>(); ResourceDescriptor versionFolder = versionFolder(resource, current.getCurrentVersion()); for (ResourceDescriptor file : listVersionFiles(versionFolder)) { ResourceItemMetadata meta = resourceService.getResourceMetadata(file); if (meta != null) { - fileEtags.put(versionFolder.getRelativePath(file), meta.getEtag()); + String relativePath = versionFolder.getRelativePath(file); + long size = meta instanceof FileMetadata fileMeta ? fileMeta.getContentLength() : 0L; + fileMetadata.put(relativePath, new FolderResourceMarker.ResourceFileMetadata(size, meta.getEtag())); } } - return fileEtags; + return fileMetadata; } private String commitMarker(ResourceDescriptor marker, ResourceDescriptor resource, String versionId, - String aggregateEtag, Map metadata, Map fileEtags, + String aggregateEtag, Map metadata, + Map fileMetadata, @Nullable FolderResourceMarker existing, String author) { long now = System.currentTimeMillis(); FolderResourceMarker document = new FolderResourceMarker(); @@ -286,7 +376,7 @@ private String commitMarker(ResourceDescriptor marker, ResourceDescriptor resour document.setState(STATE_ACTIVE); document.setCurrentVersion(versionId); document.setEtag(aggregateEtag); - document.setFiles(fileEtags); + document.setFileMetadata(fileMetadata); document.setCreatedAt(existing == null ? now : existing.getCreatedAt()); document.setUpdatedAt(now); document.setAuthor(author); @@ -313,11 +403,11 @@ private void writeReference(ResourceDescriptor resource) { * Aggregate etag derived from the per-file etags (relative path -> etag), sorted by path so the result * is deterministic regardless of insertion order. */ - private static String aggregateEtag(Map fileEtags) { + private static String aggregateEtag(Map fileMetadata) { EtagBuilder builder = new EtagBuilder(); - for (Map.Entry entry : new TreeMap<>(fileEtags).entrySet()) { + for (Map.Entry entry : new TreeMap<>(fileMetadata).entrySet()) { builder.append(entry.getKey().getBytes(StandardCharsets.UTF_8)) - .append(entry.getValue().getBytes(StandardCharsets.UTF_8)); + .append(entry.getValue().getEtag().getBytes(StandardCharsets.UTF_8)); } return builder.build(); } diff --git a/server/src/main/resources/aidial.settings.json b/server/src/main/resources/aidial.settings.json index 5b509d421..7713ac6da 100644 --- a/server/src/main/resources/aidial.settings.json +++ b/server/src/main/resources/aidial.settings.json @@ -65,6 +65,11 @@ "compressionMinSize": 256, "heartbeatPeriod": 60000 }, + "complexResource": { + "maxFiles": 100, + "maxTotalBytes": 16777216, + "maxFileSizeBytes": 1048576 + }, "complexResourceSweep": { "period": 10000, "batch": 1000, diff --git a/server/src/test/java/com/epam/aidial/core/server/SkillResourceApiTest.java b/server/src/test/java/com/epam/aidial/core/server/SkillResourceApiTest.java index ce8e76a47..b4a22b3e1 100644 --- a/server/src/test/java/com/epam/aidial/core/server/SkillResourceApiTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/SkillResourceApiTest.java @@ -4,6 +4,7 @@ import com.epam.aidial.core.server.util.ProxyUtil; import com.fasterxml.jackson.databind.JsonNode; import io.vertx.core.http.HttpMethod; +import io.vertx.core.json.JsonObject; import lombok.SneakyThrows; import org.apache.hc.client5.http.classic.methods.HttpUriRequest; import org.apache.hc.client5.http.classic.methods.HttpUriRequestBase; @@ -720,4 +721,207 @@ private static Map unzip(byte[] archive) { private record BinaryResponse(int status, byte[] body, Map headers) { } + + /** + * Boundary tests for the per-resource {@code maxFiles}/{@code maxTotalBytes}/{@code maxFileSizeBytes} limits, + * run against small overridden limits so the boundaries are cheap to exercise. + */ + public static class LimitsApiTest extends ResourceBaseTest { + + private static final String MANIFEST = """ + --- + name: My Skill + description: Does something useful + --- + """; + + @Override + protected JsonObject additionalSettingsOverrides() { + return new JsonObject(""" + { + "complexResource": { + "maxFiles": 3, + "maxTotalBytes": 200, + "maxFileSizeBytes": 100 + } + } + """); + } + + @Test + void testWholeResourcePutRejectsExceedingMaxFiles() { + Map files = new LinkedHashMap<>(); + files.put("SKILL.md", MANIFEST.getBytes(StandardCharsets.UTF_8)); + files.put("a.txt", "a".getBytes(StandardCharsets.UTF_8)); + files.put("b.txt", "b".getBytes(StandardCharsets.UTF_8)); + files.put("c.txt", "c".getBytes(StandardCharsets.UTF_8)); + + verify(uploadSkill("/too-many-files", files), 400); + assertEquals(404, downloadSkill("/too-many-files").status()); + } + + @Test + void testWholeResourcePutAcceptsExactlyMaxFiles() { + Map files = new LinkedHashMap<>(); + files.put("SKILL.md", MANIFEST.getBytes(StandardCharsets.UTF_8)); + files.put("a.txt", "a".getBytes(StandardCharsets.UTF_8)); + files.put("b.txt", "b".getBytes(StandardCharsets.UTF_8)); + + verify(uploadSkill("/exact-max-files", files), 200); + } + + @Test + void testWholeResourcePutRejectsExceedingMaxTotalBytes() { + Map files = new LinkedHashMap<>(); + files.put("SKILL.md", MANIFEST.getBytes(StandardCharsets.UTF_8)); + files.put("big.txt", "x".repeat(200).getBytes(StandardCharsets.UTF_8)); + + Response put = uploadSkill("/too-big-total", files); + verify(put, 413); + assertEquals(404, downloadSkill("/too-big-total").status()); + } + + @Test + void testWholeResourcePutRejectsExceedingMaxFileSize() { + Map files = new LinkedHashMap<>(); + files.put("SKILL.md", MANIFEST.getBytes(StandardCharsets.UTF_8)); + files.put("huge.txt", "x".repeat(150).getBytes(StandardCharsets.UTF_8)); + + Response put = uploadSkill("/too-big-file", files); + verify(put, 413); + assertEquals(404, downloadSkill("/too-big-file").status()); + } + + @Test + void testSingleFilePutRejectsExceedingMaxFileSize() { + Map files = Map.of("SKILL.md", MANIFEST.getBytes(StandardCharsets.UTF_8)); + verify(uploadSkill("/file-too-big", files), 200); + + verify(putSkillFile("/file-too-big", "big.txt", "x".repeat(150).getBytes(StandardCharsets.UTF_8)), 413); + // rejected mutation left nothing observable: the file was never added + assertEquals(404, getSkillFile("/file-too-big", "big.txt").status()); + } + + @Test + void testSingleFilePutRejectsExceedingMaxTotalBytes() { + Map files = Map.of("SKILL.md", MANIFEST.getBytes(StandardCharsets.UTF_8)); + Response created = uploadSkill("/file-total-too-big", files); + verify(created, 200); + String before = created.headers().get("etag"); + + // SKILL.md is small; adding a near-limit file tips the aggregate past maxTotalBytes(200) + verify(putSkillFile("/file-total-too-big", "big.txt", "x".repeat(190).getBytes(StandardCharsets.UTF_8)), 413); + + // rejected mutation must not have changed the resource + assertEquals(before, downloadSkill("/file-total-too-big").headers().get("etag")); + } + + @Test + void testSingleFilePutRejectsExceedingMaxFiles() { + Map files = new LinkedHashMap<>(); + files.put("SKILL.md", MANIFEST.getBytes(StandardCharsets.UTF_8)); + files.put("a.txt", "a".getBytes(StandardCharsets.UTF_8)); + files.put("b.txt", "b".getBytes(StandardCharsets.UTF_8)); + verify(uploadSkill("/file-count-too-big", files), 200); + + // adding a 4th distinct file exceeds maxFiles(3) + verify(putSkillFile("/file-count-too-big", "c.txt", "c".getBytes(StandardCharsets.UTF_8)), 400); + + // replacing an existing file is fine (file count unchanged) + verify(putSkillFile("/file-count-too-big", "a.txt", "aa".getBytes(StandardCharsets.UTF_8)), 200); + } + + @SneakyThrows + private Response uploadSkill(String skillPath, Map files, String... headers) { + String uri = "http://127.0.0.1:" + serverPort + "/v2/skills/" + bucket + skillPath; + HttpUriRequest request = new HttpUriRequestBase(HttpMethod.PUT.name(), URI.create(uri)); + + for (int i = 0; i < headers.length; i += 2) { + request.setHeader(headers[i], headers[i + 1]); + } + if (!request.containsHeader("authorization") && !request.containsHeader("api-key")) { + request.setHeader("api-key", "proxyKey1"); + } + + MultipartEntityBuilder builder = MultipartEntityBuilder.create(); + builder.setMode(HttpMultipartMode.LEGACY); + builder.setCharset(StandardCharsets.UTF_8); + for (Map.Entry entry : files.entrySet()) { + builder.addBinaryBody(entry.getKey(), entry.getValue(), ContentType.DEFAULT_BINARY, entry.getKey()); + } + request.setEntity(builder.build()); + + return client.execute(request, ResourceBaseTest::toResponse); + } + + @SneakyThrows + private BinaryResponse downloadSkill(String skillPath, String... headers) { + String uri = "http://127.0.0.1:" + serverPort + "/v2/skills/" + bucket + skillPath; + HttpUriRequest request = new HttpUriRequestBase(HttpMethod.GET.name(), URI.create(uri)); + + for (int i = 0; i < headers.length; i += 2) { + request.setHeader(headers[i], headers[i + 1]); + } + if (!request.containsHeader("authorization") && !request.containsHeader("api-key")) { + request.setHeader("api-key", "proxyKey1"); + } + + return client.execute(request, response -> { + int status = response.getCode(); + byte[] body = response.getEntity() == null ? new byte[0] : EntityUtils.toByteArray(response.getEntity()); + Map responseHeaders = new HashMap<>(); + for (Header header : response.getHeaders()) { + responseHeaders.put(header.getName(), header.getValue()); + } + return new BinaryResponse(status, body, responseHeaders); + }); + } + + @SneakyThrows + private Response putSkillFile(String skillPath, String filePath, byte[] content, String... headers) { + String uri = "http://127.0.0.1:" + serverPort + "/v2/skills/" + bucket + skillPath + "/files/" + filePath; + HttpUriRequest request = new HttpUriRequestBase(HttpMethod.PUT.name(), URI.create(uri)); + + for (int i = 0; i < headers.length; i += 2) { + request.setHeader(headers[i], headers[i + 1]); + } + if (!request.containsHeader("authorization") && !request.containsHeader("api-key")) { + request.setHeader("api-key", "proxyKey1"); + } + + MultipartEntityBuilder builder = MultipartEntityBuilder.create(); + builder.setMode(HttpMultipartMode.LEGACY); + builder.setCharset(StandardCharsets.UTF_8); + builder.addBinaryBody("file", content, ContentType.DEFAULT_BINARY, "file"); + request.setEntity(builder.build()); + + return client.execute(request, ResourceBaseTest::toResponse); + } + + @SneakyThrows + private BinaryResponse getSkillFile(String skillPath, String filePath, String... headers) { + String uri = "http://127.0.0.1:" + serverPort + "/v2/skills/" + bucket + skillPath + "/files/" + filePath; + HttpUriRequest request = new HttpUriRequestBase(HttpMethod.GET.name(), URI.create(uri)); + + for (int i = 0; i < headers.length; i += 2) { + request.setHeader(headers[i], headers[i + 1]); + } + if (!request.containsHeader("authorization") && !request.containsHeader("api-key")) { + request.setHeader("api-key", "proxyKey1"); + } + + return client.execute(request, response -> { + int status = response.getCode(); + byte[] body = response.getEntity() == null ? new byte[0] : EntityUtils.toByteArray(response.getEntity()); + Map responseHeaders = new HashMap<>(); + for (Header header : response.getHeaders()) { + responseHeaders.put(header.getName(), header.getValue()); + } + return new BinaryResponse(status, body, responseHeaders); + }); + } + + private record BinaryResponse(int status, byte[] body, Map headers) { + } + } } diff --git a/server/src/test/java/com/epam/aidial/core/server/service/resource/ComplexResourceServiceTest.java b/server/src/test/java/com/epam/aidial/core/server/service/resource/ComplexResourceServiceTest.java index 883d06a9a..de53bb77e 100644 --- a/server/src/test/java/com/epam/aidial/core/server/service/resource/ComplexResourceServiceTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/service/resource/ComplexResourceServiceTest.java @@ -7,6 +7,8 @@ import com.epam.aidial.core.storage.blobstore.BlobStorage; import com.epam.aidial.core.storage.blobstore.Storage; import com.epam.aidial.core.storage.data.ResourceItemMetadata; +import com.epam.aidial.core.storage.http.HttpException; +import com.epam.aidial.core.storage.http.HttpStatus; import com.epam.aidial.core.storage.resource.ResourceDescriptor; import com.epam.aidial.core.storage.resource.ResourceTypes; import com.epam.aidial.core.storage.service.LockService; @@ -37,6 +39,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; class ComplexResourceServiceTest { @@ -110,8 +113,12 @@ void init() throws IOException { ShareService shareService = Mockito.mock(ShareService.class); InvitationService invitationService = Mockito.mock(InvitationService.class); + ComplexResourceService.Settings complexResourceSettings = new ComplexResourceService.Settings(); + complexResourceSettings.setMaxFiles(3); + complexResourceSettings.setMaxTotalBytes(200); + complexResourceSettings.setMaxFileSizeBytes(100); complexResourceService = new ComplexResourceService( - resourceService, lockService, shareService, invitationService, blobStorage); + resourceService, lockService, shareService, invitationService, blobStorage, complexResourceSettings); } catch (Throwable e) { destroy(); throw e; @@ -257,4 +264,104 @@ void testGcObsoleteVersionsNoOpWhenDeleting() { assertFalse(complexResourceService.gcObsoleteVersions(resource, 0L)); } + + @Test + void testPutRejectsExceedingMaxFiles() { + ResourceDescriptor resource = skill("tooManyFiles"); + Map uploads = Map.of( + "SKILL.md", Buffer.buffer(MANIFEST), + "a.txt", Buffer.buffer("a"), + "b.txt", Buffer.buffer("b"), + "c.txt", Buffer.buffer("c")); + + assertThrowsBadRequest(() -> complexResourceService.put(resource, handler, uploads, EtagHeader.ANY, "user1")); + assertNull(complexResourceService.readMarkerForSweep(resource)); + } + + @Test + void testPutAcceptsExactlyMaxFiles() { + ResourceDescriptor resource = skill("exactMaxFiles"); + Map uploads = Map.of( + "SKILL.md", Buffer.buffer(MANIFEST), + "a.txt", Buffer.buffer("a"), + "b.txt", Buffer.buffer("b")); + + complexResourceService.put(resource, handler, uploads, EtagHeader.ANY, "user1"); + assertNotNull(complexResourceService.readMarkerForSweep(resource)); + } + + @Test + void testPutRejectsExceedingMaxTotalBytes() { + ResourceDescriptor resource = skill("tooBigTotal"); + Map uploads = Map.of( + "SKILL.md", Buffer.buffer(MANIFEST), + "big.txt", Buffer.buffer("x".repeat(200))); + + HttpException e = assertThrows(HttpException.class, + () -> complexResourceService.put(resource, handler, uploads, EtagHeader.ANY, "user1")); + assertEquals(HttpStatus.REQUEST_ENTITY_TOO_LARGE, e.getStatus()); + assertNull(complexResourceService.readMarkerForSweep(resource)); + } + + @Test + void testPutRejectsExceedingMaxFileSize() { + ResourceDescriptor resource = skill("tooBigFile"); + Map uploads = Map.of( + "SKILL.md", Buffer.buffer(MANIFEST), + "huge.txt", Buffer.buffer("x".repeat(150))); + + HttpException e = assertThrows(HttpException.class, + () -> complexResourceService.put(resource, handler, uploads, EtagHeader.ANY, "user1")); + assertEquals(HttpStatus.REQUEST_ENTITY_TOO_LARGE, e.getStatus()); + assertNull(complexResourceService.readMarkerForSweep(resource)); + } + + @Test + void testPutFileRejectsExceedingMaxFileSize() { + ResourceDescriptor resource = skill("fileTooBig"); + putSkill(resource, EtagHeader.ANY); + + HttpException e = assertThrows(HttpException.class, () -> complexResourceService.putFile( + resource, handler, "big.txt", "x".repeat(150).getBytes(), EtagHeader.ANY, "user1")); + assertEquals(HttpStatus.REQUEST_ENTITY_TOO_LARGE, e.getStatus()); + } + + @Test + void testPutFileRejectsExceedingMaxTotalBytes() { + ResourceDescriptor resource = skill("fileTotalTooBig"); + putSkill(resource, EtagHeader.ANY); + FolderResourceMarker before = complexResourceService.readMarkerForSweep(resource); + + // SKILL.md (68 bytes) + a 190-byte file tips the aggregate past maxTotalBytes(200). + HttpException e = assertThrows(HttpException.class, () -> complexResourceService.putFile( + resource, handler, "big.txt", "x".repeat(190).getBytes(), EtagHeader.ANY, "user1")); + assertEquals(HttpStatus.REQUEST_ENTITY_TOO_LARGE, e.getStatus()); + + // rejected mutation must not have changed the resource + FolderResourceMarker after = complexResourceService.readMarkerForSweep(resource); + assertEquals(before.getEtag(), after.getEtag()); + } + + @Test + void testPutFileRejectsExceedingMaxFiles() { + ResourceDescriptor resource = skill("fileCountTooBig"); + Map uploads = Map.of( + "SKILL.md", Buffer.buffer(MANIFEST), + "a.txt", Buffer.buffer("a"), + "b.txt", Buffer.buffer("b")); + complexResourceService.put(resource, handler, uploads, EtagHeader.ANY, "user1"); + + // adding a 4th distinct file exceeds maxFiles(3) + HttpException e = assertThrows(HttpException.class, () -> complexResourceService.putFile( + resource, handler, "c.txt", "c".getBytes(), EtagHeader.ANY, "user1")); + assertEquals(HttpStatus.BAD_REQUEST, e.getStatus()); + + // replacing an existing file is fine (file count unchanged) + complexResourceService.putFile(resource, handler, "a.txt", "aa".getBytes(), EtagHeader.ANY, "user1"); + } + + private static void assertThrowsBadRequest(Runnable action) { + HttpException e = assertThrows(HttpException.class, action::run); + assertEquals(HttpStatus.BAD_REQUEST, e.getStatus()); + } } diff --git a/server/src/test/java/com/epam/aidial/core/server/service/resource/ComplexResourceSweepServiceTest.java b/server/src/test/java/com/epam/aidial/core/server/service/resource/ComplexResourceSweepServiceTest.java index e114b19cb..f7fc3bc60 100644 --- a/server/src/test/java/com/epam/aidial/core/server/service/resource/ComplexResourceSweepServiceTest.java +++ b/server/src/test/java/com/epam/aidial/core/server/service/resource/ComplexResourceSweepServiceTest.java @@ -128,7 +128,7 @@ void init() throws IOException { ShareService shareService = Mockito.mock(ShareService.class); InvitationService invitationService = Mockito.mock(InvitationService.class); complexResourceService = new ComplexResourceService( - resourceService, lockService, shareService, invitationService, blobStorage); + resourceService, lockService, shareService, invitationService, blobStorage, new ComplexResourceService.Settings()); encryptionService = new EncryptionService(new JsonObject().put("secret", "secret").put("key", "key")); } catch (Throwable e) { destroy(); diff --git a/server/src/test/resources/aidial.settings.json b/server/src/test/resources/aidial.settings.json index 8748a4294..d2a83d8cc 100644 --- a/server/src/test/resources/aidial.settings.json +++ b/server/src/test/resources/aidial.settings.json @@ -61,6 +61,11 @@ "compressionMinSize": 256, "heartbeatPeriod": 60000 }, + "complexResource": { + "maxFiles": 100, + "maxTotalBytes": 16777216, + "maxFileSizeBytes": 1048576 + }, "complexResourceSweep": { "period": 10000, "batch": 1000,