From 75f3a0ed04672d05be29d63e0fe455bcf898c22f Mon Sep 17 00:00:00 2001
From: Nixyan <tocka@tockanest.com>
Date: Wed, 11 Mar 2026 11:48:38 -0300
Subject: [PATCH] feat: enhance security and testing for federation routes.
 Added routes for uploading files to posts and initial logic of handling it
 client-side.

- Added a new test suite for attack vectors targeting the /discover federation routes, ensuring (known) vulnerabilities are addressed.
- Implemented a proxy function to check for blacklisted servers, enhancing security measures.
- Introduced URL validation to prevent SSRF attacks by blocking internal addresses.
- Updated package.json with a new test command for the attack tests.
- Refactored server and route handling to improve type safety and error handling.
- Added new middleware for blacklist checks and URL validation to prevent unauthorized access.
---
 package.json                                  |   1 +
 playwright.config.ts                          |   8 -
 src/app/PostTestForm.tsx                      |  90 ++++
 src/app/discover/rotate/confirm/route.ts      | 132 +++--
 src/app/discover/rotate/init/route.ts         |  17 +
 src/app/discover/route.ts                     | 130 +++--
 src/app/page.tsx                              |   3 +-
 src/lib/auth.ts                               |  11 +-
 src/lib/db/schema/index.ts                    |  34 +-
 src/lib/federation/blacklist-middleware.ts    |  19 +
 src/lib/federation/url-guard.ts               |  76 +++
 src/lib/plugins/client/social.ts              |  98 +++-
 .../server/helpers/social/endpoints/index.ts  |   4 +-
 .../server/helpers/social/endpoints/posts.ts  |  63 ++-
 .../plugins/server/helpers/social/social.ts   |   4 +
 .../plugins/server/storage/minio.client.ts    |  21 +
 src/proxy.ts                                  |  49 ++
 src/server.ts                                 |   4 +-
 tests/attacks.test.ts                         | 459 ++++++++++++++++++
 tests/helpers/db.ts                           |  15 +-
 20 files changed, 1086 insertions(+), 152 deletions(-)
 create mode 100644 src/app/PostTestForm.tsx
 create mode 100644 src/lib/federation/blacklist-middleware.ts
 create mode 100644 src/lib/federation/url-guard.ts
 create mode 100644 src/lib/plugins/server/storage/minio.client.ts
 create mode 100644 src/proxy.ts
 create mode 100644 tests/attacks.test.ts

diff --git a/package.json b/package.json
index dc9f50c..612eea2 100644
--- a/package.json
+++ b/package.json
@@ -19,6 +19,7 @@
 		"keygen": "bun run src/lib/federation/keygen.ts",
 		"test:key": "cross-env NODE_ENV=test playwright test tests/key.test.ts",
 		"test:discover": "cross-env NODE_ENV=test playwright test tests/discover.test.ts",
+		"test:attacks": "cross-env NODE_ENV=test playwright test tests/attacks.test.ts",
 		"build": "next build",
 		"start": "cross-env NODE_ENV=production node src/server.ts",
 		"db:push": "drizzle-kit push",
diff --git a/playwright.config.ts b/playwright.config.ts
index 7100f36..fe473a3 100644
--- a/playwright.config.ts
+++ b/playwright.config.ts
@@ -25,13 +25,5 @@ export default defineConfig({
 			name: 'chromium',
 			use: { ...devices['Desktop Chrome'] },
 		},
-		{
-			name: 'firefox',
-			use: { ...devices['Desktop Firefox'] },
-		},
-		{
-			name: 'webkit',
-			use: { ...devices['Desktop Safari'] },
-		},
 	],
 });
diff --git a/src/app/PostTestForm.tsx b/src/app/PostTestForm.tsx
new file mode 100644
index 0000000..6fcbc8f
--- /dev/null
+++ b/src/app/PostTestForm.tsx
@@ -0,0 +1,90 @@
+"use client";
+
+import { authClient } from "@/lib/auth-client";
+import { useState } from "react";
+
+export function PostTestForm() {
+	const [text, setText] = useState("");
+	const [files, setFiles] = useState<File[]>([]);
+	const [status, setStatus] = useState<string | null>(null);
+
+	const handleSubmit = async () => {
+		setStatus("Submitting...");
+		try {
+			const content: { type: "text" | "image"; value: string | File }[] = [];
+
+			if (text.trim()) {
+				content.push({ type: "text", value: text.trim() });
+			}
+
+			for (const file of files) {
+				content.push({ type: "image", value: file });
+			}
+
+			if (content.length === 0) {
+				setStatus("Add some text or images first.");
+				return;
+			}
+
+			const result = await authClient.createPost(content);
+			setStatus(`Done: ${JSON.stringify(result)}`);
+		} catch (err) {
+			setStatus(`Error: ${err instanceof Error ? err.message : String(err)}`);
+		}
+	};
+
+	return (
+		<div style={{ padding: 32, maxWidth: 480, margin: "0 auto", fontFamily: "sans-serif" }}>
+			<h2>Test Post</h2>
+
+			<textarea
+				value={text}
+				onChange={(e) => setText(e.target.value)}
+				placeholder="Write something..."
+				rows={4}
+				style={{ width: "100%", marginBottom: 12, padding: 8, fontSize: 14 }}
+			/>
+
+			<div style={{ marginBottom: 12 }}>
+				<label style={{ display: "block", marginBottom: 4, fontWeight: 600 }}>
+					Images
+				</label>
+				<input
+					type="file"
+					accept="image/*"
+					multiple
+					onChange={(e) => setFiles(Array.from(e.target.files ?? []))}
+				/>
+				{files.length > 0 && (
+					<div style={{ marginTop: 8, fontSize: 13, color: "#666" }}>
+						{files.map((f, i) => (
+							<div key={i}>{f.name} ({(f.size / 1024).toFixed(1)} KB)</div>
+						))}
+					</div>
+				)}
+			</div>
+
+			<button
+				onClick={handleSubmit}
+				style={{
+					padding: "10px 24px",
+					fontSize: 14,
+					fontWeight: 600,
+					cursor: "pointer",
+					background: "#111",
+					color: "#fff",
+					border: "none",
+					borderRadius: 6,
+				}}
+			>
+				Create Post
+			</button>
+
+			{status && (
+				<pre style={{ marginTop: 16, padding: 12, background: "#f4f4f4", borderRadius: 6, fontSize: 13, whiteSpace: "pre-wrap" }}>
+					{status}
+				</pre>
+			)}
+		</div>
+	);
+}
diff --git a/src/app/discover/rotate/confirm/route.ts b/src/app/discover/rotate/confirm/route.ts
index 2c30ae7..2d24853 100644
--- a/src/app/discover/rotate/confirm/route.ts
+++ b/src/app/discover/rotate/confirm/route.ts
@@ -1,10 +1,10 @@
 import db from "@/lib/db";
 import { blacklistedServers, rotateChallengeTokens, serverRegistry } from "@/lib/db/schema";
 import { decryptPayload } from "@/lib/federation/keytools";
+import createDebug from "debug";
 import { eq, sql } from "drizzle-orm";
 import { NextRequest, NextResponse } from "next/server";
 import { z } from "zod";
-import createDebug from "debug";
 
 const debug = createDebug("app:discover:rotate:confirm");
 
@@ -31,8 +31,8 @@ const debug = createDebug("app:discover:rotate:confirm");
  *    - If both match: update serverRegistry with newPublicKey and delete the challenge.
  *
  * What each check proves:
- * - signedOldChallenge match → SB holds the old private key (identity proof — "they are who they say they are")
- * - signedNewChallenge match → SB holds the new private key (ownership proof — "they own the key they want to rotate to")
+ * - signedOldChallenge match → SB holds the old private key (identity proof: "they are who they say they are")
+ * - signedNewChallenge match → SB holds the new private key (ownership proof: "they own the key they want to rotate to")
  * - re-encryption with SA's public key → SB fetched SA's identity from /discover
  *
  * TODO: on success, announce the completed rotation to other known federation peers
@@ -58,76 +58,74 @@ export async function POST(request: NextRequest) {
 	}
 
 	debug("POST /discover/rotate/confirm – fetching pending challenge for %s", validated.data.serverUrl);
