ai-chatbot/scripts/storage/configure-direct-upload-cors.ts

import { config as loadEnv } from "dotenv";

import {
  GetBucketCorsCommand,
  PutBucketCorsCommand,
  S3Client,
  type CORSRule,
} from "@aws-sdk/client-s3";
import { and, eq, inArray } from "drizzle-orm";

loadEnv({ path: ".env.local" });
loadEnv();

type S3UploadConfig = {
  name?: string;
  endpoint: string;
  region: string;
  bucket: string;
  accessKeyId: string;
  secretAccessKey: string;
  provider?: "s3" | "qiniu-kodo";
};

const DEFAULT_ORIGINS = [
  "https://ai.yydd.ai",
  "https://yydd.ai",
  "https://bbyy.ai",
  "https://yybb.top",
  "https://goyingup.chat",
  "http://localhost:3000",
  "http://localhost:3001",
];

function normalizeEndpoint(endpoint: string): string {
  const trimmed = endpoint.trim().replace(/\/+$/, "");
  if (!trimmed) return "";

  if (/aliyuncs\.com/i.test(trimmed)) {
    const withoutProtocol = trimmed.replace(/^https?:\/\//i, "");
    return `https://${withoutProtocol}`;
  }

  return trimmed;
}

function createS3Client(config: S3UploadConfig) {
  const normalizedEndpoint = normalizeEndpoint(config.endpoint);
  const lowerEndpoint = normalizedEndpoint.toLowerCase();
  const isCloudS3 =
    lowerEndpoint.includes("aliyuncs.com") ||
    lowerEndpoint.includes("amazonaws.com");

  return new S3Client({
    endpoint: normalizedEndpoint,
    region: config.region || "auto",
    credentials: {
      accessKeyId: config.accessKeyId,
      secretAccessKey: config.secretAccessKey,
    },
    forcePathStyle: config.provider === "qiniu-kodo" ? true : !isCloudS3,
  });
}

function toS3Config(value: unknown): S3UploadConfig | null {
  const raw = value as Record<string, unknown> | null | undefined;
  const endpoint = raw?.endpoint ? normalizeEndpoint(String(raw.endpoint)) : "";
  const bucket = raw?.bucket ? String(raw.bucket).trim() : "";
  const accessKeyId = raw?.accessKeyId ? String(raw.accessKeyId).trim() : "";
  const secretAccessKey = raw?.secretAccessKey
    ? String(raw.secretAccessKey).trim()
    : "";

  if (
    !raw ||
    (raw.type !== "s3" && raw.type !== "qiniu-kodo") ||
    raw.isEnabled !== true ||
    !endpoint ||
    !bucket ||
    !accessKeyId ||
    !secretAccessKey
  ) {
    return null;
  }

  return {
    name: raw.name ? String(raw.name) : undefined,
    endpoint,
    region: String(raw.region || "auto").trim() || "auto",
    bucket,
    accessKeyId,
    secretAccessKey,
    provider: raw.type === "qiniu-kodo" ? "qiniu-kodo" : "s3",
  };
}

function readEnvS3Config(): S3UploadConfig | null {
  const endpoint = process.env.S3_ENDPOINT
    ? normalizeEndpoint(process.env.S3_ENDPOINT)
    : "";
  const bucket = process.env.S3_BUCKET?.trim() || "";
  const accessKeyId = process.env.S3_ACCESS_KEY_ID?.trim() || "";
  const secretAccessKey = process.env.S3_SECRET_ACCESS_KEY?.trim() || "";

  if (!endpoint || !bucket || !accessKeyId || !secretAccessKey) {
    return null;
  }

  return {
    name: "Environment S3",
    endpoint,
    region: process.env.S3_REGION?.trim() || "auto",
    bucket,
    accessKeyId,
    secretAccessKey,
  };
}

async function resolveS3Configs() {
  const [{ db }, { storageConfig }] = await Promise.all([
    import("@/lib/db"),
    import("@/lib/db/schema"),
  ]);
  const storageConfigs = await db
    .select()
    .from(storageConfig)
    .where(
      and(
        inArray(storageConfig.type, ["s3", "qiniu-kodo"]),
        eq(storageConfig.isEnabled, true)
      )
    );

  const configs = storageConfigs
    .map(toS3Config)
    .filter((config): config is S3UploadConfig => Boolean(config));
  const envConfig = readEnvS3Config();

  if (configs.length === 0 && envConfig) {
    return [envConfig];
  }

  return configs;
}

function readOrigins() {
  const raw = process.env.DIRECT_UPLOAD_CORS_ORIGINS || "";
  const origins = raw
    .split(",")
    .map((origin) => origin.trim())
    .filter(Boolean);

  return origins.length > 0 ? origins : DEFAULT_ORIGINS;
}

function normalizeRule(rule: CORSRule): CORSRule {
  return {
    AllowedHeaders: Array.from(new Set(rule.AllowedHeaders || [])),
    AllowedMethods: Array.from(new Set(rule.AllowedMethods || [])),
    AllowedOrigins: Array.from(new Set(rule.AllowedOrigins || [])),
    ExposeHeaders: Array.from(new Set(rule.ExposeHeaders || [])),
    MaxAgeSeconds: rule.MaxAgeSeconds,
  };
}

function isDirectUploadRule(rule: CORSRule, origins: string[]) {
  const methods = new Set(rule.AllowedMethods || []);
  const allowedOrigins = new Set(rule.AllowedOrigins || []);
  return (
    methods.has("PUT") &&
    methods.has("POST") &&
    origins.some((origin) => allowedOrigins.has(origin))
  );
}

function isCorsApiUnsupported(error: unknown) {
  const value = error as { Code?: string; name?: string; message?: string } | null;
  const code = String(value?.Code || value?.name || "");
  const message = String(value?.message || "");
  return code === "NotImplemented" || /not implemented/i.test(message);
}

async function main() {
  const configs = await resolveS3Configs();
  if (configs.length === 0) {
    console.log("[direct-upload-cors] skip: no S3-compatible upload targets");
    return;
  }

  const origins = readOrigins();
  const directRule: CORSRule = {
    AllowedOrigins: origins,
    AllowedMethods: ["GET", "PUT", "POST", "HEAD"],
    AllowedHeaders: ["*"],
    ExposeHeaders: ["ETag", "x-amz-request-id", "x-amz-id-2"],
    MaxAgeSeconds: 3600,
  };
  const failures: Array<{ config: S3UploadConfig; error: unknown }> = [];

  for (const config of configs) {
    const client = createS3Client(config);
    let existingRules: CORSRule[] = [];
    try {
      const current = await client.send(
        new GetBucketCorsCommand({ Bucket: config.bucket })
      );
      existingRules = current.CORSRules || [];
    } catch (error) {
      console.warn(
        "[direct-upload-cors] existing CORS not found or unreadable, replacing with direct upload rule:",
        {
          name: config.name,
          bucket: config.bucket,
          error: error instanceof Error ? error.message : String(error),
        }
      );
    }

    const mergedRules = [
      directRule,
      ...existingRules.filter((rule) => !isDirectUploadRule(rule, origins)),
    ]
      .slice(0, 100)
      .map(normalizeRule);

    try {
      await client.send(
        new PutBucketCorsCommand({
          Bucket: config.bucket,
          CORSConfiguration: {
            CORSRules: mergedRules,
          },
        })
      );

      console.log("[direct-upload-cors] configured", {
        name: config.name,
        provider: config.provider,
        bucket: config.bucket,
        origins,
        ruleCount: mergedRules.length,
      });
    } catch (error) {
      if (isCorsApiUnsupported(error)) {
        console.warn("[direct-upload-cors] bucket CORS API unsupported; leaving provider/global CORS in place", {
          name: config.name,
          provider: config.provider,
          bucket: config.bucket,
          error: error instanceof Error ? error.message : String(error),
        });
        continue;
      }

      failures.push({ config, error });
      console.error("[direct-upload-cors] configure bucket failed", {
        name: config.name,
        provider: config.provider,
        bucket: config.bucket,
        error: error instanceof Error ? error.message : String(error),
      });
    }
  }

  if (failures.length > 0) {
    throw new Error(
      `Failed to configure direct upload CORS for ${failures.length} bucket(s)`
    );
  }
}

main()
  .then(() => {
    process.exit(0);
  })
  .catch((error) => {
    console.error("[direct-upload-cors] failed:", error);
    process.exit(1);
  });