Skip to content

Use HTTP/2 to make requests from proxies #1375

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: canary
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
150 changes: 87 additions & 63 deletions typescript/fiddle-proxy/server.js
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
const cors = require('cors')
const { createProxyMiddleware } = require('http-proxy-middleware')
const assert = require('assert')
const app = require('express')()
require('dotenv').config()
const cors = require('cors');
const express = require('express');
const http2 = require('http2');
const { URL } = require('url');
require('dotenv').config();

app.use(cors())
const app = express();
app.use(cors());

// From https://nodejs.org/api/url.html#url-strings-and-url-objects:
// ┌────────────────────────────────────────────────────────────────────────────────────────────────┐
Expand Down Expand Up @@ -37,9 +39,9 @@ app.use(cors())
const API_KEY_INJECTION_ALLOWED = {
'https://api.openai.com': { Authorization: `Bearer ${process.env.OPENAI_API_KEY}` },
'https://api.anthropic.com': { 'x-api-key': process.env.ANTHROPIC_API_KEY },
'https://generativelanguage.googleapis.com': { 'x-goog-api-key': process.env.GOOGLE_API_KEY },
'https://generativelanguage.googleapis.com': { Authorization: `Bearer ${process.env.GOOGLE_API_KEY}` },
'https://openrouter.ai': { Authorization: `Bearer ${process.env.OPENROUTER_API_KEY}` },
}
};

// Consult sam@ before changing this.
for (const url of Object.keys(API_KEY_INJECTION_ALLOWED)) {
Expand All @@ -49,62 +51,84 @@ for (const url of Object.keys(API_KEY_INJECTION_ALLOWED)) {
)
}

app.use(
createProxyMiddleware({
changeOrigin: true,
pathRewrite: (path, req) => {
// Ensure the URL does not end with a slash
if (path.endsWith('/')) {
return path.slice(0, -1)
// Middleware to handle proxy requests.
app.use(async (req, res) => {
const originalUrl = req.headers['baml-original-url'];
if (!originalUrl) {
res.status(400).send('Missing baml-original-url header');
return;
}

try {
// Parse the original URL and append the request path.
const targetUrl = new URL(originalUrl);

const removeTrailingSlash = req.path.endsWith('/')
? req.path.slice(0, -1) // Remove trailing slash
: req.path;

targetUrl.pathname = `${targetUrl.pathname}${removeTrailingSlash}`;

const proxyReqHeaders = { ...req.headers }; // Clone incoming headers
delete proxyReqHeaders.host; // Remove host header for upstream requests
delete proxyReqHeaders.origin; // Remove origin header for upstream requests

// It is very important that we ONLY resolve against API_KEY_INJECTION_ALLOWED
// by using the URL origin! (i.e. NOT using str.startsWith - the latter can still
// leak API keys to malicious subdomains e.g. https://api.openai.com.evil.com)
const allowedHeaders = API_KEY_INJECTION_ALLOWED[targetUrl.origin];

if (allowedHeaders) {
// Override headers.
for (const [header, value] of Object.entries(allowedHeaders)) {
proxyReqHeaders[header.toLowerCase()] = value;
}
return path
},
router: (req) => {
// Extract the original target URL from the custom header
const originalUrl = req.headers['baml-original-url']

if (typeof originalUrl === 'string') {
return originalUrl
} else {
throw new Error('baml-original-url header is missing or invalid')
}

// Establish HTTP/2 connection
const client = http2.connect(targetUrl.origin);

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider adding error handling for http2.connect to manage cases where the connection cannot be established.

const proxyReq = client.request({
':method': req.method,
':path': `${targetUrl.pathname}${targetUrl.search}`,
...proxyReqHeaders,
});

// Pipe the request body to the upstream server.
req.pipe(proxyReq);

// Handle the response from the upstream server.
proxyReq.on('response', (headers) => {
// Set response headers
for (const [key, value] of Object.entries(headers)) {
if (key.startsWith(':')) continue; // Skip pseudo-headers
res.setHeader(key, value);
}
},
logger: console,
on: {
proxyReq: (proxyReq, req, res) => {
try {
const bamlOriginalUrl = req.headers['baml-original-url']
if (bamlOriginalUrl === undefined) {
return
}
const proxyOrigin = new URL(bamlOriginalUrl).origin
// It is very important that we ONLY resolve against API_KEY_INJECTION_ALLOWED
// by using the URL origin! (i.e. NOT using str.startsWith - the latter can still
// leak API keys to malicious subdomains e.g. https://api.openai.com.evil.com)
const headers = API_KEY_INJECTION_ALLOWED[proxyOrigin]
if (headers === undefined) {
return
}
for (const [header, value] of Object.entries(headers)) {
proxyReq.setHeader(header, value)
}
proxyReq.removeHeader('origin')
} catch (err) {
// This is not console.warn because it's not important
console.log('baml-original-url is not parsable', err)
}
},
proxyRes: (proxyRes, req, res) => {
proxyRes.headers['Access-Control-Allow-Origin'] = '*'
},
error: (error) => {
console.error('proxy error:', error)
},
},
}),
)

// Start web server on port 3000
res.setHeader('Access-Control-Allow-Origin', '*');
res.statusCode = headers[':status'];
});

proxyReq.on('data', (chunk) => {
res.write(chunk); // Forward the data to the client
});

proxyReq.on('end', () => {
res.end(); // End the response
client.close(); // Close the HTTP/2 connection
});

proxyReq.on('error', (err) => {
console.error('Proxy request error:', err);
res.status(500).send('Internal Server Error');
client.close();
});
} catch (err) {
console.error('Proxy error:', err);
res.status(500).send('Failed to process request');
}
});

// Start the server
app.listen(3000, () => {
console.log('Server is listening on port 3000')
})
console.log('Server is listening on port 3000');
});
141 changes: 74 additions & 67 deletions typescript/vscode-ext/packages/vscode/src/extension.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import { telemetry } from './plugins/language-server'
import cors from 'cors'
import { createProxyMiddleware } from 'http-proxy-middleware'
import { type LanguageClient, type ServerOptions, TransportKind } from 'vscode-languageclient/node'
import http2 from "node:http2";

