79 lines
2.4 KiB
JavaScript
79 lines
2.4 KiB
JavaScript
import ws from 'k6/ws';
|
|
import { check } from 'k6';
|
|
import crypto from "k6/crypto";
|
|
import encoding from "k6/encoding";
|
|
|
|
const hmacSecret = "secret";
|
|
|
|
const algToHash = {
|
|
HS256: "sha256",
|
|
HS384: "sha384",
|
|
HS512: "sha512"
|
|
};
|
|
|
|
function sign(data, hashAlg, secret) {
|
|
let hasher = crypto.createHMAC(hashAlg, secret);
|
|
hasher.update(data);
|
|
// Some manual base64 encoding as `Hasher.digest(encodingType)` doesn't support that encoding type yet.
|
|
return hasher.digest("base64").replace(/\//g, "_").replace(/\+/g, "-").replace(/=/g, "");
|
|
}
|
|
|
|
function encode(payload, secret, algorithm) {
|
|
algorithm = algorithm || "HS256";
|
|
let header = encoding.b64encode(JSON.stringify({ typ: "JWT", alg: algorithm }), "rawurl");
|
|
payload = encoding.b64encode(JSON.stringify(payload), "rawurl");
|
|
let sig = sign(header + "." + payload, algToHash[algorithm], secret);
|
|
return [header, payload, sig].join(".");
|
|
}
|
|
|
|
// Generate a token for each user separately using HMAC JWT
|
|
function generateToken(userId) {
|
|
const payload = {
|
|
sub: userId,
|
|
exp: Math.floor(Date.now() / 1000) + 60, // Expires in 60 seconds
|
|
};
|
|
return encode(payload, hmacSecret, "HS256")
|
|
}
|
|
|
|
export let options = {
|
|
stages: [
|
|
{ duration: "10s", target: 1000 },
|
|
{ duration: "50s", target: 1000 },
|
|
]
|
|
};
|
|
|
|
export default function () {
|
|
const url = 'ws://localhost:8000/connection/websocket';
|
|
|
|
const user = `user_${__VU}_${__ITER}`;
|
|
const token = generateToken(user);
|
|
|
|
const response = ws.connect(url, {}, function (socket) {
|
|
socket.on('open', () => {
|
|
const connectCommand = JSON.stringify({ id: 1, connect: { token: token } });
|
|
const subscribeCommand = JSON.stringify({ id: 2, subscribe: { channel: 'test' } });
|
|
socket.send(connectCommand + "\n" + subscribeCommand);
|
|
});
|
|
|
|
socket.on('message', (message) => {
|
|
// Respond to server pings.
|
|
const substrings = message.split("\n");
|
|
if (substrings.includes("{}")) {
|
|
socket.send('{}');
|
|
}
|
|
});
|
|
|
|
socket.on('close', () => {});
|
|
|
|
socket.on('error', (error) => {
|
|
console.log('WebSocket error: ', error);
|
|
});
|
|
|
|
socket.setTimeout(() => {
|
|
socket.close();
|
|
}, 60000);
|
|
});
|
|
|
|
check(response, { "status is 101": (r) => r && r.status === 101 });
|
|
}
|