Skip to content

Commit 29e74b8

Browse files
committed
chore: add memory profiling in streams examples
1 parent 1ace146 commit 29e74b8

File tree

9 files changed

+2604
-0
lines changed

9 files changed

+2604
-0
lines changed

stream_mem_stress_test/.gitignore

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
# Node modules
2+
node_modules/
3+
4+
# Large files
5+
mem_leak_data1/*
6+
mem_leak_data_default/*

stream_mem_stress_test/README.md

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# Streams - Memory Stress Test and Memory Profiling
2+
This directory contains test code for running a memory profiler when using the `encryptStream` and `decryptUnsignedMessageStream`.
3+
This directory contains everything you need to run a memory profiler on these two operations.
4+
5+
## Requirements
6+
- Node >= 12
7+
- Chrome Browser
8+
9+
## How to run the application and memory profiler
10+
1. For easier debugging open two chrome windows
11+
1. One where you can look at the profiler
12+
2. One where you can navigate through the application paths.
13+
1. On Chrome, navigate to: `chrome://inspect/#devices`
14+
1. Make sure you are in the `stream_mem_stress_test` directory.
15+
1. Start debugger and server by running `npm run start`
16+
1. On the devices page click on `inspect` for the remote target that just appeared
17+
1. Navigate to the Memory tab. You will have three options:
18+
1. Heap snapshot
19+
- Useful to focus on a specific action during runtime.
20+
1. Allocation instrumentation on timeline
21+
- Works better for our stress tests since we can see memory allocation and
22+
garbage collection during runtime
23+
1. Allocation Sampling (not useful for our test)
24+
1. Navigate to any of the provided paths and watch memory allocation and garbage collection in
25+
real time 🍿

stream_mem_stress_test/index.ts

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
import express from 'express'
2+
import { kmsDecryptStream, kmsEncryptStream } from './memory'
3+
4+
const app = express()
5+
const leak = []
6+
7+
app.get('/', (req, res) => {
8+
// Quick little Test so we can check both in browser and in terminal
9+
res.write("Test")
10+
res.send()
11+
console.log("Back in home")
12+
})
13+
14+
/** This path is here to see how a memory leak would look like in the memory profiler.
15+
* To test it you can run `npm run load-mem` to see memory get allocated but never
16+
* garbage collected
17+
*/
18+
app.get('/now', (req, res) => {
19+
let resp = JSON.stringify({ now: new Date() })
20+
leak.push(JSON.parse(resp))
21+
res.writeHead(200, { 'Content-Type': 'application/json' })
22+
res.write(resp)
23+
res.end()
24+
})
25+
26+
/** Path to read 1mb file, parameter to readFile takes any file in the current directory
27+
* so you can test with any file you'd like. We supply 2 test files for you to use.
28+
*/
29+
app.get('/readRandom_1mb', (req, res) => {
30+
res.write("Attempt to encrypt 1mb random")
31+
// you can optionaly pass in a frame size, otherwise uses the default frame size of 4096 bytes
32+
readFile('./random_1mb.txt')
33+
res.send()
34+
})
35+
36+
// Path to decrypt 1mb encrypted file
37+
app.get('/readRandom_1mbEnc', (req, res) => {
38+
res.write("Attempt to decrypt ./random_1mb.txt.encrypted")
39+
readEncryptedFile('./random_1mb.txt.encrypted')
40+
res.send()
41+
42+
})
43+
44+
// Path to encrypt 5mb file
45+
app.get('/readRandom_5mb', (req, res) => {
46+
res.write("Attempt to encrypt 5mb random")
47+
readFile('./random_5mb.txt', 1)
48+
res.send()
49+
})
50+
51+
// Path to decrypt 5mb encrypted file
52+
app.get('/readRandom_5mbEnc', (req, res) => {
53+
res.write("Attempt to decrypt ./random_5mb.txt.encrypted")
54+
readEncryptedFile('./random_5mb.txt.encrypted')
55+
res.send()
56+
57+
})
58+
59+
/** Path to encrypt 1Gb file, not included in this directory because size exceeds GitHub transfer size limit.
60+
* If you want to create a 1Gb file of random data, you can do so by running
61+
* `dd if=/dev/urandom of=rand_1gb.txt bs=1024 count=1024000` on a linux system
62+
*/
63+
app.get('/readRandom_1gb', (req, res) => {
64+
res.write("Attempt to encrypt large random")
65+
readFile('./rand_1gb.txt')
66+
res.send()
67+
})
68+
69+
// Path to decrypt 1mb encrypted file
70+
app.get('/readRandom_1gbEnc', (req, res) => {
71+
res.write("Attempt to encrypt large random")
72+
kmsDecryptStream('./rand_1gb.txt.encrypted')
73+
res.send()
74+
})
75+
76+
app.listen(3000, () => {
77+
console.log("Listening on port 3000");
78+
})
79+
80+
async function readFile(filename:string, framesize?:number) {
81+
await kmsEncryptStream(filename, framesize);
82+
}
83+
84+
async function readEncryptedFile(filename: string) {
85+
await kmsDecryptStream(filename)
86+
}

