diff --git a/app/common/src/main/java/stirling/software/common/util/GeneralUtils.java b/app/common/src/main/java/stirling/software/common/util/GeneralUtils.java index 1d7320ca84..ea753d27c6 100644 --- a/app/common/src/main/java/stirling/software/common/util/GeneralUtils.java +++ b/app/common/src/main/java/stirling/software/common/util/GeneralUtils.java @@ -870,12 +870,20 @@ public class GeneralUtils { * Internal Implementation Details * *------------------------------------------------------------------------*/ + /** + * Guards the read-modify-write cycle below. Every writer reloads the whole file, edits one key + * and writes it all back, so two unsynchronised writers would silently drop one another's keys. + */ + private final Object SETTINGS_WRITE_LOCK = new Object(); + public void saveKeyToSettings(String key, Object newValue) throws IOException { - String[] keyArray = key.split("\\."); - Path settingsPath = Path.of(InstallationPathConfig.getSettingsPath()); - YamlHelper settingsYaml = new YamlHelper(settingsPath); - settingsYaml.updateValue(Arrays.asList(keyArray), newValue); - settingsYaml.saveOverride(settingsPath); + synchronized (SETTINGS_WRITE_LOCK) { + String[] keyArray = key.split("\\."); + Path settingsPath = Path.of(InstallationPathConfig.getSettingsPath()); + YamlHelper settingsYaml = new YamlHelper(settingsPath); + settingsYaml.updateValue(Arrays.asList(keyArray), newValue); + settingsYaml.saveOverride(settingsPath); + } } /** @@ -893,19 +901,21 @@ public class GeneralUtils { return; } - Path settingsPath = Path.of(InstallationPathConfig.getSettingsPath()); - YamlHelper settingsYaml = new YamlHelper(settingsPath); + synchronized (SETTINGS_WRITE_LOCK) { + Path settingsPath = Path.of(InstallationPathConfig.getSettingsPath()); + YamlHelper settingsYaml = new YamlHelper(settingsPath); - // Apply all updates to the same YamlHelper instance - for (Map.Entry entry : settingsMap.entrySet()) { - String key = entry.getKey(); - Object value = entry.getValue(); - String[] keyArray = key.split("\\."); - settingsYaml.updateValue(Arrays.asList(keyArray), value); + // Apply all updates to the same YamlHelper instance + for (Map.Entry entry : settingsMap.entrySet()) { + String key = entry.getKey(); + Object value = entry.getValue(); + String[] keyArray = key.split("\\."); + settingsYaml.updateValue(Arrays.asList(keyArray), value); + } + + // Save only once after all updates are applied + settingsYaml.saveOverride(settingsPath); } - - // Save only once after all updates are applied - settingsYaml.saveOverride(settingsPath); } /* diff --git a/app/common/src/main/java/stirling/software/common/util/YamlHelper.java b/app/common/src/main/java/stirling/software/common/util/YamlHelper.java index 2715f50dfb..3efb08a4c3 100644 --- a/app/common/src/main/java/stirling/software/common/util/YamlHelper.java +++ b/app/common/src/main/java/stirling/software/common/util/YamlHelper.java @@ -2,8 +2,10 @@ package stirling.software.common.util; import java.io.IOException; import java.io.StringWriter; +import java.nio.file.AtomicMoveNotSupportedException; import java.nio.file.Files; import java.nio.file.Path; +import java.nio.file.StandardCopyOption; import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Arrays; @@ -349,13 +351,36 @@ public class YamlHelper { public MappingNode save(Path saveFilePath) throws IOException { if (!saveFilePath.equals(originalFilePath)) { - Files.writeString(saveFilePath, convertNodeToYaml(getUpdatedRootNode())); + writeAtomically(saveFilePath, convertNodeToYaml(getUpdatedRootNode())); } return (MappingNode) getUpdatedRootNode(); } public void saveOverride(Path saveFilePath) throws IOException { - Files.writeString(saveFilePath, convertNodeToYaml(getUpdatedRootNode())); + writeAtomically(saveFilePath, convertNodeToYaml(getUpdatedRootNode())); + } + + /** + * Write via a sibling temp file and rename. A direct write truncates first, so a crash or a + * full disk part-way through would leave settings.yml half-written and the app unable to boot. + */ + private static void writeAtomically(Path target, String content) throws IOException { + Path dir = target.getParent() != null ? target.getParent() : Path.of("."); + Path tmp = Files.createTempFile(dir, ".yaml-", ".tmp"); + try { + Files.writeString(tmp, content); + try { + Files.move( + tmp, + target, + StandardCopyOption.REPLACE_EXISTING, + StandardCopyOption.ATOMIC_MOVE); + } catch (AtomicMoveNotSupportedException e) { + Files.move(tmp, target, StandardCopyOption.REPLACE_EXISTING); + } + } finally { + Files.deleteIfExists(tmp); + } } /** diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/controller/FormDetectionController.java b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/controller/FormDetectionController.java index b3722facd6..2d6936b4d5 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/controller/FormDetectionController.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/controller/FormDetectionController.java @@ -79,6 +79,10 @@ public class FormDetectionController { boolean applyToPdf) throws IOException { + if (file.isEmpty()) { + return ResponseEntity.badRequest() + .body(Map.of("reason", "INVALID_PDF", "message", "The uploaded file is empty")); + } if (!manager.isReady()) { return ResponseEntity.status(HttpStatus.SERVICE_UNAVAILABLE) .body( @@ -105,47 +109,50 @@ public class FormDetectionController { : spec.getScoreThreshold(); byte[] pdfBytes = file.getBytes(); - List detections = new ArrayList<>(); + // Pages are consumed as they are rendered, so only one page of RGBA is ever live. + List collected = new ArrayList<>(); + List detections; try { - List pages = - rasterizer.rasterize(pdfBytes, spec.getInputSize()); - if (pages.size() > MAX_PAGES) { - return ResponseEntity.badRequest() - .body( - Map.of( - "reason", - "LIMIT", - "message", - "PDF has " - + pages.size() - + " pages; the limit is " - + MAX_PAGES)); - } - for (PageRasterizer.RasterPage page : pages) { - Yolo.Preprocessed pre = - Yolo.preprocess(page.rgba(), page.widthPx(), page.heightPx(), spec); - Map out = detector.infer(pre.chw(), spec.getInputSize()); - for (Yolo.Detection d : decodeFor(spec, out, pre, score)) { - DetectedField.RectPt rect = CoordinateMapper.toPdfPoints(d, page); - if (rect.w() <= 0 || rect.h() <= 0) { - continue; - } - detections.add( - new DetectedField( - fieldType(spec, d.classId()), - page.pageIndex(), - rect, - d.score())); - } - } - if (detections.size() > MAX_FIELDS) { + rasterizer.rasterize( + pdfBytes, + spec.getInputSize(), + MAX_PAGES, + page -> { + Yolo.Preprocessed pre = + Yolo.preprocess(page.rgba(), page.widthPx(), page.heightPx(), spec); + Map out = + detector.infer(pre.chw(), spec.getInputSize()); + for (Yolo.Detection d : decodeFor(spec, out, pre, score)) { + DetectedField.RectPt rect = CoordinateMapper.toPdfPoints(d, page); + if (rect.w() <= 0 || rect.h() <= 0) { + continue; + } + collected.add( + new DetectedField( + fieldType(spec, d.classId()), + page.pageIndex(), + rect, + d.score())); + } + }); + detections = collected; + if (collected.size() > MAX_FIELDS) { log.info( "Capping {} detections to {} highest-confidence fields", - detections.size(), + collected.size(), MAX_FIELDS); - detections.sort((a, b) -> Double.compare(b.confidence(), a.confidence())); - detections = new ArrayList<>(detections.subList(0, MAX_FIELDS)); + collected.sort((a, b) -> Double.compare(b.confidence(), a.confidence())); + detections = new ArrayList<>(collected.subList(0, MAX_FIELDS)); } + } catch (PageRasterizer.PageLimitExceededException e) { + return ResponseEntity.badRequest() + .body(Map.of("reason", "LIMIT", "message", e.getMessage())); + } catch (PageRasterizer.UnreadablePdfException e) { + // The user's file is the problem, not the engine - do not report this as a dependency + // failure, which would send an admin looking for a missing model. + log.debug("Auto Form Detection rejected an unreadable PDF: {}", e.getMessage()); + return ResponseEntity.badRequest() + .body(Map.of("reason", "INVALID_PDF", "message", e.getMessage())); } catch (IllegalStateException e) { // e.g. ONNX Runtime native unavailable for this OS/arch - report unavailable cleanly // rather than a 500. Cannot happen on a normally-built jar (all platforms bundled), but @@ -156,6 +163,8 @@ public class FormDetectionController { } if (applyToPdf) { + // PDFium is more forgiving than PDFBox, so a file can rasterize and still fail to load + // here; that is still the file's fault rather than the server's. try (PDDocument document = pdfDocumentFactory.load(file)) { FormUtils.repairMissingWidgetPageReferences(document); List defs = new ArrayList<>(); @@ -165,6 +174,16 @@ public class FormDetectionController { FormUtils.addFields(document, defs); return WebResponseUtils.pdfDocToWebResponse( document, baseName(file) + ".pdf", tempFileManager); + } catch (IOException e) { + log.debug("Auto Form Detection could not apply fields: {}", e.getMessage()); + return ResponseEntity.badRequest() + .body( + Map.of( + "reason", + "INVALID_PDF", + "message", + "The PDF could not be opened for editing; it may be" + + " corrupt or password-protected")); } } return ResponseEntity.ok(new DetectResponse(detections)); diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/inference/FormDetectionEngine.java b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/inference/FormDetectionEngine.java new file mode 100644 index 0000000000..7cca78105e --- /dev/null +++ b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/inference/FormDetectionEngine.java @@ -0,0 +1,14 @@ +package stirling.software.proprietary.formdetection.inference; + +/** + * Lets the model manager drop an engine's loaded model without knowing what the engine is. + * + *

Deliberately free of any ONNX type: the manager is an unconditional bean, and referencing + * {@code OnnxFormDetector} directly would make Spring introspect it on builds that ship no + * onnxruntime, failing startup for the whole app. + */ +public interface FormDetectionEngine { + + /** Discard any loaded model so the next inference reloads from disk. */ + void unload(); +} diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/inference/OnnxFormDetector.java b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/inference/OnnxFormDetector.java index 30a58c58cb..188cb71b04 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/inference/OnnxFormDetector.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/inference/OnnxFormDetector.java @@ -35,7 +35,7 @@ import ai.onnxruntime.OrtSession; @Service @ConditionalOnClass(name = "ai.onnxruntime.OrtEnvironment") @RequiredArgsConstructor -public class OnnxFormDetector { +public class OnnxFormDetector implements FormDetectionEngine { private final FormDetectionModelManager manager; @@ -103,6 +103,7 @@ public class OnnxFormDetector { } /** Force the next inference to reload from disk (called after install/uninstall). */ + @Override public void unload() { lock.writeLock().lock(); try { diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/render/PageRasterizer.java b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/render/PageRasterizer.java index 6341bf1695..3e4f84b5b9 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/render/PageRasterizer.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/render/PageRasterizer.java @@ -1,14 +1,14 @@ package stirling.software.proprietary.formdetection.render; import java.io.IOException; -import java.util.ArrayList; -import java.util.List; +import java.util.function.Consumer; import org.apache.pdfbox.pdmodel.PDDocument; import org.apache.pdfbox.pdmodel.PDPage; import org.apache.pdfbox.pdmodel.common.PDRectangle; import org.springframework.stereotype.Service; +import lombok.Getter; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; @@ -28,6 +28,10 @@ import stirling.software.jpdfium.model.RenderResult; *

PDFium renders the page as displayed: /Rotate baked in and the crop box anchored at (0,0). The * per-page rotation and crop-box origin needed to map detections back into unrotated user space are * not exposed by JPDFium, so they are read from PDFBox alongside the render. + * + *

Pages are handed to the caller one at a time rather than returned as a list: a rendered page + * is several megabytes of RGBA, so holding a whole document worth of them at once is enough to + * exhaust the heap on a large upload. */ @Slf4j @Service @@ -36,6 +40,26 @@ public class PageRasterizer { private final CustomPDFDocumentFactory pdfDocumentFactory; + /** The PDF cannot be opened or rendered: a bad request, not an engine failure. */ + public static class UnreadablePdfException extends RuntimeException { + public UnreadablePdfException(String message, Throwable cause) { + super(message, cause); + } + } + + /** The document has more pages than the caller allows. Thrown before anything is rendered. */ + @Getter + public static class PageLimitExceededException extends RuntimeException { + private final int pageCount; + private final int limit; + + public PageLimitExceededException(int pageCount, int limit) { + super("PDF has " + pageCount + " pages; the limit is " + limit); + this.pageCount = pageCount; + this.limit = limit; + } + } + /** * A rendered page: RGBA pixels in display space (rotation applied, crop-box origin at 0,0) plus * the geometry needed to map display-space points back to unrotated user space. @@ -65,61 +89,98 @@ public class PageRasterizer { float cropLlxPt, float cropLlyPt) {} - public List rasterize(byte[] pdfBytes, int inputSize) { - List pages = new ArrayList<>(); - try (PdfDocument doc = PdfDocument.open(pdfBytes); + /** + * Render each page in turn and pass it to {@code handler}. Only one page of pixels is reachable + * at a time, so peak memory is one bitmap rather than the whole document. + * + * @param maxPages reject the document if it has more pages than this, before rendering any + * @throws PageLimitExceededException the document exceeds {@code maxPages} + * @throws UnreadablePdfException the PDF is empty, corrupt or password-protected + */ + public void rasterize( + byte[] pdfBytes, int inputSize, int maxPages, Consumer handler) { + if (pdfBytes == null || pdfBytes.length == 0) { + throw new UnreadablePdfException("The uploaded file is empty", null); + } + try (PdfDocument doc = openForRender(pdfBytes); PDDocument boxDoc = openForGeometry(pdfBytes)) { - int count = doc.pageCount(); + int count = pageCount(doc); + // Checked before the loop: rendering first and counting after would let a huge upload + // exhaust the heap on its way to being rejected. + if (count > maxPages) { + throw new PageLimitExceededException(count, maxPages); + } for (int i = 0; i < count; i++) { - try (PdfPage page = doc.page(i)) { - PageSize size = page.size(); - float maxSide = Math.max(size.width(), size.height()); - int dpi = maxSide <= 0 ? 150 : Math.round(72f * inputSize / maxSide); - dpi = Math.max(36, Math.min(dpi, 300)); - RenderResult r = page.renderAt(dpi); - float scaleX = size.width() > 0 ? r.width() / size.width() : dpi / 72f; - float scaleY = size.height() > 0 ? r.height() / size.height() : dpi / 72f; - - int rotation = 0; - float userW = size.width(); - float userH = size.height(); - float llx = 0; - float lly = 0; - if (boxDoc != null && i < boxDoc.getNumberOfPages()) { - PDPage boxPage = boxDoc.getPage(i); - rotation = normalizeRotation(boxPage.getRotation()); - PDRectangle crop = boxPage.getCropBox(); - userW = crop.getWidth(); - userH = crop.getHeight(); - llx = crop.getLowerLeftX(); - lly = crop.getLowerLeftY(); - } else if (boxDoc == null) { - log.warn( - "Page geometry unavailable; assuming unrotated page with origin" - + " (0,0)"); - } - - pages.add( - new RasterPage( - i, - r.rgba(), - r.width(), - r.height(), - size.width(), - size.height(), - scaleX, - scaleY, - rotation, - userW, - userH, - llx, - lly)); - } + // Rendered in a helper so a PDFium failure is classified as bad input, while + // anything the handler throws propagates untouched. + handler.accept(renderPage(doc, boxDoc, i, inputSize)); } } catch (IOException e) { - throw new IllegalStateException("Failed to read PDF geometry", e); + throw new UnreadablePdfException("Failed to read the PDF", e); + } + } + + private RasterPage renderPage(PdfDocument doc, PDDocument boxDoc, int index, int inputSize) { + try (PdfPage page = doc.page(index)) { + PageSize size = page.size(); + float maxSide = Math.max(size.width(), size.height()); + int dpi = maxSide <= 0 ? 150 : Math.round(72f * inputSize / maxSide); + dpi = Math.max(36, Math.min(dpi, 300)); + RenderResult r = page.renderAt(dpi); + float scaleX = size.width() > 0 ? r.width() / size.width() : dpi / 72f; + float scaleY = size.height() > 0 ? r.height() / size.height() : dpi / 72f; + + int rotation = 0; + float userW = size.width(); + float userH = size.height(); + float llx = 0; + float lly = 0; + if (boxDoc != null && index < boxDoc.getNumberOfPages()) { + PDPage boxPage = boxDoc.getPage(index); + rotation = normalizeRotation(boxPage.getRotation()); + PDRectangle crop = boxPage.getCropBox(); + userW = crop.getWidth(); + userH = crop.getHeight(); + llx = crop.getLowerLeftX(); + lly = crop.getLowerLeftY(); + } else if (boxDoc == null) { + log.warn("Page geometry unavailable; assuming unrotated page with origin (0,0)"); + } + + return new RasterPage( + index, + r.rgba(), + r.width(), + r.height(), + size.width(), + size.height(), + scaleX, + scaleY, + rotation, + userW, + userH, + llx, + lly); + } catch (RuntimeException e) { + throw new UnreadablePdfException("Failed to render page " + (index + 1), e); + } + } + + private PdfDocument openForRender(byte[] pdfBytes) { + try { + return PdfDocument.open(pdfBytes); + } catch (Exception e) { + throw new UnreadablePdfException( + "The PDF could not be opened; it may be corrupt or password-protected", e); + } + } + + private int pageCount(PdfDocument doc) { + try { + return doc.pageCount(); + } catch (RuntimeException e) { + throw new UnreadablePdfException("The PDF has no readable page tree", e); } - return pages; } private PDDocument openForGeometry(byte[] pdfBytes) { diff --git a/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/service/FormDetectionModelManager.java b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/service/FormDetectionModelManager.java index 293bc484fe..0b8436198d 100644 --- a/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/service/FormDetectionModelManager.java +++ b/app/proprietary/src/main/java/stirling/software/proprietary/formdetection/service/FormDetectionModelManager.java @@ -25,6 +25,7 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.regex.Pattern; import org.apache.commons.lang3.StringUtils; +import org.springframework.beans.factory.ObjectProvider; import org.springframework.stereotype.Service; import jakarta.annotation.PostConstruct; @@ -38,6 +39,7 @@ import stirling.software.common.configuration.RuntimePathConfig; import stirling.software.common.model.ApplicationProperties; import stirling.software.common.util.GeneralUtils; import stirling.software.proprietary.formdetection.catalog.ModelCatalogService; +import stirling.software.proprietary.formdetection.inference.FormDetectionEngine; import stirling.software.proprietary.formdetection.model.FormDetectionStatus; import stirling.software.proprietary.formdetection.model.ModelCatalogEntry; import stirling.software.proprietary.formdetection.model.ModelStatusResponse; @@ -91,6 +93,12 @@ public class FormDetectionModelManager { private final ApplicationProperties applicationProperties; private final EndpointConfiguration endpointConfiguration; + /** + * Resolved lazily and by interface: the ONNX engine is absent from builds without onnxruntime, + * and an eager or concrete-typed dependency would fail startup there. + */ + private final ObjectProvider engineProvider; + private final AtomicBoolean installing = new AtomicBoolean(false); private volatile FormDetectionStatus state = FormDetectionStatus.NOT_INSTALLED; private volatile int progress = 0; @@ -385,8 +393,13 @@ public class FormDetectionModelManager { } } - /** Mark a verified, on-disk model as the active one and (re)enable the feature. */ - private void activate(String modelId, String expectedSha) { + /** + * Mark a verified, on-disk model as the active one and (re)enable the feature. + * + *

Synchronized because it runs on the install thread and writes settings, which the admin + * setters also do; without the shared monitor a concurrent toggle could drop one of the keys. + */ + private synchronized void activate(String modelId, String expectedSha) { clearTombstone(modelId); applicationProperties.getFormDetection().setActiveModelId(modelId); try { @@ -397,10 +410,20 @@ public class FormDetectionModelManager { activeSha = expectedSha; progress = 100; state = FormDetectionStatus.READY; + invalidateEngine(); applyEndpointState(); log.info("Auto Form Detection model '{}' installed and ready", modelId); } + /** + * Drop any model an engine still holds. Reinstalling the same id leaves the loaded id + * unchanged, so without this the engine would keep serving from the session it opened before + * the swap. + */ + private void invalidateEngine() { + engineProvider.ifAvailable(FormDetectionEngine::unload); + } + /** SHA-256 of an existing model file as lowercase hex, or {@code null} if it cannot be read. */ private String sha256OfFile(Path file) { try (InputStream in = Files.newInputStream(file)) { @@ -458,6 +481,8 @@ public class FormDetectionModelManager { state = FormDetectionStatus.NOT_INSTALLED; error = null; } + // The file is gone; releasing the engine's handle on it frees the native session too. + invalidateEngine(); applyEndpointState(); } diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/controller/FormDetectionControllerTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/controller/FormDetectionControllerTest.java index e9ac398b1d..e222f68deb 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/controller/FormDetectionControllerTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/controller/FormDetectionControllerTest.java @@ -4,7 +4,6 @@ import static org.springframework.test.web.servlet.request.MockMvcRequestBuilder import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; -import java.util.List; import java.util.Optional; import org.junit.jupiter.api.Test; @@ -40,6 +39,18 @@ class FormDetectionControllerTest { return new MockMultipartFile("file", "test.pdf", "application/pdf", "%PDF-1.4".getBytes()); } + /** A rasterizer that renders nothing, so the detector is never reached. */ + private PageRasterizer noPages() { + return Mockito.mock(PageRasterizer.class); + } + + private FormDetectionModelManager readyManager() { + FormDetectionModelManager manager = Mockito.mock(FormDetectionModelManager.class); + Mockito.when(manager.isReady()).thenReturn(true); + Mockito.when(manager.getActiveEntry()).thenReturn(Optional.of(new ModelCatalogEntry())); + return manager; + } + @Test void detectReturns503WhenModelNotReady() throws Exception { FormDetectionModelManager manager = Mockito.mock(FormDetectionModelManager.class); @@ -53,15 +64,7 @@ class FormDetectionControllerTest { @Test void detectReturnsEmptyDetectionsForBlankRender() throws Exception { - FormDetectionModelManager manager = Mockito.mock(FormDetectionModelManager.class); - Mockito.when(manager.isReady()).thenReturn(true); - Mockito.when(manager.getActiveEntry()).thenReturn(Optional.of(new ModelCatalogEntry())); - - PageRasterizer rasterizer = Mockito.mock(PageRasterizer.class); - Mockito.when(rasterizer.rasterize(Mockito.any(), Mockito.anyInt())) - .thenReturn(List.of()); // no pages -> no detections, detector never called - - mvc(manager, Mockito.mock(OnnxFormDetector.class), rasterizer) + mvc(readyManager(), Mockito.mock(OnnxFormDetector.class), noPages()) .perform(multipart("/api/v1/form/form-detection/detect").file(pdf())) .andExpect(status().isOk()) .andExpect(jsonPath("$.detections").isArray()) @@ -70,34 +73,79 @@ class FormDetectionControllerTest { @Test void detectRejectsPdfsOverThePageLimit() throws Exception { - FormDetectionModelManager manager = Mockito.mock(FormDetectionModelManager.class); - Mockito.when(manager.isReady()).thenReturn(true); - Mockito.when(manager.getActiveEntry()).thenReturn(Optional.of(new ModelCatalogEntry())); - - PageRasterizer.RasterPage blank = - new PageRasterizer.RasterPage( - 0, new byte[0], 1, 1, 1f, 1f, 1f, 1f, 0, 1f, 1f, 0f, 0f); - List tooMany = - java.util.Collections.nCopies(FormDetectionController.MAX_PAGES + 1, blank); PageRasterizer rasterizer = Mockito.mock(PageRasterizer.class); - Mockito.when(rasterizer.rasterize(Mockito.any(), Mockito.anyInt())).thenReturn(tooMany); + Mockito.doThrow( + new PageRasterizer.PageLimitExceededException( + FormDetectionController.MAX_PAGES + 1, + FormDetectionController.MAX_PAGES)) + .when(rasterizer) + .rasterize(Mockito.any(), Mockito.anyInt(), Mockito.anyInt(), Mockito.any()); - mvc(manager, Mockito.mock(OnnxFormDetector.class), rasterizer) + mvc(readyManager(), Mockito.mock(OnnxFormDetector.class), rasterizer) .perform(multipart("/api/v1/form/form-detection/detect").file(pdf())) .andExpect(status().isBadRequest()) .andExpect(jsonPath("$.reason").value("LIMIT")); } @Test - void detectToleratesOutOfRangeConfThreshold() throws Exception { - FormDetectionModelManager manager = Mockito.mock(FormDetectionModelManager.class); - Mockito.when(manager.isReady()).thenReturn(true); - Mockito.when(manager.getActiveEntry()).thenReturn(Optional.of(new ModelCatalogEntry())); + void detectPassesThePageLimitToTheRasterizerSoItIsCheckedBeforeRendering() throws Exception { + PageRasterizer rasterizer = noPages(); + mvc(readyManager(), Mockito.mock(OnnxFormDetector.class), rasterizer) + .perform(multipart("/api/v1/form/form-detection/detect").file(pdf())) + .andExpect(status().isOk()); + + Mockito.verify(rasterizer) + .rasterize( + Mockito.any(), + Mockito.anyInt(), + Mockito.eq(FormDetectionController.MAX_PAGES), + Mockito.any()); + } + + @Test + void detectRejectsAnUnreadablePdfAsBadRequestNotAsAMissingDependency() throws Exception { PageRasterizer rasterizer = Mockito.mock(PageRasterizer.class); - Mockito.when(rasterizer.rasterize(Mockito.any(), Mockito.anyInt())).thenReturn(List.of()); + Mockito.doThrow(new PageRasterizer.UnreadablePdfException("corrupt", null)) + .when(rasterizer) + .rasterize(Mockito.any(), Mockito.anyInt(), Mockito.anyInt(), Mockito.any()); - mvc(manager, Mockito.mock(OnnxFormDetector.class), rasterizer) + mvc(readyManager(), Mockito.mock(OnnxFormDetector.class), rasterizer) + .perform(multipart("/api/v1/form/form-detection/detect").file(pdf())) + .andExpect(status().isBadRequest()) + .andExpect(jsonPath("$.reason").value("INVALID_PDF")); + } + + @Test + void detectRejectsAnEmptyUploadBeforeTouchingTheEngine() throws Exception { + PageRasterizer rasterizer = Mockito.mock(PageRasterizer.class); + MockMultipartFile empty = + new MockMultipartFile("file", "empty.pdf", "application/pdf", new byte[0]); + + mvc(readyManager(), Mockito.mock(OnnxFormDetector.class), rasterizer) + .perform(multipart("/api/v1/form/form-detection/detect").file(empty)) + .andExpect(status().isBadRequest()) + .andExpect(jsonPath("$.reason").value("INVALID_PDF")); + + Mockito.verifyNoInteractions(rasterizer); + } + + @Test + void detectStillReports503WhenTheEngineItselfIsUnavailable() throws Exception { + PageRasterizer rasterizer = Mockito.mock(PageRasterizer.class); + Mockito.doThrow(new IllegalStateException("ONNX Runtime is unavailable")) + .when(rasterizer) + .rasterize(Mockito.any(), Mockito.anyInt(), Mockito.anyInt(), Mockito.any()); + + mvc(readyManager(), Mockito.mock(OnnxFormDetector.class), rasterizer) + .perform(multipart("/api/v1/form/form-detection/detect").file(pdf())) + .andExpect(status().isServiceUnavailable()) + .andExpect(jsonPath("$.reason").value("DEPENDENCY")); + } + + @Test + void detectToleratesOutOfRangeConfThreshold() throws Exception { + mvc(readyManager(), Mockito.mock(OnnxFormDetector.class), noPages()) .perform( multipart("/api/v1/form/form-detection/detect") .file(pdf()) diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/render/PageRasterizerRotationTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/render/PageRasterizerRotationTest.java index 9a5438cc3c..a76331d0aa 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/render/PageRasterizerRotationTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/render/PageRasterizerRotationTest.java @@ -1,12 +1,15 @@ package stirling.software.proprietary.formdetection.render; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assumptions.assumeTrue; import static org.mockito.Mockito.mock; import java.io.ByteArrayOutputStream; +import java.util.ArrayList; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; import org.apache.pdfbox.Loader; import org.apache.pdfbox.pdmodel.PDDocument; @@ -40,6 +43,13 @@ class PageRasterizerRotationTest { return new PageRasterizer(factory); } + /** Gather the streamed pages; these fixtures are one page, so holding them all is fine. */ + private List collect(byte[] pdf) { + List pages = new ArrayList<>(); + rasterizer().rasterize(pdf, 1216, 100, pages::add); + return pages; + } + private static byte[] pdfWithBlackRect(int rotation, float llx, float lly) throws Exception { try (PDDocument doc = new PDDocument()) { PDPage page = new PDPage(new PDRectangle(llx, lly, 200f, 300f)); @@ -107,7 +117,7 @@ class PageRasterizerRotationTest { assumeTrue(pdfiumAvailable(), "JPDFium native not available on this platform"); byte[] pdf = pdfWithBlackRect(rotation, 0f, 0f); - List pages = rasterizer().rasterize(pdf, 1216); + List pages = collect(pdf); assertEquals(1, pages.size()); PageRasterizer.RasterPage page = pages.get(0); assertEquals(rotation, page.rotationDegrees()); @@ -134,7 +144,7 @@ class PageRasterizerRotationTest { assertEquals(llx, check.getPage(0).getCropBox().getLowerLeftX(), 1e-3); } - List pages = rasterizer().rasterize(pdf, 1216); + List pages = collect(pdf); PageRasterizer.RasterPage page = pages.get(0); assertEquals(llx, page.cropLlxPt(), 1e-3); assertEquals(lly, page.cropLlyPt(), 1e-3); @@ -148,4 +158,50 @@ class PageRasterizerRotationTest { assertEquals(RECT_W, r.w(), 3.0); assertEquals(RECT_H, r.h(), 3.0); } + + @Test + void rejectsAnOverLongDocumentWithoutRenderingAnyPage() throws Exception { + assumeTrue(pdfiumAvailable(), "JPDFium native not available on this platform"); + + byte[] pdf = pdfWithPages(3); + AtomicInteger rendered = new AtomicInteger(); + + assertThrows( + PageRasterizer.PageLimitExceededException.class, + () -> rasterizer().rasterize(pdf, 1216, 2, page -> rendered.incrementAndGet())); + // The whole point of the limit: refuse before paying for any bitmap. + assertEquals(0, rendered.get()); + } + + @Test + void rejectsEmptyAndCorruptInputAsUnreadableRatherThanFailingLater() { + assertThrows( + PageRasterizer.UnreadablePdfException.class, + () -> rasterizer().rasterize(new byte[0], 1216, 10, page -> {})); + assertThrows( + PageRasterizer.UnreadablePdfException.class, + () -> rasterizer().rasterize("not a pdf".getBytes(), 1216, 10, page -> {})); + } + + @Test + void streamsPagesOneAtATime() throws Exception { + assumeTrue(pdfiumAvailable(), "JPDFium native not available on this platform"); + + byte[] pdf = pdfWithPages(3); + List seen = new ArrayList<>(); + rasterizer().rasterize(pdf, 1216, 10, page -> seen.add(page.pageIndex())); + + assertEquals(List.of(0, 1, 2), seen); + } + + private static byte[] pdfWithPages(int count) throws Exception { + try (PDDocument doc = new PDDocument()) { + for (int i = 0; i < count; i++) { + doc.addPage(new PDPage(PDRectangle.A6)); + } + ByteArrayOutputStream out = new ByteArrayOutputStream(); + doc.save(out); + return out.toByteArray(); + } + } } diff --git a/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/service/FormDetectionModelManagerTest.java b/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/service/FormDetectionModelManagerTest.java index 8e6703c80c..33b742de77 100644 --- a/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/service/FormDetectionModelManagerTest.java +++ b/app/proprietary/src/test/java/stirling/software/proprietary/formdetection/service/FormDetectionModelManagerTest.java @@ -25,6 +25,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; import org.mockito.Mockito; +import org.springframework.beans.factory.ObjectProvider; import com.sun.net.httpserver.HttpServer; @@ -32,6 +33,7 @@ import stirling.software.SPDF.config.EndpointConfiguration; import stirling.software.common.configuration.RuntimePathConfig; import stirling.software.common.model.ApplicationProperties; import stirling.software.proprietary.formdetection.catalog.ModelCatalogService; +import stirling.software.proprietary.formdetection.inference.FormDetectionEngine; import stirling.software.proprietary.formdetection.model.ModelCatalogEntry; class FormDetectionModelManagerTest { @@ -41,6 +43,12 @@ class FormDetectionModelManagerTest { private String modelSha; private int port; + /** Stands in for a build with no ONNX engine bean, which is the default packaging. */ + @SuppressWarnings("unchecked") + private static ObjectProvider noEngine() { + return Mockito.mock(ObjectProvider.class); + } + @BeforeEach void startServer() throws Exception { modelBytes = "fake-onnx-model-content-1234567890".getBytes(); @@ -92,7 +100,7 @@ class FormDetectionModelManagerTest { .thenReturn(Optional.empty()); Mockito.when(catalog.getAll()).thenReturn(List.of(entry)); // The real fetch only allows the catalog host, so stub the hop to the local test server. - return new FormDetectionModelManager(paths, catalog, props, ep) { + return new FormDetectionModelManager(paths, catalog, props, ep, noEngine()) { @Override HttpURLConnection openModelDownload(String url) throws IOException { HttpURLConnection conn = @@ -192,7 +200,8 @@ class FormDetectionModelManagerTest { paths, Mockito.mock(ModelCatalogService.class), new ApplicationProperties(), - Mockito.mock(EndpointConfiguration.class)); + Mockito.mock(EndpointConfiguration.class), + noEngine()); assertThrows( IOException.class, diff --git a/frontend/editor/src/core/components/tools/autoFormDetection/DetectionProgressPanel.tsx b/frontend/editor/src/core/components/tools/autoFormDetection/DetectionProgressPanel.tsx index ead26286b3..dbb2576c3d 100644 --- a/frontend/editor/src/core/components/tools/autoFormDetection/DetectionProgressPanel.tsx +++ b/frontend/editor/src/core/components/tools/autoFormDetection/DetectionProgressPanel.tsx @@ -56,8 +56,10 @@ export default function DetectionProgressPanel({ "Getting the model ready...", ); break; + // Each page is rendered and analysed back to back, so both stages share one band keyed on the + // page number. Giving them separate bands would send the bar backwards on every page. case "rendering": - value = 0.4 + 0.15 * (stage.page / Math.max(1, stage.pageCount)); + value = 0.4 + 0.53 * (stage.page / Math.max(1, stage.pageCount)); label = t( "autoFormDetection.progress.rendering", "Preparing page {{page}} of {{pageCount}}...", @@ -65,7 +67,7 @@ export default function DetectionProgressPanel({ ); break; case "analyzing": - value = 0.55 + 0.38 * (stage.page / Math.max(1, stage.pageCount)); + value = 0.4 + 0.53 * (stage.page / Math.max(1, stage.pageCount)); label = t( "autoFormDetection.progress.analyzing", "Analyzing page {{page}} of {{pageCount}}...", diff --git a/frontend/editor/src/core/services/formDetection/modelCache.test.ts b/frontend/editor/src/core/services/formDetection/modelCache.test.ts new file mode 100644 index 0000000000..47849fc4f1 --- /dev/null +++ b/frontend/editor/src/core/services/formDetection/modelCache.test.ts @@ -0,0 +1,43 @@ +import { afterEach, describe, expect, it } from "vitest"; + +import { + ChecksumUnsupportedError, + canVerifyChecksums, +} from "@app/services/formDetection/modelCache"; + +// Web Crypto only exists in a secure context, so a self-hosted http:// deployment has no way to +// hash the model it downloads. That has to surface as an explanation, not a TypeError. +describe("checksum capability", () => { + const realSubtle = globalThis.crypto?.subtle; + + afterEach(() => { + if (realSubtle) { + Object.defineProperty(globalThis.crypto, "subtle", { + value: realSubtle, + configurable: true, + }); + } + }); + + function withoutSubtle() { + Object.defineProperty(globalThis.crypto, "subtle", { + value: undefined, + configurable: true, + }); + } + + it("reports availability in a secure context", () => { + expect(canVerifyChecksums()).toBe(true); + }); + + it("reports unavailability when subtle crypto is missing", () => { + withoutSubtle(); + expect(canVerifyChecksums()).toBe(false); + }); + + it("names HTTPS as the fix rather than leaking a TypeError", () => { + const error = new ChecksumUnsupportedError(); + expect(error.name).toBe("ChecksumUnsupportedError"); + expect(error.message).toContain("HTTPS"); + }); +}); diff --git a/frontend/editor/src/core/services/formDetection/modelCache.ts b/frontend/editor/src/core/services/formDetection/modelCache.ts index c171a3de5c..5234faa1a8 100644 --- a/frontend/editor/src/core/services/formDetection/modelCache.ts +++ b/frontend/editor/src/core/services/formDetection/modelCache.ts @@ -41,8 +41,31 @@ async function readWithProgress( return out.buffer; } +/** + * Browsers only expose Web Crypto in a secure context, so an http:// deployment cannot hash the + * model it just downloaded. + */ +export class ChecksumUnsupportedError extends Error { + constructor() { + super( + "The model cannot be verified because Web Crypto is unavailable; serve the app over HTTPS (or localhost) to run detection in the browser", + ); + this.name = "ChecksumUnsupportedError"; + } +} + +/** Whether this context can hash the downloaded model at all. */ +export function canVerifyChecksums(): boolean { + return typeof crypto !== "undefined" && crypto.subtle != null; +} + async function verify(bytes: ArrayBuffer, expectedSha?: string): Promise { if (!expectedSha) return; + // Refusing beats skipping the check: an unverified model would go on to build a form the user + // has no way to know was never validated. + if (!canVerifyChecksums()) { + throw new ChecksumUnsupportedError(); + } const digest = await crypto.subtle.digest("SHA-256", bytes); const actual = toHex(digest); if (actual.toLowerCase() !== expectedSha.toLowerCase()) { diff --git a/frontend/editor/src/core/services/formDetection/pdfRender.ts b/frontend/editor/src/core/services/formDetection/pdfRender.ts index efcbd59f47..7648fe6bf6 100644 --- a/frontend/editor/src/core/services/formDetection/pdfRender.ts +++ b/frontend/editor/src/core/services/formDetection/pdfRender.ts @@ -26,18 +26,32 @@ function normalizeRotation(degrees: number): number { return (Math.floor(r / 90) * 90) % 360; } +/** + * Render each page and hand it straight to `onPage`. Pages are streamed rather than returned as an + * array because one page is several megabytes of RGBA, and holding a long document worth of them at + * once is enough to crash the tab. + * + * `maxPages` is checked before anything is rendered, so an over-long document costs nothing. + * Returns the page count. + */ export async function renderPages( pdfBytes: ArrayBuffer | Uint8Array, inputSize: number, - onPage?: (page: number, pageCount: number) => void, -): Promise { + maxPages: number, + onPage: (page: RasterPage, pageCount: number) => Promise | void, + onPageStart?: (page: number, pageCount: number) => void, +): Promise { const data = pdfBytes instanceof Uint8Array ? pdfBytes : new Uint8Array(pdfBytes); const pdf = await pdfWorkerManager.createDocument(data); try { - const pages: RasterPage[] = []; + if (pdf.numPages > maxPages) { + throw new Error( + `PDF has ${pdf.numPages} pages; the limit is ${maxPages}`, + ); + } for (let i = 1; i <= pdf.numPages; i++) { - onPage?.(i, pdf.numPages); + onPageStart?.(i, pdf.numPages); const page = await pdf.getPage(i); // Display-space dims (rotation applied); the crop box is page.view in user space. const base = page.getViewport({ scale: 1 }); @@ -61,23 +75,29 @@ export async function renderPages( await page.render({ canvas, canvasContext: ctx, viewport: vp }).promise; const rgba = ctx.getImageData(0, 0, canvas.width, canvas.height).data; - pages.push({ - pageIndex: i - 1, - rgba, - widthPx: canvas.width, - heightPx: canvas.height, - pageWidthPt, - pageHeightPt, - scaleX: pageWidthPt > 0 ? canvas.width / pageWidthPt : scale, - scaleY: pageHeightPt > 0 ? canvas.height / pageHeightPt : scale, - rotationDegrees: normalizeRotation(page.rotate ?? 0), - userWidthPt: urx - llx, - userHeightPt: ury - lly, - cropLlxPt: llx, - cropLlyPt: lly, - }); + await onPage( + { + pageIndex: i - 1, + rgba, + widthPx: canvas.width, + heightPx: canvas.height, + pageWidthPt, + pageHeightPt, + scaleX: pageWidthPt > 0 ? canvas.width / pageWidthPt : scale, + scaleY: pageHeightPt > 0 ? canvas.height / pageHeightPt : scale, + rotationDegrees: normalizeRotation(page.rotate ?? 0), + userWidthPt: urx - llx, + userHeightPt: ury - lly, + cropLlxPt: llx, + cropLlyPt: lly, + }, + pdf.numPages, + ); + // Drop the backing bitmap now rather than waiting on GC to notice the canvas. + canvas.width = 0; + canvas.height = 0; } - return pages; + return pdf.numPages; } finally { pdfWorkerManager.destroyDocument(pdf); } diff --git a/frontend/editor/src/core/services/formDetection/runBrowserPipeline.ts b/frontend/editor/src/core/services/formDetection/runBrowserPipeline.ts index 895e8184d6..89b72e29a5 100644 --- a/frontend/editor/src/core/services/formDetection/runBrowserPipeline.ts +++ b/frontend/editor/src/core/services/formDetection/runBrowserPipeline.ts @@ -7,7 +7,11 @@ import { FormDetectionCatalogEntry } from "@app/hooks/useFormDetectionModelStatu import { applyFields } from "@app/services/formDetection/applyFields"; import { toPdfPoints } from "@app/services/formDetection/coordinateMapping"; import { decode, decodeRfDetr } from "@app/services/formDetection/decode"; -import { loadModelBytes } from "@app/services/formDetection/modelCache"; +import { + ChecksumUnsupportedError, + canVerifyChecksums, + loadModelBytes, +} from "@app/services/formDetection/modelCache"; import { getSession, runInference, @@ -66,6 +70,11 @@ export async function runBrowserDetection( ? Math.min(1, Math.max(0, confThreshold)) : spec.scoreThreshold; + // Checked up front so `auto` can hand off before spending a ~37MB download it cannot verify. + if (activeEntry.sha256 && !canVerifyChecksums()) { + throw new ChecksumUnsupportedError(); + } + const modelBytes = await loadModelBytes( activeEntry.sha256, (loadedBytes, totalBytes) => @@ -84,50 +93,45 @@ export async function runBrowserDetection( : "text"; }; + let fields: DetectedField[] = []; + // Each page is analysed as it renders, so only one page of pixels is ever held. // pdf.js may detach the input buffer, so give each consumer its own copy. - const pages = await renderPages( + const pageCount = await renderPages( pdfBytes.slice(0), spec.inputSize, - (page, pageCount) => { - if (pageCount > MAX_PAGES) { - throw new Error( - `PDF has ${pageCount} pages; the limit is ${MAX_PAGES}`, - ); - } - onStage?.({ kind: "rendering", page, pageCount }); - }, - ); - let fields: DetectedField[] = []; - for (const [index, page] of pages.entries()) { - onStage?.({ - kind: "analyzing", - page: page.pageIndex + 1, - pageCount: pages.length, - }); - const startedAt = performance.now(); - const pre = preprocess(page.rgba, page.widthPx, page.heightPx, spec); - const out = await runInference(session, pre.chw, spec.inputSize); - // Measure the first page and bail before paying the same cost for every remaining one. The - // work already done is discarded rather than merged, so the server sees the whole document and - // the result cannot be a mix of two engines. - const pageMs = performance.now() - startedAt; - const budget = options?.pageBudgetMs; - if (index === 0 && budget && pageMs > budget && pages.length > 1) { - throw new BrowserEngineTooSlowError(pageMs); - } - for (const d of decodeFor(spec, out, pre, score)) { - const rect = toPdfPoints(d, page); - if (rect.w <= 0 || rect.h <= 0) { - continue; - } - fields.push({ - type: fieldType(d.classId), - page: page.pageIndex, - rectInPdfPoints: rect, - confidence: d.score, + MAX_PAGES, + async (page, total) => { + onStage?.({ + kind: "analyzing", + page: page.pageIndex + 1, + pageCount: total, }); - } - } + const startedAt = performance.now(); + const pre = preprocess(page.rgba, page.widthPx, page.heightPx, spec); + const out = await runInference(session, pre.chw, spec.inputSize); + // Measure the first page and bail before paying the same cost for every remaining one. The + // work already done is discarded rather than merged, so the server sees the whole document + // and the result cannot be a mix of two engines. + const pageMs = performance.now() - startedAt; + const budget = options?.pageBudgetMs; + if (page.pageIndex === 0 && budget && pageMs > budget && total > 1) { + throw new BrowserEngineTooSlowError(pageMs); + } + for (const d of decodeFor(spec, out, pre, score)) { + const rect = toPdfPoints(d, page); + if (rect.w <= 0 || rect.h <= 0) { + continue; + } + fields.push({ + type: fieldType(d.classId), + page: page.pageIndex, + rectInPdfPoints: rect, + confidence: d.score, + }); + } + }, + (page, total) => onStage?.({ kind: "rendering", page, pageCount: total }), + ); if (fields.length > MAX_FIELDS) { fields = fields .slice() @@ -137,7 +141,7 @@ export async function runBrowserDetection( onStage?.({ kind: "applying" }); const appliedPdf = await applyFields(pdfBytes.slice(0), fields); - return { fields, appliedPdf, pageCount: pages.length }; + return { fields, appliedPdf, pageCount }; } /**