TransferTribe/pages/api/chunk-upload-v2.ts

// pages/api/chunk-upload-v2.ts
// Server-side encryption approach: Upload raw chunks, encrypt on server, store encrypted

import { NextApiRequest, NextApiResponse } from 'next';
import { IncomingForm } from 'formidable';
import fs from 'fs';
import { promises as fsPromises } from 'fs';
import path from 'path';
import { v4 as uuidv4 } from 'uuid';
import { UPLOAD_DIR } from '@/lib/config';
import { prisma } from '@/lib/prisma';
import bcrypt from 'bcrypt';
import { PLAN_CONFIG } from '@/lib/subscription';
import { sendMailjetEmail } from '@/lib/mailjet';
import { deriveKeyFromPassword, encryptChunkFrame, buildHeader } from '@/lib/server-encryption';
import { getServerSession } from 'next-auth/next';
import { authOptions } from '@/lib/auth';
import crypto from 'crypto';

export const config = {
  api: {
    bodyParser: false,
  },
};

interface UploadSession {
  uploadId: string;
  totalSize: number;
  totalChunks: number;
  receivedChunks: Set<number>;
  chunkSize: number;
  sender: string;
  recipient: string;
  password?: string;
  originalFilename: string;
  message?: string;
  filenames?: string;
  createdAt: Date;
  // Encryption metadata. There is deliberately no session-wide IV: each chunk
  // frame carries its own, generated at encryption time.
  salt: Buffer;
  encryptionKey: Buffer;
  receivedBytes: number;
  maxBytes: number;
}

// Store active upload sessions
const uploadSessions = new Map<string, UploadSession>();

// Cleanup abandoned uploads every 5 minutes
const SESSION_TIMEOUT = 24 * 60 * 60 * 1000; // 24 hours
setInterval(() => {
  const now = Date.now();
  for (const [uploadId, session] of uploadSessions.entries()) {
    if (now - session.createdAt.getTime() > SESSION_TIMEOUT) {
      uploadSessions.delete(uploadId);
      const chunkDir = path.join(UPLOAD_DIR, '.tmp', uploadId);
      fsPromises.rm(chunkDir, { recursive: true }).catch(err => 
        console.warn(`Failed to clean abandoned upload ${uploadId}:`, err)
      );
    }
  }
}, 5 * 60 * 1000);

export default async function handler(req: NextApiRequest, res: NextApiResponse) {
  const authSession = await getServerSession(req, res, authOptions);
  if (!authSession?.user?.email) {
    return res.status(401).json({ success: false, message: 'Unauthorized' });
  }
  const sessionEmail = authSession.user.email;

  if (req.method === 'POST') {
    return handleChunkUpload(req, res, sessionEmail);
  } else if (req.method === 'GET') {
    return handleStatusCheck(req, res, sessionEmail);
  } else if (req.method === 'PUT') {
    return handleChunkComplete(req, res, sessionEmail);
  }

  return res.status(405).json({ success: false, message: 'Method not allowed' });
}

