使用CloudFlare Workers搭建网页截图API

export default {
  async fetch(request, env, ctx) {
    const cache = caches.default;
    const kv = env.SCREENSHOT;

    const requestUrl = new URL(request.url);
    const targetUrl = requestUrl.searchParams.get("url");

    if (!targetUrl) {
      return new Response("Missing 'url' parameter", { status: 400 });
    }

    const referer = request.headers.get("referer") || "";
    const origin = request.headers.get("origin") || "";
    const allowedDomains = env.ALLOWED_DOMAINS ? env.ALLOWED_DOMAINS.split(",") : [];

    const isAllowed = allowedDomains.some(domain => {
      const trimmedDomain = domain.trim();
      return referer.includes(trimmedDomain) || origin.includes(trimmedDomain);
    });

    if (allowedDomains.length > 0 && !isAllowed) {
      return new Response("Access denied: request not from allowed domain", { status: 403 });
    }

    const date = new Date().toISOString().split("T")[0];
    const cacheKey = targetUrl;
    const datedKey = `${targetUrl}?${date}`;

    // 工具函数：构建 Response 对象
    const buildResponse = (buffer) =>
      new Response(buffer, {
        headers: {
          "content-type": "image/png",
          "cache-control": "public, max-age=86400, immutable",
        },
      });

    // 工具函数：尝试从 KV 和 Cache 中加载已有截图
    const tryGetCachedResponse = async (key) => {
      let res = await cache.match(key);
      if (res) return res;

      const kvData = await kv.get(key, { type: "arrayBuffer" });
      if (kvData) {
        res = buildResponse(kvData);
        ctx.waitUntil(cache.put(key, res.clone()));
        return res;
      }
      return null;
    };

    // 1. 优先使用当日缓存
    let res = await tryGetCachedResponse(datedKey);
    if (res) return res;

    // 2. 若缓存不存在，则请求 Cloudflare Screenshot API
    try {
      const payload = {
        url: targetUrl,
        viewport: { width: 1200, height: 800 },
        gotoOptions: { waitUntil: "networkidle0" },
      };

      const apiRes = await fetch(
        `https://api.cloudflare.com/client/v4/accounts/${env.CF_ACCOUNT_ID}/browser-rendering/screenshot?cacheTTL=86400`,
        {
          method: "POST",
          headers: {
            Authorization: `Bearer ${env.CF_API_TOKEN}`,
            "Content-Type": "application/json",
          },
          body: JSON.stringify(payload),
        }
      );

      if (!apiRes.ok) throw new Error(`API returned ${apiRes.status}`);

      const buffer = await apiRes.arrayBuffer();
      res = buildResponse(buffer);

      // 后台缓存更新
      ctx.waitUntil(Promise.all([
        kv.put(cacheKey, buffer),
        kv.put(datedKey, buffer, { expirationTtl: 86400 }),
        cache.put(cacheKey, res.clone()),
        cache.put(datedKey, res.clone()),
      ]));

      return res;
    } catch (err) {
      console.error("Screenshot generation failed:", err);

      // 3. 回退到通用旧缓存
      res = await tryGetCachedResponse(cacheKey);
      if (res) return res;

      return new Response("Screenshot generation failed", { status: 502 });
    }
  },
};