Line data Source code
1 : import { ChannelArgs, ChannelError, ChannelMetadataProvider, RawProxyError } from '../channel_common.js';
2 : import { Logger } from "../logger/logger.js";
3 : import { HEADER_NAME_BATCH_BYTES, HEADER_NAME_BATCH_TIMEOUT, HEADER_NAME_CHANNEL_ID, HEADER_NAME_ENDPOINT, HEADER_NAME_PATH, HEADER_NAME_READ_TIMEOUT, HEADER_NAME_STREAM_ID, HEADER_NAME_TLS, HEADER_NAME_TLS_CACERTS, HEADER_NAME_TLS_CLIENT_CERT, HEADER_NAME_TLS_CLIENT_KEY } from '../native_proxy_headers.js';
4 : import { nativeProxyFetch } from '../electron_native_fetch.js';
5 :
6 1 : const LOG_CTX = 'native_grpc_client';
7 :
8 1 : export enum NativeGrpcServerStreamBatchEvent {
9 1 : StreamFailed = "StreamFailed",
10 1 : StreamFinished = "StreamFinished",
11 1 : FlushAfterClose = "FlushAfterClose",
12 1 : FlushAfterTimeout = "FlushAfterTimeout",
13 1 : FlushAfterBytes = "FlushAfterBytes",
14 1 : }
15 :
16 : export interface NativeGrpcServerStreamBatch {
17 : /// The batch event
18 : event: NativeGrpcServerStreamBatchEvent;
19 : /// The messages
20 : messages: Uint8Array[];
21 : /// The trailers (if any)
22 : metadata?: Record<string, string>;
23 : }
24 :
25 : export interface NativeGrpcProxyConfig {
26 : /// The endpoint URL
27 : proxyEndpoint: URL;
28 : }
29 :
30 : export interface NativeGrpcServerStreamOptions {
31 : readTimeoutMs?: number;
32 : }
33 :
34 1 : const DEFAULT_READ_TIMEOUT = 60_000;
35 1 : const DEFAULT_BATCH_TIMEOUT = 100;
36 1 : const DEFAULT_BATCH_BYTES = 8000000;
37 :
38 1 : function requireStringHeader(headers: Headers, key: string): string {
39 1 : return headers.get(key)!;
40 1 : }
41 1 : function requireIntegerHeader(headers: Headers, key: string): number {
42 1 : const raw = headers.get(key)!;
43 1 : return Number.parseInt(raw);
44 1 : }
45 1 : async function throwIfError(response: Response): Promise<void> {
46 1 : if (response.headers.get("dashql-error") ?? false) {
47 1 : const proxyError = await response.json() as RawProxyError;
48 1 : const statusCode = requireIntegerHeader(response.headers, "dashql-grpc-status");
49 1 : throw new ChannelError(proxyError, statusCode, response.headers, LOG_CTX);
50 1 : } else if (response.status != 200) {
51 0 : const proxyError: RawProxyError = {
52 0 : message: response.statusText
53 0 : };
54 0 : throw new ChannelError(proxyError, 500, response.headers, LOG_CTX);
55 0 : }
56 1 : }
57 :
58 1 : export class NativeGrpcServerStream implements AsyncIterator<NativeGrpcServerStreamBatch> {
59 : /// The logger
60 1 : logger: Logger;
61 : /// The endpoint
62 1 : endpoint: NativeGrpcProxyConfig;
63 : /// The channel id
64 1 : channelId: number;
65 : /// The stream id
66 1 : streamId: number;
67 : /// The metadata provider
68 1 : metadataProvider: ChannelMetadataProvider;
69 : /// Maximum time to wait for the next stream batch.
70 1 : readTimeoutMs: number;
71 : /// Reached the end of the stream?
72 1 : reachedEndOfStream: boolean;
73 :
74 1 : constructor(endpoint: NativeGrpcProxyConfig, channelId: number, streamId: number, metadata: ChannelMetadataProvider, logger: Logger, options: NativeGrpcServerStreamOptions = {}) {
75 1 : this.logger = logger;
76 1 : this.endpoint = endpoint;
77 1 : this.channelId = channelId;
78 1 : this.streamId = streamId;
79 1 : this.metadataProvider = metadata;
80 1 : this.readTimeoutMs = options.readTimeoutMs ?? DEFAULT_READ_TIMEOUT;
81 1 : this.reachedEndOfStream = false;
82 1 : }
83 :
84 : /// Read the next messages from the stream
85 1 : public async read(): Promise<NativeGrpcServerStreamBatch> {
86 1 : const url = new URL(this.endpoint.proxyEndpoint);
87 1 : url.pathname = `/grpc/channel/${this.channelId}/stream/${this.streamId}`;
88 1 : const requestMetadata = await this.metadataProvider.getRequestMetadata();
89 1 : const headers = new Headers({
90 1 : ...requestMetadata,
91 1 : [HEADER_NAME_READ_TIMEOUT]: this.readTimeoutMs.toString(),
92 1 : [HEADER_NAME_BATCH_TIMEOUT]: DEFAULT_BATCH_TIMEOUT.toString(),
93 1 : [HEADER_NAME_BATCH_BYTES]: DEFAULT_BATCH_BYTES.toString(),
94 1 : });
95 1 : const request = new Request(url, {
96 1 : method: 'GET',
97 1 : headers
98 1 : });
99 1 : const response = await nativeProxyFetch(request);
100 1 : await throwIfError(response);
101 :
102 1 : const streamBatchEvent = requireStringHeader(response.headers, "dashql-batch-event");
103 1 : const streamBatchMessages = requireIntegerHeader(response.headers, "dashql-batch-messages");
104 1 : const bodyBuffer = await response.arrayBuffer();
105 1 : const bodyBufferView = new DataView(bodyBuffer);
106 :
107 : // Unpack the messages
108 1 : let offset = 0;
109 1 : const messages = [];
110 1 : while (offset < bodyBuffer.byteLength) {
111 1 : const length = bodyBufferView.getUint32(offset, true);
112 1 : offset += 4;
113 1 : messages.push(new Uint8Array(bodyBuffer, offset, length));
114 1 : offset += length;
115 1 : }
116 :
117 : // Message count mismatch?
118 : // We treat this as an error since this means our encoded response buffer is corrupt.
119 1 : if (streamBatchMessages != messages.length) {
120 : // Fatal error, silently drop the stream
121 1 : const url = new URL(this.endpoint.proxyEndpoint);
122 1 : url.pathname = `/grpc/channel/${this.channelId}/stream/${this.streamId}`;
123 1 : const request = new Request(url, {
124 1 : method: 'DELETE'
125 1 : });
126 1 : await nativeProxyFetch(request);
127 : // XXX Log if the dropping failed
128 1 : throw new ChannelError({ message: "batch message count mismatch" }, 13, undefined, LOG_CTX);
129 1 : }
130 :
131 : // Return the batch event and all messages
132 1 : const metadata: Record<string, string> = {};
133 1 : response.headers.forEach((value, key) => {
134 1 : metadata[key] = value;
135 1 : });
136 :
137 : // Return the batch event and all messages
138 1 : return {
139 1 : event: streamBatchEvent as NativeGrpcServerStreamBatchEvent,
140 1 : messages: messages,
141 1 : metadata
142 1 : };
143 1 : }
144 :
145 : /// Get the next batch
146 1 : public async next(): Promise<IteratorResult<NativeGrpcServerStreamBatch>> {
147 1 : if (this.reachedEndOfStream) {
148 1 : return { value: undefined, done: true };
149 : }
150 1 : const batch = await this.read();
151 1 : switch (batch.event) {
152 1 : case NativeGrpcServerStreamBatchEvent.FlushAfterBytes:
153 1 : case NativeGrpcServerStreamBatchEvent.FlushAfterTimeout:
154 1 : return { value: batch, done: false };
155 :
156 1 : case NativeGrpcServerStreamBatchEvent.StreamFinished:
157 1 : case NativeGrpcServerStreamBatchEvent.FlushAfterClose:
158 1 : this.reachedEndOfStream = true;
159 1 : return { value: batch, done: false };
160 :
161 1 : case NativeGrpcServerStreamBatchEvent.StreamFailed:
162 1 : this.reachedEndOfStream = true;
163 1 : throw new ChannelError({ message: "stream failed" }, 400, new Headers(batch.metadata));
164 1 : }
165 1 : }
166 :
167 1 : public async destroy(): Promise<void> {
168 1 : if (this.reachedEndOfStream) return;
169 1 : this.reachedEndOfStream = true;
170 1 : const url = new URL(this.endpoint.proxyEndpoint);
171 1 : url.pathname = `/grpc/channel/${this.channelId}/stream/${this.streamId}`;
172 1 : const response = await nativeProxyFetch(new Request(url, { method: 'DELETE' }));
173 1 : await throwIfError(response);
174 1 : }
175 1 : }
176 :
177 1 : export class NativeGrpcServerStreamMessageIterator implements AsyncIterator<Uint8Array> {
178 : /// The logger
179 1 : protected logger: Logger;
180 : /// The batch iterator
181 1 : protected batchIterator: AsyncIterator<NativeGrpcServerStreamBatch>;
182 : /// The current batch
183 1 : protected currentBatch: NativeGrpcServerStreamBatch | null;
184 : /// The next index in the current batch
185 1 : protected nextInCurrentBatch: number;
186 : /// The combined metadata
187 1 : public metadata: Map<string, string>;
188 :
189 1 : constructor(batchIterator: AsyncIterator<NativeGrpcServerStreamBatch>, logger: Logger) {
190 1 : this.logger = logger;
191 1 : this.batchIterator = batchIterator;
192 1 : this.currentBatch = null;
193 1 : this.nextInCurrentBatch = 0;
194 1 : this.metadata = new Map();
195 1 : }
196 :
197 : /// Get the bytes from the next message in the gRPC stream
198 1 : async next(): Promise<IteratorResult<Uint8Array>> {
199 1 : while (true) {
200 : // Fast path, we still have a buffered message
201 1 : if (this.currentBatch !== null && this.nextInCurrentBatch < this.currentBatch.messages.length) {
202 1 : const mId = this.nextInCurrentBatch;
203 1 : this.nextInCurrentBatch += 1;
204 1 : return { value: this.currentBatch.messages[mId], done: false };
205 1 : }
206 1 : this.currentBatch = null;
207 1 : this.nextInCurrentBatch = 0;
208 :
209 : // Otherwise, get a new batch
210 1 : const result = await this.batchIterator.next();
211 1 : if (result.done) {
212 1 : return { value: undefined, done: true, }
213 1 : } else {
214 1 : this.currentBatch = result.value;
215 1 : this.nextInCurrentBatch = 0;
216 1 : if (result.value.metadata) {
217 1 : for (const key in result.value.metadata) {
218 1 : this.metadata.set(key, result.value.metadata[key]);
219 : }
220 : }
221 1 : }
222 1 : }
223 1 : }
224 1 : }
225 :
226 : interface StartServerStreamArgs {
227 : /// The path
228 : path: string;
229 : /// The request body
230 : body: Uint8Array;
231 : /// Maximum time to wait for the next stream batch.
232 : readTimeoutMs?: number;
233 : }
234 :
235 1 : export class NativeGrpcChannel {
236 : /// The logger
237 1 : logger: Logger;
238 : /// The endpoint
239 1 : endpoint: NativeGrpcProxyConfig;
240 : /// The metadata provider
241 1 : metadataProvider: ChannelMetadataProvider;
242 : /// The channel id
243 1 : channelId: number;
244 :
245 1 : constructor(endpoint: NativeGrpcProxyConfig, channelId: number, metadataProvider: ChannelMetadataProvider, logger: Logger) {
246 1 : this.logger = logger;
247 1 : this.endpoint = endpoint;
248 1 : this.metadataProvider = metadataProvider;
249 1 : this.channelId = channelId;
250 1 : }
251 :
252 : /// Call a server streaming.
253 1 : public async startServerStream(args: StartServerStreamArgs): Promise<NativeGrpcServerStream> {
254 1 : const url = new URL(this.endpoint.proxyEndpoint);
255 1 : url.pathname = `/grpc/channel/${this.channelId}/streams`;
256 :
257 : // Resolve the additional metadata
258 1 : const additionalMetadata = await this.metadataProvider.getRequestMetadata();
259 :
260 : // Collect the request headers
261 1 : const headers = new Headers(additionalMetadata);
262 1 : headers.set(HEADER_NAME_CHANNEL_ID, this.channelId.toString());
263 1 : headers.set(HEADER_NAME_PATH, args.path);
264 :
265 : // Send the request
266 1 : const request = new Request(url, {
267 1 : method: 'POST',
268 1 : headers,
269 1 : body: args.body as BodyInit,
270 1 : });
271 1 : const response = await nativeProxyFetch(request);
272 1 : await throwIfError(response);
273 :
274 1 : const streamId = requireIntegerHeader(response.headers, HEADER_NAME_STREAM_ID);
275 1 : return new NativeGrpcServerStream(this.endpoint, this.channelId, streamId, this.metadataProvider, this.logger, {
276 1 : readTimeoutMs: args.readTimeoutMs,
277 1 : });
278 1 : }
279 :
280 : // Destroy a channel
281 1 : public async close(): Promise<void> {
282 1 : const url = new URL(this.endpoint.proxyEndpoint);
283 1 : url.pathname = `/grpc/channel/${this.channelId}`;
284 :
285 1 : const headers = new Headers();
286 1 : headers.set(HEADER_NAME_CHANNEL_ID, this.channelId.toString());
287 :
288 : // Send the request
289 1 : const request = new Request(url, {
290 1 : method: 'DELETE',
291 1 : headers,
292 1 : body: "",
293 1 : });
294 1 : const response = await nativeProxyFetch(request);
295 1 : await throwIfError(response);
296 1 : }
297 1 : }
298 :
299 1 : export class NativeGrpcClient {
300 : /// The logger
301 1 : logger: Logger;
302 : /// The endpoint
303 1 : proxy: NativeGrpcProxyConfig;
304 :
305 1 : constructor(proxy: NativeGrpcProxyConfig, logger: Logger) {
306 1 : this.logger = logger;
307 1 : this.proxy = proxy;
308 1 : }
309 :
310 : /// Create a gRPC channel
311 1 : public async connect(args: ChannelArgs, metadataProvider: ChannelMetadataProvider): Promise<NativeGrpcChannel> {
312 1 : const url = new URL(this.proxy.proxyEndpoint);
313 1 : url.pathname = `/grpc/channels`;
314 :
315 : // Resolve additional metadata (such as the authorization header)
316 1 : const additionalMetadata = await metadataProvider.getRequestMetadata()
317 :
318 : // Set system headers
319 1 : const headers = new Headers(additionalMetadata);
320 1 : headers.set(HEADER_NAME_ENDPOINT, args.endpoint);
321 1 : if (args.tls) {
322 1 : headers.set(HEADER_NAME_TLS, "1");
323 1 : if (args.tls.keyPath && args.tls.keyPath !== "") {
324 1 : headers.set(HEADER_NAME_TLS_CLIENT_KEY, args.tls.keyPath);
325 : }
326 1 : if (args.tls.pubPath && args.tls.pubPath !== "") {
327 1 : headers.set(HEADER_NAME_TLS_CLIENT_CERT, args.tls.pubPath);
328 : }
329 1 : if (args.tls.caPath && args.tls.caPath !== "") {
330 1 : headers.set(HEADER_NAME_TLS_CACERTS, args.tls.caPath);
331 : }
332 1 : }
333 1 : const request = new Request(url, {
334 1 : method: 'POST',
335 1 : headers
336 1 : });
337 1 : const response = await nativeProxyFetch(request);
338 1 : await throwIfError(response);
339 :
340 1 : const channelId = requireIntegerHeader(response.headers, HEADER_NAME_CHANNEL_ID);
341 1 : return new NativeGrpcChannel(this.proxy, channelId, metadataProvider, this.logger);
342 1 : }
343 1 : }
|
