MyGoNavi/frontend/src/utils/connectionReadOnly.ts

import type { ConnectionConfig } from "../types";
import { convertMongoShellToJsonCommand } from "./mongodb";
import { resolveSqlDialect } from "./sqlDialect";
import { findSqlStatementRanges } from "./sqlStatementSelection";

export type ConnectionProtectionKey =
  | "restrictDataEdit"
  | "restrictStructureEdit"
  | "restrictScriptExecution"
  | "restrictDataImport";

export type ConnectionProtectionConfig = NonNullable<
  ConnectionConfig["protection"]
>;

type ConnectionReadOnlyLike = Pick<
  ConnectionConfig,
  "type" | "driver" | "oceanBaseProtocol" | "readOnly" | "protection"
> | null | undefined;

export const CONNECTION_PROTECTION_KEYS: ConnectionProtectionKey[] = [
  "restrictDataEdit",
  "restrictStructureEdit",
  "restrictScriptExecution",
  "restrictDataImport",
];

const EMPTY_CONNECTION_PROTECTION: ConnectionProtectionConfig = {
  restrictDataEdit: false,
  restrictStructureEdit: false,
  restrictScriptExecution: false,
  restrictDataImport: false,
};

const FULL_CONNECTION_PROTECTION: ConnectionProtectionConfig = {
  restrictDataEdit: true,
  restrictStructureEdit: true,
  restrictScriptExecution: true,
  restrictDataImport: true,
};

const CONNECTION_READ_ONLY_TYPES = new Set([
  "mysql",
  "goldendb",
  "mariadb",
  "oceanbase",
  "diros",
  "starrocks",
  "sphinx",
  "postgres",
  "kingbase",
  "highgo",
  "vastbase",
  "opengauss",
  "gaussdb",
  "sqlserver",
  "iris",
  "sqlite",
  "duckdb",
  "oracle",
  "dameng",
  "tdengine",
  "clickhouse",
  "trino",
  "mongodb",
]);

const SQL_READ_ONLY_KEYWORDS = new Set([
  "select",
  "with",
  "show",
  "describe",
  "desc",
  "explain",
  "pragma",
  "values",
  "consume",
]);

const SQL_MUTATING_WITH_KEYWORDS = /\b(insert|update|delete|replace|merge|upsert)\b/i;
const SQL_SELECT_INTO_PATTERN = /^\s*select\b[\s\S]*\binto\b/i;

const MONGO_READ_ONLY_COMMANDS = new Set([
  "aggregate",
  "buildinfo",
  "collstats",
  "connectionstatus",
  "count",
  "countdocuments",
  "dbstats",
  "distinct",
  "explain",
  "find",
  "findone",
  "getparameter",
  "hello",
  "hostinfo",
  "ismaster",
  "listcollections",
  "listdatabases",
  "listindexes",
  "ping",
  "serverstatus",
]);

const MONGO_WRITE_COMMANDS = new Set([
  "bulkwrite",
  "collmod",
  "create",
  "createindexes",
  "delete",
  "drop",
  "dropdatabase",
  "dropindexes",
  "findandmodify",
  "insert",
  "mapreduce",
  "renamecollection",
  "update",
]);

const MONGO_META_KEYS = new Set([
  "$db",
  "$readpreference",
  "api",
  "apideprecationerrors",
  "apistrict",
  "comment",
  "let",
  "lsid",
  "maxtimems",
  "ordered",
  "readconcern",
  "writeconcern",
]);

const stripLeadingSqlComments = (statement: string): string => {
  let text = String(statement || "").trim();
  while (text) {
    if (text.startsWith("--")) {
      const next = text.indexOf("\n");
      text = next >= 0 ? text.slice(next + 1).trimStart() : "";
      continue;
    }
    if (text.startsWith("#")) {
      const next = text.indexOf("\n");
      text = next >= 0 ? text.slice(next + 1).trimStart() : "";
      continue;
    }
    if (text.startsWith("/*")) {
      const next = text.indexOf("*/");
      text = next >= 0 ? text.slice(next + 2).trimStart() : "";
      continue;
    }
    break;
  }
  return text;
};

const extractLeadingSqlKeyword = (statement: string): string => {
  const text = stripLeadingSqlComments(statement);
  const match = text.match(/^[A-Za-z_][A-Za-z0-9_]*/);
  return match ? match[0].toLowerCase() : "";
};

const resolveConnectionReadOnlyType = (
  config: ConnectionReadOnlyLike,
): string => {
  if (!config) return "";
  return String(
    resolveSqlDialect(String(config.type || ""), String(config.driver || ""), {
      oceanBaseProtocol: config.oceanBaseProtocol,
    }),
  )
    .trim()
    .toLowerCase();
};

const isReadOnlySqlStatement = (statement: string): boolean => {
  const text = stripLeadingSqlComments(statement);
  if (!text) return true;
  const keyword = extractLeadingSqlKeyword(text);
  if (!keyword || !SQL_READ_ONLY_KEYWORDS.has(keyword)) {
    return false;
  }
  if (keyword === "select") {
    return !SQL_SELECT_INTO_PATTERN.test(text);
  }
  if (keyword === "with") {
    return !SQL_SELECT_INTO_PATTERN.test(text) &&
      !SQL_MUTATING_WITH_KEYWORDS.test(text);
  }
  return true;
};

const normalizeMongoCommandText = (statement: string): string => {
  const trimmed = String(statement || "").trim();
  if (!trimmed) return "";
  if (trimmed.startsWith("{")) {
    return trimmed;
  }
  const converted = convertMongoShellToJsonCommand(trimmed);
  if (converted.recognized && converted.command) {
    return converted.command;
  }
  return "";
};

