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 fd118305d..8dbeec18c 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 @@ -669,24 +669,19 @@ private boolean doOutputStreamEncrypt07() throws Exception { 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!"); + try (FileOutputStream tempFileOutputStream = new FileOutputStream(tempXlsx)) { + try { + writeWorkbookHolder.getWorkbook().write(tempFileOutputStream); + } finally { + writeWorkbookHolder.getWorkbook().close(); } - throw e; } - } - try (POIFSFileSystem fileSystem = openFileSystemAndEncrypt(tempXlsx)) { - fileSystem.writeFilesystem(writeWorkbookHolder.getOutputStream()); + try (POIFSFileSystem fileSystem = openFileSystemAndEncrypt(tempXlsx)) { + fileSystem.writeFilesystem(writeWorkbookHolder.getOutputStream()); + } } finally { - if (!tempXlsx.delete()) { + if (tempXlsx.exists() && !tempXlsx.delete()) { throw new ExcelGenerateException("Can not delete temp File!"); } } 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..affc72e05 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,127 @@ void xlsPasswordWrite_isActuallyEncrypted() throws Exception { .doReadSync(); Assertions.assertEquals(10, dataList.size()); } + + /** + * 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); + } + } }