let client: LanguageClient

Expand Down Expand Up @@ -191,77 +192,83 @@ export function activate(context: vscode.ExtensionContext) {
return addr.port
}

app.use(
createProxyMiddleware({
changeOrigin: true,
pathRewrite: (path, req) => {
console.log('reqmethod', req.method)
app.use(async (req, res) => {
const originalUrl = req.headers['baml-original-url'];
if (typeof originalUrl !== 'string') {
console.log('baml-original-url header is missing or invalid')
throw new Error('baml-original-url header is missing or invalid')
}

// Remove the path in the case of images. Since we request things differently for image GET requests, where we add the url to localhost:4500/actual-url.png
// to prevent caching issues with Rust reqwest.
// But for normal completion POST requests, we always call localhost:4500.
// The original url is always in baml-original-url header.
try {
// Parse the original URL and append the request path.
const targetUrl = new URL(originalUrl);

// Check for file extensions and set path to empty string.
if (/\.[a-zA-Z0-9]+$/.test(path) && req.method === 'GET') {
return ''
}
// Remove trailing slash
if (path.endsWith('/')) {
return path.slice(0, -1)
}
return path
},
router: (req) => {
// Extract the original target URL from the custom header
let originalUrl = req.headers['baml-original-url']
if (typeof originalUrl === 'string') {
// For some reason, Node doesn't like deleting headers in the proxyReq function.
delete req.headers['baml-original-url']
delete req.headers['origin']

// Ensure the URL does not end with a slash
console.log('originalUrl1', originalUrl)
if (originalUrl.endsWith('/')) {
originalUrl = originalUrl.slice(0, -1)
}
console.log('returning original url', originalUrl)
// return new URL(originalUrl).toString()
let pathRewrite = req.path;

// Remove the path in the case of images. Since we request things differently for image GET requests, where we add the url to localhost:4500/actual-url.png
// to prevent caching issues with Rust reqwest.
// But for normal completion POST requests, we always call localhost:4500.
// The original url is always in baml-original-url header.

// Check for file extensions and set path to empty string.
if (/\.[a-zA-Z0-9]+$/.test(req.path) && req.method === 'GET') {
pathRewrite = '';
}

// Remove trailing slash
if (req.path.endsWith('/')) {
pathRewrite = req.path.slice(0, -1)
}

targetUrl.pathname = `${targetUrl.pathname}${pathRewrite}`;

const proxyReqHeaders = { ...req.headers }; // Clone incoming headers
delete proxyReqHeaders.host; // Remove host header for upstream requests
delete proxyReqHeaders.origin; // Remove origin header for upstream requests
delete req.headers['baml-original-url']; // Remove the custom header

// Establish HTTP/2 connection
const client = http2.connect(targetUrl.origin);

return originalUrl
} else {
console.log('baml-original-url header is missing or invalid')
throw new Error('baml-original-url header is missing or invalid')
const proxyReq = client.request({
':method': req.method,
':path': `${targetUrl.pathname}${targetUrl.search}`,
...proxyReqHeaders,
});

// Pipe the request body to the upstream server.
req.pipe(proxyReq);

// Handle the response from the upstream server.
proxyReq.on('response', (headers) => {
// Set response headers
for (const [key, value] of Object.entries(headers)) {
if (key.startsWith(':')) continue; // Skip pseudo-headers
res.setHeader(key, value as any);
}
},
logger: console,
on: {
proxyReq: (proxyReq, req, res) => {
// const bamlOriginalUrl = req.headers['baml-original-url']
// if (typeof bamlOriginalUrl === 'string') {
// const targetUrl = new URL(bamlOriginalUrl)
// // Copy all original headers except those we want to modify/remove
// Object.entries(req.headers).forEach(([key, value]) => {
// if (key !== 'host' && key !== 'origin' && key !== 'baml-original-url') {
// proxyReq.setHeader(key, value)
// }
// })
// // Set the correct origin and host headers
// proxyReq.setHeader('origin', targetUrl.origin)
// proxyReq.setHeader('host', targetUrl.host)
// }
},
proxyRes: (proxyRes, req, res) => {
proxyRes.headers['Access-Control-Allow-Origin'] = '*'
},
error: (error, req, res) => {
console.error('proxy error:', error)

res.end(JSON.stringify({ error: error }))
},
},
}),
)
res.setHeader('Access-Control-Allow-Origin', '*');
res.statusCode = headers[':status'] as number;
});

proxyReq.on('data', (chunk) => {
res.write(chunk); // Forward the data to the client
});

proxyReq.on('end', () => {
res.end(); // End the response
client.close(); // Close the HTTP/2 connection
});

proxyReq.on('error', (err) => {
console.error('Proxy request error:', err);
res.status(500).send('Internal Server Error');
client.close();
});
} catch (err) {
console.error('Proxy error:', err);
res.status(500).send('Failed to process request');
}
});

const bamlPlaygroundCommand = vscode.commands.registerCommand(
'baml.openBamlPanel',
Expand Down
Loading