Harden Zip CLI input validation

Author: krotname

src/main/java/algorithms/sprint0/Zip.java

@@ -8,15 +8,18 @@
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.io.OutputStreamWriter;
+import java.io.StringReader;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 
 import static algorithms.sprint0.Utils.printList;
-import static algorithms.sprint0.Utils.readList;
 
 public class Zip {
 
+    private static final int MAX_LIST_SIZE = 100_000;
+    private static final int MAX_INPUT_LINE_LENGTH = 1_200_001;
+
     static List<Integer> zip(List<Integer> a, List<Integer> b, int n) {
         if (n < 0) {
             throw new IllegalArgumentException("n >= 0 required");
@@ -34,14 +37,49 @@ static List<Integer> zip(List<Integer> a, List<Integer> b, int n) {
     public static void main(String[] args) throws IOException {
         try (BufferedReader reader = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8));
              BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(System.out, StandardCharsets.UTF_8))) {
-            String sizeLine = reader.readLine();
-            if (sizeLine == null) {
-                throw new EOFException("Missing list size");
+            try {
+                process(reader, writer);
+            } catch (IllegalArgumentException | EOFException exception) {
+                System.err.println("Invalid input: " + exception.getMessage());
+            }
+        }
+    }
+
+    static void process(BufferedReader reader, BufferedWriter writer) throws IOException {
+        String sizeLine = readBoundedLine(reader);
+        if (sizeLine == null) {
+            throw new EOFException("Missing list size");
+        }
+        int n = parseInt(sizeLine.trim());
+        if (n < 0 || n > MAX_LIST_SIZE) {
+            throw new IllegalArgumentException("List size must be between 0 and " + MAX_LIST_SIZE);
+        }
+        List<Integer> a = parseList(readBoundedLine(reader));
+        List<Integer> b = parseList(readBoundedLine(reader));
+        if (a.size() < n || b.size() < n) {
+            throw new IllegalArgumentException("Each list must contain at least n integers");
+        }
+        printList(zip(a, b, n), writer);
+    }
+
+    private static String readBoundedLine(BufferedReader reader) throws IOException {
+        StringBuilder line = new StringBuilder();
+        int character;
+        while ((character = reader.read()) != -1 && character != '\n') {
+            if (line.length() == MAX_INPUT_LINE_LENGTH) {
+                throw new IllegalArgumentException("Input line is too long");
+            }
+            if (character != '\r') {
+                line.append((char) character);
             }
-            int n = parseInt(sizeLine.trim());
-            List<Integer> a = readList(reader);
-            List<Integer> b = readList(reader);
-            printList(zip(a, b, n), writer);
         }
+        return character == -1 && line.length() == 0 ? null : line.toString();
+    }
+
+    private static List<Integer> parseList(String line) throws IOException {
+        if (line == null) {
+            throw new EOFException("Missing integer list");
+        }
+        return Utils.readList(new BufferedReader(new StringReader(line)));
     }
 }

src/test/java/algorithms/sprint0/ZipTest.java

@@ -4,6 +4,11 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.io.BufferedReader;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.StringReader;
+import java.io.StringWriter;
 import java.util.Arrays;
 import java.util.List;
 import org.junit.jupiter.api.Tag;
@@ -48,4 +53,32 @@ void rejectsNullLists() {
         assertThrows(NullPointerException.class, () -> Zip.zip(null, List.of(2), 1));
         assertThrows(NullPointerException.class, () -> Zip.zip(List.of(1), null, 1));
     }
+
+    @Test
+    void processAcceptsValidInput() throws IOException {
+        StringWriter output = new StringWriter();
+        BufferedWriter writer = new BufferedWriter(output);
+
+        Zip.process(new BufferedReader(new StringReader("3\n1 5 6\n7 8 9\n")), writer);
+        writer.flush();
+
+        assertEquals("1 7 5 8 6 9 ", output.toString());
+    }
+
+    @Test
+    void processRejectsMissingOrShortLists() {
+        assertThrows(IOException.class, () -> process("3\n1 2 3\n"));
+        assertThrows(IllegalArgumentException.class, () -> process("3\n1 2\n4 5 6\n"));
+    }
+
+    @Test
+    void processRejectsUnboundedSizesAndLines() {
+        assertThrows(IllegalArgumentException.class, () -> process("100001\n1\n2\n"));
+        String oversizedLine = "1".repeat(1_200_002);
+        assertThrows(IllegalArgumentException.class, () -> process("1\n" + oversizedLine + "\n2\n"));
+    }
+
+    private static void process(String input) throws IOException {
+        Zip.process(new BufferedReader(new StringReader(input)), new BufferedWriter(new StringWriter()));
+    }
 }