Handle gRPC compressed payloads + headers are sent to formatters + im…

…prove Protobuf/gRPC tests coverage
httptoolkit · Sep 26, 2024 · 0849ad0 · 0849ad0
1 parent 93372c5
commit 0849ad0
Show file tree

Hide file tree

Showing 11 changed files with 299 additions and 93 deletions.
diff --git a/src/components/editor/content-viewer.tsx b/src/components/editor/content-viewer.tsx
@@ -5,6 +5,7 @@ import { observer } from 'mobx-react';
 import { SchemaObject } from 'openapi3-ts';
 import * as portals from 'react-reverse-portal';
 
+import { Headers } from '../../types';
 import { styled } from '../../styles';
 import { ObservablePromise, isObservablePromise } from '../../util/observable';
 import { asError, unreachableCheck } from '../../util/error';
@@ -22,7 +23,7 @@ interface ContentViewerProps {
     children: Buffer | string;
     schema?: SchemaObject;
     expanded: boolean;
-    rawContentType?: string;
+    headers?: Headers;
     contentType: ViewableContentType;
     editorNode: portals.HtmlPortalNode<typeof SelfSizedEditor | typeof ContainerSizedEditor>;
     cache: Map<Symbol, unknown>;
@@ -199,7 +200,7 @@ export class ContentViewer extends React.Component<ContentViewerProps> {
             return <FormatterContainer expanded={this.props.expanded}>
                 <formatterConfig.Component
                     content={this.contentBuffer}
-                    rawContentType={this.props.rawContentType}
+                    headers={this.props.headers}
                 />
             </FormatterContainer>;
         }

diff --git a/src/components/send/sent-response-body.tsx b/src/components/send/sent-response-body.tsx
@@ -121,7 +121,7 @@ export class SentResponseBodyCard extends React.Component<ExpandableCardProps &
                     <ContentViewer
                         contentId={message.id}
                         editorNode={this.props.editorNode}
-                        rawContentType={lastHeader(message.headers['content-type'])}
+                        headers={message.headers}
                         contentType={decodedContentType}
                         expanded={!!expanded}
                         cache={message.cache}

diff --git a/src/components/view/http/http-body-card.tsx b/src/components/view/http/http-body-card.tsx
@@ -121,7 +121,7 @@ export class HttpBodyCard extends React.Component<ExpandableCardProps & {
                     <ContentViewer
                         contentId={`${message.id}-${direction}`}
                         editorNode={this.props.editorNode}
-                        rawContentType={lastHeader(message.headers['content-type'])}
+                        headers={message.headers}
                         contentType={decodedContentType}
                         schema={apiBodySchema}
                         expanded={!!expanded}

diff --git a/src/model/events/body-formatting.ts b/src/model/events/body-formatting.ts
@@ -1,3 +1,4 @@
+import { Headers } from '../../types';
 import { styled } from '../../styles';
 
 import { ViewableContentType } from '../events/content-types';
@@ -13,12 +14,12 @@ export interface EditorFormatter {
     language: string;
     cacheKey: Symbol;
     isEditApplicable: boolean; // Can you apply this manually during editing to format an input?
-    render(content: Buffer): string | ObservablePromise<string>;
+    render(content: Buffer, headers?: Headers): string | ObservablePromise<string>;
 }
 
 type FormatComponentProps = {
     content: Buffer;
-    rawContentType: string | undefined;
+    headers?: Headers;
 };
 
 type FormatComponent = React.ComponentType<FormatComponentProps>;
@@ -35,17 +36,17 @@ export function isEditorFormatter(input: any): input is EditorFormatter {
 }
 
 const buildAsyncRenderer = (formatKey: WorkerFormatterKey) =>
-    (input: Buffer) => observablePromise(
-        formatBufferAsync(input, formatKey)
+    (input: Buffer, headers?: Headers) => observablePromise(
+        formatBufferAsync(input, formatKey, headers)
     );
 
 export const Formatters: { [key in ViewableContentType]: Formatter } = {
     raw: {
         language: 'text',
         cacheKey: Symbol('raw'),
         isEditApplicable: false,
-        render: (input: Buffer) => {
-            if (input.byteLength < 2000) {
+        render: (input: Buffer, headers?: Headers) => {
+            if (input.byteLength < 2_000) {
                 try {
                     // For short-ish inputs, we return synchronously - conveniently this avoids
                     // showing the loading spinner that churns the layout in short content cases.
@@ -55,7 +56,7 @@ export const Formatters: { [key in ViewableContentType]: Formatter } = {
                 }
             } else {
                 return observablePromise(
-                    formatBufferAsync(input, 'raw')
+                    formatBufferAsync(input, 'raw', headers)
                 );
             }
         }
@@ -64,7 +65,7 @@ export const Formatters: { [key in ViewableContentType]: Formatter } = {
         language: 'text',
         cacheKey: Symbol('text'),
         isEditApplicable: false,
-        render: (input: Buffer) => {
+        render: (input: Buffer, headers?: Headers) => {
             return bufferToString(input);
         }
     },
@@ -102,24 +103,24 @@ export const Formatters: { [key in ViewableContentType]: Formatter } = {
         language: 'json',
         cacheKey: Symbol('json'),
         isEditApplicable: true,
-        render: (input: Buffer) => {
-            if (input.byteLength < 10000) {
+        render: (input: Buffer, headers?: Headers) => {
+            if (input.byteLength < 10_000) {
                 const inputAsString = bufferToString(input);
 
                 try {
                     // For short-ish inputs, we return synchronously - conveniently this avoids
                     // showing the loading spinner that churns the layout in short content cases.
                     return JSON.stringify(
                         JSON.parse(inputAsString),
-                    null, 2);
+                        null, 2);
                     // ^ Same logic as in UI-worker-formatter
                 } catch (e) {
                     // Fallback to showing the raw un-formatted JSON:
                     return inputAsString;
                 }
             } else {
                 return observablePromise(
-                    formatBufferAsync(input, 'json')
+                    formatBufferAsync(input, 'json', headers)
                 );
             }
         }

diff --git a/src/model/events/content-types.ts b/src/model/events/content-types.ts
@@ -21,7 +21,7 @@ export const getBaseContentType = (mimeType: string | undefined) => {
         return type + '/' + combinedSubTypes;
     }
 
-    // Otherwise, wr collect a list of types from most specific to most generic: [svg, xml] for image/svg+xml
+    // Otherwise, we collect a list of types from most specific to most generic: [svg, xml] for image/svg+xml
     // and then look through in order to see if there are any matches here:
     const subTypes = combinedSubTypes.split('+');
     const possibleTypes = subTypes.map(st => type + '/' + st);
@@ -112,6 +112,9 @@ const mimeTypeToContentTypeMap: { [mimeType: string]: ViewableContentType } = {
     'application/x-protobuffer': 'protobuf', // Commonly seen in Google apps
 
     'application/grpc+proto': 'grpc-proto', // Used in GRPC requests (protobuf but with special headers)
+    'application/grpc+protobuf': 'grpc-proto',
+    'application/grpc-proto': 'grpc-proto',
+    'application/grpc-protobuf': 'grpc-proto',
 
     'application/octet-stream': 'raw'
 } as const;
@@ -180,10 +183,6 @@ export function getCompatibleTypes(
         types.add('xml');
     }
 
-    if (!types.has('grpc-proto') && rawContentType === 'application/grpc') {
-        types.add('grpc-proto')
-    }
-
     if (
         body &&
         isProbablyProtobuf(body) &&
@@ -205,7 +204,7 @@ export function getCompatibleTypes(
         body &&
         body.length > 0 &&
         body.length % 4 === 0 && // Multiple of 4 bytes
-        body.length < 1000 * 100 && // < 100 KB of content
+        body.length < 100_000 && // < 100 KB of content
         body.every(isValidBase64Byte)
     ) {
         types.add('base64');

diff --git a/src/services/ui-worker-api.ts b/src/services/ui-worker-api.ts
@@ -21,7 +21,7 @@ import type {
     ParseCertResponse
 } from './ui-worker';
 
-import { Omit } from '../types';
+import { Headers, Omit } from '../types';
 import type { ApiMetadata, ApiSpec } from '../model/api/api-interfaces';
 import { WorkerFormatterKey } from './ui-worker-formatters';
 
@@ -149,10 +149,11 @@ export async function parseCert(buffer: ArrayBuffer) {
     })).result;
 }
 
-export async function formatBufferAsync(buffer: ArrayBuffer, format: WorkerFormatterKey) {
+export async function formatBufferAsync(buffer: ArrayBuffer, format: WorkerFormatterKey, headers?: Headers) {
     return (await callApi<FormatRequest, FormatResponse>({
         type: 'format',
         buffer,
-        format
+        format,
+        headers,
     })).formatted;
 }
diff --git a/src/services/ui-worker-formatters.ts b/src/services/ui-worker-formatters.ts
@@ -5,6 +5,7 @@ import {
 } from 'js-beautify/js/lib/beautifier';
 import * as beautifyXml from 'xml-beautifier';
 
+import { Headers } from '../types';
 import { bufferToHex, bufferToString, getReadableSize } from '../util/buffer';
 import { parseRawProtobuf, extractProtobufFromGrpc } from '../util/protobuf';
 
@@ -13,10 +14,25 @@ const FIVE_MB = 1024 * 1024 * 5;
 
 export type WorkerFormatterKey = keyof typeof WorkerFormatters;
 
-export function formatBuffer(buffer: ArrayBuffer, format: WorkerFormatterKey): string {
-    return WorkerFormatters[format](Buffer.from(buffer));
+export function formatBuffer(buffer: ArrayBuffer, format: WorkerFormatterKey, headers?: Headers): string {
+    return WorkerFormatters[format](Buffer.from(buffer), headers);
 }
 
+const prettyProtobufView = (data: any) => JSON.stringify(data, (_key, value) => {
+    // Buffers have toJSON defined, so arrive here in JSONified form:
+    if (value.type === 'Buffer' && Array.isArray(value.data)) {
+        const buffer = Buffer.from(value.data);
+
+        return {
+            "Type": `Buffer (${getReadableSize(buffer)})`,
+            "As string": bufferToString(buffer, 'detect-encoding'),
+            "As hex": bufferToHex(buffer)
+        }
+    } else {
+        return value;
+    }
+}, 2);
+
 // A subset of all possible formatters (those allowed by body-formatting), which require
 // non-trivial processing, and therefore need to be processed async.
 const WorkerFormatters = {
@@ -76,44 +92,15 @@ const WorkerFormatters = {
         });
     },
     protobuf: (content: Buffer) => {
-        const data = parseRawProtobuf(content, {
-            prefix: ''
-        });
-
-        return JSON.stringify(data, (_key, value) => {
-            // Buffers have toJSON defined, so arrive here in JSONified form:
-            if (value.type === 'Buffer' && Array.isArray(value.data)) {
-                const buffer = Buffer.from(value.data);
-
-                return {
-                    "Type": `Buffer (${getReadableSize(buffer)})`,
-                    "As string": bufferToString(buffer, 'detect-encoding'),
-                    "As hex": bufferToHex(buffer)
-                }
-            } else {
-                return value;
-            }
-        }, 2);
+        const data = parseRawProtobuf(content, { prefix: '' });
+        return prettyProtobufView(data);
     },
-    'grpc-proto': (content: Buffer) => {
-        const protobufMessages = extractProtobufFromGrpc(content);
+    'grpc-proto': (content: Buffer, headers?: Headers) => {
+        const protobufMessages = extractProtobufFromGrpc(content, headers ?? {});
 
         let data = protobufMessages.map((msg) => parseRawProtobuf(msg, { prefix: '' }));
         if (data.length === 1) data = data[0];
 
-        return JSON.stringify(data, (_key, value) => {
-            // Buffers have toJSON defined, so arrive here in JSONified form:
-            if (value.type === 'Buffer' && Array.isArray(value.data)) {
-                const buffer = Buffer.from(value.data);
-
-                return {
-                    "Type": `Buffer (${getReadableSize(buffer)})`,
-                    "As string": bufferToString(buffer, 'detect-encoding'),
-                    "As hex": bufferToHex(buffer)
-                }
-            } else {
-                return value;
-            }
-        }, 2);
+        return prettyProtobufView(data);
     }
 } as const;
diff --git a/src/services/ui-worker.ts b/src/services/ui-worker.ts
@@ -13,6 +13,7 @@ import {
 } from 'http-encoding';
 import { OpenAPIObject } from 'openapi-directory';
 
+import { Headers } from '../types';
 import { ApiMetadata, ApiSpec } from '../model/api/api-interfaces';
 import { buildOpenApiMetadata, buildOpenRpcMetadata } from '../model/api/build-api-metadata';
 import { parseCert, ParsedCertificate, validatePKCS12, ValidationResult } from '../model/crypto';
@@ -91,6 +92,7 @@ export interface FormatRequest extends Message {
     type: 'format';
     buffer: ArrayBuffer;
     format: WorkerFormatterKey;
+    headers?: Headers;
 }
 
 export interface FormatResponse extends Message {
@@ -217,7 +219,7 @@ ctx.addEventListener('message', async (event: { data: BackgroundRequest }) => {
                 break;
 
             case 'format':
-                const formatted = formatBuffer(event.data.buffer, event.data.format);
+                const formatted = formatBuffer(event.data.buffer, event.data.format, event.data.headers);
                 ctx.postMessage({ id: event.data.id, formatted });
                 break;