-	const [challenge] = await db.select().from(rotateChallengeTokens)
-		.where(eq(rotateChallengeTokens.serverUrl, validated.data.serverUrl));
 
-	if (!challenge) {
-		debug("POST /discover/rotate/confirm – no pending challenge found");
-		return NextResponse.json({ error: "No pending rotation challenge found for this server." }, { status: 404 });
-	}
+	// transaction to ensure that the challenge is deleted and the server registry is updated atomically and that there's no race condition.
+	return await db.transaction(async (tx) => {
+		const [challenge] = await tx.select().from(rotateChallengeTokens)
+			.where(eq(rotateChallengeTokens.serverUrl, validated.data.serverUrl))
+			.for("update");
 
-	if (challenge.expiresAt < new Date()) {
-		debug("POST /discover/rotate/confirm – challenge expired at %s", challenge.expiresAt.toISOString());
-		await db.delete(rotateChallengeTokens).where(eq(rotateChallengeTokens.id, challenge.id));
-		return NextResponse.json({ error: "Challenge token has expired." }, { status: 400 });
-	}
+		if (!challenge) {
+			debug("POST /discover/rotate/confirm – no pending challenge found");
+			return NextResponse.json({ error: "No pending rotation challenge found for this server." }, { status: 404 });
+		}
 
-	if (challenge.attemptsLeft <= 0) {
-		debug("POST /discover/rotate/confirm – no attempts left, blacklisting %s", challenge.serverUrl);
-		await db.insert(blacklistedServers).values({
-			id: crypto.randomUUID(),
-			serverUrl: challenge.serverUrl,
-			reason: "Too many failed attempts to confirm key rotation challenge",
-			createdAt: new Date(),
-		});
-		await db.delete(rotateChallengeTokens).where(eq(rotateChallengeTokens.id, challenge.id));
-		return NextResponse.json({ error: "Your server has been blacklisted. Please contact support to unblacklist your server." }, { status: 403 });
-	}
+		if (challenge.expiresAt < new Date()) {
+			debug("POST /discover/rotate/confirm – challenge expired at %s", challenge.expiresAt.toISOString());
+			await tx.delete(rotateChallengeTokens).where(eq(rotateChallengeTokens.id, challenge.id));
+			return NextResponse.json({ error: "Challenge token has expired." }, { status: 400 });
+		}
 
-	debug("POST /discover/rotate/confirm – %d attempt(s) left, decrypting challenges", challenge.attemptsLeft);
-	let decryptedOld: string;
-	let decryptedNew: string;
-	try {
-		decryptedOld = decryptPayload(validated.data.signedOldChallenge, process.env.FEDERATION_PRIVATE_KEY!);
-		decryptedNew = decryptPayload(validated.data.signedNewChallenge, process.env.FEDERATION_PRIVATE_KEY!);
-	} catch {
-		debug("POST /discover/rotate/confirm – decryption failed, decrementing attempts");
-		await db.update(rotateChallengeTokens).set({
-			attemptsLeft: sql`${rotateChallengeTokens.attemptsLeft} - 1`,
-		}).where(eq(rotateChallengeTokens.id, challenge.id))
-		return NextResponse.json({
-			error: `Failed to decrypt one or both challenges. You have ${challenge.attemptsLeft - 1} attempts left before your server is blacklisted.`,
-		}, { status: 400 });
-	}
+		if (challenge.attemptsLeft <= 0) {
+			debug("POST /discover/rotate/confirm – no attempts left, blacklisting %s", challenge.serverUrl);
+			await tx.insert(blacklistedServers).values({
+				id: crypto.randomUUID(),
+				serverUrl: challenge.serverUrl,
+				reason: "Too many failed attempts to confirm key rotation challenge",
+				createdAt: new Date(),
+			});
+			await tx.delete(rotateChallengeTokens).where(eq(rotateChallengeTokens.id, challenge.id));
+			return NextResponse.json({ error: "Your server has been blacklisted. Please contact support to unblacklist your server." }, { status: 403 });
+		}
 
-	// Both plaintexts must match their stored tokens.
-	// A mismatch on oldKeyChallenge means the requester does not hold the registered private key.
-	// A mismatch on newKeyChallenge means the requester does not actually own the new key.
-	if (decryptedOld !== challenge.oldKeyToken || decryptedNew !== challenge.newKeyToken) {
-		debug("POST /discover/rotate/confirm – token mismatch (old=%s, new=%s), decrementing attempts",
-			decryptedOld === challenge.oldKeyToken ? "ok" : "MISMATCH",
-			decryptedNew === challenge.newKeyToken ? "ok" : "MISMATCH",
-		);
-		await db.update(rotateChallengeTokens).set({
-			attemptsLeft: sql`${rotateChallengeTokens.attemptsLeft} - 1`,
-		}).where(eq(rotateChallengeTokens.id, challenge.id));
-		return NextResponse.json({
-			error: `Challenge mismatch. You have ${challenge.attemptsLeft - 1} attempts left before your server is blacklisted.`,
-		}, { status: 400 });
-	}
+		debug("POST /discover/rotate/confirm – %d attempt(s) left, decrypting challenges", challenge.attemptsLeft);
+		let decryptedOld: string;
+		let decryptedNew: string;
+		try {
+			decryptedOld = decryptPayload(validated.data.signedOldChallenge, process.env.FEDERATION_PRIVATE_KEY!);
+			decryptedNew = decryptPayload(validated.data.signedNewChallenge, process.env.FEDERATION_PRIVATE_KEY!);
+		} catch {
+			debug("POST /discover/rotate/confirm – decryption failed, decrementing attempts");
+			await tx.update(rotateChallengeTokens).set({
+				attemptsLeft: sql`${rotateChallengeTokens.attemptsLeft} - 1`,
+			}).where(eq(rotateChallengeTokens.id, challenge.id))
+			return NextResponse.json({
+				error: `Failed to decrypt one or both challenges. You have ${challenge.attemptsLeft - 1} attempts left before your server is blacklisted.`,
+			}, { status: 400 });
+		}
 
-	// Both challenges passed:
-	// — SA holds the old private key (they are who they claim to be)
-	// — SA holds the new private key (they own the key they want to rotate to)
-	// — SA knows our public key (they fetched our identity to re-encrypt)
-	debug("POST /discover/rotate/confirm – both challenges passed, rotating key for %s", challenge.serverUrl);
-	await db.update(serverRegistry).set({
-		publicKey: challenge.newPublicKey,
-		updatedAt: new Date(),
-	}).where(eq(serverRegistry.url, challenge.serverUrl));
+		if (decryptedOld !== challenge.oldKeyToken || decryptedNew !== challenge.newKeyToken) {
+			debug("POST /discover/rotate/confirm – token mismatch (old=%s, new=%s), decrementing attempts",
+				decryptedOld === challenge.oldKeyToken ? "ok" : "MISMATCH",
+				decryptedNew === challenge.newKeyToken ? "ok" : "MISMATCH",
+			);
+			await tx.update(rotateChallengeTokens).set({
+				attemptsLeft: sql`${rotateChallengeTokens.attemptsLeft} - 1`,
+			}).where(eq(rotateChallengeTokens.id, challenge.id));
+			return NextResponse.json({
+				error: `Challenge mismatch. You have ${challenge.attemptsLeft - 1} attempts left before your server is blacklisted.`,
+			}, { status: 400 });
+		}
 
-	await db.delete(rotateChallengeTokens).where(eq(rotateChallengeTokens.id, challenge.id));
+		debug("POST /discover/rotate/confirm – both challenges passed, rotating key for %s", challenge.serverUrl);
+		await tx.update(serverRegistry).set({
+			publicKey: challenge.newPublicKey,
+			updatedAt: new Date(),
+		}).where(eq(serverRegistry.url, challenge.serverUrl));
 
-	debug("POST /discover/rotate/confirm – key rotation complete for %s", challenge.serverUrl);
-	return NextResponse.json({ message: "Key rotation confirmed successfully." });
+		await tx.delete(rotateChallengeTokens).where(eq(rotateChallengeTokens.id, challenge.id));
+
+		debug("POST /discover/rotate/confirm – key rotation complete for %s", challenge.serverUrl);
+		return NextResponse.json({ message: "Key rotation confirmed successfully." });
+	});
 }