async function handleChunkUpload(
  req: NextApiRequest,
  res: NextApiResponse,
  sessionEmail: string
) {
  const form = new IncomingForm({
    uploadDir: path.join(UPLOAD_DIR, '.tmp'),
    keepExtensions: true,
  });

  try {
    await fsPromises.mkdir(path.join(UPLOAD_DIR, '.tmp'), { recursive: true });
  } catch (err) {
    console.error('Failed to create temp directory:', err);
  }

  return new Promise((resolve) => {
    form.parse(req, async (err: Error | null, fields: any, files: any) => {
      if (err) {
        return resolve(res.status(400).json({ success: false, message: 'Parse error: ' + err.message }));
      }

      try {
        const chunkFile = Array.isArray(files.chunk) ? files.chunk[0] : files.chunk;
        const uploadId = Array.isArray(fields.uploadId) ? fields.uploadId[0] : fields.uploadId;
        const chunkIndex = parseInt(Array.isArray(fields.chunkIndex) ? fields.chunkIndex[0] : fields.chunkIndex);
        const totalChunks = parseInt(Array.isArray(fields.totalChunks) ? fields.totalChunks[0] : fields.totalChunks);
        const chunkSize = parseInt(Array.isArray(fields.chunkSize) ? fields.chunkSize[0] : fields.chunkSize);

        if (!chunkFile || !uploadId || isNaN(chunkIndex) || isNaN(totalChunks)) {
          return resolve(res.status(400).json({ success: false, message: 'Missing chunk metadata' }));
        }

        // Get or create upload session
        let session = uploadSessions.get(uploadId);
        if (session && session.sender !== sessionEmail) {
          return resolve(res.status(403).json({ success: false, message: 'Forbidden' }));
        }
        if (!session) {
          const sender = sessionEmail; // never trust a sender field from the body
          const recipient = Array.isArray(fields.email) ? fields.email[0] : fields.email;
          const password = Array.isArray(fields.password) ? fields.password[0] : fields.password;
          const originalFilename = Array.isArray(fields.originalFilename) ? fields.originalFilename[0] : fields.originalFilename;
          const message = Array.isArray(fields.message) ? fields.message[0] : fields.message;
          const filenames = Array.isArray(fields.filenames) ? fields.filenames[0] : fields.filenames;

          if (!recipient) {
            return resolve(res.status(400).json({ success: false, message: 'Missing required fields' }));
          }

          // A password is mandatory: it is the key material. Without it every
          // file would be encrypted under a key derived from the empty string.
          if (!password) {
            return resolve(res.status(400).json({ success: false, message: 'A password is required' }));
          }

          // Verify user exists
          const user = await prisma.user.findUnique({ where: { email: sender } });
          if (!user) {
            return resolve(res.status(404).json({ success: false, message: 'User not found' }));
          }

          // Generate salt and derive encryption key
          const salt = crypto.randomBytes(16);
          const encryptionKey = await deriveKeyFromPassword(password, salt);

          const plan = (user.plan || 'free') as 'free' | 'rookie' | 'pro';
          const limits = PLAN_CONFIG[plan];

          session = {
            uploadId,
            totalSize: chunkSize * totalChunks,
            totalChunks,
            receivedChunks: new Set(),
            chunkSize,
            sender,
            recipient,
            password,
            originalFilename,
            message,
            filenames,
            createdAt: new Date(),
            salt,
            encryptionKey,
            receivedBytes: 0,
            maxBytes: limits.maxFileSize,
          };

          // The advertised size is client-supplied, so this is only an early
          // reject; the real enforcement is the running receivedBytes check.
          if (limits.maxFileSize !== Infinity && session.totalSize > limits.maxFileSize) {
            return resolve(res.status(413).json({
              success: false,
              message: `File exceeds max size for your plan (${plan})`
            }));
          }

          uploadSessions.set(uploadId, session);
        }

        // Read and encrypt chunk into a self-contained frame (own IV + own tag)
        const chunkData = await fsPromises.readFile(chunkFile.filepath);

        if (
          session.maxBytes !== Infinity &&
          session.receivedBytes + chunkData.length > session.maxBytes
        ) {
          await fsPromises.unlink(chunkFile.filepath).catch(() => {});
          return resolve(res.status(413).json({
            success: false,
            message: 'File exceeds max size for your plan',
          }));
        }

        const frame = encryptChunkFrame(chunkData, session.encryptionKey);

        // Save encrypted chunk to temp location
        const chunkDir = path.join(UPLOAD_DIR, '.tmp', uploadId);
        await fsPromises.mkdir(chunkDir, { recursive: true });
        const encryptedChunkPath = path.join(chunkDir, `chunk-${chunkIndex}.enc`);

        await fsPromises.writeFile(encryptedChunkPath, frame);

        // Clean up formidable temp file
        try {
          await fsPromises.unlink(chunkFile.filepath);
        } catch (err) {
          console.warn('Failed to clean up temp file:', err);
        }

        // Guard against a retried chunk double-counting toward the quota.
        if (!session.receivedChunks.has(chunkIndex)) {
          session.receivedBytes += chunkData.length;
        }
        session.receivedChunks.add(chunkIndex);

        return resolve(res.status(200).json({
          success: true,
          uploadId,
          chunkIndex,
          receivedChunks: Array.from(session.receivedChunks),
        }));
      } catch (err: any) {
        console.error('Chunk upload error:', err);
        return resolve(res.status(500).json({ success: false, message: 'Upload error: ' + err.message }));
      }
    });
  });
}