const resolveMongoCommandKey = (command: Record<string, unknown>): string => {
  const keys = Object.keys(command).map((key) => key.trim());
  const effectiveKeys = keys.filter((key) => {
    const normalized = key.toLowerCase();
    return normalized !== "" && !MONGO_META_KEYS.has(normalized);
  });
  for (const key of effectiveKeys) {
    if (MONGO_WRITE_COMMANDS.has(key.toLowerCase())) {
      return key;
    }
  }
  for (const key of effectiveKeys) {
    if (MONGO_READ_ONLY_COMMANDS.has(key.toLowerCase())) {
      return key;
    }
  }
  return effectiveKeys[0] || "";
};

const isReadOnlyMongoStatement = (statement: string): boolean => {
  const commandText = normalizeMongoCommandText(statement);
  if (!commandText) return false;
  try {
    const parsed = JSON.parse(commandText) as Record<string, unknown>;
    if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
      return false;
    }
    const commandKey = resolveMongoCommandKey(parsed).toLowerCase();
    if (!commandKey) return false;
    if (MONGO_WRITE_COMMANDS.has(commandKey)) {
      return false;
    }
    return MONGO_READ_ONLY_COMMANDS.has(commandKey);
  } catch {
    return false;
  }
};

const isConnectionReadOnlyStatement = (
  config: ConnectionReadOnlyLike,
  statement: string,
): boolean => {
  const dialect = resolveConnectionReadOnlyType(config);
  if (dialect === "mongodb") {
    return isReadOnlyMongoStatement(statement);
  }
  return isReadOnlySqlStatement(statement);
};

export const supportsConnectionReadOnlyMode = (
  config: ConnectionReadOnlyLike,
): boolean => {
  return CONNECTION_READ_ONLY_TYPES.has(resolveConnectionReadOnlyType(config));
};

export const normalizeConnectionProtectionConfig = (
  value: unknown,
): ConnectionProtectionConfig => {
  const raw =
    value && typeof value === "object"
      ? (value as Record<string, unknown>)
      : {};
  return {
    restrictDataEdit: raw.restrictDataEdit === true,
    restrictStructureEdit: raw.restrictStructureEdit === true,
    restrictScriptExecution: raw.restrictScriptExecution === true,
    restrictDataImport: raw.restrictDataImport === true,
  };
};

export const createEmptyConnectionProtectionConfig =
  (): ConnectionProtectionConfig => ({ ...EMPTY_CONNECTION_PROTECTION });

export const deriveLegacyConnectionReadOnlyFlag = (
  protection: unknown,
): boolean => {
  const normalized = normalizeConnectionProtectionConfig(protection);
  return CONNECTION_PROTECTION_KEYS.every((key) => normalized[key] === true);
};

export const resolveConnectionProtectionConfig = (
  config: ConnectionReadOnlyLike,
): ConnectionProtectionConfig => {
  if (!supportsConnectionReadOnlyMode(config)) {
    return createEmptyConnectionProtectionConfig();
  }
  const normalized = normalizeConnectionProtectionConfig(config?.protection);
  const hasExplicitRestriction = CONNECTION_PROTECTION_KEYS.some(
    (key) => normalized[key] === true,
  );
  if (hasExplicitRestriction) {
    return normalized;
  }
  if (config?.readOnly === true) {
    return { ...FULL_CONNECTION_PROTECTION };
  }
  return createEmptyConnectionProtectionConfig();
};

export const isConnectionProtectionEnabled = (
  config: ConnectionReadOnlyLike,
  key: ConnectionProtectionKey,
): boolean => {
  return resolveConnectionProtectionConfig(config)[key] === true;
};

export const getConnectionProtectionEnabledCount = (
  config: ConnectionReadOnlyLike,
): number => {
  const protection = resolveConnectionProtectionConfig(config);
  return CONNECTION_PROTECTION_KEYS.filter(
    (key) => protection[key] === true,
  ).length;
};

export const hasAnyConnectionProtection = (
  config: ConnectionReadOnlyLike,
): boolean => {
  return getConnectionProtectionEnabledCount(config) > 0;
};

export const isConnectionDataEditRestricted = (
  config: ConnectionReadOnlyLike,
): boolean => {
  return isConnectionProtectionEnabled(config, "restrictDataEdit");
};

export const isConnectionStructureEditRestricted = (
  config: ConnectionReadOnlyLike,
): boolean => {
  return isConnectionProtectionEnabled(config, "restrictStructureEdit");
};

export const isConnectionScriptExecutionRestricted = (
  config: ConnectionReadOnlyLike,
): boolean => {
  return isConnectionProtectionEnabled(config, "restrictScriptExecution");
};

export const isConnectionDataImportRestricted = (
  config: ConnectionReadOnlyLike,
): boolean => {
  return isConnectionProtectionEnabled(config, "restrictDataImport");
};

export const isConnectionForcedReadOnly = (
  config: ConnectionReadOnlyLike,
): boolean => {
  return deriveLegacyConnectionReadOnlyFlag(
    resolveConnectionProtectionConfig(config),
  );
};

export const findConnectionMutatingStatements = (
  config: ConnectionReadOnlyLike,
  sql: string,
): string[] => {
  if (!isConnectionScriptExecutionRestricted(config)) {
    return [];
  }
  return findSqlStatementRanges(String(sql || ""))
    .map((range) => range.text.trim())
    .filter((statement) => statement.length > 0)
    .filter((statement) => !isConnectionReadOnlyStatement(config, statement));
};