DragonWind/lib/upload-manager.js

"use strict";
const crypto = require("crypto");
const { getSignedUrl } = require("@aws-sdk/s3-request-presigner");
const { PutObjectCommand } = require("@aws-sdk/client-s3");

/**
 * Dragon Wind Upload Manager
 * Tracks upload sessions for both HTTP (presigned S3) and UDP (relay) modes.
 */
class UploadManager {
  constructor(getS3Client, getConfig) {
    // Accept getter functions so config stays dynamic
    this._getS3 = getS3Client;
    this._getConfig = getConfig;
    this.sessions = new Map();
  }

  get config() { return this._getConfig(); }
  get s3() { return this._getS3(); }

  createSession({ filename, size, mode = "http", prefix = "" }) {
    if (!["http", "udp"].includes(mode)) throw new Error("mode must be 'http' or 'udp'");
    const sessionId = crypto.randomBytes(16).toString("hex");
    const key = prefix ? `${prefix.replace(/[-\/]+$/, "")}--${filename}` : filename;
    const session = {
      sessionId, filename, size, mode, key, prefix,
      status: "pending",
      createdAt: Date.now(),
      updatedAt: Date.now(),
      expiresAt: Date.now() + 2 * 60 * 60 * 1000, // 2h
      presignedUrl: null,
      uploadedBytes: 0,
      s3Key: key,
      error: null,
    };
    this.sessions.set(sessionId, session);
    // Auto-cleanup
    setTimeout(() => this.sessions.delete(sessionId), 2 * 60 * 60 * 1000);
    return session;
  }

  getSession(sessionId) { return this.sessions.get(sessionId) || null; }

  async getPresignedUrl(sessionId) {
    const session = this.sessions.get(sessionId);
    if (!session) throw new Error(`Session not found: ${sessionId}`);
    if (!this.s3) throw new Error("S3 client not initialized — configure S3 in Admin settings");
    const cfg = this.config;
    const cmd = new PutObjectCommand({ Bucket: cfg.bucket, Key: session.key });
    const url = await getSignedUrl(this.s3, cmd, { expiresIn: 3600 });
    session.presignedUrl = url;
    session.updatedAt = Date.now();
    return { url, key: session.key, bucket: cfg.bucket, sessionId };
  }

  async initializeUdpSession(sessionId) {
    const session = this.sessions.get(sessionId);
    if (!session) throw new Error(`Session not found: ${sessionId}`);
    const cfg = this.config;
    if (!cfg.relayUrl) throw new Error("UDP relay not configured — set relay URL in Admin settings");
    session.status = "udp_pending";
    session.updatedAt = Date.now();
    return {
      sessionId,
      relayUrl: cfg.relayUrl,
      udpPort: cfg.udpPort || 5000,
      key: session.key,
      bucket: cfg.bucket,
      filename: session.filename,
      size: session.size,
    };
  }

  completeUpload(sessionId, success = true, error = null) {
    const session = this.sessions.get(sessionId);
    if (!session) throw new Error(`Session not found: ${sessionId}`);
    session.status = success ? "completed" : "failed";
    session.error = error;
    session.updatedAt = Date.now();
    return session;
  }

  getStats() {
    const all = Array.from(this.sessions.values());
    return {
      total: all.length,
      pending: all.filter((s) => s.status === "pending").length,
      completed: all.filter((s) => s.status === "completed").length,
      failed: all.filter((s) => s.status === "failed").length,
      http: all.filter((s) => s.mode === "http").length,
      udp: all.filter((s) => s.mode === "udp").length,
    };
  }

  async getRelayHealth(relayUrl) {
    const url = relayUrl || this.config.relayUrl;
    if (!url) return { healthy: false, error: "Relay URL not configured" };
    try {
      const controller = new AbortController();
      const t = setTimeout(() => controller.abort(), 5000);
      const r = await fetch(`${url}/health`, { signal: controller.signal });
      clearTimeout(t);
      const data = await r.json().catch(() => ({}));
      return { healthy: r.ok, status: r.status, ...data };
    } catch (err) {
      return { healthy: false, error: err.message };
    }
  }
}

module.exports = UploadManager;