More race patching

#239
jdabtieu · Dec 6, 2023 · 9b5093c · 9b5093c
1 parent 4b0237b
commit 9b5093c
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 18 deletions.
diff --git a/src/application.py b/src/application.py
@@ -321,14 +321,10 @@ def register():
     except ValueError:
         if db.execute("SELECT COUNT(*) AS cnt FROM users WHERE username=?", username)[0]["cnt"] > 0:
             flash('Username already exists', 'danger')
-            return render_template("auth/register.html",
-                                   site_key=app.config['HCAPTCHA_SITE']), 400
         elif db.execute("SELECT COUNT(*) AS cnt FROM users WHERE email=?", email)[0]["cnt"] > 0:
             flash('Email already exists', 'danger')
-            return render_template("auth/register.html",
-                                   site_key=app.config['HCAPTCHA_SITE']), 400
-        else:
-            abort(500)
+        return render_template("auth/register.html",
+                               site_key=app.config['HCAPTCHA_SITE']), 400
 
     if not app.config['TESTING']:
         token = create_jwt({'email': email}, app.config['SECRET_KEY'])
@@ -731,19 +727,16 @@ def create_problem():
     if not hints:
         hints = ""
 
-    # Ensure problem does not already exist
-    problem_info = db.execute("SELECT * FROM problems WHERE id=:problem_id OR name=:name",
-                              problem_id=problem_id, name=name)
-    if len(problem_info) != 0:
+    # Create & ensure problem doesn't already exist
+    try:
+        db.execute(("INSERT INTO problems (id, name, point_value, category, flag, draft, "
+                    "flag_hint, instanced) VALUES (:id, :name, :point_value, :category, "
+                    ":flag, :draft, :fhint, :inst)"),
+                   id=problem_id, name=name, point_value=point_value, category=category,
+                   flag=flag, draft=draft, fhint=flag_hint, inst=instanced)
+    except ValueError:
         flash('A problem with this name or ID already exists', 'danger')
-        return render_template("problem/create.html"), 409
-
-    # Modify problems table
-    db.execute(("INSERT INTO problems (id, name, point_value, category, flag, draft, "
-                "flag_hint, instanced) VALUES (:id, :name, :point_value, :category, "
-                ":flag, :draft, :fhint, :inst)"),
-               id=problem_id, name=name, point_value=point_value, category=category,
-               flag=flag, draft=draft, fhint=flag_hint, inst=instanced)
+        return render_template("problem/create.html"), 400
 
     # Check if file exists & upload if it does
     file = request.files["file"]

diff --git a/src/tests/test_problems.py b/src/tests/test_problems.py
@@ -35,6 +35,38 @@ def test_problem(client, database):
     })
     assert result.status_code == 302
 
+    result = client.post('/problems/create', data={
+        'id': 'hello world testing',
+        'name': 'hello world',
+        'description': 'a short fun problem',
+        'hints': 'try looking at the title',
+        'point_value': 1,
+        'category': 'general',
+        'flag': 'ctf{hello}',
+        'flag_hint': 'ctf{...}',
+        'instanced': True,
+        'file': ('test_upload.txt', 'test_upload.txt'),
+        'draft': True
+    })
+    assert result.status_code == 400
+    assert b'Invalid problem ID' in result.data
+
+    result = client.post('/problems/create', data={
+        'id': 'helloworldtesting',
+        'name': 'hello world',
+        'description': 'a short fun problem',
+        'hints': 'try looking at the title',
+        'point_value': 1,
+        'category': 'general',
+        'flag': 'a very nefariout ',
+        'flag_hint': 'ctf{...}',
+        'instanced': True,
+        'file': ('test_upload.txt', 'test_upload.txt'),
+        'draft': True
+    })
+    assert result.status_code == 400
+    assert b'already exists' in result.data
+
     # TODO Assert the instancer interface exists
 
     result = client.post('/problem/helloworldtesting',