Migrate stream to resource for stability (#6160)

2026-09-14 12:42:08 +02:00 · 2026-04-23 15:56:31 +01:00
parent c294e9b2cb
commit 177c776658
109 changed files with 1390 additions and 855 deletions
@@ -4,8 +4,11 @@ import static org.junit.jupiter.api.Assertions.*;
 import static org.mockito.Mockito.*;

 import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.stream.Stream;

 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -13,6 +16,8 @@ import org.junit.jupiter.api.io.TempDir;
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.springframework.core.io.ByteArrayResource;
+import org.springframework.core.io.Resource;
 import org.springframework.http.MediaType;
 import org.springframework.test.util.ReflectionTestUtils;
 import org.springframework.web.multipart.MultipartFile;
@@ -185,4 +190,76 @@ class FileStorageTest {
        // Assert
        assertFalse(result);
    }
+
+    @Test
+    void storeFromResource_happyPath_persistsContent() throws IOException {
+        // Arrange
+        byte[] payload = "resource-body-bytes".getBytes(StandardCharsets.UTF_8);
+        Resource resource = new ByteArrayResource(payload);
+
+        // Act
+        String fileId = fileStorage.storeFromResource(resource, "whatever.pdf");
+
+        // Assert
+        assertNotNull(fileId);
+        assertTrue(Files.exists(tempDir.resolve(fileId)));
+        assertArrayEquals(payload, fileStorage.retrieveBytes(fileId));
+    }
+
+    @Test
+    void storeFromResource_failureCleansUpPartialFile() throws IOException {
+        // Arrange: a Resource whose getInputStream returns a stream that throws after
+        // emitting a few bytes. The finally block in storeFromResource should remove
+        // the partial file on disk.
+        byte[] head = "partial".getBytes(StandardCharsets.UTF_8);
+        Resource flakyResource =
+                new ByteArrayResource(head) {
+                    @Override
+                    public InputStream getInputStream() {
+                        return new InputStream() {
+                            private int position = 0;
+
+                            @Override
+                            public int read() throws IOException {
+                                if (position < head.length) {
+                                    return head[position++] & 0xFF;
+                                }
+                                throw new IOException("simulated mid-copy read failure");
+                            }
+
+                            @Override
+                            public int read(byte[] b, int off, int len) throws IOException {
+                                if (position >= head.length) {
+                                    throw new IOException("simulated mid-copy read failure");
+                                }
+                                int toCopy = Math.min(len, head.length - position);
+                                System.arraycopy(head, position, b, off, toCopy);
+                                position += toCopy;
+                                return toCopy;
+                            }
+                        };
+                    }
+                };
+
+        // Snapshot dir contents before the call so we can detect any lingering file.
+        long filesBefore;
+        try (Stream<Path> s = Files.list(tempDir)) {
+            filesBefore = s.count();
+        }
+
+        // Act + Assert: IOException must propagate out — not be swallowed.
+        assertThrows(
+                IOException.class, () -> fileStorage.storeFromResource(flakyResource, "n.pdf"));
+
+        // Assert: no partial file lingers under the storage directory — the finally
+        // branch's deleteIfExists must have cleaned it up.
+        long filesAfter;
+        try (Stream<Path> s = Files.list(tempDir)) {
+            filesAfter = s.count();
+        }
+        assertEquals(
+                filesBefore,
+                filesAfter,
+                "partial file must be cleaned up by storeFromResource finally block");
+    }
 }
@@ -8,9 +8,11 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;

+import java.nio.charset.StandardCharsets;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeoutException;
@@ -24,7 +26,12 @@ import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.core.io.ByteArrayResource;
+import org.springframework.core.io.Resource;
+import org.springframework.http.ContentDisposition;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;

 import jakarta.servlet.http.HttpServletRequest;
@@ -180,6 +187,46 @@ class JobExecutorServiceTest {
        verify(spy).runJobGeneric(eq(false), any(Supplier.class), eq(customTimeout));
    }

