-
Notifications
You must be signed in to change notification settings - Fork 290
/
jupyterRequestCreator.node.ts
121 lines (104 loc) · 4.99 KB
/
jupyterRequestCreator.node.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
import { IJupyterRequestCreator } from '../types';
import * as nodeFetch from 'node-fetch';
import { ClassType } from '../../../platform/ioc/types';
import * as WebSocketIsomorphic from 'isomorphic-ws';
import { traceError } from '../../../platform/logging';
import { noop } from '../../../platform/common/utils/misc';
import { KernelSocketWrapper } from '../../common/kernelSocketWrapper';
import { IKernelSocket } from '../../types';
import { injectable } from 'inversify';
/* eslint-disable @typescript-eslint/no-explicit-any */
const JupyterWebSockets = new Map<string, WebSocketIsomorphic & IKernelSocket>(); // NOSONAR
// Function for creating node Request object that prevents jupyterlab services from writing its own
// authorization header.
/* eslint-disable @typescript-eslint/no-explicit-any */
@injectable()
export class JupyterRequestCreator implements IJupyterRequestCreator {
public getRequestCtor(getAuthHeader?: () => any) {
class AuthorizingRequest extends nodeFetch.Request {
constructor(input: nodeFetch.RequestInfo, init?: nodeFetch.RequestInit) {
super(input, init);
// Add all of the authorization parts onto the headers.
const origHeaders = this.headers;
const authorizationHeader = getAuthHeader!();
const keys = Object.keys(authorizationHeader);
keys.forEach((k) => origHeaders.append(k, authorizationHeader[k].toString()));
origHeaders.set('Content-Type', 'application/json');
// Rewrite the 'append' method for the headers to disallow 'authorization' after this point
const origAppend = origHeaders.append.bind(origHeaders);
origHeaders.append = (k, v) => {
if (k.toLowerCase() !== 'authorization') {
origAppend(k, v);
}
};
}
}
return (getAuthHeader ? AuthorizingRequest : nodeFetch.Request) as any;
}
public getWebsocketCtor(
cookieString?: string,
allowUnauthorized?: boolean,
getAuthHeaders?: () => any
): ClassType<WebSocket> {
class JupyterWebSocket extends KernelSocketWrapper(WebSocketIsomorphic) {
private kernelId: string | undefined;
private timer: NodeJS.Timeout | number;
constructor(url: string, protocols?: string | string[] | undefined) {
let co: WebSocketIsomorphic.ClientOptions = {};
let co_headers: { [key: string]: string } | undefined;
if (allowUnauthorized) {
co = { ...co, rejectUnauthorized: false };
}
if (cookieString) {
co_headers = { Cookie: cookieString };
}
// Auth headers have to be refetched every time we create a connection. They may have expired
// since the last connection.
if (getAuthHeaders) {
const authorizationHeader = getAuthHeaders();
co_headers = co_headers ? { ...co_headers, ...authorizationHeader } : authorizationHeader;
}
if (co_headers) {
co = { ...co, headers: co_headers };
}
super(url, protocols, co);
let timer: NodeJS.Timeout | undefined = undefined;
// Parse the url for the kernel id
const parsed = /.*\/kernels\/(.*)\/.*/.exec(url);
if (parsed && parsed.length > 1) {
this.kernelId = parsed[1];
}
if (this.kernelId) {
JupyterWebSockets.set(this.kernelId, this);
this.on('close', () => {
if (timer && this.timer !== timer) {
clearInterval(timer as any);
}
if (JupyterWebSockets.get(this.kernelId!) === this) {
JupyterWebSockets.delete(this.kernelId!);
}
});
} else {
traceError('KernelId not extracted from Kernel WebSocket URL');
}
// Ping the websocket connection every 30 seconds to make sure it stays alive
timer = this.timer = setInterval(() => this.ping(noop), 30_000);
}
}
return JupyterWebSocket as any;
}
public getWebsocket(id: string): IKernelSocket | undefined {
return JupyterWebSockets.get(id);
}
public getFetchMethod(): (input: RequestInfo, init?: RequestInit) => Promise<Response> {
return nodeFetch.default as any;
}
public getHeadersCtor(): ClassType<Headers> {
return nodeFetch.Headers as any;
}
public getRequestInit(): RequestInit {
return { cache: 'no-store', credentials: 'same-origin' };
}
}