AeroFetch/src/main/lib/verifiedDownload.ts

import { net } from 'electron'
import { createWriteStream, type WriteStream } from 'fs'
import { unlink } from 'fs/promises'
import { createHash } from 'crypto'

/** Result of streaming a verified download to disk. */
export interface VerifiedDownloadResult {
  ok: boolean
  /** absolute path to the written file, when ok */
  filePath?: string
  error?: string
}

/** Live byte progress of a streaming download (main → renderer). */
export interface VerifiedDownloadProgress {
  /** bytes received so far */
  received: number
  /** total bytes, when the server reports Content-Length */
  total?: number
  /** 0..1 fraction, when total is known */
  fraction?: number
}

/** Context-specific message overrides; the defaults are generic and suit any caller. */
export interface VerifiedDownloadMessages {
  /** shown when the streamed bytes don't match expectedSha */
  checksumMismatch?: string
  /** shown when the caller-armed cancel fires */
  canceled?: string
}

export interface StreamVerifiedFileOptions {
  url: string
  filePath: string
  /** published SHA-256 (lowercase hex) to verify against, or null to skip verification */
  expectedSha: string | null
  /** per-hop trust gate: a URL is fetched or followed only if this returns true for it */
  isTrusted: (url: string) => boolean
  onProgress: (p: VerifiedDownloadProgress) => void
  onCancelReady: (cancel: () => void) => void
  onSettled: (cancel: () => void) => void
  /** stall budget: how long with no bytes received before abandoning (default 60s) */
  idleTimeoutMs?: number
  messages?: VerifiedDownloadMessages
}

/** Default stall (no-bytes) budget before a download is abandoned. */
const DEFAULT_IDLE_TIMEOUT_MS = 60_000

/**
 * Stream one HTTPS resource into a file with full download discipline: per-hop redirect
 * re-validation against a caller-supplied trust gate, streaming SHA-256 verification,
 * idle-timeout, content-length truncation detection, write-stream backpressure, and a
 * single teardown point that aborts the request and removes the partial file on every
 * failure path.
 *
 * Extracted from the app updater (CL4) so the same audited loop backs both the installer
 * download and the first-run ffmpeg fetch; the differences (which hosts are trusted, the
 * checksum source, a couple of user-facing strings) are parameters, not forks. net.request
 * (not fetch) is used so the host can be re-validated on EVERY redirect hop — undici's
 * fetch hides the Location header under redirect:'manual', so it can't gate hops.
 *
 * The caller owns the cancel slot: `onCancelReady` hands it this download's cancel
 * function once armed, and `onSettled` hands the same function back on completion so the
 * caller can release the slot only if this download still owns it.
 */
export function streamVerifiedFile(
  opts: StreamVerifiedFileOptions
): Promise<VerifiedDownloadResult> {
  const { url, filePath, expectedSha, isTrusted, onProgress, onCancelReady, onSettled } = opts
  const idleTimeoutMs = opts.idleTimeoutMs ?? DEFAULT_IDLE_TIMEOUT_MS
  const canceledMsg = opts.messages?.canceled ?? 'Download canceled.'
  const checksumMsg =
    opts.messages?.checksumMismatch ??
    'The download failed its checksum check — it may be corrupt or tampered with.'

  return new Promise<VerifiedDownloadResult>((resolve) => {
    let settled = false
    let fileStream: WriteStream | null = null
    let idle: NodeJS.Timeout | null = null

    // Single teardown point. On failure it also aborts the request, so a write error
    // (disk full, etc.) can't leave Electron pulling bytes into a dead stream. abort()
    // is idempotent, so the call sites below don't repeat it.
    const finish = (result: VerifiedDownloadResult): void => {
      if (settled) return
      settled = true
      onSettled(requestCancel)
      if (idle) clearTimeout(idle)
      if (!result.ok) {
        request.abort()
        // Close the handle before unlinking (Windows won't delete an open file) and
        // never leave a partial/aborted file lying around.
        if (fileStream && !fileStream.destroyed) fileStream.destroy()
        unlink(filePath).catch(() => {})
      }
      resolve(result)
    }

    // Expose this download's abort so the caller's Cancel routes through the same
    // teardown — request.abort() + partial cleanup.
    const requestCancel = (): void => finish({ ok: false, error: canceledMsg })
    onCancelReady(requestCancel)

    // 'manual' means a hop only proceeds if we call followRedirect().
    const request = net.request({ url, redirect: 'manual' })

    const armIdle = (): void => {
      if (idle) clearTimeout(idle)
      idle = setTimeout(() => {
        finish({ ok: false, error: 'Download stalled — please try again.' })
      }, idleTimeoutMs)
    }

    request.on('redirect', (_status, _method, redirectUrl) => {
      if (!isTrusted(redirectUrl)) {
        finish({ ok: false, error: 'Refused to follow a redirect to an untrusted location.' })
        return
      }
      request.followRedirect()
    })

    request.on('response', (response) => {
      const status = response.statusCode
      if (status < 200 || status >= 300) {
        finish({ ok: false, error: `Download failed (HTTP ${status}).` })
        return
      }

      const lenHeader = response.headers['content-length']
      const total = Number(Array.isArray(lenHeader) ? lenHeader[0] : lenHeader) || undefined

      const stream = createWriteStream(filePath)
      fileStream = stream
      // Hash the bytes as they stream by, so verification needs no second pass.
      const hash = createHash('sha256')
      // Electron's IncomingMessage is event-based (no .pipe). pause()/resume() exist at
      // runtime but aren't in its type, so feature-detect them to apply backpressure —
      // without it a large download can buffer in memory.
      const flow = response as unknown as { pause?(): void; resume?(): void }
      let received = 0
      armIdle()

      response.on('data', (chunk: Buffer) => {
        armIdle()
        received += chunk.length
        hash.update(chunk)
        if (!stream.write(chunk) && flow.pause && flow.resume) {
          flow.pause()
          stream.once('drain', () => flow.resume?.())
        }
        onProgress({ received, total, fraction: total ? received / total : undefined })
      })

      response.on('end', () => {
        // Body fully received — "stalled" is no longer meaningful past this point.
        if (idle) clearTimeout(idle)
        stream.end(() => {
          // A truncated download (connection dropped mid-stream) must never be used.
          if (total !== undefined && received !== total) {
            finish({ ok: false, error: 'The download was incomplete — please try again.' })
            return
          }
          // Verify the bytes we wrote match the published SHA-256.
          if (expectedSha && hash.digest('hex') !== expectedSha) {
            finish({ ok: false, error: checksumMsg })
            return
          }
          finish({ ok: true, filePath })
        })
      })

      // A mid-stream abort emits neither 'end' nor always 'error'; catch it so the
      // promise can't hang.
      response.on('aborted', () => finish({ ok: false, error: 'The download was interrupted.' }))
      response.on('error', (e: Error) => finish({ ok: false, error: e.message }))
      stream.on('error', (e: Error) => finish({ ok: false, error: e.message }))
    })

    request.on('error', (e: Error) => finish({ ok: false, error: e.message }))
    request.end()
  })
}