From 5599ed001b20bdb616f39a97e7701253f20d1b94 Mon Sep 17 00:00:00 2001
From: Cameron Custer <73217097+cameroncuster@users.noreply.github.com>
Date: Thu, 23 Jul 2026 21:31:52 +0000
Subject: [PATCH] fix(security): validate problem source hosts

---
 e2e/data.spec.ts                  |  5 +++--
 src/lib/services/problem.ts       |  8 +-------
 src/lib/services/problemSource.ts |  6 ++++++
 tests/problem-source.test.ts      | 24 ++++++++++++++++++++++++
 4 files changed, 34 insertions(+), 9 deletions(-)
 create mode 100644 src/lib/services/problemSource.ts
 create mode 100644 tests/problem-source.test.ts

diff --git a/e2e/data.spec.ts b/e2e/data.spec.ts
index d81c90a..cb6c128 100644
--- a/e2e/data.spec.ts
+++ b/e2e/data.spec.ts
@@ -2,6 +2,7 @@ import { test, expect, type Page } from '@playwright/test';
 import { collect, expectClean, waitForShell } from './support/harness.ts';
 import { PROBLEMS, CONTESTS, LEADERBOARD } from './support/fixtures.ts';
 import { setScenario } from './support/scenario.ts';
+import { getProblemSource } from '../src/lib/services/problemSource.ts';
 
 test.beforeEach(async () => {
   await setScenario('data');
@@ -67,7 +68,7 @@ test.describe('problems data', () => {
 
     // First click -> codeforces only.
     await sourceFilter.click();
-    const codeforces = PROBLEMS.filter((p) => !p.url.includes('kattis.com'));
+    const codeforces = PROBLEMS.filter((p) => getProblemSource(p.url) === 'codeforces');
     await expect(rows).toHaveCount(codeforces.length);
     for (const p of codeforces) {
       await expect(page.getByRole('link', { name: p.name, exact: true })).toBeVisible();
@@ -75,7 +76,7 @@ test.describe('problems data', () => {
 
     // Second click -> kattis only.
     await sourceFilter.click();
-    const kattis = PROBLEMS.filter((p) => p.url.includes('kattis.com'));
+    const kattis = PROBLEMS.filter((p) => getProblemSource(p.url) === 'kattis');
     await expect(rows).toHaveCount(kattis.length);
     for (const p of kattis) {
       await expect(page.getByRole('link', { name: p.name, exact: true })).toBeVisible();
diff --git a/src/lib/services/problem.ts b/src/lib/services/problem.ts
index 57b9de0..482f5c4 100644
--- a/src/lib/services/problem.ts
+++ b/src/lib/services/problem.ts
@@ -3,6 +3,7 @@
  */
 import { supabase } from './database';
 import { user } from './auth';
+import { getProblemSource } from './problemSource';
 import { get } from 'svelte/store';
 
 /**
@@ -52,13 +53,6 @@ export type UserSolvedProblem = {
   solved_at: string;
 };
 
-/**
- * Determine the problem source based on URL
- */
-export function getProblemSource(url: string): 'codeforces' | 'kattis' {
-  return url.includes('kattis.com') ? 'kattis' : 'codeforces';
-}
-
 /**
  * Columns selected when reading problems for display. Kept explicit so we only
  * fetch what the UI needs instead of every row column.
diff --git a/src/lib/services/problemSource.ts b/src/lib/services/problemSource.ts
new file mode 100644
index 0000000..366a972
--- /dev/null
+++ b/src/lib/services/problemSource.ts
@@ -0,0 +1,6 @@
+import { parseKattisProblemId } from './kattisUrl.ts';
+
+/** Determine the problem source using the exact-host Kattis URL validator. */
+export function getProblemSource(url: string): 'codeforces' | 'kattis' {
+  return parseKattisProblemId(url) ? 'kattis' : 'codeforces';
+}
diff --git a/tests/problem-source.test.ts b/tests/problem-source.test.ts
new file mode 100644
index 0000000..7a11bf3
--- /dev/null
+++ b/tests/problem-source.test.ts
@@ -0,0 +1,24 @@
+import { test } from 'node:test';
+import assert from 'node:assert/strict';
+import { getProblemSource } from '../src/lib/services/problemSource.ts';
+
+test('classifies canonical Kattis problem URLs', () => {
+  assert.equal(getProblemSource('https://open.kattis.com/problems/hello'), 'kattis');
+  assert.equal(getProblemSource('https://kattis.com/problems/hello'), 'kattis');
+});
+
+test('classifies ordinary Codeforces problem URLs', () => {
+  assert.equal(getProblemSource('https://codeforces.com/problemset/problem/1/A'), 'codeforces');
+});
+
+test('does not classify URLs by an untrusted hostname substring', () => {
+  for (const url of [
+    'https://open.kattis.com.evil.example/problems/hello',
+    'https://evil.example/open.kattis.com/problems/hello',
+    'https://open.kattis.com@evil.example/problems/hello',
+    'https://evil.example/?next=https://open.kattis.com/problems/hello',
+    'not-a-url-containing-kattis.com'
+  ]) {
+    assert.equal(getProblemSource(url), 'codeforces');
+  }
+});