+    @Test
+    void shouldPersistResponseEntityResourceBodyViaFileStorage() throws Exception {
+        // Given: an async job whose result is a ResponseEntity<Resource> — the new
+        // branch added by the stream-to-Resource migration. The executor must route
+        // the body through FileStorage.storeFromResource and then record the result
+        // via TaskManager.setFileResult with the filename/content-type extracted
+        // from the response headers.
+        byte[] payload = "resource-bytes".getBytes(StandardCharsets.UTF_8);
+        Resource resource = new ByteArrayResource(payload);
+
+        HttpHeaders headers = new HttpHeaders();
+        headers.setContentType(MediaType.APPLICATION_PDF);
+        headers.setContentDisposition(
+                ContentDisposition.formData().name("attachment").filename("result.pdf").build());
+
+        Supplier<Object> work = () -> new ResponseEntity<>(resource, headers, HttpStatus.OK);
+
+        when(fileStorage.storeFromResource(any(Resource.class), anyString()))
+                .thenReturn("stored-file-id");
+
+        // When: run the job asynchronously — processJobResult runs on the executor.
+        ResponseEntity<?> response = jobExecutorService.runJobGeneric(true, work);
+
+        // Then: the immediate return must be the JobResponse envelope.
+        assertEquals(HttpStatus.OK, response.getStatusCode());
+        assertInstanceOf(JobResponse.class, response.getBody());
+
+        // Wait for async processing and verify the Resource branch was taken —
+        // FileStorage.storeFromResource was invoked with the same Resource instance,
+        // and TaskManager.setFileResult recorded the extracted filename + content-type.
+        verify(fileStorage, timeout(5000)).storeFromResource(eq(resource), eq("result.pdf"));
+        verify(taskManager, timeout(5000))
+                .setFileResult(
+                        anyString(),
+                        eq("stored-file-id"),
+                        eq("result.pdf"),
+                        eq(MediaType.APPLICATION_PDF_VALUE));
+        verify(taskManager, timeout(5000)).setComplete(anyString());
+    }
+
    @Test
    void shouldHandleTimeout() throws Exception {
        // Given
@@ -25,12 +25,12 @@ import org.junit.jupiter.api.io.TempDir;
 import org.mockito.Mock;
 import org.mockito.MockedStatic;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.core.io.Resource;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;
 import org.springframework.mock.web.MockMultipartFile;
 import org.springframework.web.multipart.MultipartFile;
-import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

 import io.github.pixee.security.ZipSecurity;

@@ -83,9 +83,11 @@ class PDFToFileTest {
        pdfToFile = new PDFToFile(mockTempFileManager, mockRuntimePathConfig);
    }

-    private static byte[] drain(ResponseEntity<StreamingResponseBody> response) throws IOException {
+    private static byte[] drain(ResponseEntity<Resource> response) throws IOException {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
-        response.getBody().writeTo(baos);
+        try (java.io.InputStream __in = response.getBody().getInputStream()) {
+            __in.transferTo(baos);
+        }
        return baos.toByteArray();
    }

@@ -100,7 +102,7 @@ class PDFToFileTest {
                        "This is not a PDF".getBytes());

        // Execute
-        ResponseEntity<StreamingResponseBody> response = pdfToFile.processPdfToMarkdown(nonPdfFile);
+        ResponseEntity<Resource> response = pdfToFile.processPdfToMarkdown(nonPdfFile);

        // Verify
        assertEquals(HttpStatus.BAD_REQUEST, response.getStatusCode());
@@ -117,7 +119,7 @@ class PDFToFileTest {
                        "This is not a PDF".getBytes());

        // Execute
-        ResponseEntity<StreamingResponseBody> response = pdfToFile.processPdfToHtml(nonPdfFile);
+        ResponseEntity<Resource> response = pdfToFile.processPdfToHtml(nonPdfFile);

        // Verify
        assertEquals(HttpStatus.BAD_REQUEST, response.getStatusCode());
@@ -135,7 +137,7 @@ class PDFToFileTest {
                        "This is not a PDF".getBytes());

        // Execute
-        ResponseEntity<StreamingResponseBody> response =
+        ResponseEntity<Resource> response =
                pdfToFile.processPdfToOfficeFormat(nonPdfFile, "docx", "draw_pdf_import");

        // Verify
@@ -154,7 +156,7 @@ class PDFToFileTest {
                        "Fake PDF content".getBytes());

        // Execute with invalid format
-        ResponseEntity<StreamingResponseBody> response =
+        ResponseEntity<Resource> response =
                pdfToFile.processPdfToOfficeFormat(pdfFile, "invalid_format", "draw_pdf_import");

        // Verify
@@ -205,8 +207,7 @@ class PDFToFileTest {
                            });

            // Execute the method
-            ResponseEntity<StreamingResponseBody> response =
-                    pdfToFile.processPdfToMarkdown(pdfFile);
+            ResponseEntity<Resource> response = pdfToFile.processPdfToMarkdown(pdfFile);

            // Verify - should now return a ZIP file instead of plain markdown
            assertEquals(HttpStatus.OK, response.getStatusCode());
@@ -298,8 +299,7 @@ class PDFToFileTest {
                            });

            // Execute the method
