From 2c1cccf25a82425678094758727afb23b9a347f8 Mon Sep 17 00:00:00 2001
From: Liam <liam@fascinated.cc>
Date: Sun, 2 Jun 2024 10:40:46 +0100
Subject: [PATCH] block file some common headers

---
 .../backend/common/FileHeaderChecker.java     | 77 +++++++++++++++++++
 .../backend/controller/PasteController.java   |  2 +-
 .../backend/service/PasteService.java         |  6 ++
 3 files changed, 84 insertions(+), 1 deletion(-)
 create mode 100644 src/main/java/cc/fascinated/backend/common/FileHeaderChecker.java

diff --git a/src/main/java/cc/fascinated/backend/common/FileHeaderChecker.java b/src/main/java/cc/fascinated/backend/common/FileHeaderChecker.java
new file mode 100644
index 0000000..a22b109
--- /dev/null
+++ b/src/main/java/cc/fascinated/backend/common/FileHeaderChecker.java
@@ -0,0 +1,77 @@
+package cc.fascinated.backend.common;
+
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * @author Fascinated (fascinated7)
+ */
+public class FileHeaderChecker {
+    // Define known file headers
+    private static final Map<String, byte[]> FILE_HEADERS = new HashMap<>();
+
+    static {
+        FILE_HEADERS.put("PNG", new byte[]{(byte) 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A});
+        FILE_HEADERS.put("JPEG", new byte[]{(byte) 0xFF, (byte) 0xD8, (byte) 0xFF});
+        FILE_HEADERS.put("GIF", new byte[]{0x47, 0x49, 0x46, 0x38});
+        FILE_HEADERS.put("PDF", new byte[]{0x25, 0x50, 0x44, 0x46});
+        FILE_HEADERS.put("ZIP", new byte[]{0x50, 0x4B, 0x03, 0x04});
+        FILE_HEADERS.put("RAR", new byte[]{0x52, 0x61, 0x72, 0x21});
+        FILE_HEADERS.put("7Z", new byte[]{0x37, 0x7A, (byte) 0xBC, (byte) 0xAF, (byte) 0x27, 0x1C});
+        FILE_HEADERS.put("TAR", new byte[]{0x75, 0x73, 0x74, 0x61, 0x72});
+        FILE_HEADERS.put("BZ2", new byte[]{0x42, 0x5A, 0x68});
+        FILE_HEADERS.put("GZ", new byte[]{0x1F, (byte) 0x8B});
+        FILE_HEADERS.put("EXE", new byte[]{0x4D, 0x5A});
+    }
+
+    /**
+     * Converts a string to a byte array using UTF-8 encoding.
+     *
+     * @param input the string to convert
+     * @return the byte array representation of the string
+     */
+    private static byte[] stringToByteArray(String input) {
+        return input.getBytes(StandardCharsets.UTF_8);
+    }
+
+    /**
+     * Checks if the byte array starts with the specified header.
+     *
+     * @param byteArray the byte array to check
+     * @param header the header to look for
+     * @return true if the byte array starts with the header, false otherwise
+     */
+    private static boolean startsWith(byte[] byteArray, byte[] header) {
+        if (byteArray.length < header.length) {
+            return false;
+        }
+
+        for (int i = 0; i < header.length; i++) {
+            if (byteArray[i] != header[i]) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    /**
+     * Checks if the given string contains a known file header.
+     *
+     * @param input the string to check
+     * @return true if the string contains a known file header, false otherwise
+     */
+    public static boolean containsFileHeader(String input) {
+        System.out.println("Checking for file headers in: " + input);
+        byte[] byteArray = stringToByteArray(input);
+
+        for (byte[] header : FILE_HEADERS.values()) {
+            if (startsWith(byteArray, header)) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+}
diff --git a/src/main/java/cc/fascinated/backend/controller/PasteController.java b/src/main/java/cc/fascinated/backend/controller/PasteController.java
index c8ed17c..3dd8cad 100644
--- a/src/main/java/cc/fascinated/backend/controller/PasteController.java
+++ b/src/main/java/cc/fascinated/backend/controller/PasteController.java
@@ -20,7 +20,7 @@ public class PasteController {
         this.pasteService = pasteService;
     }
 
-    @PostMapping(value = "/upload", consumes = MediaType.TEXT_PLAIN_VALUE)
+    @PostMapping(value = "/upload")
     public ResponseEntity<?> uploadPaste(@RequestBody String content) {
         String id = pasteService.createPaste(content);
         return ResponseEntity.ok(Map.of("id", id));
diff --git a/src/main/java/cc/fascinated/backend/service/PasteService.java b/src/main/java/cc/fascinated/backend/service/PasteService.java
index 984b77f..3574f4e 100644
--- a/src/main/java/cc/fascinated/backend/service/PasteService.java
+++ b/src/main/java/cc/fascinated/backend/service/PasteService.java
@@ -1,5 +1,6 @@
 package cc.fascinated.backend.service;
 
+import cc.fascinated.backend.common.FileHeaderChecker;
 import cc.fascinated.backend.exception.impl.BadRequestException;
 import cc.fascinated.backend.exception.impl.ResourceNotFoundException;
 import cc.fascinated.backend.model.Paste;
@@ -54,6 +55,11 @@ public class PasteService {
             throw new BadRequestException("The paste content is too large, the limit is " + uploadSizeLimit + " characters");
         }
 
+        // Ensure the paste content does not contain a file header.
+        if (FileHeaderChecker.containsFileHeader(content)) {
+            throw new BadRequestException("The paste content contains a file header");
+        }
+
         // Save the paste to the database.
         Paste paste = pasteRepository.save(new Paste(
                 RandomStringUtils.randomAlphabetic(idLength),