Checked quota limit during tar/zip archive generation

src/main/java/it/inaf/ia2/transfer/controller/PutFileController.java

@@ -73,7 +73,7 @@ public class PutFileController extends FileController {
                 // if MultipartFile provides file size it is possible to check
                 // quota limit before reading the stream
                 if (remainingQuota != null && file != null && file.getSize() > remainingQuota) {
-                    throw new InsufficientStorageException(fileInfo.getVirtualPath());
+                    throw new InsufficientStorageException("QuotaExceeded Path: " + fileInfo.getVirtualPath());
                 }
                 
                 if (file != null) {
@@ -127,7 +127,7 @@ public class PutFileController extends FileController {
             // Quota limit is checked again to handle cases where MultipartFile is not used
             if (remainingQuota != null && fileSize > remainingQuota) {
                 file.delete();
-                throw new InsufficientStorageException(fileInfo.getVirtualPath());
+                throw new InsufficientStorageException("QuotaExceeded Path: " + fileInfo.getVirtualPath());
             }
 
             String md5Checksum = makeMD5Checksum(file);

src/main/java/it/inaf/ia2/transfer/exception/InsufficientStorageException.java

@@ -11,8 +11,8 @@ import org.springframework.web.bind.annotation.ResponseStatus;
 @ResponseStatus(HttpStatus.INSUFFICIENT_STORAGE)
 public class InsufficientStorageException extends JobException {
 
-    public InsufficientStorageException(String path) {
+    public InsufficientStorageException(String errorDetail) {
         super(Type.FATAL, "Quota Exceeded");
-        setErrorDetail("QuotaExceeded Path: " + path);
+        setErrorDetail(errorDetail);
     }
 }

src/main/java/it/inaf/ia2/transfer/service/ArchiveService.java

@@ -6,6 +6,7 @@
 package it.inaf.ia2.transfer.service;
 
 import it.inaf.ia2.transfer.auth.TokenPrincipal;
+import it.inaf.ia2.transfer.exception.InsufficientStorageException;
 import it.inaf.ia2.transfer.exception.JobException;
 import it.inaf.ia2.transfer.exception.JobException.Type;
 import it.inaf.ia2.transfer.persistence.FileDAO;
@@ -38,6 +39,7 @@ import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.stereotype.Service;
 import org.springframework.util.FileSystemUtils;
+import org.springframework.util.unit.DataSize;
 import org.springframework.web.client.RestTemplate;
 
 @Service
@@ -63,10 +65,15 @@ public class ArchiveService {
     @Value("${upload_location_id}")
     private int uploadLocationId;
 
+    // Directory containing temporary files generated by jobs.
     @Value("${generated.dir}")
     private String generatedDirString;
     private File generatedDir;
 
+    // Maximum size of the working directory for each registered user
+    @Value("${generated.dir.max-size}")
+    private DataSize generatedDirMaxSize;
+
     @PostConstruct
     public void init() {
         this.generatedDir = new File(generatedDirString);
@@ -138,7 +145,10 @@ public class ArchiveService {
             }
         }
 
+        checkQuotaLimit(parentDir);
+
         File archiveFile = parentDir.toPath().resolve(job.getJobId() + "." + job.getType().getExtension()).toFile();
+
         if (!archiveFile.createNewFile()) {
             LOG.error("Unable to create file " + archiveFile.getAbsolutePath());
             throw new JobException(Type.FATAL, "Internal Fault")
@@ -148,6 +158,18 @@ public class ArchiveService {
         return archiveFile;
     }
 
+    /**
+     * If used working space exceeds quota limit throws an
+     * InsufficientStorageException.
+     */
+    private void checkQuotaLimit(File parentDir) throws IOException {
+        long usedSpace = Files.walk(parentDir.toPath()).mapToLong(p -> p.toFile().length()).sum();
+
+        if (usedSpace > generatedDirMaxSize.toBytes()) {
+            throw new InsufficientStorageException("Archive size limit exceeded.");
+        }
+    }
+
     public File getArchiveParentDir(Principal principal) {
         return generatedDir.toPath().resolve(principal.getName()).toFile();
     }
@@ -173,20 +195,23 @@ public class ArchiveService {
         return commonParent;
     }
 
-    private static abstract class ArchiveHandler<O extends OutputStream, E> implements AutoCloseable {
+    private abstract class ArchiveHandler<O extends OutputStream, E> implements AutoCloseable {
 
         private final O os;
+        private final File parentDir;
 
-        ArchiveHandler(O os) {
+        ArchiveHandler(O os, File parentDir) {
             this.os = os;
+            this.parentDir = parentDir;
         }
 
         public abstract E getEntry(File file, String path);
 
-        public abstract void putNextEntry(E entry) throws IOException;
+        protected abstract void putNextEntry(E entry) throws IOException;
 
-        public void putNextEntry(File file, String path) throws IOException {
+        public final void putNextEntry(File file, String path) throws IOException {
             putNextEntry(getEntry(file, path));
+            checkQuotaLimit(parentDir);
         }
 
         public final O getOutputStream() {
@@ -202,7 +227,7 @@ public class ArchiveService {
     private class TarArchiveHandler extends ArchiveHandler<TarOutputStream, TarEntry> {
 
         TarArchiveHandler(File archiveFile) throws IOException {
-            super(new TarOutputStream(new BufferedOutputStream(new FileOutputStream(archiveFile))));
+            super(new TarOutputStream(new BufferedOutputStream(new FileOutputStream(archiveFile))), archiveFile.getParentFile());
         }
 
         @Override
@@ -219,7 +244,7 @@ public class ArchiveService {
     private class ZipArchiveHandler extends ArchiveHandler<ZipOutputStream, ZipEntry> {
 
         ZipArchiveHandler(File archiveFile) throws IOException {
-            super(new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(archiveFile))));
+            super(new ZipOutputStream(new BufferedOutputStream(new FileOutputStream(archiveFile))), archiveFile.getParentFile());
         }
 
         @Override

src/main/resources/application.properties

@@ -6,6 +6,7 @@ file-catalog.datasource.username=postgres
 file-catalog.datasource.password=
 
 generated.dir=/tmp/vospace/gen
+generated.dir.max-size=10GB
 
 gms_base_url=https://sso.ia2.inaf.it/gms
 jwks_uri=https://sso.ia2.inaf.it/rap-ia2/auth/oidc/jwks

src/test/java/it/inaf/ia2/transfer/service/ArchiveServiceTest.java

@@ -6,6 +6,7 @@
 package it.inaf.ia2.transfer.service;
 
 import it.inaf.ia2.transfer.auth.TokenPrincipal;
+import it.inaf.ia2.transfer.exception.InsufficientStorageException;
 import it.inaf.ia2.transfer.persistence.FileDAO;
 import it.inaf.ia2.transfer.persistence.JobDAO;
 import it.inaf.ia2.transfer.persistence.LocationDAO;
@@ -13,6 +14,7 @@ import it.inaf.ia2.transfer.persistence.model.FileInfo;
 import java.io.ByteArrayInputStream;
 import java.io.File;
 import java.io.FileInputStream;
+import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.UncheckedIOException;
@@ -25,50 +27,53 @@ import java.util.function.Function;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
 import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.Assertions;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.extension.ExtendWith;
 import org.kamranzafar.jtar.TarEntry;
 import org.kamranzafar.jtar.TarInputStream;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
-import org.mockito.InjectMocks;
-import org.mockito.Mock;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
-import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.mock.mockito.MockBean;
+import org.springframework.boot.test.util.TestPropertyValues;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.client.ClientHttpResponse;
-import org.springframework.test.util.ReflectionTestUtils;
+import org.springframework.test.context.ContextConfiguration;
 import org.springframework.util.FileSystemUtils;
 import org.springframework.web.client.RequestCallback;
 import org.springframework.web.client.ResponseExtractor;
 import org.springframework.web.client.RestTemplate;
 
-@ExtendWith(MockitoExtension.class)
+@SpringBootTest
+@ContextConfiguration(initializers = ArchiveServiceTest.TestPropertiesInitializer.class)
 public class ArchiveServiceTest {
 
-    @Mock
+    @MockBean
     private JobDAO jobDAO;
 
-    @Mock
+    @MockBean
     private FileDAO fileDAO;
 
-    @Mock
+    @MockBean
     private LocationDAO locationDAO;
 
-    @Mock
+    @MockBean
     private RestTemplate restTemplate;
 
-    @Mock
+    @MockBean
     private AuthorizationService authorizationService;
 
-    @InjectMocks
+    @Autowired
     private ArchiveService archiveService;
 
     private static File tmpDir;
@@ -83,11 +88,6 @@ public class ArchiveServiceTest {
         FileSystemUtils.deleteRecursively(tmpDir);
     }
 
-    @BeforeEach
-    public void setUp() {
-        ReflectionTestUtils.setField(archiveService, "generatedDir", tmpDir);
-    }
-
     @Test
     public void testTarGeneration() throws Exception {
 
@@ -130,6 +130,32 @@ public class ArchiveServiceTest {
         });
     }
 
+    @Test
+    public void testArchiveQuotaExceeded() throws Exception {
+
+        ArchiveJob job = new ArchiveJob();
+        job.setPrincipal(new TokenPrincipal("user2", "token2"));
+        job.setJobId("job2");
+        job.setType(ArchiveJob.Type.ZIP);
+        job.setVosPaths(Arrays.asList("/ignore"));
+
+        File user2Dir = tmpDir.toPath().resolve("user2").toFile();
+        user2Dir.mkdir();
+
+        File fillQuotaFile = user2Dir.toPath().resolve("fillQuotaFile").toFile();
+
+        // create a file bigger than test quota limit (20 KB)
+        try (FileInputStream fis = new FileInputStream("/dev/zero");
+                FileOutputStream fos = new FileOutputStream(fillQuotaFile)) {
+            byte[] junk = fis.readNBytes(20 * 1024);
+            fos.write(junk);
+        }
+
+        Assertions.assertThrows(InsufficientStorageException.class, () -> {
+            archiveService.createArchive(job);
+        });
+    }
+
     private static abstract class TestArchiveHandler<I extends InputStream, E> {
 
         private final I is;
@@ -163,7 +189,7 @@ public class ArchiveServiceTest {
         File file7 = createFile(tmpParent, "portal-file");
 
         ArchiveJob job = new ArchiveJob();
-        job.setPrincipal(new TokenPrincipal("user123", "token123"));
+        job.setPrincipal(new TokenPrincipal("user1", "token1"));
         job.setJobId("abcdef");
         job.setType(type);
         job.setVosPaths(Arrays.asList(parent + "/dir1", parent + "/dir2", parent + "/file6"));
@@ -200,7 +226,7 @@ public class ArchiveServiceTest {
 
         archiveService.createArchive(job);
 
-        File result = tmpDir.toPath().resolve("user123").resolve("abcdef." + extension).toFile();
+        File result = tmpDir.toPath().resolve("user1").resolve("abcdef." + extension).toFile();
 
         assertTrue(result.exists());
 
@@ -265,4 +291,20 @@ public class ArchiveServiceTest {
         }
         throw new IllegalStateException("Files have to be created");
     }
+
+    /**
+     * @TestPropertySource annotation can't be used in this test because we need
+     * to set the generated.dir property dynamically (since the test directory
+     * is generated by the @BeforeAll method), so this inner class is used to
+     * perform test property initialization.
+     */
+    static class TestPropertiesInitializer implements ApplicationContextInitializer<ConfigurableApplicationContext> {
+
+        @Override
+        public void initialize(ConfigurableApplicationContext configurableApplicationContext) {
+            TestPropertyValues.of("generated.dir=" + tmpDir.getAbsolutePath(),
+                    "generated.dir.max-size=20KB", "upload_location_id=3")
+                    .applyTo(configurableApplicationContext.getEnvironment());
+        }
+    }
 }