-            ResponseEntity<StreamingResponseBody> response =
-                    pdfToFile.processPdfToMarkdown(pdfFile);
+            ResponseEntity<Resource> response = pdfToFile.processPdfToMarkdown(pdfFile);

            // Verify
            assertEquals(HttpStatus.OK, response.getStatusCode());
@@ -377,7 +377,7 @@ class PDFToFileTest {
                            });

            // Execute the method
-            ResponseEntity<StreamingResponseBody> response = pdfToFile.processPdfToHtml(pdfFile);
+            ResponseEntity<Resource> response = pdfToFile.processPdfToHtml(pdfFile);

            // Verify
            assertEquals(HttpStatus.OK, response.getStatusCode());
@@ -463,7 +463,7 @@ class PDFToFileTest {
                            });

            // Execute the method with docx format
-            ResponseEntity<StreamingResponseBody> response =
+            ResponseEntity<Resource> response =
                    pdfToFile.processPdfToOfficeFormat(pdfFile, "docx", "draw_pdf_import");

            // Verify
@@ -535,7 +535,7 @@ class PDFToFileTest {
                            });

            // Execute the method with ODP format
-            ResponseEntity<StreamingResponseBody> response =
+            ResponseEntity<Resource> response =
                    pdfToFile.processPdfToOfficeFormat(pdfFile, "odp", "draw_pdf_import");

            // Verify
@@ -620,7 +620,7 @@ class PDFToFileTest {
                            });

            // Execute the method with text format
-            ResponseEntity<StreamingResponseBody> response =
+            ResponseEntity<Resource> response =
                    pdfToFile.processPdfToOfficeFormat(pdfFile, "txt:Text", "draw_pdf_import");

            // Verify
@@ -679,7 +679,7 @@ class PDFToFileTest {
                            });

            // Execute the method
-            ResponseEntity<StreamingResponseBody> response =
+            ResponseEntity<Resource> response =
                    pdfToFile.processPdfToOfficeFormat(pdfFile, "docx", "draw_pdf_import");

            // Verify
@@ -726,7 +726,7 @@ class PDFToFileTest {
                                return mockExecutorResult;
                            });

-            ResponseEntity<StreamingResponseBody> response =
+            ResponseEntity<Resource> response =
                    pdfToFileWithUno.processPdfToOfficeFormat(pdfFile, "docx", "writer_pdf_import");

            assertEquals(HttpStatus.OK, response.getStatusCode());
@@ -790,7 +790,7 @@ class PDFToFileTest {
                                return mockExecutorResult;
                            });

-            ResponseEntity<StreamingResponseBody> response =
+            ResponseEntity<Resource> response =
                    pdfToFileWithUno.processPdfToOfficeFormat(pdfFile, "docx", "writer_pdf_import");

            assertEquals(HttpStatus.OK, response.getStatusCode());
@@ -3,17 +3,39 @@ package stirling.software.common.util;
 import static org.junit.jupiter.api.Assertions.*;

 import java.io.ByteArrayOutputStream;
+import java.io.File;
 import java.io.IOException;
+import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.concurrent.atomic.AtomicBoolean;

+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
 import org.springframework.http.ResponseEntity;

