Add syntax highlighting html rewriter

packages/astro-md/middleware.ts

@@ -6,6 +6,61 @@ import { myRemark } from "../my-remark";
 import { remarkObsidian } from "../remark-obsidian";
 import remarkFrontmatter from "remark-frontmatter";
 
+const decode = (str: string) =>
+  str.replace(/&#x(\d+);/g, (_, hex) => String.fromCharCode(parseInt(hex, 16)));
+
+const theme = "nord";
+
+class SyntaxHighlightRewriter implements HTMLRewriterElementContentHandlers {
+  private lang: string = "";
+  private code: string = "";
+  cache: KVNamespace;
+
+  constructor(private runtime: Runtime["runtime"]) {
+    this.cache = runtime.env.KV_HIGHLIGHT;
+  }
+
+  element(element: Element) {
+    this.lang = element.getAttribute("data-lang") ?? "";
+    this.code = "";
+  }
+  async text(text: Text) {
+    this.code += decode(text.text);
+    if (text.lastInTextNode) {
+      const hashBuffer = await crypto.subtle.digest(
+        "SHA-1",
+        new TextEncoder().encode(`${theme}-${this.lang}-${this.code}`)
+      );
+      const hashArray = Array.from(new Uint8Array(hashBuffer));
+      const hash = hashArray
+        .map((b) => b.toString(16).padStart(2, "0"))
+        .join("");
+
+      const cachedResult = await this.cache.get(hash);
+      if (cachedResult) {
+        text.replace(cachedResult, { html: true });
+        return;
+      }
+
+      const res = await fetch(
+        `https://highlight.val.just-be.dev?lang=${this.lang}&theme=${theme}`,
+        {
+          method: "POST",
+          body: this.code,
+          headers: {
+            Accept: "text/html",
+          },
+        }
+      );
+      const data = await res.text();
+      this.runtime.ctx.waitUntil(this.cache.put(hash, data));
+      text.replace(data, { html: true });
+    } else {
+      text.remove();
+    }
+  }
+}
+
 export const onRequest = defineMiddleware(async (context, next) => {
   const fileResolver = async (path: string) => {
     return (await context.locals.runtime.env.KV_MAPPINGS.get(path)) ?? path;
@@ -18,5 +73,15 @@ export const onRequest = defineMiddleware(async (context, next) => {
     ],
   });
   context.locals.render = processor.render;
-  return next();
+  const res = await next();
+  if (import.meta.env.DEV && typeof HTMLRewriter === "undefined") {
+    // @ts-expect-error Isn't defined on globalThis, but that's fine
+    globalThis.HTMLRewriter = (
+      await import("@worker-tools/html-rewriter/base64")
+    ).HTMLRewriter;
+  }
+
+  return new HTMLRewriter()
+    .on("code[data-lang]", new SyntaxHighlightRewriter(context.locals.runtime))
+    .transform(res);
 });

packages/astro-md/package.json

@@ -12,5 +12,8 @@
   "dependencies": {
     "@astrojs/markdown-remark": "catalog:",
     "remark-frontmatter": "^5.0.0"
+  },
+  "devDependencies": {
+    "@worker-tools/html-rewriter": "0.1.0-pre.19"
   }
 }

worker-configuration.d.ts

@@ -1,8 +1,8 @@
-// Generated by Wrangler on Wed Aug 21 2024 20:59:22 GMT-0400 (Eastern Daylight Time)
-// by running `wrangler types`
+// Generated by Wrangler by running `wrangler types`
 
 interface Env {
 	KV_MAPPINGS: KVNamespace;
+	KV_HIGHLIGHT: KVNamespace;
 	NODE_VERSION: "v20.16.0";
 	PNPM_VERSION: "v9.7.0";
 	SITE: string;

wrangler.toml

@@ -11,6 +11,10 @@ PNPM_VERSION = "v9.7.0"
 binding = "KV_MAPPINGS"
 id = "36e4319084c4498fa08e5bf00f36331a"
 
+[[kv_namespaces]]
+binding = "KV_HIGHLIGHT"
+id = "9503a7c1e6f948cb8f1754c08e046179"
+
 [[r2_buckets]]
 binding = "R2_BUCKET"
 bucket_name = "just-be-dev"