diff --git a/src/app/discover/rotate/init/route.ts b/src/app/discover/rotate/init/route.ts
index 2c80841..b70b945 100644
--- a/src/app/discover/rotate/init/route.ts
+++ b/src/app/discover/rotate/init/route.ts
@@ -60,6 +60,23 @@ export async function POST(request: NextRequest) {
 		return NextResponse.json({ error: "Your server is already registered with this public key." }, { status: 400 });
 	}
 
+	// Check for existing pending challenges, only one active challenge per server is allowed.
+	// This got removed by accident on a previous commit.
+	const [existing] = await db.select().from(rotateChallengeTokens)
+		.where(eq(rotateChallengeTokens.serverUrl, validated.data.url.toString()));
+
+	if (existing) {
+		if (existing.expiresAt > new Date()) {
+			debug("POST /discover/rotate/init – active challenge already exists, rejecting");
+			return NextResponse.json(
+				{ error: "A rotation challenge is already pending for this server." },
+				{ status: 409 },
+			);
+		}
+		debug("POST /discover/rotate/init – deleting expired challenge");
+		await db.delete(rotateChallengeTokens).where(eq(rotateChallengeTokens.id, existing.id));
+	}
+
 	// Issue two independent challenges:
 	//
 	// oldKeyChallenge — encrypted with the SA's CURRENT registered public key.
diff --git a/src/app/discover/route.ts b/src/app/discover/route.ts
index 4beacdc..8d50fb3 100644
--- a/src/app/discover/route.ts
+++ b/src/app/discover/route.ts
@@ -1,8 +1,9 @@
 import db from "@/lib/db";
 import { serverRegistry } from "@/lib/db/schema";
 import { decryptPayload } from "@/lib/federation/keytools";
+import { assertSafeUrl, UrlGuardError } from "@/lib/federation/url-guard";
 import createDebug from "debug";
-import { eq } from "drizzle-orm";
+import { desc, eq } from "drizzle-orm";
 import { NextRequest, NextResponse } from "next/server";
 import forge from "node-forge";
 import { z } from "zod";
