Merge pull request #2083 from quadratichq/ayush/1899

feat: request proxying for JS & fix auth header override for both JS and Python requests

quadratic-client/src/app/web-workers/javascriptWebWorker/javascriptClientMessages.ts

@@ -22,15 +22,31 @@ export interface JavascriptClientState {
   version?: string;
 }
 
+export interface JavascriptClientInit {
+  type: 'javascriptClientInit';
+  version: string;
+}
+
 export interface ClientJavascriptCoreChannel {
   type: 'clientJavascriptCoreChannel';
+  env: ImportMetaEnv;
 }
 
-export interface JavascriptClientInit {
-  type: 'javascriptClientInit';
-  version: string;
+export interface ClientJavascriptGetJwt {
+  type: 'clientJavascriptGetJwt';
+  id: number;
+  jwt: string;
+}
+
+export interface JavascriptClientGetJwt {
+  type: 'javascriptClientGetJwt';
+  id: number;
 }
 
-export type JavascriptClientMessage = JavascriptClientLoadError | JavascriptClientState | JavascriptClientInit;
+export type JavascriptClientMessage =
+  | JavascriptClientLoadError
+  | JavascriptClientState
+  | JavascriptClientInit
+  | JavascriptClientGetJwt;
 
-export type ClientJavascriptMessage = ClientJavascriptCoreChannel;
+export type ClientJavascriptMessage = ClientJavascriptCoreChannel | ClientJavascriptGetJwt;

quadratic-client/src/app/web-workers/javascriptWebWorker/javascriptWebWorker.ts

@@ -1,8 +1,9 @@
 import { events } from '@/app/events/events';
+import { authClient } from '@/auth/auth';
 import mixpanel from 'mixpanel-browser';
 import { LanguageState } from '../languageTypes';
 import { quadraticCore } from '../quadraticCore/quadraticCore';
-import { ClientJavascriptMessage, JavascriptClientMessage } from './javascriptClientMessages';
+import { ClientJavascriptMessage, JavascriptClientGetJwt, JavascriptClientMessage } from './javascriptClientMessages';
 
 class JavascriptWebWorker {
   state: LanguageState = 'loading';
@@ -30,6 +31,14 @@ class JavascriptWebWorker {
         events.emit('javascriptState', message.data.state, message.data.current, message.data.awaitingExecution);
         break;
 
+      case 'javascriptClientGetJwt':
+        authClient.getTokenOrRedirect().then((jwt) => {
+          const data = message.data as JavascriptClientGetJwt;
+          this.send({ type: 'clientJavascriptGetJwt', id: data.id, jwt });
+        });
+
+        break;
+
       default:
         throw new Error(`Unhandled message type ${message.type}`);
     }
@@ -40,7 +49,7 @@ class JavascriptWebWorker {
     this.worker.onmessage = this.handleMessage;
 
     const JavascriptCoreChannel = new MessageChannel();
-    this.send({ type: 'clientJavascriptCoreChannel' }, JavascriptCoreChannel.port1);
+    this.send({ type: 'clientJavascriptCoreChannel', env: import.meta.env }, JavascriptCoreChannel.port1);
     quadraticCore.sendJavascriptInit(JavascriptCoreChannel.port2);
   }
 

quadratic-client/src/app/web-workers/javascriptWebWorker/worker/javascript/getJavascriptFetchOverride.ts

@@ -0,0 +1,28 @@
+export function getJavascriptFetchOverride(proxyUrl: string, jwt: string) {
+  return `
+self['fetch'] = new Proxy(fetch, {
+  apply: function (target, thisArg, args) {
+    const [url, config] = args;
+
+    const newConfig = config || {};
+    const headers = newConfig.headers || {};
+    const newHeaders = {};
+
+    // Prefix all original request headers with X-Proxy-
+    for (const key in headers) {
+      newHeaders[\`X-Proxy-\${key}\`] = headers[key];
+    }
+
+    // Set the original request URL on X-Proxy-Url header
+    newHeaders['X-Proxy-Url'] = url.toString();
+
+    // Set the authorization header for the proxy server
+    newHeaders['Authorization'] = 'Bearer ${jwt}';
+
+    newConfig.headers = newHeaders;
+
+    return target.call(thisArg, '${proxyUrl}', newConfig);
+  },
+});\n
+`;
+}

quadratic-client/src/app/web-workers/javascriptWebWorker/worker/javascript/getJavascriptXHROverride.ts

@@ -0,0 +1,50 @@
+export function getJavascriptXHROverride(proxyUrl: string, jwt: string) {
+  return `
+self['XMLHttpRequest'] = new Proxy(XMLHttpRequest, {
+  construct: function (target, args) {
+    const xhr = new target();
+
+    xhr.open = new Proxy(xhr.open, {
+      apply: function (target, thisArg, args) {
+        Object.defineProperty(xhr, '__url', { value: args[1].toString(), writable: true });
+        args[1] = '${proxyUrl}';
+        return target.apply(thisArg, args);
+      },
+    });
+
+    xhr.setRequestHeader = new Proxy(xhr.setRequestHeader, {
+      apply: function (target, thisArg, args) {
+        // apply quadratic-authorization header as the only authorization header
+        // this is required for authentication with the proxy server
+        if (args[0] === 'Quadratic-Authorization') {
+          args[0] = 'Authorization';
+        } else {
+          // apply all headers on the original request prefixed with X-Proxy-
+          args[0] = \`X-Proxy-\${args[0]}\`;
+        }
+        return target.apply(thisArg, args);
+      },
+    });
+
+    xhr.onreadystatechange = function () {
+      if (xhr.readyState === XMLHttpRequest.OPENED) {
+        // this applies the quadratic-authorization header as the only authorization header
+        // this is required for authentication with the proxy server
+        xhr.setRequestHeader('Quadratic-Authorization', 'Bearer ${jwt}');
+
+        // this applies the original request URL as the x-proxy-url header
+        // this will get prefixed with X-Proxy due to above setRequestHeader override
+        xhr.setRequestHeader('Url', xhr.__url);
+      }
+      // After completion of XHR request
+      if (xhr.readyState === 4) {
+        if (xhr.status === 401) {
+        }
+      }
+    };
+
+    return xhr;
+  },
+});\n
+`;
+}

quadratic-client/src/app/web-workers/javascriptWebWorker/worker/javascript/javascript.ts

@@ -34,7 +34,7 @@ export class Javascript {
     this.api = new JavascriptAPI(this);
   }
 
-  init = async () => {
+  private init = async () => {
     await esbuild.initialize({
       wasmURL: '/esbuild.wasm',
       // this would create another worker to run the actual code. I don't
@@ -97,7 +97,9 @@ export class Javascript {
     this.row = message.y;
 
     try {
-      const code = prepareJavascriptCode(transformedCode, message.x, message.y, this.withLineNumbers);
+      const proxyUrl = `${javascriptClient.env.VITE_QUADRATIC_CONNECTION_URL}/proxy`;
+      const jwt = await javascriptClient.getJwt();
+      const code = prepareJavascriptCode(transformedCode, message.x, message.y, this.withLineNumbers, proxyUrl, jwt);
       const runner = new Worker(URL.createObjectURL(new Blob([code], { type: 'application/javascript' })), {
         type: 'module',
         name: 'javascriptWorker',

quadratic-client/src/app/web-workers/javascriptWebWorker/worker/javascript/javascriptCompile.ts

@@ -5,10 +5,11 @@
 //! numbers to all return statements via a caught thrown error (the only way to
 //! get line numbers in JS).
 
+import { getJavascriptFetchOverride } from '@/app/web-workers/javascriptWebWorker/worker/javascript/getJavascriptFetchOverride';
+import { getJavascriptXHROverride } from '@/app/web-workers/javascriptWebWorker/worker/javascript/getJavascriptXHROverride';
 import * as esbuild from 'esbuild-wasm';
 import { LINE_NUMBER_VAR } from './javascript';
 import { javascriptLibrary } from './runner/generateJavascriptForRunner';
-
 export interface JavascriptTransformedCode {
   imports: string;
   code: string;
@@ -66,10 +67,16 @@ export function prepareJavascriptCode(
   transform: JavascriptTransformedCode,
   x: number,
   y: number,
-  withLineNumbers: boolean
+  withLineNumbers: boolean,
+  proxyUrl: string,
+  jwt: string
 ): string {
   const code = withLineNumbers ? javascriptAddLineNumberVars(transform) : transform.code;
+  const javascriptXHROverride = getJavascriptXHROverride(proxyUrl, jwt);
+  const javascriptFetchOverride = getJavascriptFetchOverride(proxyUrl, jwt);
   const compiledCode =
+    javascriptXHROverride +
+    javascriptFetchOverride +
     transform.imports +
     (withLineNumbers ? `let ${LINE_NUMBER_VAR} = 0;` : '') +
     javascriptLibrary.replace('{x:0,y:0}', `{x:${x},y:${y}}`) + // replace the pos() with the correct x,y coordinates

quadratic-client/src/app/web-workers/javascriptWebWorker/worker/javascriptClient.ts

@@ -1,12 +1,20 @@
 import { debugWebWorkers, debugWebWorkersMessages } from '@/app/debugFlags';
 import { LanguageState } from '@/app/web-workers/languageTypes';
 import { CodeRun } from '../../CodeRun';
-import type { ClientJavascriptMessage, JavascriptClientMessage } from '../javascriptClientMessages';
+import type {
+  ClientJavascriptGetJwt,
+  ClientJavascriptMessage,
+  JavascriptClientMessage,
+} from '../javascriptClientMessages';
 import { javascriptCore } from './javascriptCore';
 
 declare var self: WorkerGlobalScope & typeof globalThis;
 
 class JavascriptClient {
+  private id = 0;
+  private waitingForResponse: Record<number, Function> = {};
+  env: Record<string, string> = {};
+
   start() {
     self.onmessage = this.handleMessage;
     if (debugWebWorkers) console.log('[javascriptClient] initialized.');
@@ -25,11 +33,21 @@ class JavascriptClient {
 
     switch (e.data.type) {
       case 'clientJavascriptCoreChannel':
+        this.env = e.data.env;
         javascriptCore.init(e.ports[0]);
         break;
 
       default:
-        console.warn('[coreClient] Unhandled message type', e.data);
+        if (e.data.id !== undefined) {
+          if (this.waitingForResponse[e.data.id]) {
+            this.waitingForResponse[e.data.id](e.data);
+            delete this.waitingForResponse[e.data.id];
+          } else {
+            console.warn('No resolve for message in javascriptClient', e.data.type);
+          }
+        } else {
+          console.warn('[javascriptClient] Unhandled message type', e.data);
+        }
     }
   };
 
@@ -50,6 +68,14 @@ class JavascriptClient {
   sendInit(version: string) {
     this.send({ type: 'javascriptClientInit', version });
   }
+
+  getJwt(): Promise<string> {
+    return new Promise((resolve) => {
+      const id = this.id++;
+      this.waitingForResponse[id] = (message: ClientJavascriptGetJwt) => resolve(message.jwt);
+      this.send({ type: 'javascriptClientGetJwt', id });
+    });
+  }
 }
 
 export const javascriptClient = new JavascriptClient();

quadratic-client/src/app/web-workers/pythonWebWorker/worker/python.ts

@@ -51,7 +51,7 @@ class Python {
     }
   };
 
-  init = async () => {
+  private init = async () => {
     const jwt = await pythonClient.getJwt();
 
     // patch XMLHttpRequest to send requests to the proxy
@@ -77,12 +77,31 @@ class Python {
           },
         });
 
+        xhr.setRequestHeader = new Proxy(xhr.setRequestHeader, {
+          apply: function (target, thisArg, args: [string, string]) {
+            // apply quadratic-authorization header as the only authorization header
+            // this is required for authentication with the proxy server
+            if (args[0] === 'Quadratic-Authorization') {
+              args[0] = 'Authorization';
+            } else {
+              // apply all headers on the original request prefixed with X-Proxy
+              args[0] = `X-Proxy-${args[0]}`;
+            }
+            return target.apply(thisArg, args);
+          },
+        });
+
         xhr.onreadystatechange = function () {
           if (xhr.readyState === XMLHttpRequest.OPENED) {
-            xhr.setRequestHeader('Proxy', (xhr as any).__url);
-            xhr.setRequestHeader('Authorization', `Bearer ${jwt}`);
+            // this applies the quadratic-authorization header as the only authorization header
+            // this is required for authentication with the proxy server
+            xhr.setRequestHeader('Quadratic-Authorization', `Bearer ${jwt}`);
+
+            // this applies the original request URL as the x-proxy-url header
+            // this will get prefixed with X-Proxy due to above setRequestHeader override
+            xhr.setRequestHeader('Url', (xhr as any).__url);
           }
-          // After complition of XHR request
+          // After completion of XHR request
           if (xhr.readyState === 4) {
             if (xhr.status === 401) {
             }

quadratic-client/src/app/web-workers/pythonWebWorker/worker/pythonClient.ts

@@ -47,7 +47,7 @@ class PythonClient {
             this.waitingForResponse[e.data.id](e.data);
             delete this.waitingForResponse[e.data.id];
           } else {
-            console.warn('No resolve for message in pythonClient', e.data.id);
+            console.warn('No resolve for message in pythonClient', e.data.type);
           }
         } else {
           console.warn('[pythonClient] Unhandled message type', e.data);

quadratic-connection/src/proxy.rs

@@ -13,7 +13,9 @@ use crate::error::{proxy_error, ConnectionError, Result};
 use crate::state::State;
 
 const REQUEST_TIMEOUT_SEC: u64 = 15;
-const PROXY_HEADER: &str = "proxy";
+const AUTHORIZATION_HEADER: &str = "authorization";
+const PROXY_URL_HEADER: &str = "x-proxy-url";
+const PROXY_HEADER_PREFIX: &str = "x-proxy-";
 
 pub(crate) async fn axum_to_reqwest(
     url: &str,
@@ -24,12 +26,20 @@ pub(crate) async fn axum_to_reqwest(
     let method = Method::from_bytes(method_bytes).map_err(proxy_error)?;
 
     let mut headers = reqwest::header::HeaderMap::with_capacity(req.headers().len());
-    let headers_to_ignore = ["host", PROXY_HEADER, "authorization"];
+    let headers_to_ignore = ["host", AUTHORIZATION_HEADER, PROXY_URL_HEADER];
 
     for (name, value) in req
         .headers()
         .into_iter()
         .filter(|(name, _)| !headers_to_ignore.contains(&name.as_str()))
+        .map(|(name, value)| {
+            let name = name.as_str();
+            if let Some(name) = name.strip_prefix(PROXY_HEADER_PREFIX) {
+                (name, value)
+            } else {
+                (name, value)
+            }
+        })
     {
         let name = reqwest::header::HeaderName::from_bytes(name.as_ref()).map_err(proxy_error)?;
         let value =
@@ -74,7 +84,7 @@ pub(crate) async fn proxy(
 
     let headers = req.headers().clone();
     let url = headers
-        .get(PROXY_HEADER)
+        .get(PROXY_URL_HEADER)
         .ok_or_else(|| ConnectionError::Proxy("No proxy header found".to_string()))?
         .to_str()
         .map_err(proxy_error)?;
@@ -103,7 +113,7 @@ mod tests {
         let mut request = Request::new(Body::empty());
         request
             .headers_mut()
-            .insert(PROXY_HEADER, HeaderValue::from_static(URL));
+            .insert(PROXY_URL_HEADER, HeaderValue::from_static(URL));
         let data = proxy(state, request).await.unwrap();
         let response = data.into_response();
 
@@ -122,7 +132,7 @@ mod tests {
             .insert(ACCEPT, HeaderValue::from_static(accept));
         request
             .headers_mut()
-            .insert(PROXY_HEADER, HeaderValue::from_static(URL));
+            .insert(PROXY_URL_HEADER, HeaderValue::from_static(URL));
 
         let result = axum_to_reqwest(URL, request, state.client.clone())
             .await
@@ -133,7 +143,7 @@ mod tests {
         assert_eq!(result.method(), Method::POST);
         assert_eq!(result.url().to_string(), URL);
 
-        // PROXY_HEADER doesn't get copied over
+        // PROXY_URL_HEADER doesn't get copied over
         assert_eq!(result.headers().len(), 1);
         assert_eq!(
             result.headers().get(reqwest::header::ACCEPT).unwrap(),