diff --git a/fesod-sheet/src/main/java/org/apache/fesod/sheet/context/WorkbookEncryption.java b/fesod-sheet/src/main/java/org/apache/fesod/sheet/context/WorkbookEncryption.java new file mode 100644 index 000000000..1feadb3cb --- /dev/null +++ b/fesod-sheet/src/main/java/org/apache/fesod/sheet/context/WorkbookEncryption.java @@ -0,0 +1,158 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.fesod.sheet.context; + +import java.io.File; +import java.io.FileOutputStream; +import java.io.OutputStream; +import java.util.UUID; +import org.apache.fesod.common.util.StringUtils; +import org.apache.fesod.sheet.exception.ExcelGenerateException; +import org.apache.fesod.sheet.support.ExcelTypeEnum; +import org.apache.fesod.sheet.util.FileUtils; +import org.apache.fesod.sheet.write.metadata.holder.WriteWorkbookHolder; +import org.apache.poi.hssf.record.crypto.Biff8EncryptionKey; +import org.apache.poi.openxml4j.opc.OPCPackage; +import org.apache.poi.openxml4j.opc.PackageAccess; +import org.apache.poi.poifs.crypt.EncryptionInfo; +import org.apache.poi.poifs.crypt.EncryptionMode; +import org.apache.poi.poifs.crypt.Encryptor; +import org.apache.poi.poifs.filesystem.POIFSFileSystem; + +/** + * Password protection applied when a workbook is finished: xlsx is encrypted into the output stream or the output + * file, and the xls password is cleared afterwards. + */ +final class WorkbookEncryption { + + private enum Target { + NONE, + OUTPUT_STREAM, + FILE + } + + private final WriteWorkbookHolder writeWorkbookHolder; + private final Target target; + + WorkbookEncryption(WriteWorkbookHolder writeWorkbookHolder) { + this.writeWorkbookHolder = writeWorkbookHolder; + if (!hasPasswordFor(ExcelTypeEnum.XLSX)) { + target = Target.NONE; + } else if (writeWorkbookHolder.getFile() == null) { + target = Target.OUTPUT_STREAM; + } else { + target = Target.FILE; + } + } + + /** + * Clears encryption settings for older Excel formats. + */ + void clearEncrypt03() { + if (!hasPasswordFor(ExcelTypeEnum.XLS)) { + return; + } + Biff8EncryptionKey.setCurrentUserPassword(null); + } + + /** + * Checks whether the workbook is password-protected and written as the given type. + */ + private boolean hasPasswordFor(ExcelTypeEnum excelType) { + return !StringUtils.isEmpty(writeWorkbookHolder.getPassword()) + && excelType.equals(writeWorkbookHolder.getExcelType()); + } + + boolean encryptsOutputStream() { + return target == Target.OUTPUT_STREAM; + } + + boolean encryptsFile() { + return target == Target.FILE; + } + + /** + * Encrypts the output stream for newer Excel formats. + * + * @throws Exception If an error occurs during encryption. + */ + void doOutputStreamEncrypt07() throws Exception { + File tempXlsx = FileUtils.createTmpFile(UUID.randomUUID() + ".xlsx"); + try { + try (FileOutputStream tempFileOutputStream = new FileOutputStream(tempXlsx)) { + try { + writeWorkbookHolder.getWorkbook().write(tempFileOutputStream); + } finally { + writeWorkbookHolder.getWorkbook().close(); + } + } + try (POIFSFileSystem fileSystem = openFileSystemAndEncrypt(tempXlsx)) { + fileSystem.writeFilesystem(writeWorkbookHolder.getOutputStream()); + } + } finally { + if (tempXlsx.exists() && !tempXlsx.delete()) { + throw new ExcelGenerateException("Can not delete temp File!"); + } + } + } + + /** + * To encrypt + */ + void doFileEncrypt07() throws Exception { + try { + // Use try-with-resources to automatically close resources, encrypt and write the file + try (POIFSFileSystem fileSystem = openFileSystemAndEncrypt(writeWorkbookHolder.getFile()); + FileOutputStream fileOutputStream = new FileOutputStream(writeWorkbookHolder.getFile())) { + fileSystem.writeFilesystem(fileOutputStream); + } + } catch (Throwable t) { + // The workbook was written to the file before encryption, so the file still holds the unprotected + // workbook and must not be left behind. + File file = writeWorkbookHolder.getFile(); + if (file.exists() && !file.delete()) { + throw new ExcelGenerateException("Can not delete unencrypted file: " + file.getAbsolutePath(), t); + } + throw t; + } + } + + /** + * Opens a file system and encrypts the given file. + * + * This method creates a new POIFSFileSystem instance, sets up an Encryptor with a standard encryption mode, + * and confirms the password for encryption. It then opens the provided file in read-write mode, saves its content + * into an encrypted output stream, and finally returns the encrypted file system. + * + * @param file The file to be encrypted. + * @return An encrypted POIFSFileSystem object. + * @throws Exception If any error occurs during the encryption process or file handling. + */ + private POIFSFileSystem openFileSystemAndEncrypt(File file) throws Exception { + POIFSFileSystem fileSystem = new POIFSFileSystem(); + Encryptor encryptor = new EncryptionInfo(EncryptionMode.standard).getEncryptor(); + encryptor.confirmPassword(writeWorkbookHolder.getPassword()); + try (OPCPackage opcPackage = OPCPackage.open(file, PackageAccess.READ_WRITE); + OutputStream outputStream = encryptor.getDataStream(fileSystem)) { + opcPackage.save(outputStream); + } + return fileSystem; + } +} diff --git a/fesod-sheet/src/main/java/org/apache/fesod/sheet/context/WriteContextImpl.java b/fesod-sheet/src/main/java/org/apache/fesod/sheet/context/WriteContextImpl.java index 17fcc7aad..d372b9aa1 100644 --- a/fesod-sheet/src/main/java/org/apache/fesod/sheet/context/WriteContextImpl.java +++ b/fesod-sheet/src/main/java/org/apache/fesod/sheet/context/WriteContextImpl.java @@ -25,13 +25,11 @@ package org.apache.fesod.sheet.context; -import java.io.File; -import java.io.FileOutputStream; +import java.io.IOException; import java.io.OutputStream; import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.UUID; import java.util.stream.Collectors; import lombok.extern.slf4j.Slf4j; import org.apache.commons.collections4.MapUtils; @@ -44,10 +42,9 @@ import org.apache.fesod.sheet.metadata.Head; import org.apache.fesod.sheet.metadata.data.WriteCellData; import org.apache.fesod.sheet.metadata.property.ExcelContentProperty; -import org.apache.fesod.sheet.support.ExcelTypeEnum; import org.apache.fesod.sheet.util.ClassUtils; import org.apache.fesod.sheet.util.DateUtils; -import org.apache.fesod.sheet.util.FileUtils; +import org.apache.fesod.sheet.util.ExceptionCollector; import org.apache.fesod.sheet.util.NumberDataFormatterUtils; import org.apache.fesod.sheet.util.NumberUtils; import org.apache.fesod.sheet.util.WorkBookUtil; @@ -64,13 +61,6 @@ import org.apache.fesod.sheet.write.metadata.holder.WriteTableHolder; import org.apache.fesod.sheet.write.metadata.holder.WriteWorkbookHolder; import org.apache.fesod.sheet.write.property.ExcelWriteHeadProperty; -import org.apache.poi.hssf.record.crypto.Biff8EncryptionKey; -import org.apache.poi.openxml4j.opc.OPCPackage; -import org.apache.poi.openxml4j.opc.PackageAccess; -import org.apache.poi.poifs.crypt.EncryptionInfo; -import org.apache.poi.poifs.crypt.EncryptionMode; -import org.apache.poi.poifs.crypt.Encryptor; -import org.apache.poi.poifs.filesystem.POIFSFileSystem; import org.apache.poi.ss.usermodel.Cell; import org.apache.poi.ss.usermodel.Row; import org.apache.poi.ss.usermodel.Sheet; @@ -495,9 +485,19 @@ public void finish(boolean onException) { finished = true; // executes the callback after all sheets has been fully written. - boolean shouldSkip = onException && !writeWorkbookHolder.getWriteExcelOnException(); - Map writeSheetHolderMap = - shouldSkip ? null : writeWorkbookHolder.getHasBeenInitializedSheetIndexMap(); + if (!onException || writeWorkbookHolder.getWriteExcelOnException()) { + afterSheetsDispose(); + } + + WriteHandlerUtils.afterWorkbookDispose(writeWorkbookHolder.getWorkbookWriteHandlerContext()); + writeAndCloseWorkbook(onException); + } + + /** + * Runs the after-sheet-dispose handlers for every initialized sheet. + */ + private void afterSheetsDispose() { + Map writeSheetHolderMap = writeWorkbookHolder.getHasBeenInitializedSheetIndexMap(); if (MapUtils.isNotEmpty(writeSheetHolderMap)) { if (MapUtils.size(writeSheetHolderMap) == 1) { SheetWriteHandlerContext sheetWriteHandlerContext = @@ -524,86 +524,74 @@ public void finish(boolean onException) { } } } + } - WriteHandlerUtils.afterWorkbookDispose(writeWorkbookHolder.getWorkbookWriteHandlerContext()); - if (writeWorkbookHolder == null) { - return; - } - Throwable throwable = null; - boolean isOutputStreamEncrypt = false; - // Determine if you need to write excel - boolean writeExcel = !onException; - if (writeWorkbookHolder.getWriteExcelOnException()) { - writeExcel = Boolean.TRUE; - } - // No data is written if an exception is thrown + /** + * Writes the workbook, encrypted if a password is set, then closes the workbook and its streams. If encrypting + * into the output stream fails, the unencrypted workbook is not written in its place. Each close step runs even + * if an earlier step failed, and the first failure is thrown at the end. + * + * @param onException Indicates whether the finish is triggered by an exception. + */ + private void writeAndCloseWorkbook(boolean onException) { + WorkbookEncryption encryption = new WorkbookEncryption(writeWorkbookHolder); + ExceptionCollector collector = new ExceptionCollector(); + boolean writeExcel = !onException || writeWorkbookHolder.getWriteExcelOnException(); if (writeExcel) { - try { - isOutputStreamEncrypt = doOutputStreamEncrypt07(); - } catch (Throwable t) { - throwable = recordFailure(throwable, t); - } - } - if (!isOutputStreamEncrypt) { - try { - if (writeExcel) { - writeWorkbookHolder.getWorkbook().write(writeWorkbookHolder.getOutputStream()); - } - writeWorkbookHolder.getWorkbook().close(); - } catch (Throwable t) { - throwable = recordFailure(throwable, t); - } - } - try { - Workbook workbook = writeWorkbookHolder.getWorkbook(); - if (workbook instanceof SXSSFWorkbook) { - ((SXSSFWorkbook) workbook).dispose(); - } - } catch (Throwable t) { - throwable = recordFailure(throwable, t); - } - try { - if (writeWorkbookHolder.getAutoCloseStream() && writeWorkbookHolder.getOutputStream() != null) { - writeWorkbookHolder.getOutputStream().close(); - } - } catch (Throwable t) { - throwable = recordFailure(throwable, t); - } - if (writeExcel && !isOutputStreamEncrypt) { - try { - doFileEncrypt07(); - } catch (Throwable t) { - throwable = recordFailure(throwable, t); + if (encryption.encryptsOutputStream()) { + collector.run(encryption::doOutputStreamEncrypt07); + } else { + collector.run(this::writeWorkbook); } } - try { - if (writeWorkbookHolder.getTempTemplateInputStream() != null) { - writeWorkbookHolder.getTempTemplateInputStream().close(); - } - } catch (Throwable t) { - throwable = recordFailure(throwable, t); + collector.run(this::disposeSxssfWorkbook); + collector.run(this::closeOutputStream); + if (writeExcel && encryption.encryptsFile()) { + collector.run(encryption::doFileEncrypt07); } - clearEncrypt03(); + collector.run(this::closeTemplateInputStream); + encryption.clearEncrypt03(); removeThreadLocalCache(); - if (throwable != null) { - throw new ExcelGenerateException("Can not close IO.", throwable); - } + collector.throwIfPresent(t -> new ExcelGenerateException("Can not close IO.", t)); if (log.isDebugEnabled()) { log.debug("Finished write."); } } /** - * Returns the first recorded failure, attaching any later one to it as a suppressed exception. + * Writes the unencrypted workbook to the output stream if requested, then closes the workbook. */ - private static Throwable recordFailure(Throwable recorded, Throwable t) { - if (recorded == null) { - return t; + private void writeWorkbook() throws IOException { + writeWorkbookHolder.getWorkbook().write(writeWorkbookHolder.getOutputStream()); + writeWorkbookHolder.getWorkbook().close(); + } + + /** + * Deletes the temporary files backing a streaming workbook. + */ + private void disposeSxssfWorkbook() { + Workbook workbook = writeWorkbookHolder.getWorkbook(); + if (workbook instanceof SXSSFWorkbook) { + ((SXSSFWorkbook) workbook).dispose(); + } + } + + /** + * Closes the output stream, unless the caller asked to keep it open. + */ + private void closeOutputStream() throws IOException { + if (writeWorkbookHolder.getAutoCloseStream() && writeWorkbookHolder.getOutputStream() != null) { + writeWorkbookHolder.getOutputStream().close(); } - if (t != null && recorded != t) { - recorded.addSuppressed(t); + } + + /** + * Closes the temporary copy of the template input stream, if one was made. + */ + private void closeTemplateInputStream() throws IOException { + if (writeWorkbookHolder.getTempTemplateInputStream() != null) { + writeWorkbookHolder.getTempTemplateInputStream().close(); } - return recorded; } /** @@ -655,96 +643,4 @@ public OutputStream getOutputStream() { public Workbook getWorkbook() { return writeWorkbookHolder.getWorkbook(); } - - /** - * Clears encryption settings for older Excel formats. - */ - private void clearEncrypt03() { - if (StringUtils.isEmpty(writeWorkbookHolder.getPassword()) - || !ExcelTypeEnum.XLS.equals(writeWorkbookHolder.getExcelType())) { - return; - } - Biff8EncryptionKey.setCurrentUserPassword(null); - } - - /** - * Encrypts the output stream for newer Excel formats. - * - * @return True if encryption is successful, otherwise false. - * @throws Exception If an error occurs during encryption. - */ - private boolean doOutputStreamEncrypt07() throws Exception { - if (StringUtils.isEmpty(writeWorkbookHolder.getPassword()) - || !ExcelTypeEnum.XLSX.equals(writeWorkbookHolder.getExcelType())) { - return false; - } - if (writeWorkbookHolder.getFile() != null) { - return false; - } - File tempXlsx = FileUtils.createTmpFile(UUID.randomUUID() + ".xlsx"); - FileOutputStream tempFileOutputStream = new FileOutputStream(tempXlsx); - try { - writeWorkbookHolder.getWorkbook().write(tempFileOutputStream); - } finally { - try { - writeWorkbookHolder.getWorkbook().close(); - tempFileOutputStream.close(); - } catch (Exception e) { - if (!tempXlsx.delete()) { - throw new ExcelGenerateException("Can not delete temp File!"); - } - throw e; - } - } - try (POIFSFileSystem fileSystem = openFileSystemAndEncrypt(tempXlsx)) { - fileSystem.writeFilesystem(writeWorkbookHolder.getOutputStream()); - } finally { - if (!tempXlsx.delete()) { - throw new ExcelGenerateException("Can not delete temp File!"); - } - } - return true; - } - - /** - * To encrypt - */ - private void doFileEncrypt07() throws Exception { - // Check if the password is empty or the file type is not xlsx, if so, return directly - if (StringUtils.isEmpty(writeWorkbookHolder.getPassword()) - || !ExcelTypeEnum.XLSX.equals(writeWorkbookHolder.getExcelType())) { - return; - } - // Check if the file is null, if so, return directly - if (writeWorkbookHolder.getFile() == null) { - return; - } - // Use try-with-resources to automatically close resources, encrypt and write the file - try (POIFSFileSystem fileSystem = openFileSystemAndEncrypt(writeWorkbookHolder.getFile()); - FileOutputStream fileOutputStream = new FileOutputStream(writeWorkbookHolder.getFile())) { - fileSystem.writeFilesystem(fileOutputStream); - } - } - - /** - * Opens a file system and encrypts the given file. - * - * This method creates a new POIFSFileSystem instance, sets up an Encryptor with a standard encryption mode, - * and confirms the password for encryption. It then opens the provided file in read-write mode, saves its content - * into an encrypted output stream, and finally returns the encrypted file system. - * - * @param file The file to be encrypted. - * @return An encrypted POIFSFileSystem object. - * @throws Exception If any error occurs during the encryption process or file handling. - */ - private POIFSFileSystem openFileSystemAndEncrypt(File file) throws Exception { - POIFSFileSystem fileSystem = new POIFSFileSystem(); - Encryptor encryptor = new EncryptionInfo(EncryptionMode.standard).getEncryptor(); - encryptor.confirmPassword(writeWorkbookHolder.getPassword()); - try (OPCPackage opcPackage = OPCPackage.open(file, PackageAccess.READ_WRITE); - OutputStream outputStream = encryptor.getDataStream(fileSystem)) { - opcPackage.save(outputStream); - } - return fileSystem; - } } diff --git a/fesod-sheet/src/main/java/org/apache/fesod/sheet/util/ExceptionCollector.java b/fesod-sheet/src/main/java/org/apache/fesod/sheet/util/ExceptionCollector.java new file mode 100644 index 000000000..ac4b21fe2 --- /dev/null +++ b/fesod-sheet/src/main/java/org/apache/fesod/sheet/util/ExceptionCollector.java @@ -0,0 +1,74 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.fesod.sheet.util; + +import java.util.function.Function; + +/** + * Collects failures from multiple actions, retaining the first and adding subsequent failures as suppressed exceptions. + */ +public final class ExceptionCollector { + + private Throwable failure; + + /** + * Runs the action and records any thrown failure, allowing subsequent actions to run. + */ + public void run(ThrowableRunnable runnable) { + try { + runnable.run(); + } catch (Throwable t) { + failure = recordFailure(failure, t); + } + } + + public boolean hasFailure() { + return failure != null; + } + + /** + * Throws the collected failure transformed into the target exception type. + */ + public void throwIfPresent(Function transformer) throws E { + if (failure == null) { + return; + } + E transformed = transformer.apply(failure); + if (transformed != null) { + throw transformed; + } + throw new IllegalStateException("transformer must not return null", failure); + } + + private static Throwable recordFailure(Throwable recorded, Throwable t) { + if (recorded == null) { + return t; + } + if (recorded != t) { + recorded.addSuppressed(t); + } + return recorded; + } + + @FunctionalInterface + public interface ThrowableRunnable { + void run() throws Throwable; + } +} diff --git a/fesod-sheet/src/test/java/org/apache/fesod/sheet/readwrite/EncryptDataTest.java b/fesod-sheet/src/test/java/org/apache/fesod/sheet/readwrite/EncryptDataTest.java index 1eb375ee8..2222db2f4 100644 --- a/fesod-sheet/src/test/java/org/apache/fesod/sheet/readwrite/EncryptDataTest.java +++ b/fesod-sheet/src/test/java/org/apache/fesod/sheet/readwrite/EncryptDataTest.java @@ -25,10 +25,16 @@ package org.apache.fesod.sheet.readwrite; +import java.io.ByteArrayOutputStream; import java.io.File; +import java.io.FileNotFoundException; +import java.io.IOException; +import java.io.OutputStream; import java.nio.file.Files; import java.util.List; +import org.apache.commons.io.output.BrokenOutputStream; import org.apache.fesod.sheet.FesodSheet; +import org.apache.fesod.sheet.exception.ExcelGenerateException; import org.apache.fesod.sheet.read.builder.ExcelReaderBuilder; import org.apache.fesod.sheet.support.ExcelTypeEnum; import org.apache.fesod.sheet.testkit.Tags; @@ -38,9 +44,14 @@ import org.apache.fesod.sheet.testkit.listeners.CollectingReadListener; import org.apache.fesod.sheet.testkit.models.SimpleData; import org.apache.fesod.sheet.testkit.params.ExcelFormatSource; +import org.apache.fesod.sheet.util.FileUtils; import org.apache.fesod.sheet.write.builder.ExcelWriterBuilder; +import org.apache.fesod.sheet.write.handler.WorkbookWriteHandler; +import org.apache.fesod.sheet.write.handler.context.WorkbookWriteHandlerContext; import org.apache.poi.EncryptedDocumentException; +import org.apache.poi.xssf.streaming.SXSSFWorkbook; import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Assumptions; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; @@ -138,4 +149,176 @@ void xlsPasswordWrite_isActuallyEncrypted() throws Exception { .doReadSync(); Assertions.assertEquals(10, dataList.size()); } + + /** + * Verifies that when encrypting an XLSX stream write fails, nothing is written to the stream rather than + * the unencrypted workbook. + */ + @Test + void xlsxStreamPasswordWrite_encryptionFails_writesNothing() throws Exception { + // A temp directory under a regular file cannot be created, so the encryption step fails + File regularFile = createTempFile("enc-blocker", ExcelFormat.XLSX); + String originalPrefix = FileUtils.getTempFilePrefix(); + FileUtils.setTempFilePrefix(regularFile.getAbsolutePath() + File.separator + "sub" + File.separator); + try { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + Assertions.assertThrows(ExcelGenerateException.class, () -> FesodSheet.write(out, SimpleData.class) + .excelType(ExcelTypeEnum.XLSX) + .password(PASSWORD) + .sheet() + .doWrite(TestDataBuilder.simpleData(10))); + Assertions.assertEquals(0, out.size()); + } finally { + FileUtils.setTempFilePrefix(originalPrefix); + } + } + + /** + * Verifies that when writing the unencrypted temp file for an XLSX stream write fails, the temp file is deleted. + */ + @Test + void xlsxStreamPasswordWrite_tempFileWriteFails_deletesTempFile() { + File tempFileDir = new File(tempDir, "fesod-temp"); + WorkbookWriteHandler failingWorkbook = swapWorkbook(new SXSSFWorkbook() { + @Override + public void write(OutputStream stream) throws IOException { + super.write(stream); + throw new IOException("write failed"); + } + }); + Assertions.assertThrows( + ExcelGenerateException.class, + () -> writeWithPassword(tempFileDir, new ByteArrayOutputStream(), failingWorkbook)); + Assertions.assertArrayEquals(new String[0], tempFileDir.list()); + } + + /** + * Verifies that when closing the workbook after writing the temp file fails, the temp file is deleted. + */ + @Test + void xlsxStreamPasswordWrite_workbookCloseFails_deletesTempFile() { + File tempFileDir = new File(tempDir, "fesod-temp"); + WorkbookWriteHandler failingWorkbook = swapWorkbook(new SXSSFWorkbook() { + @Override + public void close() throws IOException { + super.close(); + throw new IOException("close failed"); + } + }); + Assertions.assertThrows( + ExcelGenerateException.class, + () -> writeWithPassword(tempFileDir, new ByteArrayOutputStream(), failingWorkbook)); + Assertions.assertArrayEquals(new String[0], tempFileDir.list()); + } + + /** + * Verifies that when closing the workbook after writing the temp file fails, the temp file stream is closed. + */ + @Test + void xlsxStreamPasswordWrite_workbookCloseFails_closesTempFileStream() { + File tempFileDir = new File(tempDir, "fesod-temp"); + OutputStream[] tempFileStream = new OutputStream[1]; + WorkbookWriteHandler failingWorkbook = swapWorkbook(new SXSSFWorkbook() { + @Override + public void write(OutputStream stream) throws IOException { + if (tempFileStream[0] == null) { + tempFileStream[0] = stream; + } + super.write(stream); + } + + @Override + public void close() throws IOException { + super.close(); + throw new IOException("close failed"); + } + }); + Assertions.assertThrows( + ExcelGenerateException.class, + () -> writeWithPassword(tempFileDir, new ByteArrayOutputStream(), failingWorkbook)); + Assertions.assertThrows(IOException.class, () -> tempFileStream[0].write(0)); + } + + /** + * Verifies that when writing the encrypted workbook to the caller's stream fails, the temp file is deleted. + */ + @Test + void xlsxStreamPasswordWrite_encryptedWriteFails_deletesTempFile() { + File tempFileDir = new File(tempDir, "fesod-temp"); + Assertions.assertThrows( + ExcelGenerateException.class, () -> writeWithPassword(tempFileDir, BrokenOutputStream.INSTANCE, null)); + Assertions.assertArrayEquals(new String[0], tempFileDir.list()); + } + + /** + * Verifies that when the temp file cannot be opened, the open error is reported rather than a failed delete. + */ + @Test + void xlsxStreamPasswordWrite_tempFileOpenFails_reportsOpenError() { + File tempFileDir = new File(tempDir, "fesod-temp"); + Assertions.assertTrue(tempFileDir.mkdirs()); + Assertions.assertTrue(tempFileDir.setWritable(false)); + try { + // Root and some file systems ignore the flag, so the open would not fail there + Assumptions.assumeFalse(tempFileDir.canWrite()); + ExcelGenerateException e = Assertions.assertThrows( + ExcelGenerateException.class, + () -> writeWithPassword(tempFileDir, new ByteArrayOutputStream(), null)); + Assertions.assertInstanceOf(FileNotFoundException.class, e.getCause()); + } finally { + Assertions.assertTrue(tempFileDir.setWritable(true)); + } + } + + private static WorkbookWriteHandler swapWorkbook(SXSSFWorkbook workbook) { + return new WorkbookWriteHandler() { + @Override + public void afterWorkbookDispose(WorkbookWriteHandlerContext context) { + workbook.createSheet("s").createRow(0).createCell(0).setCellValue("secret"); + context.getWriteWorkbookHolder().setWorkbook(workbook); + } + }; + } + + private void writeWithPassword(File tempFileDir, OutputStream outputStream, WorkbookWriteHandler writeHandler) { + String originalPrefix = FileUtils.getTempFilePrefix(); + FileUtils.setTempFilePrefix(tempFileDir.getAbsolutePath() + File.separator); + try { + ExcelWriterBuilder writerBuilder = FesodSheet.write(outputStream, SimpleData.class) + .excelType(ExcelTypeEnum.XLSX) + .password(PASSWORD); + if (writeHandler != null) { + writerBuilder.registerWriteHandler(writeHandler); + } + writerBuilder.sheet("s").doWrite(TestDataBuilder.simpleData(10)); + } finally { + FileUtils.setTempFilePrefix(originalPrefix); + } + } + + /** + * Verifies that when encrypting an XLSX file write fails, the unencrypted file is not left behind. + */ + @Test + void xlsxFilePasswordWrite_encryptionFails_leavesNoFile() throws Exception { + File file = createTempFile("enc-readonly", ExcelFormat.XLSX); + // Root ignores file permissions, so the failure forced below cannot happen + Assumptions.assumeTrue(file.setReadOnly() && !file.canWrite(), "file permissions are not enforced"); + Assertions.assertTrue(file.setWritable(true)); + // Once the write stream is open, making the file read-only makes the encrypt-in-place step fail + WorkbookWriteHandler makeReadOnly = new WorkbookWriteHandler() { + @Override + public void afterWorkbookDispose(WorkbookWriteHandlerContext context) { + Assertions.assertTrue(file.setReadOnly()); + } + }; + + Assertions.assertThrows(ExcelGenerateException.class, () -> FesodSheet.write(file, SimpleData.class) + .excelType(ExcelTypeEnum.XLSX) + .password(PASSWORD) + .registerWriteHandler(makeReadOnly) + .sheet() + .doWrite(TestDataBuilder.simpleData(10))); + Assertions.assertFalse(file.exists()); + } }