@@ -11,7 +12,10 @@ const debug = createDebug("app:discover");
 
 export async function GET() {
 	debug("GET /discover – fetching healthy peers");
-	const peers = await db.select().from(serverRegistry).where(eq(serverRegistry.isHealthy, true));
+	const peers = await db.select({
+		url: serverRegistry.url,
+		isHealthy: serverRegistry.isHealthy,
+	}).from(serverRegistry).where(eq(serverRegistry.isHealthy, true)).orderBy(desc(serverRegistry.lastSeen));
 	debug("GET /discover – found %d peer(s)", peers.length);
 
 	return NextResponse.json({
@@ -50,28 +54,40 @@ const publicKeySchema = z.string().superRefine((key, ctx) => {
 	}
 });
 
-const schema = z.discriminatedUnion("method", [
-	z.object({
-		method: z.literal("DISCOVER"),
-		publicKey: publicKeySchema,
-		signature: z.string().refine((signature) => {
-			try {
-				const sig = decryptPayload(signature, process.env.FEDERATION_PRIVATE_KEY!);
-				const data = JSON.parse(sig);
-				return data.publicKey != null && data.url != null;
-			} catch {
-				return false;
-			}
-		}, { message: "Invalid signature" }),
-	}),
-	z.object({
-		method: z.literal("REGISTER"),
-		url: z.url(),
-		publicKey: publicKeySchema,
-	})
-]);
+function fingerprintKey(pem: string): string {
+	const md = forge.md.sha256.create();
+	md.update(pem, "utf8");
+	return md.digest().toHex();
+}
 
-async function discoverServer(validated: Extract<z.infer<typeof schema>, { method: "DISCOVER" }>) {
+const discoverSchema = z.object({
+	method: z.literal("DISCOVER"),
+	publicKey: publicKeySchema,
+	signature: z.string(),
+}).superRefine((data, ctx) => {
+	try {
+		const decrypted = decryptPayload(data.signature, process.env.FEDERATION_PRIVATE_KEY!);
+		const parsed = JSON.parse(decrypted);
+		// The signature contains a SHA-256 fingerprint of the public key
+		// (since the full PEM exceeds RSA-OAEP's size limit) plus a url.
+		if (parsed.publicKeyFingerprint !== fingerprintKey(data.publicKey)) {
+			ctx.addIssue({ code: "custom", message: "Signature does not match the provided public key" });
+		}
+		if (!parsed.url) {
+			ctx.addIssue({ code: "custom", message: "Signature is missing the url field" });
+		}
+	} catch {
+		ctx.addIssue({ code: "custom", message: "Invalid signature" });
+	}
+});
+
+const registerSchema = z.object({
+	method: z.literal("REGISTER"),
+	url: z.url(),
+	publicKey: publicKeySchema,
+});
+
+async function discoverServer(validated: z.infer<typeof discoverSchema>) {
 	debug("DISCOVER – looking up server by public key");
 	const server = await db.select().from(serverRegistry).where(eq(serverRegistry.publicKey, validated.publicKey));
 	if (server.length === 0) {
@@ -79,6 +95,16 @@ async function discoverServer(validated: Extract<z.infer<typeof schema>, { metho
 		return NextResponse.json({ error: "Server not found" }, { status: 404 });
 	}
 
+	try {
+		assertSafeUrl(server[0].url);
+	} catch (err) {
+		debug("DISCOVER – stored URL failed SSRF check: %s", server[0].url);
+		if (err instanceof UrlGuardError) {
+			return NextResponse.json({ error: "Stored server URL is blocked" }, { status: 400 });
+		}
+		throw err;
+	}
+
 	const confirmations = {
 		sameKeyOnServer: false,
 		sameKeyOnFetch: false,
@@ -86,16 +112,38 @@ async function discoverServer(validated: Extract<z.infer<typeof schema>, { metho
 
 	if (server[0].publicKey === validated.publicKey) confirmations.sameKeyOnServer = true;
 	debug("DISCOVER – fetching public key from federation server %s", server[0].url);
-	const federationResponse = await (await fetch(server[0].url + "/discover")).json();
-	if (federationResponse.publicKey === validated.publicKey) confirmations.sameKeyOnFetch = true;
+	try {
+		const federationResponse = await (await fetch(server[0].url + "/discover")).json();
+		if (federationResponse.publicKey === validated.publicKey) confirmations.sameKeyOnFetch = true;
+	} catch (err) {
+		debug("DISCOVER – fetch to %s failed: %o", server[0].url, err);
+		return NextResponse.json({ error: "Failed to reach the federation server" }, { status: 502 });
+	}
 
 	debug("DISCOVER – confirmations: %o", confirmations);
 	return NextResponse.json(confirmations);
 }
 
-async function registerServer(validated: Extract<z.infer<typeof schema>, { method: "REGISTER" }>) {
+async function registerServer(validated: z.infer<typeof registerSchema>) {
+	try {
+		assertSafeUrl(validated.url);
+	} catch (err) {
+		debug("REGISTER – URL failed SSRF check: %s", validated.url);
+		if (err instanceof UrlGuardError) {
+			return NextResponse.json({ error: err.message }, { status: 400 });
+		}
+		throw err;
+	}
+
 	debug("REGISTER – fetching /discover from %s to validate server", validated.url);
-	const response = await (await fetch(validated.url + "/discover")).json();
+	let response: { publicKey?: string };
+	try {
+		response = await (await fetch(validated.url + "/discover")).json();
+	} catch (err) {
+		debug("REGISTER – fetch to %s failed: %o", validated.url, err);
+		return NextResponse.json({ error: "Failed to reach the server" }, { status: 502 });
+	}
+
 	if (!response.publicKey) {
 		debug("REGISTER – remote server returned no public key");
 		return NextResponse.json({ error: "Invalid server" }, { status: 400 });
@@ -129,19 +177,23 @@ export async function POST(request: NextRequest) {
 	const body = await request.json();
 	debug("POST /discover – method: %s", body?.method);
 
-	const validated = schema.safeParse(body);
-
-	if (!validated.success) {
-		debug("POST /discover – validation failed: %o", validated.error.message);
-		return NextResponse.json({ error: validated.error.message }, { status: 400 });
+	if (body?.method === "DISCOVER") {
+		const validated = discoverSchema.safeParse(body);
+		if (!validated.success) {
+			debug("POST /discover – DISCOVER validation failed: %o", validated.error.message);
+			return NextResponse.json({ error: validated.error.message }, { status: 400 });
+		}
+		return await discoverServer(validated.data);
 	}
 
-	switch (validated.data.method) {
-		case "DISCOVER":
-			return await discoverServer(validated.data);
-		case "REGISTER":
-			return await registerServer(validated.data);
-		default:
-			return NextResponse.json({ error: "Invalid method" }, { status: 400 });
+	if (body?.method === "REGISTER") {
+		const validated = registerSchema.safeParse(body);
+		if (!validated.success) {
+			debug("POST /discover – REGISTER validation failed: %o", validated.error.message);
+			return NextResponse.json({ error: validated.error.message }, { status: 400 });
+		}
+		return await registerServer(validated.data);
 	}
+
+	return NextResponse.json({ error: "Invalid method" }, { status: 400 });
 }
\ No newline at end of file
diff --git a/src/app/page.tsx b/src/app/page.tsx
index 0255261..fa01183 100644
--- a/src/app/page.tsx
+++ b/src/app/page.tsx
@@ -3,6 +3,7 @@
 import { auth } from "@/lib/auth";
 import { headers } from "next/headers";
 import { redirect } from "next/navigation";
+import { PostTestForm } from "./PostTestForm";
 
 export default async function Home() {
 
@@ -12,6 +13,6 @@ export default async function Home() {
 	if (!session) redirect(`/auth`);
 
 	return (
-		<></>
+		<PostTestForm />
 	);
 }
diff --git a/src/lib/auth.ts b/src/lib/auth.ts
index bb9ab60..ee914f3 100644
--- a/src/lib/auth.ts
+++ b/src/lib/auth.ts
@@ -7,6 +7,7 @@ import { bearer, haveIBeenPwned, openAPI, testUtils, twoFactor, username } from
 import db from "./db";
 import * as schema from "./db/schema";
 import EmailService from "./mail";
+import minioClient from "./plugins/server/storage/minio.client";
 
 const isTest = process.env.NODE_ENV === "test";
 const emailService: EmailService | undefined = isTest ? undefined : new EmailService();
@@ -16,13 +17,14 @@ if (!federationKeysExist) {
 	throw new Error("FEDERATION_PUBLIC_KEY and FEDERATION_PRIVATE_KEY must be set, please run `bun run keygen` to generate them.");
 }
 
-export const auth = betterAuth({
+const bAuth = betterAuth({
 	secret: process.env.BETTER_AUTH_SECRET!,
 	baseURL: process.env.BETTER_AUTH_URL ?? (process.env.NODE_ENV === "test" ? "http://localhost:3000" : undefined),
 	experimental: {
 		joins: true
 	},
 	emailAndPassword: {
+		autoSignIn: false,
 		enabled: true,
 	},
 	emailVerification: {
@@ -81,4 +83,9 @@ export const auth = betterAuth({
 			}
 		}
 	}
-});
\ No newline at end of file
+});
+
+export const auth: typeof bAuth & { minio: typeof minioClient } = {
+	...bAuth,
+	minio: minioClient
+}
\ No newline at end of file
diff --git a/src/lib/db/schema/index.ts b/src/lib/db/schema/index.ts
index 8cee988..4b8f545 100644
--- a/src/lib/db/schema/index.ts
+++ b/src/lib/db/schema/index.ts
@@ -102,23 +102,6 @@ export const twoFactor = pgTable(
   ],
 );
 
-export const serverRegistry = pgTable(
-  "server_registry",
-  {
-    id: text("id").primaryKey(),
-    url: text("url").notNull().unique(),
-    publicKey: text("public_key").notNull().unique(),
-    lastSeen: timestamp("last_seen").notNull(),
-    createdAt: timestamp("created_at").notNull(),
-    updatedAt: timestamp("updated_at").notNull(),
-    isHealthy: boolean("is_healthy").notNull(),
-  },
-  (table) => [
-    uniqueIndex("serverRegistry_publicKey_uidx").on(table.publicKey),
-    index("serverRegistry_lastSeen_idx").on(table.lastSeen),
-  ],
-);
-
 export const posts = pgTable("posts", {
   id: text("id").primaryKey(),
   content: jsonb("content").notNull(),
@@ -175,6 +158,23 @@ export const blocks = pgTable("blocks", {
   createdAt: timestamp("created_at").notNull(),
 });
 
+export const serverRegistry = pgTable(
+  "server_registry",
+  {
+    id: text("id").primaryKey(),
+    url: text("url").notNull().unique(),
+    publicKey: text("public_key").notNull().unique(),
+    lastSeen: timestamp("last_seen").notNull(),
+    createdAt: timestamp("created_at").notNull(),
+    updatedAt: timestamp("updated_at").notNull(),
+    isHealthy: boolean("is_healthy").notNull(),
+  },
+  (table) => [
+    uniqueIndex("serverRegistry_publicKey_uidx").on(table.publicKey),
+    index("serverRegistry_lastSeen_idx").on(table.lastSeen),
+  ],
+);
+
 export const rotateChallengeTokens = pgTable(
   "rotate_challenge_tokens",
   {
diff --git a/src/lib/federation/blacklist-middleware.ts b/src/lib/federation/blacklist-middleware.ts
new file mode 100644
index 0000000..19c6422
--- /dev/null
+++ b/src/lib/federation/blacklist-middleware.ts
@@ -0,0 +1,19 @@
+import db from "@/lib/db";
+import { blacklistedServers } from "@/lib/db/schema";
+import createDebug from "debug";
+import { eq } from "drizzle-orm";
+
+const debug = createDebug("federation:blacklist");
+
+/**
+ * Check if a server URL is blacklisted.
+ * Exported so route handlers can call it with body-extracted URLs.
+ */
+export async function isBlacklisted(serverUrl: string): Promise<boolean> {
+	const [row] = await db
+		.select({ id: blacklistedServers.id })
+		.from(blacklistedServers)
+		.where(eq(blacklistedServers.serverUrl, serverUrl))
+		.limit(1);
+	return !!row;
+}
\ No newline at end of file
diff --git a/src/lib/federation/url-guard.ts b/src/lib/federation/url-guard.ts
new file mode 100644
index 0000000..c4875a5
--- /dev/null
+++ b/src/lib/federation/url-guard.ts
@@ -0,0 +1,76 @@
+import createDebug from "debug";
+
+const debug = createDebug("federation:url-guard");
+
+const BLOCKED_HOSTNAMES = new Set([
+	"localhost",
+	"0.0.0.0",
+	"[::1]",
+	"[::0]",
+]);
+
+function isPrivateIPv4(hostname: string): boolean {
+	const parts = hostname.split(".").map(Number);
+	if (parts.length !== 4 || parts.some((p) => isNaN(p))) return false;
+
+	const [a, b] = parts;
+	if (a === 127) return true;                       // 127.0.0.0/8
+	if (a === 10) return true;                        // 10.0.0.0/8
+	if (a === 172 && b >= 16 && b <= 31) return true; // 172.16.0.0/12
+	if (a === 192 && b === 168) return true;           // 192.168.0.0/16
+	if (a === 169 && b === 254) return true;           // 169.254.0.0/16 (link-local / AWS metadata)
+	if (a === 0) return true;                          // 0.0.0.0/8
+
+	return false;
+}
+
+function isPrivateIPv6(hostname: string): boolean {
+	const bare = hostname.replace(/^\[|\]$/g, "").toLowerCase();
+	if (bare === "::1" || bare === "::0" || bare === "::") return true;
+	if (bare.startsWith("fc") || bare.startsWith("fd")) return true;  // ULA
+	if (bare.startsWith("fe80")) return true;                          // link-local
+	return false;
+}
+
+/**
+ * Throws if the URL points to a private/internal address or uses a
+ * non-HTTP(S) protocol. Call before any server-side fetch to prevent SSRF.
+ */
+export function assertSafeUrl(url: string): void {
+	let parsed: URL;
+	try {
+		parsed = new URL(url);
+	} catch {
+		throw new UrlGuardError(`Invalid URL: ${url}`);
+	}
+
+	if (parsed.protocol !== "http:" && parsed.protocol !== "https:") {
+		throw new UrlGuardError(`Blocked protocol: ${parsed.protocol}`);
+	}
+
+	const hostname = parsed.hostname;
+
+	if (BLOCKED_HOSTNAMES.has(hostname)) {
+		debug("blocked hostname: %s", hostname);
+		throw new UrlGuardError(`Blocked internal address: ${hostname}`);
+	}
+
+	if (isPrivateIPv4(hostname)) {
+		debug("blocked private IPv4: %s", hostname);
+		throw new UrlGuardError(`Blocked internal address: ${hostname}`);
+	}
+
+	if (hostname.startsWith("[") || isPrivateIPv6(hostname)) {
+		if (isPrivateIPv6(hostname)) {
+			debug("blocked private IPv6: %s", hostname);
+			throw new UrlGuardError(`Blocked internal address: ${hostname}`);
+		}
+	}
+}
+
+export class UrlGuardError extends Error {
+	constructor(message: string) {
+		super(message);
+		this.name = "UrlGuardError";
+	}
+}
diff --git a/src/lib/plugins/client/social.ts b/src/lib/plugins/client/social.ts
index cdf7076..8c2fc65 100644
--- a/src/lib/plugins/client/social.ts
+++ b/src/lib/plugins/client/social.ts
@@ -1,24 +1,104 @@
 import type { BetterAuthClientPlugin } from "better-auth/client";
 import { z } from "zod";
-import { postContentSchema } from "../server/helpers/social/social";
 import type { sipherSocial } from "../server/social";
 
+const clientPostContentSchmema = z.array(
+	z.object(
+		{
+			type: z.enum(["text", "image", "video", "audio", "link"]),
+			// value could be a string, a file, a url, etc.
+			value: z.union([z.string(), z.instanceof(File), z.url()], { error: "Value must be a string, a file or a URL" }),
+		}
+	)
+)
+
 type SipherSocialPlugin = typeof sipherSocial;
 
+export const MAX_IMAGE_SIZE = 10 * 1024 * 1024; // 10MB
+export const MAX_VIDEO_SIZE = 100 * 1024 * 1024; // 100MB
+export const MAX_AUDIO_SIZE = 100 * 1024 * 1024; // 100MB
+export const MAX_IMAGE_COUNT = 5; // 5 images per post
+export const MAX_AUDIO_COUNT = 1; // 1 audio per post
+export const MAX_VIDEO_COUNT = 2; // 2 videos per post
+
 export const sipherSocialClientPlugin = () => {
 	return {
 		id: "sipher-social",
 		$InferServerPlugin: {} as ReturnType<SipherSocialPlugin>,
 		getActions($fetch, $store, options) {
 			return {
-				createPost: async (content: z.infer<typeof postContentSchema>) => {
-					const response = await $fetch("/social/posts", {
-						method: "POST",
-						body: {
-							content,
-						},
-					});
-					return response;
+				createPost: async (content: z.infer<typeof clientPostContentSchmema>) => {
+
+					// Allow only these combinations of content:
+					// 1. Text only
+					// 2. Text and images
+					// 3. Text, images and videos
+					// 4. Text and audio
+					// No other combinations are allowed
+					// Check the content types and throw an error if the combination is not allowed
+					const contentTypes = content.map((block) => block.type);
+					if (contentTypes.length > 1) {
+						if (contentTypes.includes("image") && contentTypes.includes("audio")) {
+							throw new Error("Images and audios cannot be combined under the same post.")
+						} else if (contentTypes.includes("video") && contentTypes.includes("audio")) {
+							throw new Error("Videos and audios cannot be combined under the same post.")
+						}
+					}
+					// Check if the content amount per type is under the allowed limits
+					const imageCount = content.filter((block) => block.type === "image").length;
+					const videoCount = content.filter((block) => block.type === "video").length;
+					const audioCount = content.filter((block) => block.type === "audio").length;
+					if (imageCount > MAX_IMAGE_COUNT) throw new Error("Maximum number of images per post exceeded");
+					if (videoCount > MAX_VIDEO_COUNT) throw new Error("Maximum number of videos per post exceeded");
+					if (audioCount > MAX_AUDIO_COUNT) throw new Error("Maximum number of audios per post exceeded");
+
+					const resolvedContent: { type: string; value?: string; url?: string; size?: number; index?: number }[] = [];
+					let mediaIndex = 0;
+
+					for (const block of content) {
+						if (block.type === "text" || block.type === "link") {
+							resolvedContent.push({ type: block.type, value: block.value as string });
+							continue;
+						}
+
+						const file = block.value as File;
+
+						const { data, error } = await $fetch<{
+							presignedUrl: string;
+							objectUrl: string;
+							objectKey: string;
+						}>("/social/posts/files", {
+							method: "POST",
+							body: {
+								fileName: file.name,
+								mimeType: file.type,
+								size: file.size,
+							},
+						});
+
+						if (error || !data) {
+							throw new Error("Failed to get upload URL");
+						}
+
+						const uploadRes = await fetch(data.presignedUrl, {
+							method: "PUT",
+							body: file,
+							headers: { "Content-Type": file.type },
+						});
+
+						if (!uploadRes.ok) {
+							throw new Error(`Failed to upload ${file.name}`);
+						}
+
+						resolvedContent.push({
+							type: block.type,
+							url: data.objectUrl,
+							size: file.size,
+							index: mediaIndex++,
+						});
+					}
+
+					console.log("Resolved content:", resolvedContent);
 				}
 			}
 		},
diff --git a/src/lib/plugins/server/helpers/social/endpoints/index.ts b/src/lib/plugins/server/helpers/social/endpoints/index.ts
index b2b9262..fadbb6d 100644
--- a/src/lib/plugins/server/helpers/social/endpoints/index.ts
+++ b/src/lib/plugins/server/helpers/social/endpoints/index.ts
@@ -1,7 +1,7 @@
 import { createBlock, deleteBlock, getBlocks } from "./blocks";
 import { followUser, getFollowers, getFollows, unfollowUser } from "./follows";
 import { createMute, deleteMute, getMutes } from "./mutes";
-import { createPost, getPost } from "./posts";
+import { createPost, getPost, uploadFile } from "./posts";
 
-export { createBlock, createMute, createPost, deleteBlock, deleteMute, followUser, getBlocks, getFollowers, getFollows, getMutes, getPost, unfollowUser };
+export { createBlock, createMute, createPost, deleteBlock, deleteMute, followUser, getBlocks, getFollowers, getFollows, getMutes, getPost, unfollowUser, uploadFile };
 
diff --git a/src/lib/plugins/server/helpers/social/endpoints/posts.ts b/src/lib/plugins/server/helpers/social/endpoints/posts.ts
index 64f8d8e..50d6f03 100644
--- a/src/lib/plugins/server/helpers/social/endpoints/posts.ts
+++ b/src/lib/plugins/server/helpers/social/endpoints/posts.ts
@@ -1,3 +1,6 @@
+import db from "@/lib/db";
+import { posts } from "@/lib/db/schema";
+import minioClient from "@/plugins/server/storage/minio.client";
 import { createAuthEndpoint, getSessionFromCtx } from "better-auth/api";
 import { z } from "zod";
 import { postContentSchema } from "../social";
@@ -7,14 +10,26 @@ export const createPost = createAuthEndpoint("/social/posts", {
 	body: postContentSchema,
 }, async (context) => {
 	const content = context.body;
-	const user = getSessionFromCtx(context)
+	const user = await getSessionFromCtx(context)
 
 	if (!user) {
 		return context.json({ error: "Unauthorized" }, { status: 401 });
 	}
 
-	console.log(content);
-	return context.json({ message: "Hello, world!" }, { status: 200 });
+
+
+	// Create post
+	const post = await db.insert(posts).values({
+		id: crypto.randomUUID(),
+		content: content,
+		authorId: user.user.id,
+		published: new Date(),
+		isLocal: true,
+		createdAt: new Date(),
+	}).returning({ id: posts.id });
+
+	return context.json({ id: post[0].id }, { status: 200 });
+
 });
 
 export const getPost = createAuthEndpoint("/social/posts/:id", {
@@ -22,4 +37,44 @@ export const getPost = createAuthEndpoint("/social/posts/:id", {
 	params: z.object({
 		id: z.string(),
 	}),
-}, async (context) => { })
\ No newline at end of file
+}, async (context) => { });
+
+const ALLOWED_MIME_TYPES = [
+	"image/jpeg", "image/png", "image/gif", "image/webp",
+	"video/mp4", "video/webm",
+	"audio/mpeg", "audio/ogg", "audio/wav",
+];
+
+const PRESIGN_EXPIRY_SECONDS = 5 * 60;
+
+export const uploadFile = createAuthEndpoint("/social/posts/files", {
+	method: "POST",
+	body: z.object({
+		fileName: z.string().min(1),
+		mimeType: z.string().refine((v) => ALLOWED_MIME_TYPES.includes(v), {
+			message: "Unsupported file type",
+		}),
+		size: z.number().positive(),
+	}),
+}, async (context) => {
+	const user = await getSessionFromCtx(context);
+
+	if (!user) {
+		return context.json({ error: "Unauthorized" }, { status: 401 });
+	}
+
+	const { fileName, mimeType, size } = context.body;
+	const ext = fileName.split(".").pop() ?? "bin";
+	const objectKey = `tmp/${user.user.id}/${crypto.randomUUID()}.${ext}`;
+
+	const presignedUrl = await minioClient.presignedPutObject(
+		process.env.MINIO_BUCKET!,
+		objectKey,
+		PRESIGN_EXPIRY_SECONDS,
+	);
+
+	const protocol = process.env.MINIO_USE_SSL === "true" ? "https" : "http";
+	const objectUrl = `${protocol}://${process.env.MINIO_ENDPOINT}:${process.env.MINIO_PORT}/${process.env.MINIO_BUCKET}/${objectKey}`;
+
+	return context.json({ presignedUrl, objectUrl, objectKey }, { status: 200 });
+})
\ No newline at end of file
diff --git a/src/lib/plugins/server/helpers/social/social.ts b/src/lib/plugins/server/helpers/social/social.ts
index fb3bddc..7a7a872 100644
--- a/src/lib/plugins/server/helpers/social/social.ts
+++ b/src/lib/plugins/server/helpers/social/social.ts
@@ -10,14 +10,18 @@ const postContentBlockSchema = z.discriminatedUnion("type", [
 		type: z.literal("image"),
 		url: z.url("Image must be a valid URL"),
 		index: z.number().min(0, "Index must be a positive number"),
+		size: z.number().min(0, "Size must be a positive number"),
 	}),
 	z.object({
 		type: z.literal("video"),
 		url: z.url("Video must be a valid URL"),
+		size: z.number().min(0, "Size must be a positive number"),
+		index: z.number().min(0, "Index must be a positive number"),
 	}),
 	z.object({
 		type: z.literal("audio"),
 		url: z.url("Audio must be a valid URL"),
+		size: z.number().min(0, "Size must be a positive number"),
 	}),
 	z.object({
 		type: z.literal("link"),
diff --git a/src/lib/plugins/server/storage/minio.client.ts b/src/lib/plugins/server/storage/minio.client.ts
new file mode 100644
index 0000000..c4f9610
--- /dev/null
+++ b/src/lib/plugins/server/storage/minio.client.ts
@@ -0,0 +1,21 @@
+import * as Minio from "minio";
+
+const MINIO_ENDPOINT_ENV = process.env.MINIO_ENDPOINT;
+const MINIO_PORT_ENV = process.env.MINIO_PORT;
+const MINIO_USE_SSL_ENV = process.env.MINIO_USE_SSL;
+const MINIO_ACCESS_KEY_ENV = process.env.MINIO_ACCESS_KEY;
+const MINIO_SECRET_KEY_ENV = process.env.MINIO_SECRET_KEY;
+
+if (!MINIO_ENDPOINT_ENV || !MINIO_PORT_ENV || !MINIO_USE_SSL_ENV || !MINIO_ACCESS_KEY_ENV || !MINIO_SECRET_KEY_ENV) {
+	throw new Error("Missing Minio environment variables");
+}
+
+const minioClient = new Minio.Client({
+	endPoint: MINIO_ENDPOINT_ENV,
+	port: parseInt(MINIO_PORT_ENV),
+	useSSL: MINIO_USE_SSL_ENV === "true",
+	accessKey: MINIO_ACCESS_KEY_ENV,
+	secretKey: MINIO_SECRET_KEY_ENV,
+})
+
+export default minioClient;
\ No newline at end of file
diff --git a/src/proxy.ts b/src/proxy.ts
new file mode 100644
index 0000000..1bc8116
--- /dev/null
+++ b/src/proxy.ts
@@ -0,0 +1,49 @@
+import { isBlacklisted } from '@/lib/federation/blacklist-middleware'
+import { NextRequest, NextResponse } from 'next/server'
+
+// Under no circumstances a blacklisted server should be able to access anything from this federation server.
+// This is a security measure to isolate the federation server from potentially malicious servers.
+// This could and should be revised in the future.
+export async function proxy(request: NextRequest) {
+
+	// If coming from self, skip
+	if (
+		request.headers.get("x-federation-origin") === process.env.BETTER_AUTH_URL ||
+		request.headers.get("origin") === process.env.BETTER_AUTH_URL
+	) {
+		return NextResponse.next()
+	}
+
+	const candidates: string[] = []
+
+	const origin = request.headers.get("origin")
+	if (origin) candidates.push(origin)
+
+	const federationOrigin = request.headers.get("x-federation-origin")
+	if (federationOrigin) candidates.push(federationOrigin)
+
+	const method = request.method.toUpperCase()
+	if (["POST", "PUT", "PATCH"].includes(method)) {
+		const contentType = request.headers.get("content-type") ?? ""
+		if (contentType.includes("application/json")) {
+			try {
+				const body = await request.clone().json()
+				if (typeof body.url === "string") candidates.push(body.url)
+				if (typeof body.serverUrl === "string") candidates.push(body.serverUrl)
+			} catch {
+				// Invalid JSON, let the route handler deal with it
+			}
+		}
+	}
+
+	for (const url of candidates) {
+		if (await isBlacklisted(url)) {
+			return NextResponse.json(
+				{ error: "Your server has been blacklisted." },
+				{ status: 403 },
+			)
+		}
+	}
+
+	return NextResponse.next()
+}
\ No newline at end of file
diff --git a/src/server.ts b/src/server.ts
index 7a35249..b4d6961 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -1,5 +1,5 @@
 import { config } from 'dotenv'
-import { createServer } from 'http'
+import { createServer, type IncomingMessage, type ServerResponse } from 'http'
 import next from 'next'
 
 config({ path: '.env.local' })
@@ -9,7 +9,7 @@ const app = next({ dev })
 const handle = app.getRequestHandler()
 
 app.prepare().then(() => {
-	createServer((req, res) => {
+	createServer(async (req: IncomingMessage, res: ServerResponse) => {
 		handle(req, res)
 	}).listen(port)
 
diff --git a/tests/attacks.test.ts b/tests/attacks.test.ts
new file mode 100644
index 0000000..1f55dd3
--- /dev/null
+++ b/tests/attacks.test.ts
@@ -0,0 +1,459 @@
+/**
+ * Attack-vector tests for the /discover federation routes.
+ *
+ * These tests verify that the security fixes are working correctly.
+ * Each describe block targets a vulnerability that was found during audit
+ * and has now been patched.
+ * 
+ * @author Tocka
+ * These tests were generated by AI (Claude Opus 4.6).
+ * I did review it and made some changes but it's still mostly AI generated so take this with a grain of salt.
+ * 
+ * The AI did expose some actual vulnerabilities so it's not all bad.
+ * After fixing the vulnerabilitie, the tests were updated to match the new behavior.
+ * -----
+ */
+import { expect, test } from "@playwright/test"
+import forge from "node-forge"
+import http from "node:http"
+import {
+	clearTables,
+	generateKeypair,
+	getBlacklistedServer,
+	getChallengesByServerUrl,
+	seedChallenge,
+	seedServer,
+} from "./helpers/db"
+
+const BASE = "http://localhost:3000"
+
+function encryptPayload(payload: string, recipientPublicKeyPem: string) {
+	const pub = forge.pki.publicKeyFromPem(recipientPublicKeyPem)
+	return forge.util.encode64(
+		pub.encrypt(forge.util.encodeUtf8(payload), "RSA-OAEP"),
+	)
+}
+
+function fingerprintKey(pem: string): string {
+	const md = forge.md.sha256.create()
+	md.update(pem, "utf8")
+	return md.digest().toHex()
+}
+
+function generate4096Keypair() {
+	const kp = forge.pki.rsa.generateKeyPair(4096)
+	return {
+		publicKey: forge.pki.publicKeyToPem(kp.publicKey),
+		privateKey: forge.pki.privateKeyToPem(kp.privateKey),
+	}
+}
+
+function createTrapServer(fakePublicKey: string) {
+	const hits: { method: string; url: string }[] = []
+	const server = http.createServer((req, res) => {
+		hits.push({ method: req.method!, url: req.url! })
+		res.writeHead(200, { "Content-Type": "application/json" })
+		res.end(JSON.stringify({ publicKey: fakePublicKey }))
+	})
+
+	return {
+		hits,
+		start: () => new Promise<number>((resolve) => {
+			server.listen(0, "127.0.0.1", () => {
+				resolve((server.address() as { port: number }).port)
+			})
+		}),
+		stop: () => new Promise<void>((resolve) => { server.close(() => resolve()) }),
+	}
+}
+
+test.beforeEach(async () => { await clearTables() })
+test.afterEach(async () => { await clearTables() })
+
+// ---------------------------------------------------------------------------
+// 1. SSRF Protection — internal URLs are now blocked
+// ---------------------------------------------------------------------------
+test.describe("SSRF protection", () => {
+	test("REGISTER rejects loopback URLs", async ({ request }) => {
+		const { publicKey: fakePub } = generateKeypair()
+		const trap = createTrapServer(fakePub)
+		const port = await trap.start()
+
+		try {
+			const res = await request.post(`${BASE}/discover`, {
+				data: {
+					method: "REGISTER",
+					url: `http://127.0.0.1:${port}`,
+					publicKey: fakePub,
+				},
+			})
+
+			expect(res.status()).toBe(400)
+			const body = await res.json()
+			expect(body.error).toMatch(/blocked/i)
+
+			// The trap server should NOT have been hit
+			expect(trap.hits.length).toBe(0)
+		} finally {
+			await trap.stop()
+		}
+	})
+
+	test("REGISTER rejects RFC-1918 and link-local URLs", async ({ request }) => {
+		const internalUrls = [
+			"http://10.0.0.1:8080",
+			"http://192.168.1.1:8080",
+			"http://169.254.169.254",
+		]
+
+		for (const url of internalUrls) {
+			const { publicKey } = generateKeypair()
+			const res = await request.post(`${BASE}/discover`, {
+				data: { method: "REGISTER", url, publicKey },
+			})
+
+			expect(res.status()).toBe(400)
+			const body = await res.json()
+			expect(body.error).toMatch(/blocked/i)
+		}
+	})
+
+	test("DISCOVER rejects stored internal URLs", async ({ request }) => {
+		const { publicKey: maliciousPub } = generateKeypair()
+		await seedServer("http://127.0.0.1:9999", maliciousPub)
+
+		// Build a valid signature using the fingerprint approach
+		const signaturePayload = JSON.stringify({
+			publicKeyFingerprint: fingerprintKey(maliciousPub),
+			url: "http://127.0.0.1:9999",
+		})
+		const signature = encryptPayload(signaturePayload, process.env.FEDERATION_PUBLIC_KEY!)
+
+		const res = await request.post(`${BASE}/discover`, {
+			data: {
+				method: "DISCOVER",
+				publicKey: maliciousPub,
+				signature,
+			},
+		})
+
+		// Should be blocked rather than fetching the internal URL
+		expect(res.status()).toBe(400)
+		const body = await res.json()
+		expect(body.error).toMatch(/blocked/i)
+	})
+})
+
+// ---------------------------------------------------------------------------
+// 2. Blacklist enforcement — blocks all federation routes
+// ---------------------------------------------------------------------------
+test.describe("Blacklist enforcement (fixed)", () => {
+	async function blacklistServer(serverUrl: string, request: any) {
+		// Seed a challenge with 1 attempt left
+		await seedChallenge({
+			serverUrl,
+			attemptsLeft: 1,
+			expiresAt: new Date(Date.now() + 1000 * 60 * 5),
+		})
+
+		// First attempt: mismatch, decrements to 0
+		await request.post(`${BASE}/discover/rotate/confirm`, {
+			data: {
+				serverUrl,
+				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+			},
+		})
+
+		// Second attempt: attemptsLeft=0, triggers blacklist
+		await request.post(`${BASE}/discover/rotate/confirm`, {
+			data: {
+				serverUrl,
+				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+			},
+		})
+
+		const bl = await getBlacklistedServer(serverUrl)
+		expect(bl).toBeDefined()
+	}
+
+	test("blacklisted server is rejected by rotate/init", async ({ request }) => {
+		const { publicKey: oldPub } = generateKeypair()
+		const serverUrl = "https://blacklisted-server.example"
+		await seedServer(serverUrl, oldPub)
+		await blacklistServer(serverUrl, request as any)
+
+		const { publicKey: newPub } = generate4096Keypair()
+		const initRes = await request.post(`${BASE}/discover/rotate/init`, {
+			data: { url: serverUrl, newPublicKey: newPub },
+		})
+		expect(initRes.status()).toBe(403)
+		const body = await initRes.json()
+		expect(body.error).toMatch(/blacklisted/i)
+	})
+
+	test("blacklisted server is rejected by rotate/confirm", async ({ request }) => {
+		const serverUrl = "https://blacklisted-confirm.example"
+		const { publicKey } = generateKeypair()
+		await seedServer(serverUrl, publicKey)
+		await blacklistServer(serverUrl, request as any)
+
+		const confirmRes = await request.post(`${BASE}/discover/rotate/confirm`, {
+			data: {
+				serverUrl,
+				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+			},
+		})
+		expect(confirmRes.status()).toBe(403)
+		const body = await confirmRes.json()
+		expect(body.error).toMatch(/blacklisted/i)
+	})
+})
+
+// ---------------------------------------------------------------------------
+// 3. Race condition fixed — transaction + FOR UPDATE
+// ---------------------------------------------------------------------------
+test.describe("Race condition fixed on rotate/confirm", () => {
+	test("concurrent requests are serialised by the row lock", async () => {
+		const serverUrl = "https://race-target.example"
+		const { publicKey } = generateKeypair()
+		await seedServer(serverUrl, publicKey)
+
+		await seedChallenge({
+			serverUrl,
+			attemptsLeft: 1,
+			expiresAt: new Date(Date.now() + 1000 * 60 * 5),
+		})
+
+		const payload = JSON.stringify({
+			serverUrl,
+			signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+			signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+		})
+
+		const fire = () =>
+			fetch(`${BASE}/discover/rotate/confirm`, {
+				method: "POST",
+				headers: { "Content-Type": "application/json" },
+				body: payload,
+			})
+
+		const results = await Promise.all(Array.from({ length: 15 }, fire))
+		const statuses = results.map((r) => r.status)
+
+		const mismatch = statuses.filter((s) => s === 400).length
+		const blacklisted = statuses.filter((s) => s === 403).length
+		const notFound = statuses.filter((s) => s === 404).length
+
+		// With the transaction lock, exactly 1 request should process the
+		// mismatch (400), then the next sees attemptsLeft=0 and blacklists (403),
+		// and the rest find no challenge (404) or hit the blacklist check.
+		// The key point: no more than 1 mismatch (400) should get through.
+		expect(mismatch).toBeLessThanOrEqual(1)
+		expect(mismatch + blacklisted + notFound).toBe(statuses.length)
+	})
+})
+
+// ---------------------------------------------------------------------------
+// 4. Challenge deduplication — no flooding, no attempt reset
+// ---------------------------------------------------------------------------
+test.describe("Challenge deduplication (fixed)", () => {
+	test("second init is rejected while a challenge is pending", async ({ request }) => {
+		const serverUrl = "https://dedup-target.example"
+		const { publicKey } = generateKeypair()
+		await seedServer(serverUrl, publicKey)
+
+		const { publicKey: newPub1 } = generate4096Keypair()
+		const { publicKey: newPub2 } = generate4096Keypair()
+
+		const res1 = await request.post(`${BASE}/discover/rotate/init`, {
+			data: { url: serverUrl, newPublicKey: newPub1 },
+		})
+		expect(res1.status()).toBe(200)
+
+		// Second init while the first is still active → 409
+		const res2 = await request.post(`${BASE}/discover/rotate/init`, {
+			data: { url: serverUrl, newPublicKey: newPub2 },
+		})
+		expect(res2.status()).toBe(409)
+		const body = await res2.json()
+		expect(body.error).toMatch(/already pending/i)
+
+		// Only one challenge exists
+		const challenges = await getChallengesByServerUrl(serverUrl)
+		expect(challenges.length).toBe(1)
+	})
+
+	test("init succeeds after the previous challenge expires", async ({ request }) => {
+		const serverUrl = "https://dedup-expire.example"
+		const { publicKey } = generateKeypair()
+		await seedServer(serverUrl, publicKey)
+
+		// Seed an already-expired challenge directly
+		await seedChallenge({
+			serverUrl,
+			expiresAt: new Date(Date.now() - 1000),
+		})
+
+		const { publicKey: newPub } = generate4096Keypair()
+		const res = await request.post(`${BASE}/discover/rotate/init`, {
+			data: { url: serverUrl, newPublicKey: newPub },
+		})
+		expect(res.status()).toBe(200)
+
+		// Old challenge was replaced, only new one exists
+		const challenges = await getChallengesByServerUrl(serverUrl)
+		expect(challenges.length).toBe(1)
+		expect(challenges[0].newPublicKey).toBe(newPub)
+	})
+
+	test("blacklisted server cannot reset attempts via new init", async ({ request }) => {
+		const serverUrl = "https://reset-blocked.example"
+		const { publicKey } = generateKeypair()
+		await seedServer(serverUrl, publicKey)
+
+		// Exhaust attempts → get blacklisted
+		await seedChallenge({
+			serverUrl,
+			attemptsLeft: 1,
+			expiresAt: new Date(Date.now() + 1000 * 60 * 5),
+		})
+		await request.post(`${BASE}/discover/rotate/confirm`, {
+			data: {
+				serverUrl,
+				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+			},
+		})
+		await request.post(`${BASE}/discover/rotate/confirm`, {
+			data: {
+				serverUrl,
+				signedOldChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+				signedNewChallenge: encryptPayload("wrong", process.env.FEDERATION_PUBLIC_KEY!),
+			},
+		})
+
+		const bl = await getBlacklistedServer(serverUrl)
+		expect(bl).toBeDefined()
+
+		// Try init → blocked by blacklist check
+		const { publicKey: freshPub } = generate4096Keypair()
+		const initRes = await request.post(`${BASE}/discover/rotate/init`, {
+			data: { url: serverUrl, newPublicKey: freshPub },
+		})
+		expect(initRes.status()).toBe(403)
+	})
+})
+
+// ---------------------------------------------------------------------------
+// 5. Signature validation — field values must match the request
+// ---------------------------------------------------------------------------
+test.describe("Signature validation (fixed)", () => {
+	test("signature with mismatched publicKey fingerprint is rejected", async ({ request }) => {
+		const { publicKey: peerPub } = generateKeypair()
+		await seedServer("https://sig-test.example", peerPub)
+
+		// Encrypt a signature where the fingerprint doesn't match
+		const badSignature = encryptPayload(
+			JSON.stringify({ publicKeyFingerprint: "wrong-fingerprint", url: "https://sig-test.example" }),
+			process.env.FEDERATION_PUBLIC_KEY!,
+		)
+
+		const res = await request.post(`${BASE}/discover`, {
+			data: {
+				method: "DISCOVER",
+				publicKey: peerPub,
+				signature: badSignature,
+			},
+		})
+
+		expect(res.status()).toBe(400)
+	})
+
+	test("signature with placeholder values is rejected", async ({ request }) => {
+		const { publicKey: peerPub } = generateKeypair()
+		await seedServer("https://sig-test2.example", peerPub)
+
+		// The old bypass: { publicKey: "x", url: "y" }, now invalid
+		const forgerySignature = encryptPayload(
+			JSON.stringify({ publicKey: "x", url: "y" }),
+			process.env.FEDERATION_PUBLIC_KEY!,
+		)
+
+		const res = await request.post(`${BASE}/discover`, {
+			data: {
+				method: "DISCOVER",
+				publicKey: peerPub,
+				signature: forgerySignature,
+			},
+		})
+
+		expect(res.status()).toBe(400)
+	})
+
+	test("signature with correct fingerprint passes validation", async ({ request }) => {
+		const { publicKey: peerPub } = generateKeypair()
+		const trap = createTrapServer(peerPub)
+		const port = await trap.start()
+		const peerUrl = `http://127.0.0.1:${port}`
+
+		try {
+			await seedServer(peerUrl, peerPub)
+
+			const validSignature = encryptPayload(
+				JSON.stringify({
+					publicKeyFingerprint: fingerprintKey(peerPub),
+					url: peerUrl,
+				}),
+				process.env.FEDERATION_PUBLIC_KEY!,
+			)
+
+			const res = await request.post(`${BASE}/discover`, {
+				data: {
+					method: "DISCOVER",
+					publicKey: peerPub,
+					signature: validSignature,
+				},
+			})
+
+			// The signature is valid, but the stored URL is internal → blocked by SSRF guard
+			expect(res.status()).toBe(400)
+			const body = await res.json()
+			expect(body.error).toMatch(/blocked/i)
+		} finally {
+			await trap.stop()
+		}
+	})
+})
+
+// ---------------------------------------------------------------------------
+// 6. Information disclosure: only url + isHealthy in peer list
+// ---------------------------------------------------------------------------
+test.describe("Information disclosure", () => {
+	test("GET /discover only returns url and isHealthy for peers", async ({ request }) => {
+		const { publicKey: peerPub1 } = generateKeypair()
+		const { publicKey: peerPub2 } = generateKeypair()
+		await seedServer("https://peer-one.example", peerPub1)
+		await seedServer("https://peer-two.example", peerPub2)
+
+		const res = await request.get(`${BASE}/discover`)
+		expect(res.status()).toBe(200)
+
+		const body = await res.json()
+		expect(body.peers).toBeInstanceOf(Array)
+		expect(body.peers.length).toBeGreaterThanOrEqual(2)
+
+		for (const peer of body.peers) {
+			expect(peer.url).toBeDefined()
+			expect(peer.isHealthy).toBeDefined()
+			// Internal fields must NOT be exposed
+			expect(peer.id).toBeUndefined()
+			expect(peer.createdAt).toBeUndefined()
+			expect(peer.updatedAt).toBeUndefined()
+			expect(peer.lastSeen).toBeUndefined()
+			expect(peer.isHealthy).toBeUndefined()
+		}
+	})
+})
diff --git a/tests/helpers/db.ts b/tests/helpers/db.ts
index aaebd2c..84438b4 100644
--- a/tests/helpers/db.ts
+++ b/tests/helpers/db.ts
@@ -1,6 +1,6 @@
 // tests/helpers/db.ts
 import db from "@/lib/db";
-import { rotateChallengeTokens, serverRegistry } from "@/lib/db/schema";
+import { blacklistedServers, rotateChallengeTokens, serverRegistry } from "@/lib/db/schema";
 import { eq } from "drizzle-orm";
 import forge from "node-forge";
 
@@ -65,9 +65,22 @@ export async function insertServerEcho(url: string, publicKey: string) {
 	}).onConflictDoNothing()
 }
 
+export async function getBlacklistedServer(serverUrl: string) {
+	return (await db.select().from(blacklistedServers).where(eq(blacklistedServers.serverUrl, serverUrl)))[0]
+}
+
+export async function getChallengesByServerUrl(serverUrl: string) {
+	return await db.select().from(rotateChallengeTokens).where(eq(rotateChallengeTokens.serverUrl, serverUrl))
+}
+
+export async function clearBlacklist() {
+	return await db.delete(blacklistedServers)
+}
+
 export async function clearTables() {
 	return await Promise.all([
 		clearRotateChallengeTokens(),
+		clearBlacklist(),
 		clearServerRegistry(),
 	])
 }
\ No newline at end of file