photoai/src/main/mediaProtocol.ts

import { protocol } from 'electron'
import { readFile } from 'node:fs/promises'
import { extname } from 'node:path'
import { MEDIA_SCHEME } from '@shared/constants'
import { profileStore } from './profileStore'
import { presetStore } from './presetStore'
import { logger } from './logger'

const MIME: Record<string, string> = {
  '.jpg': 'image/jpeg',
  '.jpeg': 'image/jpeg',
  '.png': 'image/png',
  '.webp': 'image/webp'
}

/**
 * 로컬 참조 이미지를 UI 렌더러에 안전하게 표시하기 위한 커스텀 프로토콜.
 *
 * URL 형식: photoai-media://img/?p=<encodeURIComponent(절대경로)>
 *
 * 보안: 렌더러가 임의 파일을 읽지 못하도록, **현재 등록된 참조 이미지 경로**에
 * 한해서만 파일을 제공한다. (그 외 경로는 403)
 */

/** app.whenReady 이전에 호출해야 함 — 스킴을 권한 있는 표준 스킴으로 등록 */
export function registerMediaScheme(): void {
  protocol.registerSchemesAsPrivileged([
    {
      scheme: MEDIA_SCHEME,
      privileges: {
        standard: true,
        secure: true,
        supportFetchAPI: true,
        stream: true
      }
    }
  ])
}

/** app.whenReady 이후에 호출 — 실제 요청 핸들러 등록 */
export function handleMediaProtocol(): void {
  protocol.handle(MEDIA_SCHEME, async (request) => {
    try {
      // request.url = photoai-media://img/?p=<encodeURIComponent(절대경로)>
      // searchParams의 자동 디코딩과의 이중 디코딩을 피하려고 raw 문자열을 직접 파싱한다.
      const marker = 'p='
      const i = request.url.indexOf(marker)
      if (i < 0) return new Response('missing path', { status: 400 })
      const filePath = decodeURIComponent(request.url.slice(i + marker.length))

      // 등록된 참조 이미지(활성 프로필 또는 프리셋)에 한해 제공 — 임의 파일 읽기 차단
      const allowed =
        (await profileStore.isReferenceImage(filePath)) ||
        (await presetStore.isReferenceImage(filePath))
      if (!allowed) return new Response('forbidden', { status: 403 })

      // net.fetch(file://) 대신 fs로 직접 읽어 바이트 반환 → 한글/공백 경로에도 안전
      const data = await readFile(filePath)
      const mime = MIME[extname(filePath).toLowerCase()] ?? 'application/octet-stream'
      return new Response(new Uint8Array(data), {
        status: 200,
        headers: { 'content-type': mime, 'cache-control': 'no-cache' }
      })
    } catch (err) {
      logger.error('media 프로토콜 처리 실패', { message: (err as Error).message })
      return new Response('not found', { status: 404 })
    }
  })
}