Fix critical and high security issues

- exec_script: relay now enforces admin role before forwarding to agent - relay CORS: restrict allow_origins via ALLOWED_ORIGINS env var (docker-compose passes app URL) - session-code: replace Math.random() with crypto.randomInt, add per-key rate limit (10 req/min) - sessions GET: fix IDOR — users can only read their own sessions (admins see all) - signal API: validate session ownership on create; enforce ownerUserId on all subsequent actions Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-10 23:38:03 -07:00
parent e27d4cfa58
commit 27673daa63
5 changed files with 86 additions and 23 deletions
--- a/app/api/agent/session-code/route.ts
+++ b/app/api/agent/session-code/route.ts
@@ -2,16 +2,34 @@ import { db } from '@/lib/db'
 import { machines, sessionCodes } from '@/lib/db/schema'
 import { eq, and } from 'drizzle-orm'
 import { NextRequest, NextResponse } from 'next/server'
+import { randomInt } from 'crypto'

 function generateSessionCode(): string {
  const chars = 'ABCDEFGHJKLMNPQRSTUVWXYZ23456789'
  let code = ''
  for (let i = 0; i < 6; i++) {
-    code += chars.charAt(Math.floor(Math.random() * chars.length))
+    code += chars.charAt(randomInt(chars.length))
  }
  return code
 }

