Added checkpoint downloader

2026-03-09 04:29:48 +00:00 · 2025-02-22 16:48:15 -07:00
parent 6e19e7449e
commit a280f78c69
7 changed files with 303 additions and 2 deletions
--- a/ui/src/app/api/files/[...filePath]/route.ts
+++ b/ui/src/app/api/files/[...filePath]/route.ts
@@ -0,0 +1,105 @@
+import { NextRequest, NextResponse } from 'next/server';
+import fs from 'fs';
+import path from 'path';
+import { getDatasetsRoot, getTrainingFolder } from '@/server/settings';
+
+export async function GET(request: NextRequest, { params }: { params: { filePath: string } }) {
+  const { filePath } = await params;
+  try {
+    // Decode the path
+    const decodedFilePath = decodeURIComponent(filePath);
+
+    // Get allowed directories
+    const datasetRoot = await getDatasetsRoot();
+    const trainingRoot = await getTrainingFolder();
+    const allowedDirs = [datasetRoot, trainingRoot];
+
+    // Security check: Ensure path is in allowed directory
+    const isAllowed = allowedDirs.some(allowedDir => decodedFilePath.startsWith(allowedDir)) && !decodedFilePath.includes('..');
+
+    if (!isAllowed) {
+      console.warn(`Access denied: ${decodedFilePath} not in ${allowedDirs.join(', ')}`);
+      return new NextResponse('Access denied', { status: 403 });
+    }
+
+    // Check if file exists
+    if (!fs.existsSync(decodedFilePath)) {
+      console.warn(`File not found: ${decodedFilePath}`);
+      return new NextResponse('File not found', { status: 404 });
+    }
+
+    // Get file info
+    const stat = fs.statSync(decodedFilePath);
+    if (!stat.isFile()) {
+      return new NextResponse('Not a file', { status: 400 });
+    }
+
+    // Get filename for Content-Disposition
+    const filename = path.basename(decodedFilePath);
+
+    // Determine content type
+    const ext = path.extname(decodedFilePath).toLowerCase();
+    const contentTypeMap: { [key: string]: string } = {
+      '.jpg': 'image/jpeg',
+      '.jpeg': 'image/jpeg',
+      '.png': 'image/png',
+      '.gif': 'image/gif',
+      '.webp': 'image/webp',
+      '.svg': 'image/svg+xml',
+      '.bmp': 'image/bmp',
+      '.safetensors': 'application/octet-stream',
+    };
+
+    const contentType = contentTypeMap[ext] || 'application/octet-stream';
+
+    // Get range header for partial content support
+    const range = request.headers.get('range');
+
+    // Common headers for better download handling
+    const commonHeaders = {
+      'Content-Type': contentType,
+      'Accept-Ranges': 'bytes',
+      'Cache-Control': 'public, max-age=86400',
+      'Content-Disposition': `attachment; filename="${encodeURIComponent(filename)}"`,
+      'X-Content-Type-Options': 'nosniff'
+    };
+
+    if (range) {
+      // Parse range header
+      const parts = range.replace(/bytes=/, '').split('-');
+      const start = parseInt(parts[0], 10);
+      const end = parts[1] ? parseInt(parts[1], 10) : Math.min(start + 10 * 1024 * 1024, stat.size - 1); // 10MB chunks
+      const chunkSize = (end - start) + 1;
+
+      const fileStream = fs.createReadStream(decodedFilePath, { 
+        start, 
+        end,
+        highWaterMark: 64 * 1024 // 64KB buffer
+      });
+
+      return new NextResponse(fileStream as any, {
+        status: 206,
+        headers: {
+          ...commonHeaders,
+          'Content-Range': `bytes ${start}-${end}/${stat.size}`,
+          'Content-Length': String(chunkSize)
+        },
+      });
+    } else {
+      // For full file download, read directly without streaming wrapper
+      const fileStream = fs.createReadStream(decodedFilePath, {
+        highWaterMark: 64 * 1024 // 64KB buffer
+      });
+
+      return new NextResponse(fileStream as any, {
+        headers: {
+          ...commonHeaders,
+          'Content-Length': String(stat.size)
+        },
+      });
+    }
+  } catch (error) {
+    console.error('Error serving file:', error);
+    return new NextResponse('Internal Server Error', { status: 500 });
+  }
+}