From 3d7fdd0f358b9e5879cd413677f3820201b4f646 Mon Sep 17 00:00:00 2001
From: Anthony Stirling <77850077+Frooodle@users.noreply.github.com>
Date: Mon, 27 Mar 2023 23:24:48 +0100
Subject: [PATCH] Bug fixes for merge and split (#80)

---
 .../SPDF/controller/MergeController.java      | 12 ++----
 .../SPDF/controller/SplitPDFController.java   | 38 ++++++++++---------
 src/main/resources/templates/merge-pdfs.html  | 10 ++---
 3 files changed, 28 insertions(+), 32 deletions(-)

diff --git a/src/main/java/stirling/software/SPDF/controller/MergeController.java b/src/main/java/stirling/software/SPDF/controller/MergeController.java
index da672b27..b7fb02d5 100644
--- a/src/main/java/stirling/software/SPDF/controller/MergeController.java
+++ b/src/main/java/stirling/software/SPDF/controller/MergeController.java
@@ -21,6 +21,8 @@ import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.multipart.MultipartFile;
 
+import stirling.software.SPDF.utils.PdfUtils;
+
 @Controller
 public class MergeController {
 
@@ -33,7 +35,7 @@ public class MergeController {
     }
 
     @PostMapping("/merge-pdfs")
-    public ResponseEntity<InputStreamResource> mergePdfs(@RequestParam("fileInput") MultipartFile[] files) throws IOException {
+    public ResponseEntity<byte[]> mergePdfs(@RequestParam("fileInput") MultipartFile[] files) throws IOException {
         // Read the input PDF files into PDDocument objects
         List<PDDocument> documents = new ArrayList<>();
 
@@ -43,15 +45,9 @@ public class MergeController {
         }
 
         PDDocument mergedDoc = mergeDocuments(documents);
-        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
-        mergedDoc.save(byteArrayOutputStream);
-        mergedDoc.close();
-
-        // Create an InputStreamResource from the merged PDF
-        InputStreamResource resource = new InputStreamResource(new ByteArrayInputStream(byteArrayOutputStream.toByteArray()));
 
         // Return the merged PDF as a response
-        return ResponseEntity.ok().contentType(MediaType.APPLICATION_PDF).body(resource);
+        return PdfUtils.pdfDocToWebResponse(mergedDoc, files[0].getOriginalFilename().replaceFirst("[.][^.]+$", "")+ "_merged.pdf");
     }
 
     private PDDocument mergeDocuments(List<PDDocument> documents) throws IOException {
diff --git a/src/main/java/stirling/software/SPDF/controller/SplitPDFController.java b/src/main/java/stirling/software/SPDF/controller/SplitPDFController.java
index 0eb0429e..e456d89b 100644
--- a/src/main/java/stirling/software/SPDF/controller/SplitPDFController.java
+++ b/src/main/java/stirling/software/SPDF/controller/SplitPDFController.java
@@ -31,7 +31,8 @@ import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.multipart.MultipartFile;
-
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipOutputStream;
 @Controller
 public class SplitPDFController {
 
@@ -105,32 +106,33 @@ public class SplitPDFController {
         // closing the original document
         document.close();
 
-        // create the zip file
         Path zipFile = Files.createTempFile("split_documents", ".zip");
-        URI uri = URI.create("jar:file:" + zipFile.toUri().getPath());
-        Map<String, String> env = new HashMap<>();
-        env.put("create", "true");
-        FileSystem zipfs = FileSystems.newFileSystem(uri, env);
 
-        // loop through the split documents and write them to the zip file
-        for (int i = 0; i < splitDocumentsBoas.size(); i++) {
-            String fileName = "split_document_" + (i + 1) + ".pdf";
-            ByteArrayOutputStream baos = splitDocumentsBoas.get(i);
-            byte[] pdf = baos.toByteArray();
-            Path pathInZipfile = zipfs.getPath(fileName);
-            try (OutputStream os = Files.newOutputStream(pathInZipfile)) {
-                os.write(pdf);
+        try (ZipOutputStream zipOut = new ZipOutputStream(Files.newOutputStream(zipFile))) {
+            // loop through the split documents and write them to the zip file
+            for (int i = 0; i < splitDocumentsBoas.size(); i++) {
+                String fileName = "split_document_" + (i + 1) + ".pdf";
+                ByteArrayOutputStream baos = splitDocumentsBoas.get(i);
+                byte[] pdf = baos.toByteArray();
+
+                // Add PDF file to the zip
+                ZipEntry pdfEntry = new ZipEntry(fileName);
+                zipOut.putNextEntry(pdfEntry);
+                zipOut.write(pdf);
+                zipOut.closeEntry();
+
                 logger.info("Wrote split document {} to zip file", fileName);
-            } catch (Exception e) {
-                logger.error("Failed writing to zip", e);
-                throw e;
             }
+        } catch (Exception e) {
+            logger.error("Failed writing to zip", e);
+            throw e;
         }
-        zipfs.close();
+
         logger.info("Successfully created zip file with split documents: {}", zipFile.toString());
         byte[] data = Files.readAllBytes(zipFile);
         ByteArrayResource resource = new ByteArrayResource(data);
         Files.delete(zipFile);
+
         // return the Resource in the response
         return ResponseEntity.ok().header(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=" + file.getOriginalFilename().replaceFirst("[.][^.]+$", "") + "_split.zip").contentType(MediaType.APPLICATION_OCTET_STREAM)
                 .contentLength(resource.contentLength()).body(resource);
diff --git a/src/main/resources/templates/merge-pdfs.html b/src/main/resources/templates/merge-pdfs.html
index 9f594be0..1daa099b 100644
--- a/src/main/resources/templates/merge-pdfs.html
+++ b/src/main/resources/templates/merge-pdfs.html
@@ -16,10 +16,8 @@
                         <form action="merge-pdfs" method="post" enctype="multipart/form-data">
                             <div class="form-group">
                                 <label th:text="#{multiPdfDropPrompt}"></label>
-                                <div class="custom-file">
-                                    <input type="file" class="custom-file-input" id="fileInput" name="fileInput" accept="application/pdf" multiple required> 
-                                    <label class="custom-file-label" th:text="#{pdfPrompt}"></label>
-                                </div>
+                                <div th:replace="~{fragments/common :: fileSelector(name='fileInput', multiple=true, accept='application/pdf')}"></div>
+
                             </div>
                             <div class="form-group">
                                 <ul id="selectedFiles" class="list-group"></ul>
@@ -61,7 +59,7 @@
 
                         </style>
                         <script>
-                        document.getElementById("fileInput").addEventListener("change", function() {
+                        document.getElementById("fileInput-input").addEventListener("change", function() {
                         	  var files = this.files;
                         	  var list = document.getElementById("selectedFiles");
                         	  list.innerHTML = "";
@@ -121,7 +119,7 @@
                         	        }
                         	      }
                         	    }
-                        	    document.getElementById("fileInput").files = dataTransfer.files;
+                        	    document.getElementById("fileInput-input").files = dataTransfer.files;
                         	  }
                         	});