stream_mem_stress_test/memory.ts

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
import {
2+
KmsKeyringNode,
3+
buildClient,
4+
CommitmentPolicy
5+
} from '@aws-crypto/client-node'
6+
import {AlgorithmSuiteIdentifier} from '@aws-crypto/material-management'
7+
8+
/* This builds client with the REQUIRE_ENCRYPT_REQUIRE_DECRYPT commitment policy.
9+
* REQUIRE_ENCRYPT_REQUIRE_DECRYPT enforces that this client only encrypts using committing algorithm suites
10+
* and enforces that this client will only decrypt encrypted messages
11+
* that were created with a committing algorithm suite.
12+
* This is the default if you build `buildClient()`.
13+
*/
14+
const { encryptStream, decryptUnsignedMessageStream } = buildClient(
15+
CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
16+
)
17+
18+
import { finished } from 'stream'
19+
import { createReadStream, createWriteStream } from 'fs'
20+
import { promisify } from 'util'
21+
const finishedAsync = promisify(finished)
22+
23+
/* A KMS CMK is required to generate the data key.
24+
* You need kms:GenerateDataKey permission on the CMK in generatorKeyId.
25+
* This key is public, DO NOT USE in production environment.
26+
*/
27+
const generatorKeyId =
28+
'arn:aws:kms:us-west-2:658956600833:alias/EncryptDecrypt'
29+
30+
// configure keyring with CMK(s) you wish to work with
31+
const keyring = new KmsKeyringNode({ generatorKeyId })
32+
33+
/**
34+
* Encryption Context is very useful if you want to assert things about the encrypted data
35+
* See: https://docs.aws.amazon.com/encryption-sdk/latest/developer-guide/concepts.html#encryption-context
36+
*/
37+
const context = {
38+
stage: 'demo',
39+
purpose: 'streaming memory stress test',
40+
origin: 'us-west-2',
41+
}
42+
43+
/**
44+
* kmsEncryptStream will use the encryptStream function and create a pipeline to encrypt a stream of data
45+
* from a file and write it to destination `./{filename}.encrypted`
46+
* @param filename string of file name you wish to encrypt
47+
* @param framesize optional parameter to determine frame size; default is 4096 bytes
48+
*/
49+
export async function kmsEncryptStream(filename:string, framesize?:number) {
50+
const readable = createReadStream(filename)
51+
const encFile = filename + '.encrypted'
52+
const writeable = createWriteStream(encFile)
53+
54+
// pipeline of read stream
55+
readable.pipe(
56+
encryptStream(keyring, {
57+
/**
58+
* Since we are streaming, and assuming that the encryption and decryption contexts
59+
* are equally trusted, using an unsigned algorithm suite is faster and avoids
60+
* the possibility of processing plaintext before the signature is verified.
61+
*/
62+
suiteId:
63+
AlgorithmSuiteIdentifier.ALG_AES256_GCM_IV12_TAG16_HKDF_SHA512_COMMIT_KEY,
64+
encryptionContext: context,
65+
frameLength: framesize
66+
})
67+
).pipe(writeable.on('finish', () => {
68+
// output from encryptStream will get piped to writeable
69+
console.log(`The new file name is ${encFile}.`);
70+
}))
71+
72+
await finishedAsync(writeable)
73+
console.log("Finished Encrypting");
74+
75+
}
76+
/**
77+
* kmsDecryptStream will take a filename and create a decryption stream to decrypt contents of file
78+
* and write it to destination `./{filename}.decrypted
79+
* @param filename string of file name you wish to encrypt
80+
*/
81+
export async function kmsDecryptStream(filename: string) {
82+
const readable = createReadStream(filename)
83+
const decFile = filename + '.decrypted'
84+
const writeable = createWriteStream(decFile)
85+
86+
readable.pipe(
87+
/**
88+
* decryptUnsignedMessageStream is recommended when streaming if you don't need
89+
* digital signatures.
90+
*/
91+
decryptUnsignedMessageStream(keyring)
92+
).pipe(writeable.on('finish', () => {
93+
console.log(`The new file name is ${decFile}.`);
94+
}))
95+
96+
await finishedAsync(writeable)
97+
console.log("Finished Decrypting")
98+
}
99+

0 commit comments

Comments
 (0)