+// Simple in-process rate limiter: max 10 requests per minute per access key
+const rateLimitMap = new Map<string, { count: number; resetAt: number }>()
+const RATE_LIMIT = 10
+const RATE_WINDOW_MS = 60_000
+
+function isRateLimited(key: string): boolean {
+  const now = Date.now()
+  const entry = rateLimitMap.get(key)
+  if (!entry || now >= entry.resetAt) {
+    rateLimitMap.set(key, { count: 1, resetAt: now + RATE_WINDOW_MS })
+    return false
+  }
+  if (entry.count >= RATE_LIMIT) return true
+  entry.count++
+  return false
+}
+
 export async function POST(request: NextRequest) {
  try {
    const { accessKey } = await request.json()
@@ -20,6 +38,10 @@ export async function POST(request: NextRequest) {
      return NextResponse.json({ error: 'Access key required' }, { status: 400 })
    }

+    if (isRateLimited(accessKey)) {
+      return NextResponse.json({ error: 'Too many requests' }, { status: 429 })
+    }
+
    const machineResult = await db
      .select()
      .from(machines)
--- a/app/api/sessions/[id]/route.ts
+++ b/app/api/sessions/[id]/route.ts
@@ -8,17 +8,22 @@ export async function GET(
  _request: NextRequest,
  { params }: { params: Promise<{ id: string }> }
 ) {
-  const session = await auth()
-  if (!session?.user?.id) {
+  const authSession = await auth()
+  if (!authSession?.user?.id) {
    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
  }

  const { id } = await params
+  const isAdmin = (authSession.user as { role?: string }).role === 'admin'

  const result = await db
    .select()
    .from(sessions)
-    .where(eq(sessions.id, id))
+    .where(
+      isAdmin
+        ? eq(sessions.id, id)
+        : and(eq(sessions.id, id), eq(sessions.viewerUserId, authSession.user.id))
+    )
    .limit(1)

  if (!result[0]) {
@@ -32,8 +37,8 @@ export async function PATCH(
  request: NextRequest,
  { params }: { params: Promise<{ id: string }> }
 ) {
-  const session = await auth()
-  if (!session?.user?.id) {
+  const authSession = await auth()
+  if (!authSession?.user?.id) {
    return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
  }

@@ -47,7 +52,7 @@ export async function PATCH(
  await db
    .update(sessions)
    .set(updates)
-    .where(and(eq(sessions.id, id), eq(sessions.viewerUserId, session.user.id)))
+    .where(and(eq(sessions.id, id), eq(sessions.viewerUserId, authSession.user.id)))

  return NextResponse.json({ success: true })
 }
--- a/app/api/signal/route.ts
+++ b/app/api/signal/route.ts
@@ -1,4 +1,7 @@
 import { auth } from '@/auth'
+import { db } from '@/lib/db'
+import { sessions } from '@/lib/db/schema'
+import { and, eq } from 'drizzle-orm'
 import { NextRequest, NextResponse } from 'next/server'

 // In-memory signaling store (use Redis for multi-instance deployments)
@@ -8,6 +11,7 @@ const signalingStore = new Map<string, {
  viewerCandidates: RTCIceCandidateInit[]
  hostCandidates: RTCIceCandidateInit[]
  createdAt: number
+  ownerUserId: string
 }>()

 function cleanupOldSessions() {
@@ -19,48 +23,59 @@ function cleanupOldSessions() {

 export async function POST(request: NextRequest) {
  try {
-    const session = await auth()
-    if (!session?.user) {
+    const authSession = await auth()
+    if (!authSession?.user?.id) {
      return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
    }

+    const userId = authSession.user.id
    const { action, sessionId, data } = await request.json()
    cleanupOldSessions()

    switch (action) {
-      case 'create-session':
-        signalingStore.set(sessionId, { viewerCandidates: [], hostCandidates: [], createdAt: Date.now() })
+      case 'create-session': {
+        // Verify the sessionId corresponds to a real session owned by this user
+        const dbSession = await db
+          .select({ id: sessions.id })
+          .from(sessions)
+          .where(and(eq(sessions.id, sessionId), eq(sessions.viewerUserId, userId)))
+          .limit(1)
+        if (!dbSession[0]) {
+          return NextResponse.json({ error: 'Session not found' }, { status: 404 })
+        }
+        signalingStore.set(sessionId, { viewerCandidates: [], hostCandidates: [], createdAt: Date.now(), ownerUserId: userId })
        return NextResponse.json({ success: true })
+      }

      case 'send-offer': {
        const s = signalingStore.get(sessionId)
-        if (!s) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
+        if (!s || s.ownerUserId !== userId) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
        s.offer = data.offer
        return NextResponse.json({ success: true })
      }

      case 'get-offer': {
        const s = signalingStore.get(sessionId)
-        if (!s) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
+        if (!s || s.ownerUserId !== userId) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
        return NextResponse.json({ offer: s.offer ?? null })
      }

      case 'send-answer': {
        const s = signalingStore.get(sessionId)
-        if (!s) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
+        if (!s || s.ownerUserId !== userId) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
        s.answer = data.answer
        return NextResponse.json({ success: true })
      }

      case 'get-answer': {
        const s = signalingStore.get(sessionId)
-        if (!s) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
+        if (!s || s.ownerUserId !== userId) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
        return NextResponse.json({ answer: s.answer ?? null })
      }

      case 'send-ice-candidate': {
        const s = signalingStore.get(sessionId)
-        if (!s) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
+        if (!s || s.ownerUserId !== userId) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
        if (data.role === 'viewer') s.viewerCandidates.push(data.candidate)
        else s.hostCandidates.push(data.candidate)
        return NextResponse.json({ success: true })
@@ -68,16 +83,19 @@ export async function POST(request: NextRequest) {

      case 'get-ice-candidates': {
        const s = signalingStore.get(sessionId)
-        if (!s) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
+        if (!s || s.ownerUserId !== userId) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
        const candidates = data.role === 'viewer' ? s.hostCandidates : s.viewerCandidates
        if (data.role === 'viewer') s.hostCandidates = []
        else s.viewerCandidates = []
        return NextResponse.json({ candidates })
      }

-      case 'close-session':
+      case 'close-session': {
+        const s = signalingStore.get(sessionId)
+        if (s && s.ownerUserId !== userId) return NextResponse.json({ error: 'Session not found' }, { status: 404 })
        signalingStore.delete(sessionId)
        return NextResponse.json({ success: true })
+      }

      default:
        return NextResponse.json({ error: 'Unknown action' }, { status: 400 })
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -52,6 +52,7 @@ services:
      - "8765:8765"
    environment:
      DATABASE_URL: postgresql://${POSTGRES_USER:-remotelink}:${POSTGRES_PASSWORD}@postgres:5432/${POSTGRES_DB:-remotelink}
+      ALLOWED_ORIGINS: ${NEXT_PUBLIC_APP_URL:-http://localhost:3000}
    depends_on:
      postgres:
        condition: service_healthy
--- a/relay/main.py
+++ b/relay/main.py
@@ -22,6 +22,8 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(mess
 log = logging.getLogger("relay")

 DATABASE_URL = os.environ["DATABASE_URL"]
+# Comma-separated list of allowed origins for CORS, e.g. "https://app.example.com"
+ALLOWED_ORIGINS = [o.strip() for o in os.environ.get("ALLOWED_ORIGINS", "").split(",") if o.strip()]

 # ── In-memory connection registry ────────────────────────────────────────────
 # machine_id (str) → WebSocket
@@ -30,6 +32,8 @@ agents: dict[str, WebSocket] = {}
 viewers: dict[str, WebSocket] = {}
 # session_id → machine_id
 session_to_machine: dict[str, str] = {}
+# session_id → viewer's user role ("admin" | "user")
+viewer_roles: dict[str, str] = {}

 db_pool: Optional[asyncpg.Pool] = None

@@ -47,7 +51,7 @@ app = FastAPI(lifespan=lifespan)

 app.add_middleware(
    CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=ALLOWED_ORIGINS if ALLOWED_ORIGINS else ["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
@@ -68,9 +72,10 @@ async def validate_agent(machine_id: str, access_key: str) -> Optional[dict]:
 async def validate_viewer(session_id: str, viewer_token: str) -> Optional[dict]:
    async with db_pool.acquire() as conn:
        row = await conn.fetchrow(
-            """SELECT id, machine_id, machine_name
-               FROM sessions
-               WHERE id = $1 AND viewer_token = $2 AND ended_at IS NULL""",
+            """SELECT s.id, s.machine_id, s.machine_name, COALESCE(u.role, 'user') AS viewer_role
+               FROM sessions s
+               LEFT JOIN users u ON u.id = s.viewer_user_id
+               WHERE s.id = $1 AND s.viewer_token = $2 AND s.ended_at IS NULL""",
            session_id, viewer_token,
        )
    return dict(row) if row else None
@@ -174,10 +179,13 @@ async def viewer_endpoint(
        await websocket.close(code=4002, reason="No machine associated with session")
        return

+    viewer_role = session.get("viewer_role", "user")
+
    await websocket.accept()
    viewers[session_id] = websocket
    session_to_machine[session_id] = machine_id
-    log.info(f"Viewer connected to session {session_id} (machine {machine_id})")
+    viewer_roles[session_id] = viewer_role
+    log.info(f"Viewer connected to session {session_id} (machine {machine_id}, role {viewer_role})")

    if machine_id in agents:
        # Tell agent to start streaming for this session
@@ -195,6 +203,14 @@ async def viewer_endpoint(
            try:
                event = json.loads(text)
                event["session_id"] = session_id
+                # exec_script is restricted to admin viewers only
+                if event.get("type") == "exec_script":
+                    if viewer_roles.get(session_id) != "admin":
+                        await send_json(websocket, {
+                            "type": "error",
+                            "message": "exec_script requires admin role",
+                        })
+                        continue
                # Forward control events to the agent
                if machine_id in agents:
                    await send_json(agents[machine_id], event)
@@ -207,6 +223,7 @@ async def viewer_endpoint(
    finally:
        viewers.pop(session_id, None)
        session_to_machine.pop(session_id, None)
+        viewer_roles.pop(session_id, None)
        if machine_id in agents:
            await send_json(agents[machine_id], {"type": "stop_stream", "session_id": session_id})
        log.info(f"Viewer disconnected from session {session_id}")