+import stirling.software.common.model.ApplicationProperties;
+
 class WebResponseUtilsTest {

+    @TempDir Path tempDir;
+
+    private TempFileManager tempFileManager;
+
+    @BeforeEach
+    void setUpTempFileManager() {
+        TempFileRegistry registry = new TempFileRegistry();
+        ApplicationProperties applicationProperties = new ApplicationProperties();
+        applicationProperties.getSystem().getTempFileManagement().setBaseTmpDir(tempDir.toString());
+        applicationProperties.getSystem().getTempFileManagement().setPrefix("wru-test-");
+        tempFileManager = new TempFileManager(registry, applicationProperties);
+    }
+
    @Test
    void testBytesToWebResponse_defaultMediaType() throws IOException {
        byte[] data = "test content".getBytes(StandardCharsets.UTF_8);
@@ -91,4 +113,87 @@ class WebResponseUtilsTest {
        assertEquals(HttpStatus.OK, response.getStatusCode());
        assertEquals(0, response.getHeaders().getContentLength());
    }
+
+    @Test
+    void managedTempFileResource_happyPathDeletesBackingFile() throws IOException {
+        // Create a real managed temp file on disk and populate it.
+        TempFile tempFile = tempFileManager.createManagedTempFile(".pdf");
+        byte[] payload = "fully-readable-pdf-body".getBytes(StandardCharsets.UTF_8);
+        Files.write(tempFile.getPath(), payload);
+        File backing = tempFile.getFile();
+        assertTrue(backing.exists(), "precondition: backing file should exist");
+
+        WebResponseUtils.ManagedTempFileResource resource =
+                new WebResponseUtils.ManagedTempFileResource(tempFile);
+
+        byte[] readBack;
+        try (InputStream in = resource.getInputStream()) {
+            readBack = in.readAllBytes();
+        }
+
+        assertArrayEquals(payload, readBack, "stream should deliver the original bytes");
+        assertFalse(
+                backing.exists(),
+                "backing temp file must be deleted once the response stream is closed");
+    }
+
+    @Test
+    void closingInputStream_propagatesReadFailure() throws IOException {
+        // ManagedTempFileResource is final, so we can't swap in a mock underlying stream.
+        // Instead: open the resource's stream, close the inner FileInputStream early
+        // (by closing the outer stream once), then confirm reading the now-closed stream
+        // throws — exercising ClosingInputStream's read() catch/log/rethrow path. Finally
+        // confirm the temp file is still cleaned up on close().
+        TempFile tempFile = tempFileManager.createManagedTempFile(".pdf");
+        Files.write(tempFile.getPath(), "some-bytes".getBytes(StandardCharsets.UTF_8));
+        File backing = tempFile.getFile();
+        assertTrue(backing.exists());
+
+        WebResponseUtils.ManagedTempFileResource resource =
+                new WebResponseUtils.ManagedTempFileResource(tempFile);
+
+        InputStream in = resource.getInputStream();
+        // Pre-close the underlying stream to guarantee read() throws.
+        in.close();
+        // After close, the temp file has been deleted already.
+        assertFalse(
+                backing.exists(),
+                "backing file is deleted on first close — precondition for read assertion");
+
+        // Any attempt to read from the already-closed stream must throw IOException
+        // (not silently return -1). This exercises the ClosingInputStream.read() path
+        // that logs and rethrows.
+        IOException ioex = assertThrows(IOException.class, in::read);
+        assertNotNull(ioex);
+
+        // close() is idempotent — a second call must not throw.
+        assertDoesNotThrow(in::close);
+    }
+
+    @Test
+    void managedTempFileResource_openFailureCleansUp() throws IOException {
+        // Arrange: build a TempFile whose backing file is deleted out from under it so
+        // super.getInputStream() throws on open. Instrument close() with an AtomicBoolean
+        // hook to confirm the cleanup path ran.
+        AtomicBoolean closed = new AtomicBoolean(false);
+        TempFile spying =
+                new TempFile(tempFileManager, ".pdf") {
+                    @Override
+                    public void close() {
+                        closed.set(true);
+                        super.close();
+                    }
+                };
+        assertTrue(
+                spying.getFile().delete(),
+                "precondition: delete backing so super.getInputStream() fails");
+
+        WebResponseUtils.ManagedTempFileResource resource =
+                new WebResponseUtils.ManagedTempFileResource(spying);
+
+        assertThrows(IOException.class, resource::getInputStream);
+        assertTrue(
+                closed.get(),
+                "tempFile.close() must run when super.getInputStream() fails on open");
+    }
 }