Skip to content

Commit

Permalink
Heap & asm pooling
Browse files Browse the repository at this point in the history
  • Loading branch information
twiss committed Nov 13, 2020
1 parent 2a7a339 commit 6e5fe3f
Show file tree
Hide file tree
Showing 15 changed files with 271 additions and 180 deletions.
67 changes: 48 additions & 19 deletions src/aes/aes.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,31 +2,62 @@ import { AES_asm, AES_mode } from './aes.asm';
import { _heap_init, _heap_write, is_bytes } from '../other/utils';
import { IllegalArgumentError, SecurityError } from '../other/errors';

const heap_pool: Uint8Array[] = [];
const asm_pool: AES_asm[] = [];

export class AES {
public readonly heap: Uint8Array;
public readonly asm: AES_asm;
public heap?: Uint8Array;
public asm?: AES_asm;
private readonly mode: string;
public padding: boolean; // TODO: This should be `private readonly`, hacking for AES-CFB?!
public pos: number = 0;
public len: number = 0;

private key: Uint8Array;
private iv: Uint8Array | undefined;

constructor(key: Uint8Array, iv: Uint8Array | undefined, padding = true, mode: AES_mode, heap?: Uint8Array, asm?: AES_asm) {
this.mode = mode;

// The AES "worker"
this.heap = heap ? heap : _heap_init().subarray(AES_asm.HEAP_DATA);
this.asm = asm ? asm : new AES_asm(null, this.heap.buffer);

// The AES object state
this.pos = 0;
this.len = 0;

this.key = key;
this.iv = iv;
this.padding = padding;

// The AES "worker"
this.acquire_asm(heap, asm);
}

acquire_asm(heap?: Uint8Array, asm?: AES_asm): { heap: Uint8Array, asm: AES_asm } {
if (this.heap === undefined || this.asm === undefined) {
this.heap = heap || heap_pool.pop() || _heap_init().subarray(AES_asm.HEAP_DATA);
this.asm = asm || asm_pool.pop() || new AES_asm(null, this.heap.buffer);
this.reset(this.key, this.iv);
}
return { heap: this.heap, asm: this.asm };
}

release_asm() {
if (this.heap !== undefined && this.asm !== undefined) {
heap_pool.push(this.heap);
asm_pool.push(this.asm);
}
this.heap = undefined;
this.asm = undefined;
}

protected reset(key: Uint8Array, iv: Uint8Array | undefined) {
const { asm } = this.acquire_asm();

// Key
const keylen = key.length;
if (keylen !== 16 && keylen !== 24 && keylen !== 32) throw new IllegalArgumentError('illegal key size');

const keyview = new DataView(key.buffer, key.byteOffset, key.byteLength);
this.asm.set_key(
asm.set_key(
keylen >> 2,
keyview.getUint32(0),
keyview.getUint32(4),
Expand All @@ -44,19 +75,16 @@ export class AES {

let ivview = new DataView(iv.buffer, iv.byteOffset, iv.byteLength);

this.asm.set_iv(ivview.getUint32(0), ivview.getUint32(4), ivview.getUint32(8), ivview.getUint32(12));
asm.set_iv(ivview.getUint32(0), ivview.getUint32(4), ivview.getUint32(8), ivview.getUint32(12));
} else {
this.asm.set_iv(0, 0, 0, 0);
asm.set_iv(0, 0, 0, 0);
}

this.padding = padding;
}

AES_Encrypt_process(data: Uint8Array): Uint8Array {
if (!is_bytes(data)) throw new TypeError("data isn't of expected type");

let asm = this.asm;
let heap = this.heap;
let { heap, asm } = this.acquire_asm();
let amode = AES_asm.ENC[this.mode];
let hpos = AES_asm.HEAP_DATA;
let pos = this.pos;
Expand Down Expand Up @@ -96,8 +124,7 @@ export class AES {
}

AES_Encrypt_finish(): Uint8Array {
let asm = this.asm;
let heap = this.heap;
let { heap, asm } = this.acquire_asm();
let amode = AES_asm.ENC[this.mode];
let hpos = AES_asm.HEAP_DATA;
let pos = this.pos;
Expand Down Expand Up @@ -128,14 +155,15 @@ export class AES {
this.pos = 0;
this.len = 0;

this.release_asm();

return result;
}

AES_Decrypt_process(data: Uint8Array): Uint8Array {
if (!is_bytes(data)) throw new TypeError("data isn't of expected type");

let asm = this.asm;
let heap = this.heap;
let { heap, asm } = this.acquire_asm();
let amode = AES_asm.DEC[this.mode];
let hpos = AES_asm.HEAP_DATA;
let pos = this.pos;
Expand Down Expand Up @@ -181,8 +209,7 @@ export class AES {
}

AES_Decrypt_finish(): Uint8Array {
let asm = this.asm;
let heap = this.heap;
let { heap, asm } = this.acquire_asm();
let amode = AES_asm.DEC[this.mode];
let hpos = AES_asm.HEAP_DATA;
let pos = this.pos;
Expand Down Expand Up @@ -221,6 +248,8 @@ export class AES {
this.pos = 0;
this.len = 0;

this.release_asm();

return result;
}
}
2 changes: 1 addition & 1 deletion src/aes/cbc.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import { AES } from './aes';
import { joinBytes } from '../other/utils';

export class AES_CBC {
private aes: AES;
public aes: AES;

static encrypt(data: Uint8Array, key: Uint8Array, padding: boolean = true, iv?: Uint8Array): Uint8Array {
return new AES_CBC(key, iv, padding).encrypt(data);
Expand Down
30 changes: 14 additions & 16 deletions src/aes/ccm.ts
Original file line number Diff line number Diff line change
Expand Up @@ -141,17 +141,17 @@ export class AES_CCM {
data.set(adata, 18);
}

let { asm, heap } = this.aes.acquire_asm();
this._cbc_mac_process(data);
this.aes.asm.get_state(AES_asm.HEAP_DATA);
asm.get_state(AES_asm.HEAP_DATA);

const iv = new Uint8Array(this.aes.heap.subarray(0, 16));
const iv = new Uint8Array(heap.subarray(0, 16));
const ivview = new DataView(iv.buffer, iv.byteOffset, iv.byteLength);
this.aes.asm.set_iv(ivview.getUint32(0), ivview.getUint32(4), ivview.getUint32(8), ivview.getUint32(12));
asm.set_iv(ivview.getUint32(0), ivview.getUint32(4), ivview.getUint32(8), ivview.getUint32(12));
}

_cbc_mac_process(data: Uint8Array): void {
const heap = this.aes.heap;
const asm = this.aes.asm;
let { asm, heap } = this.aes.acquire_asm();
let dpos = 0;
let dlen = data.length || 0;
let wlen = 0;
Expand All @@ -167,8 +167,7 @@ export class AES_CCM {
}

AES_CCM_Encrypt_process(data: Uint8Array): Uint8Array {
const asm = this.aes.asm;
const heap = this.aes.heap;
let { asm, heap } = this.aes.acquire_asm();

let dpos = 0;
let dlen = data.length || 0;
Expand Down Expand Up @@ -216,8 +215,7 @@ export class AES_CCM {
}

AES_CCM_Encrypt_finish(): Uint8Array {
const asm = this.aes.asm;
const heap = this.aes.heap;
let { asm, heap } = this.aes.acquire_asm();
const tagSize = this.tagSize;
const pos = this.aes.pos;
const len = this.aes.len;
Expand Down Expand Up @@ -246,8 +244,7 @@ export class AES_CCM {
AES_CCM_Decrypt_process(data: Uint8Array): Uint8Array {
let dpos = 0;
let dlen = data.length || 0;
const asm = this.aes.asm;
const heap = this.aes.heap;
let { asm, heap } = this.aes.acquire_asm();
let counter = this.counter;
const tagSize = this.tagSize;
let pos = this.aes.pos;
Expand Down Expand Up @@ -290,8 +287,7 @@ export class AES_CCM {
}

AES_CCM_Decrypt_finish(): Uint8Array {
const asm = this.aes.asm;
const heap = this.aes.heap;
let { asm, heap } = this.aes.acquire_asm();
const tagSize = this.tagSize;
const pos = this.aes.pos;
const len = this.aes.len;
Expand Down Expand Up @@ -327,8 +323,10 @@ export class AES_CCM {
private AES_CTR_set_options(nonce: Uint8Array, counter: number, size: number): void {
if (size < 8 || size > 48) throw new IllegalArgumentError('illegal counter size');

let { asm } = this.aes.acquire_asm();

const mask = Math.pow(2, size) - 1;
this.aes.asm.set_mask(0, 0, (mask / 0x100000000) | 0, mask | 0);
asm.set_mask(0, 0, (mask / 0x100000000) | 0, mask | 0);

const len = nonce.length;
if (!len || len > 16) throw new IllegalArgumentError('illegal nonce size');
Expand All @@ -338,12 +336,12 @@ export class AES_CCM {
const view = new DataView(new ArrayBuffer(16));
new Uint8Array(view.buffer).set(nonce);

this.aes.asm.set_nonce(view.getUint32(0), view.getUint32(4), view.getUint32(8), view.getUint32(12));
asm.set_nonce(view.getUint32(0), view.getUint32(4), view.getUint32(8), view.getUint32(12));

if (counter < 0 || counter >= Math.pow(2, size)) throw new IllegalArgumentError('illegal counter value');

this.counter = counter;

this.aes.asm.set_counter(0, 0, (counter / 0x100000000) | 0, counter | 0);
asm.set_counter(0, 0, (counter / 0x100000000) | 0, counter | 0);
}
}
4 changes: 2 additions & 2 deletions src/aes/cmac.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,9 @@ export class AES_CMAC {

process(data: Uint8Array): this {
if (this.bufferLength + data.length > 16) {
this.cbc.encrypt(this.buffer.subarray(0, this.bufferLength));
this.cbc.aes.AES_Encrypt_process(this.buffer.subarray(0, this.bufferLength));
const offset = ((this.bufferLength + data.length - 1) & ~15) - this.bufferLength;
this.cbc.encrypt(data.subarray(0, offset));
this.cbc.aes.AES_Encrypt_process(data.subarray(0, offset));
this.buffer.set(data.subarray(offset));
this.bufferLength = data.length - offset;
} else {
Expand Down
10 changes: 6 additions & 4 deletions src/aes/ctr.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,14 +35,16 @@ export class AES_CTR {
}

private AES_CTR_set_options(nonce: Uint8Array, counter?: number, size?: number): void {
let { asm } = this.aes.acquire_asm();

if (size !== undefined) {
if (size < 8 || size > 48) throw new IllegalArgumentError('illegal counter size');

let mask = Math.pow(2, size) - 1;
this.aes.asm.set_mask(0, 0, (mask / 0x100000000) | 0, mask | 0);
asm.set_mask(0, 0, (mask / 0x100000000) | 0, mask | 0);
} else {
size = 48;
this.aes.asm.set_mask(0, 0, 0xffff, 0xffffffff);
asm.set_mask(0, 0, 0xffff, 0xffffffff);
}

if (nonce !== undefined) {
Expand All @@ -52,15 +54,15 @@ export class AES_CTR {
let view = new DataView(new ArrayBuffer(16));
new Uint8Array(view.buffer).set(nonce);

this.aes.asm.set_nonce(view.getUint32(0), view.getUint32(4), view.getUint32(8), view.getUint32(12));
asm.set_nonce(view.getUint32(0), view.getUint32(4), view.getUint32(8), view.getUint32(12));
} else {
throw new Error('nonce is required');
}

if (counter !== undefined) {
if (counter < 0 || counter >= Math.pow(2, size)) throw new IllegalArgumentError('illegal counter value');

this.aes.asm.set_counter(0, 0, (counter / 0x100000000) | 0, counter | 0);
asm.set_counter(0, 0, (counter / 0x100000000) | 0, counter | 0);
}
}
}
Loading

0 comments on commit 6e5fe3f

Please sign in to comment.