Discourse AI custom tool: BigQuery SQL query with embedded RSA crypto for MiniRacer sandbox

bigquery_sql_tool.js

/* eslint-disable no-undef, no-unused-vars */

// ---- Base64 ----

const B64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

function b64Dec(s) {
  s = s.replace(/[\s=]/g, "");
  const o = [];
  let buf = 0, bits = 0;
  for (let i = 0; i < s.length; i++) {
    buf = (buf << 6) | B64.indexOf(s[i]);
    bits += 6;
    if (bits >= 8) { bits -= 8; o.push((buf >> bits) & 0xff); }
  }
  return o;
}

function b64Enc(b) {
  let s = "";
  for (let i = 0; i < b.length; i += 3) {
    const x = b[i], y = b[i + 1] || 0, z = b[i + 2] || 0;
    s += B64[x >> 2] + B64[((x & 3) << 4) | (y >> 4)];
    if (i + 1 < b.length) s += B64[((y & 15) << 2) | (z >> 6)];
    if (i + 2 < b.length) s += B64[z & 63];
  }
  return s.replace(/\+/g, "-").replace(/\//g, "_");
}

function str2b(s) {
  const b = [];
  for (let i = 0; i < s.length; i++) b.push(s.charCodeAt(i));
  return b;
}

// ---- SHA-256 ----

const SK = [
  0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,
  0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
  0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,
  0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
  0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,
  0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
  0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,
  0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2
];

function rr(x, n) { return ((x >>> n) | (x << (32 - n))) >>> 0; }

function sha256(msg) {
  const p = msg.slice();
  const bl = msg.length * 8;
  p.push(0x80);
  while (p.length % 64 !== 56) p.push(0);
  p.push(0, 0, 0, 0, (bl >>> 24) & 0xff, (bl >>> 16) & 0xff, (bl >>> 8) & 0xff, bl & 0xff);

  let h0 = 0x6a09e667, h1 = 0xbb67ae85, h2 = 0x3c6ef372, h3 = 0xa54ff53a;
  let h4 = 0x510e527f, h5 = 0x9b05688c, h6 = 0x1f83d9ab, h7 = 0x5be0cd19;

  for (let off = 0; off < p.length; off += 64) {
    const w = Array(64);
    for (let i = 0; i < 16; i++)
      w[i] = ((p[off+i*4] << 24) | (p[off+i*4+1] << 16) |
              (p[off+i*4+2] << 8) | p[off+i*4+3]) >>> 0;
    for (let i = 16; i < 64; i++)
      w[i] = (w[i-16] + (rr(w[i-15],7) ^ rr(w[i-15],18) ^ (w[i-15]>>>3)) +
              w[i-7] + (rr(w[i-2],17) ^ rr(w[i-2],19) ^ (w[i-2]>>>10))) >>> 0;

    let a = h0, b = h1, c = h2, d = h3, e = h4, f = h5, g = h6, hh = h7;
    for (let i = 0; i < 64; i++) {
      const t1 = (hh + (rr(e,6) ^ rr(e,11) ^ rr(e,25)) +
                  ((e & f) ^ (~e & g)) + SK[i] + w[i]) >>> 0;
      const t2 = ((rr(a,2) ^ rr(a,13) ^ rr(a,22)) +
                  ((a & b) ^ (a & c) ^ (b & c))) >>> 0;
      hh = g; g = f; f = e; e = (d + t1) >>> 0;
      d = c; c = b; b = a; a = (t1 + t2) >>> 0;
    }
    h0 = (h0+a)>>>0; h1 = (h1+b)>>>0; h2 = (h2+c)>>>0; h3 = (h3+d)>>>0;
    h4 = (h4+e)>>>0; h5 = (h5+f)>>>0; h6 = (h6+g)>>>0; h7 = (h7+hh)>>>0;
  }

  const out = [];
  for (const v of [h0, h1, h2, h3, h4, h5, h6, h7])
    out.push((v >>> 24) & 0xff, (v >>> 16) & 0xff, (v >>> 8) & 0xff, v & 0xff);
  return out;
}

// ---- ASN.1 DER parser (PKCS8 → RSA key) ----

function parsePKCS8(pem) {
  const der = b64Dec(pem.replace(/-----[^-]+-----/g, "").replace(/\s/g, ""));
  let p = 0;

  // helpers — flat, no nested closures (avoids V8 closure/pos interaction bug)
  function readLen() {
    var l = der[p++];
    if (l & 0x80) { var nb = l & 0x7f; l = 0; for (var i = 0; i < nb; i++) l = (l << 8) | der[p++]; }
    return l;
  }
  function readSeq()  { p++; readLen(); }                                          // skip tag + length header
  function readInt()  { p++; var l = readLen(); var v = der.slice(p, p + l); p += l; return v; }
  function readSkip() { p++; var l = readLen(); p += l; }                          // skip tag + length + content

  readSeq();            // PrivateKeyInfo SEQUENCE
  readInt();            // version (0)
  readSkip();           // algorithm (rsaEncryption OID + NULL)
  p++; readLen();       // OCTET STRING wrapper
  readSeq();            // RSAPrivateKey SEQUENCE
  readInt();            // version (0)
  var n = readInt();    // modulus
  readInt();            // publicExponent (skip)
  var d = readInt();    // privateExponent
  return { n: n, d: d };
}

// ---- RSA PKCS1v15-SHA256 sign (uses native BigInt) ----

function b2bi(bytes) {
  if (bytes.length === 0) return 0n;
  let h = "0x";
  for (let i = 0; i < bytes.length; i++) h += bytes[i].toString(16).padStart(2, "0");
  return BigInt(h);
}

function bi2b(n, len) {
  let h = n.toString(16);
  if (h.length & 1) h = "0" + h;
  while (h.length < len * 2) h = "00" + h;
  const o = [];
  for (let i = 0; i < h.length; i += 2) o.push(parseInt(h.substr(i, 2), 16));
  return o;
}

function modpow(base, exp, mod) {
  let r = 1n;
  base = base % mod;
  while (exp > 0n) {
    if (exp & 1n) r = r * base % mod;
    exp >>= 1n;
    base = base * base % mod;
  }
  return r;
}

// DER-encoded DigestInfo prefix for SHA-256
const DI_PREFIX = [
  0x30,0x31,0x30,0x0d,0x06,0x09,0x60,0x86,0x48,0x01,
  0x65,0x03,0x04,0x02,0x01,0x05,0x00,0x04,0x20
];

function rsaSign(hash, key) {
  const nBI = b2bi(key.n), dBI = b2bi(key.d);
  const kLen = Math.ceil(nBI.toString(16).length / 2);
  const di = DI_PREFIX.concat(hash);

  // PKCS1v15: 0x00 0x01 [0xff padding] 0x00 [DigestInfo]
  const em = [0x00, 0x01];
  for (let i = 0; i < kLen - di.length - 3; i++) em.push(0xff);
  em.push(0x00, ...di);

  return bi2b(modpow(b2bi(em), dBI, nBI), kLen);
}

// ---- JWT + OAuth2 ----

function signJWT(cred) {
  const now = Math.floor(Date.now() / 1000);
  const hdr = b64Enc(str2b(JSON.stringify({ alg: "RS256", typ: "JWT" })));
  const pay = b64Enc(str2b(JSON.stringify({
    iss: cred.client_email,
    scope: "https://www.googleapis.com/auth/bigquery.readonly",
    aud: "https://oauth2.googleapis.com/token",
    iat: now,
    exp: now + 3600
  })));
  const msg = hdr + "." + pay;
  const sig = rsaSign(sha256(str2b(msg)), parsePKCS8(cred.private_key));
  return msg + "." + b64Enc(sig);
}

function getToken(cred) {
  const jwt = signJWT(cred);
  const r = http.post("https://oauth2.googleapis.com/token", {
    headers: { "Content-Type": "application/x-www-form-urlencoded" },
    body: "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer&assertion=" + jwt
  });
  if (r.status !== 200) throw new Error("OAuth2 failed (" + r.status + "): " + r.body);
  return JSON.parse(r.body).access_token;
}

// ---- BigQuery ----

function parseBqError(status, body) {
  var result = { error: true, status: status };
  try {
    var parsed = JSON.parse(body);
    var err = parsed.error || {};
    result.message = err.message || body;
    if (err.errors && err.errors.length > 0) {
      result.details = err.errors.map(function(e) {
        var d = { reason: e.reason, message: e.message };
        if (e.location) d.location = e.location;
        return d;
      });
    }
    // Classify for the model
    if (status === 400) {
      result.category = "INVALID_QUERY";
      result.hint = "Check SQL syntax, table/column names, and types. Use bigquery_schema tool to verify the schema.";
    } else if (status === 403) {
      result.category = "ACCESS_DENIED";
      result.hint = "The service account lacks permission for this dataset or table.";
    } else if (status === 404) {
      result.category = "NOT_FOUND";
      result.hint = "Table or dataset does not exist. Use bigquery_list_datasets and bigquery_schema to verify names.";
    } else if (status === 409) {
      result.category = "ALREADY_EXISTS";
      result.hint = "Resource conflict — this shouldn't happen for SELECT queries.";
    } else if (status === 429 || (err.errors && err.errors[0] && err.errors[0].reason === "rateLimitExceeded")) {
      result.category = "RATE_LIMITED";
      result.hint = "Too many requests. Wait a moment and try again.";
    } else {
      result.category = "API_ERROR";
      result.hint = "Unexpected BigQuery error. Check the message for details.";
    }
  } catch (e) {
    result.message = body ? body.substring(0, 500) : "Empty response from BigQuery";
    result.category = "UNKNOWN";
    result.hint = "Could not parse the error response.";
  }
  return result;
}

function bqQuery(token, project, sql) {
  const url = "https://bigquery.googleapis.com/bigquery/v2/projects/" +
              encodeURIComponent(project) + "/queries";
  const r = http.post(url, {
    headers: { Authorization: "Bearer " + token, "Content-Type": "application/json" },
    body: JSON.stringify({
      query: sql, useLegacySql: false, maxResults: 100, timeoutMs: 30000
    })
  });
  if (r.status !== 200) return parseBqError(r.status, r.body);
  return JSON.parse(r.body);
}

// ---- Tool ----

let lastSQL = "";

function invoke(params) {
  var sql = (params.sql || "").trim();
  var clean = sql.replace(/--[^\n]*/g, "").replace(/\/\*[\s\S]*?\*\//g, "").trim();
  if (!sql) return { error: true, category: "INVALID_QUERY", message: "No SQL provided", hint: "Provide a SELECT or WITH...SELECT query in the sql parameter." };
  if (!/^(SELECT|WITH)\b/i.test(clean))
    return { error: true, category: "INVALID_QUERY", message: "Only SELECT (and WITH...SELECT) statements are allowed", hint: "Rewrite as a read-only SELECT query. DML/DDL is not permitted." };
  lastSQL = sql;

  var cred, token, res;
  try {
    cred = JSON.parse(secrets.get("json_cred"));
  } catch (e) {
    return { error: true, category: "CONFIG_ERROR", message: "Failed to parse json_cred secret", hint: "The json_cred secret is missing or contains invalid JSON." };
  }

  try {
    token = getToken(cred);
  } catch (e) {
    return { error: true, category: "AUTH_ERROR", message: "OAuth2 authentication failed: " + e.message, hint: "Check that the service account credentials are valid and have BigQuery access." };
  }

  res = bqQuery(token, cred.project_id, sql);

  // bqQuery returns an error object on failure
  if (res.error === true) {
    res.sql = sql;
    return res;
  }

  if (!res.jobComplete)
    return { error: true, category: "TIMEOUT", message: "Query did not complete within 30 seconds", sql: sql, hint: "Simplify the query, add filters, reduce JOINs, or add a LIMIT clause." };

  var cols = (res.schema && res.schema.fields || []).map(function(f) { return f.name; });
  var rows = (res.rows || []).map(function(row) {
    var o = {};
    row.f.forEach(function(c, i) { o[cols[i]] = c.v; });
    return o;
  });

  var out = { columns: cols, rows: rows, total_rows: res.totalRows || "0" };

  // Bound response size for LLM context
  if (JSON.stringify(out).length > 50000) {
    out.rows = rows.slice(0, 20);
    out.note = "Truncated to 20 rows (response too large for LLM context). Add a LIMIT or select fewer columns.";
    if (JSON.stringify(out).length > 50000) {
      out.rows = rows.slice(0, 5);
      out.note = "Truncated to 5 rows (very large row data). Select fewer columns or use aggregation.";
    }
  }

  return out;
}

function details() {
  return lastSQL ? "<code>" + lastSQL.substring(0, 200) + "</code>" : "";
}