Skip to content

Commit

Permalink
Merge pull request #1900 from cjihrig/encoding
Browse files Browse the repository at this point in the history
fix: support multi-byte characters in WebSocket handler
  • Loading branch information
k8s-ci-robot authored Oct 1, 2024
2 parents 5e0bc07 + f77bffe commit 0906eed
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 23 deletions.
54 changes: 34 additions & 20 deletions src/web-socket-handler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -51,18 +51,11 @@ export class WebSocketHandler implements WebSocketInterface {

public static handleStandardInput(
ws: WebSocket.WebSocket,
stdin: stream.Readable | any,
stdin: stream.Readable,
streamNum: number = 0,
): boolean {
stdin.on('data', (data) => {
const buff = Buffer.alloc(data.length + 1);
buff.writeInt8(streamNum, 0);
if (data instanceof Buffer) {
data.copy(buff, 1);
} else {
buff.write(data, 1);
}
ws.send(buff);
ws.send(copyChunkForWebSocket(streamNum, data, stdin.readableEncoding));
});

stdin.on('end', () => {
Expand All @@ -78,16 +71,9 @@ export class WebSocketHandler implements WebSocketInterface {
createWS: () => Promise<WebSocket.WebSocket>,
streamNum: number = 0,
retryCount: number = 3,
encoding?: BufferEncoding | null,
): Promise<WebSocket.WebSocket | null> {
const buff = Buffer.alloc(data.length + 1);

buff.writeInt8(streamNum, 0);
if (data instanceof Buffer) {
data.copy(buff, 1);
} else {
buff.write(data, 1);
}

const buff = copyChunkForWebSocket(streamNum, data, encoding);
let i = 0;
for (; i < retryCount; ++i) {
if (ws !== null && ws.readyState === WebSocket.OPEN) {
Expand All @@ -109,7 +95,7 @@ export class WebSocketHandler implements WebSocketInterface {

public static restartableHandleStandardInput(
createWS: () => Promise<WebSocket.WebSocket>,
stdin: stream.Readable | any,
stdin: stream.Readable,
streamNum: number = 0,
retryCount: number = 3,
): () => WebSocket.WebSocket | null {
Expand All @@ -122,7 +108,14 @@ export class WebSocketHandler implements WebSocketInterface {

stdin.on('data', (data) => {
queue = queue.then(async () => {
ws = await WebSocketHandler.processData(data, ws, createWS, streamNum, retryCount);
ws = await WebSocketHandler.processData(
data,
ws,
createWS,
streamNum,
retryCount,
stdin.readableEncoding,
);
});
});

Expand Down Expand Up @@ -201,3 +194,24 @@ export class WebSocketHandler implements WebSocketInterface {
});
}
}

function copyChunkForWebSocket(
streamNum: number,
chunk: string | Buffer,
encoding?: BufferEncoding | null,
): Buffer {
let buff: Buffer;

if (chunk instanceof Buffer) {
buff = Buffer.alloc(chunk.length + 1);
chunk.copy(buff, 1);
} else {
encoding ??= 'utf-8';
const size = Buffer.byteLength(chunk, encoding);
buff = Buffer.alloc(size + 1);
buff.write(chunk, 1, size, encoding);
}

buff.writeInt8(streamNum, 0);
return buff;
}
23 changes: 20 additions & 3 deletions src/web-socket-handler_test.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import { Readable } from 'node:stream';
import { promisify } from 'util';
import { expect } from 'chai';
import WebSocket = require('isomorphic-ws');
Expand Down Expand Up @@ -303,14 +304,30 @@ describe('WebSocket', () => {
expect(datum).to.equal(fill);
}
});
it('handles multi-byte characters', () => {
return new Promise((resolve) => {
const stream = new Readable({ read() {} });
const mockWs = {
close() {},
send(data) {
expect(data).to.deep.equal(Buffer.from([0x0f, 0xe2, 0x98, 0x83]));
resolve(undefined);
},
} as WebSocket.WebSocket;

stream.setEncoding('utf8');
stream.push('☃');
WebSocketHandler.handleStandardInput(mockWs, stream, 0x0f);
});
});
});

describe('Restartable Handle Standard Input', () => {
it('should throw on negative retry', () => {
const p = new Promise<WebSocket.WebSocket>(() => {});
expect(() => WebSocketHandler.restartableHandleStandardInput(() => p, null, 0, -1)).to.throw(
"retryCount can't be lower than 0.",
);
expect(() =>
WebSocketHandler.restartableHandleStandardInput(() => p, new Readable({ read() {} }), 0, -1),
).to.throw("retryCount can't be lower than 0.");
});

it('should retry n times', () => {
Expand Down

0 comments on commit 0906eed

Please sign in to comment.