From e8ef435efffd7f3415e05a3f6a73284b489e1b0c Mon Sep 17 00:00:00 2001
From: Samuel Colvin <s@muelcolvin.com>
Date: Thu, 16 Jan 2025 12:24:05 +0000
Subject: [PATCH] add traces proxy including CORS

---
 src/index.ts  | 131 ++++++++++++++++++++++++++++++++------------------
 wrangler.toml |   4 ++
 2 files changed, 89 insertions(+), 46 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 5a6ebd4..63bfba3 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -9,58 +9,97 @@ export default {
     // console.log(Object.fromEntries(request.headers))
     const { pathname } = new URL(request.url)
     if (request.method === 'GET' && pathname === '/') {
-      return new Response(`See https://github.com/pydantic/logfire-logs-proxy for details (commit ${env.GITHUB_SHA}).`)
-    } else if (pathname !== '/v1/logs' || request.method !== 'POST') {
-      return new Response('Only POST requests to `/v1/logs` are supported', { status: 404 })
+      return new Response(index_html(env), { headers: { 'Content-Type': 'text/html' } })
+    } else if (pathname === '/v1/logs' && request.method === 'POST') {
+      return await logProxy(request)
+    } else if (pathname === '/v1/traces' && request.method === 'POST') {
+      return await traceProxy(request)
+    } else if (pathname === '/v1/traces' && request.method === 'OPTIONS') {
+      return tracePreflight()
+    } else {
+      return new Response(`404: '${request.method} ${pathname}' not found`, { status: 404 })
     }
+  },
+} satisfies ExportedHandler<Env>
 
-    const auth = request.headers.get('Authorization')
-    if (!auth) {
-      return new Response('No "Authorization" header', { status: 401 })
-    }
+const index_html = (env: Env) => `
+<h1>logfire-logs-proxy</h1>
+<p>
+  See <a href="https://github.com/pydantic/logfire-logs-proxy">github.com/pydantic/logfire-logs-proxy</a>
+  for details (commit <code>${env.GITHUB_SHA}</code>).
+</p>
+`
 
-    let body: ArrayBuffer
-    try {
-      body = await getBody(request)
-    } catch (e) {
-      console.log('Error parsing request body:', e)
-      return new Response(`Error collecting request body: ${e}`, { status: 400 })
-    }
+async function logProxy(request: Request): Promise<Response> {
+  const auth = request.headers.get('Authorization')
+  if (!auth) {
+    return new Response('No "Authorization" header', { status: 401 })
+  }
 
-    let logRequest
-    try {
-      logRequest = decodeLogs(body)
-    } catch (e) {
-      console.log('Error parsing protobuf:', e)
-      return new Response(`Error parsing protobuf: ${e}`, { status: 400 })
-    }
+  let body: ArrayBuffer
+  try {
+    body = await getBody(request)
+  } catch (e) {
+    console.log('Error parsing request body:', e)
+    return new Response(`Error collecting request body: ${e}`, { status: 400 })
+  }
 
-    const traceRequest = convert(logRequest)
-    if (!traceRequest || !traceRequest.resourceSpans) {
-      return new Response('no data to proxy', { status: 202 })
-    }
+  let logRequest
+  try {
+    logRequest = decodeLogs(body)
+  } catch (e) {
+    console.log('Error parsing protobuf:', e)
+    return new Response(`Error parsing protobuf: ${e}`, { status: 400 })
+  }
 
-    console.log('Sending trace to logfire')
-    // console.log('Sending trace to logfire', JSON.stringify(traceRequest.resourceSpans, null, 2))
-    const response = await fetch('https://logfire-api.pydantic.dev/v1/traces', {
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/x-protobuf',
-        Authorization: auth,
-        'User-Agent': `logfire-logs-proxy ${request.headers.get('User-Agent')}`,
-      },
-      body: encodeTraces(traceRequest),
-    })
-    if (response.ok) {
-      console.log('Successfully sent trace to logfire')
-      return response
-    } else {
-      const text = await response.text()
-      console.warn('Unexpected response:', { status: response.status, text })
-      return new Response(text, response)
-    }
-  },
-} satisfies ExportedHandler<Env>
+  const traceRequest = convert(logRequest)
+  if (!traceRequest || !traceRequest.resourceSpans) {
+    return new Response('no data to proxy', { status: 202 })
+  }
+
+  console.log('Sending trace to logfire')
+  // console.log('Sending trace to logfire', JSON.stringify(traceRequest.resourceSpans, null, 2))
+  const response = await fetch('https://logfire-api.pydantic.dev/v1/traces', {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/x-protobuf',
+      Authorization: auth,
+      'User-Agent': `logfire-logs-proxy ${request.headers.get('User-Agent')}`,
+    },
+    body: encodeTraces(traceRequest),
+  })
+  if (response.ok) {
+    console.log('Successfully sent trace to logfire')
+    return response
+  } else {
+    const text = await response.text()
+    console.warn('Unexpected response:', { status: response.status, text })
+    return new Response(text, response)
+  }
+}
+
+async function traceProxy(request: Request): Promise<Response> {
+  const response = await fetch('https://logfire-api.pydantic.dev/v1/traces', request)
+  // add CORS headers
+  return new Response(response.body, {
+    status: response.status,
+    statusText: response.statusText,
+    headers: {
+      ...Object.fromEntries(response.headers),
+      'Access-Control-Allow-Origin': '*',
+    },
+  })
+}
+
+const tracePreflight = () =>
+  new Response(null, {
+    status: 204,
+    headers: {
+      'Access-Control-Allow-Origin': '*',
+      'Access-Control-Allow-Methods': 'POST',
+      'Access-Control-Allow-Headers': 'Authorization, Content-Type',
+    },
+  })
 
 async function getBody(request: Request): Promise<ArrayBuffer> {
   if (request.body === null) {
diff --git a/wrangler.toml b/wrangler.toml
index edd23a4..e9d4a65 100644
--- a/wrangler.toml
+++ b/wrangler.toml
@@ -8,3 +8,7 @@ enabled = true
 
 [vars]
 GITHUB_SHA = "unknown"
+
+[env.staging]
+name = "logfire-logs-proxy-staging"
+vars = { GITHUB_SHA = "staging" }