From b8bbd504b03252cb53c204c330f33a6075b57272 Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Sun, 4 Jun 2023 16:55:28 +0800
Subject: [PATCH] core: add preferredPrefix option for problem_import

---
 packages/hydrooj/src/handler/import.ts |  9 ++++++---
 packages/hydrooj/src/model/problem.ts  | 16 +++++++++++++---
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/packages/hydrooj/src/handler/import.ts b/packages/hydrooj/src/handler/import.ts
index 23e831a08..a9f771ac1 100644
--- a/packages/hydrooj/src/handler/import.ts
+++ b/packages/hydrooj/src/handler/import.ts
@@ -2,17 +2,20 @@ import { Context } from '../context';
 import { ValidationError } from '../error';
 import { PERM, PRIV } from '../model/builtin';
 import problem from '../model/problem';
-import { Handler } from '../service/server';
+import { Handler, param, Types } from '../service/server';
 
 class ProblemImportHydroHandler extends Handler {
     async get() {
         this.response.template = 'problem_import.html';
     }
 
-    async post({ domainId, keepUser }) {
+    @param('keepUser', Types.Boolean)
+    @param('preferredPrefix', Types.String, true)
+    async post(domainId: string, keepUser: boolean, preferredPrefix?: string) {
         if (keepUser) this.checkPriv(PRIV.PRIV_EDIT_SYSTEM);
         if (!this.request.files.file) throw new ValidationError('file');
-        await problem.import(domainId, this.request.files.file.filepath, keepUser ? null : this.user._id);
+        if (preferredPrefix && !/^[a-zA-Z]+$/.test(preferredPrefix)) throw new ValidationError('preferredPrefix');
+        await problem.import(domainId, this.request.files.file.filepath, keepUser ? null : this.user._id, preferredPrefix);
         this.response.redirect = this.url('problem_main');
     }
 }
diff --git a/packages/hydrooj/src/model/problem.ts b/packages/hydrooj/src/model/problem.ts
index 1546b50cb..be237202b 100644
--- a/packages/hydrooj/src/model/problem.ts
+++ b/packages/hydrooj/src/model/problem.ts
@@ -407,7 +407,7 @@ export class ProblemModel {
         return true;
     }
 
-    static async import(domainId: string, filepath: string, operator = 1) {
+    static async import(domainId: string, filepath: string, operator = 1, preferredPrefix?: string) {
         let tmpdir = '';
         let del = false;
         if (filepath.endsWith('.zip')) {
@@ -441,9 +441,19 @@ export class ProblemModel {
                 const pdoc: ProblemDoc = yaml.load(content) as any;
                 if (!pdoc) continue;
                 let pid = pdoc.pid;
+
+                const isValidPid = async (id: string) => {
+                    if (!(/^[A-Za-z]+[0-9A-Za-z]*$/.test(id))) return false;
+                    if (await ProblemModel.get(domainId, id)) return false;
+                    return true;
+                };
+
                 if (pid) {
-                    if (!(/^[A-Za-z]+[0-9A-Za-z]*$/.test(pid))) pid = undefined;
-                    else if (await ProblemModel.get(domainId, pid)) pid = undefined;
+                    if (preferredPrefix) {
+                        const newPid = pid.replace(/^[A-Za-z]+/, preferredPrefix);
+                        if (isValidPid(newPid)) pid = newPid;
+                    }
+                    if (!isValidPid(pid)) pid = undefined;
                 }
                 const overrideContent = findOverrideContent(path.join(tmpdir, i));
                 const docId = await ProblemModel.add(