feat: custom file IDs

2 files changed, 34 insertions, 12 deletions

diff --git a/public/index.php b/public/index.php
index 0493cc4..f5809a5 100644
--- a/public/index.php
+++ b/public/index.php
@@ -339,6 +339,14 @@ $privacy_exists = is_file($_SERVER['DOCUMENT_ROOT'] . '/static/PRIVACY.txt');
                             <p class="remove-script">Details:</p>
                             <hr class="remove-script">
                             <table class="vertical left" id="form-upload-options">
+                                <?php if (FILE_CUSTOM_ID): ?>
+                                    <tr>
+                                        <th>File ID:</th>
+                                        <td><input type="text" name="id" placeholder="Leave empty for a random ID"
+                                                maxlength="<?= FILE_CUSTOM_ID_LENGTH ?>">
+                                        </td>
+                                    </tr>
+                                <?php endif; ?>
                                 <tr>
                                     <th>Title:</th>
                                     <td>
diff --git a/public/upload.php b/public/upload.php
index 08de5a1..21219ad 100644
--- a/public/upload.php
+++ b/public/upload.php
@@ -128,20 +128,34 @@ try {
 
     $db = new PDO(DB_URL, DB_USER, DB_PASS);
 
-    $file_id_length = FILE_ID_LENGTH;
-    $file_id_gen_attempts = 0;
-    $sql = 'SELECT id FROM files WHERE id = ? AND extension = ?';
-    do {
-        $file_id = FILE_ID_PREFIX . generate_random_char_sequence(FILE_ID_CHARACTERS, $file_id_length);
-        if ($file_id_gen_attempts > 20) {
-            $file_id_length++;
-            $file_id_gen_attempts = 0;
+    if (FILE_CUSTOM_ID && isset($_POST['id']) && !empty(trim($_POST['id']))) {
+        $file_id = $_POST['id'];
+        if (!preg_match(FILE_CUSTOM_ID_REGEX, $file_id) || strlen($file_id) > FILE_CUSTOM_ID_LENGTH) {
+            throw new RuntimeException('Invalid file ID.');
         }
-        $file_id_gen_attempts++;
 
-        $stmt = $db->prepare($sql);
-        $stmt->execute([$file_id, $file_data['extension']]);
-    } while ($stmt->rowCount() > 0);
+        $stmt = $db->prepare('SELECT id FROM files WHERE id = ?');
+        $stmt->execute([$file_id]);
+        if ($stmt->rowCount() > 0) {
+            throw new RuntimeException('File ID has already been taken.');
+        }
+    } else {
+        $file_id_length = FILE_ID_LENGTH;
+        $file_id_gen_attempts = 0;
+        $sql = 'SELECT id FROM files WHERE id = ? AND extension = ?';
+        do {
+            $file_id = FILE_ID_PREFIX . generate_random_char_sequence(FILE_ID_CHARACTERS, $file_id_length);
+            if ($file_id_gen_attempts > 20) {
+                $file_id_length++;
+                $file_id_gen_attempts = 0;
+            }
+            $file_id_gen_attempts++;
+
+            $stmt = $db->prepare($sql);
+            $stmt->execute([$file_id, $file_data['extension']]);
+        } while ($stmt->rowCount() > 0);
+    }
+
     $file_data['id'] = $file_id;
 
     if (isset($url)) {