async function handleStatusCheck(
  req: NextApiRequest,
  res: NextApiResponse,
  sessionEmail: string
) {
  const { uploadId } = req.query;

  if (!uploadId || typeof uploadId !== 'string') {
    return res.status(400).json({ success: false, message: 'Missing uploadId' });
  }

  const session = uploadSessions.get(uploadId);
  if (!session || session.sender !== sessionEmail) {
    return res.status(404).json({ success: false, message: 'Upload session not found' });
  }

  return res.status(200).json({
    success: true,
    uploadId,
    totalChunks: session.totalChunks,
    receivedChunks: Array.from(session.receivedChunks),
    isComplete: session.receivedChunks.size === session.totalChunks,
    receivedBytes: session.receivedBytes,
    totalBytes: session.totalSize,
  });
}

async function handleChunkComplete(
  req: NextApiRequest,
  res: NextApiResponse,
  sessionEmail: string
) {
  const { uploadId } = req.query;

  if (!uploadId || typeof uploadId !== 'string') {
    return res.status(400).json({ success: false, message: 'Missing uploadId' });
  }

  const session = uploadSessions.get(uploadId);
  if (!session || session.sender !== sessionEmail) {
    return res.status(404).json({ success: false, message: 'Upload session not found' });
  }

  // Check if all chunks received
  if (session.receivedChunks.size !== session.totalChunks) {
    return res.status(400).json({
      success: false,
      message: `Missing chunks. Received ${session.receivedChunks.size}/${session.totalChunks}`,
    });
  }

  try {
    // Reassemble encrypted chunks into final encrypted file
    console.log(`[${uploadId}] Starting reassembly of ${session.totalChunks} encrypted chunks`);
    
    const chunkDir = path.join(UPLOAD_DIR, '.tmp', uploadId);
    const finalFilename = uuidv4() + '.enc';
    const finalPath = path.join(UPLOAD_DIR, finalFilename);

    const writeStream = fs.createWriteStream(finalPath);

    // Honour backpressure: createWriteStream queues in memory whenever write()
    // returns false, so a multi-GB reassembly that ignores the return value
    // grows the internal buffer instead of flushing to disk.
    const write = (buf: Buffer): Promise<void> =>
      new Promise((resolve, reject) => {
        if (writeStream.write(buf)) return resolve();
        writeStream.once('drain', resolve);
        writeStream.once('error', reject);
      });

    // Write the format header: magic (8) + salt (16). The per-chunk IVs live
    // in the frames themselves, so nothing needs storing in the DB.
    await write(buildHeader(session.salt));

    // Assemble encrypted frames in order, with progress tracking
    let chunksAssembled = 0;
    for (let i = 0; i < session.totalChunks; i++) {
      const chunkPath = path.join(chunkDir, `chunk-${i}.enc`);

      // Retry reading with exponential backoff
      let chunkData: Buffer | null = null;
      for (let attempt = 0; attempt < 3; attempt++) {
        try {
          chunkData = await fsPromises.readFile(chunkPath);
          break;
        } catch (err: any) {
          if (attempt < 2 && err.code === 'EACCES') {
            await new Promise(resolve => setTimeout(resolve, Math.pow(2, attempt) * 100));
          } else {
            throw err;
          }
        }
      }

      // A silently skipped chunk would corrupt the file, so fail loudly.
      if (!chunkData) {
        throw new Error(`Chunk ${i} missing during reassembly`);
      }

      await write(chunkData);
      chunksAssembled++;

      if (chunksAssembled % 10 === 0 || chunksAssembled === session.totalChunks) {
        console.log(`[${uploadId}] Reassembly progress: ${chunksAssembled}/${session.totalChunks} chunks`);
      }
    }

    console.log(`[${uploadId}] Waiting for write stream to finish...`);
    
    await new Promise<void>((resolve, reject) => {
      writeStream.on('finish', resolve);
      writeStream.on('error', reject);
      writeStream.end();
    });

    console.log(`[${uploadId}] Reassembly complete, validating file...`);

    // Sizes recorded and metered are PLAINTEXT bytes actually received, not the
    // on-disk size, which is inflated by the header and per-frame IVs and tags.
    const plaintextSize = session.receivedBytes;

    // Verify user and check monthly limits
    const user = await prisma.user.findUnique({ where: { email: session.sender } });
    if (!user) {
      await fsPromises.unlink(finalPath);
      return res.status(404).json({ success: false, message: 'User not found' });
    }

    const plan = (user.plan || 'free') as 'free' | 'rookie' | 'pro';
    const limits = PLAN_CONFIG[plan];

    const startOfMonth = new Date();
    startOfMonth.setDate(1);
    startOfMonth.setHours(0, 0, 0, 0);

    const transfersThisMonth = await prisma.transfer.findMany({
      where: {
        senderEmail: session.sender,
        createdAt: { gte: startOfMonth },
      },
    });

    const totalSizeSentThisMonth = transfersThisMonth.reduce((acc, t) => acc + BigInt(t.totalSize), BigInt(0));
    if (
      (limits.maxTransfersPerMonth !== Infinity && transfersThisMonth.length >= limits.maxTransfersPerMonth) ||
      (limits.maxTransferSizePerMonth !== Infinity && totalSizeSentThisMonth + BigInt(plaintextSize) > BigInt(limits.maxTransferSizePerMonth))
    ) {
      await fsPromises.unlink(finalPath);
      return res.status(429).json({ success: false, message: 'Plan limits exceeded' });
    }

    // Hash password for verification (don't encrypt again)
    const hash = session.password ? await bcrypt.hash(session.password, 10) : null;

    // Parse filenames
    let totalFiles = 1;
    if (session.filenames) {
      try {
        const filesArray = JSON.parse(session.filenames);
        if (Array.isArray(filesArray)) {
          totalFiles = filesArray.length;
        }
      } catch {}
    }

    const expiresAt = new Date(Date.now() + limits.maxExpiryMs);

    // Create transfer record
    const transfer = await prisma.transfer.create({
      data: {
        senderEmail: session.sender,
        recipientEmail: session.recipient,
        passwordHash: hash,
        downloadUrl: uuidv4(),
        expiresAt,
        filenames: session.filenames || null,
        message: session.message || null,
        totalFiles,
        totalSize: plaintextSize,
        userId: user.id,
      },
    });

    // Create file record with encryption metadata
    // Store salt and IV in the file path (prepended to encrypted data)
    // So we don't need separate DB fields
    await prisma.file.create({
      data: {
        name: session.originalFilename || 'file',
        path: finalPath,
        size: plaintextSize,
        transferId: transfer.id,
      },
    });

    // Send email
    const link = `${process.env.NEXT_PUBLIC_APP_URL}/download/${transfer.downloadUrl}`;

    await sendMailjetEmail({
      to: session.recipient,
      subject: 'You\'ve received an encrypted file',
      text: `
${session.sender} sent you an encrypted file via TransferTribe.

Link: ${link}

${session.message ? `Message:\n${session.message}\n\n` : ''}Note: You'll need the password they shared with you to decrypt it.
`.trim(),
    });

    // Clean up session and temp files
    uploadSessions.delete(uploadId);
    try {
      await fsPromises.rm(chunkDir, { recursive: true });
    } catch (err) {
      console.warn('Failed to clean up temp directory:', err);
    }

    return res.status(200).json({
      success: true,
      message: 'Upload complete',
      downloadUrl: transfer.downloadUrl,
    });
  } catch (err: any) {
    console.error('Chunk completion error:', err);
    return res.status(500).json({ success: false, message: 'Completion error: ' + err.message });
  }
}