Improved ShardedKey

We now use direct allocated ByteBuffers under the hood which allows us to more strict control over the memory and its lifetime

Author: Wavesonics

app/src/main/kotlin/com/darkrockstudios/app/securecamera/security/SharedKey.kt

@@ -1,29 +1,37 @@
 package com.darkrockstudios.app.securecamera.security
 
 import dev.whyoleg.cryptography.random.CryptographyRandom
+import java.nio.ByteBuffer
 import kotlin.random.Random
 
 class ShardedKey(
 	key: ByteArray
 ) {
-	private val part1: ByteArray
-	private val part2: ByteArray
+	private val part1: ByteBuffer
+	private val part2: ByteBuffer
 	private val keySize: Int = key.size
 
 	init {
-		// Randomly size the storage array so it is not the exact length of an AES key
-		part1 = ByteArray(keySize + CryptographyRandom.nextInt(3, 155))
-		CryptographyRandom.nextBytes(part1)
+		// Randomly size the storage buffer so it is not the exact length of an AES key
+		val part1Size = keySize + CryptographyRandom.nextInt(3, 155)
+		part1 = ByteBuffer.allocateDirect(part1Size)
+		val part1Array = ByteArray(part1Size)
+		CryptographyRandom.nextBytes(part1Array)
+		part1.put(part1Array)
+		part1.flip()
 
 		// Best effort to ensure the two key parts don't live next to each other in memory.
-		// Because this is the JVM, Buckets and Garbage collection and what not will move
-		// things around, so they may end up next to each other now or later, but at least
-		// we tried. This will be released once the constructor returns (or when ever the GC
-		// runs next)
-		val spacer = ByteArray(Random.nextInt(978, 2893))
+		// This will be released once the constructor returns (or when ever the GC runs next)
+		val spacerSize = Random.nextInt(978, 2893)
+		val spacer = ByteBuffer.allocateDirect(spacerSize)
+		val spacerArray = ByteArray(spacerSize)
+		CryptographyRandom.nextBytes(spacerArray)
+		spacer.put(spacerArray)
+		spacer.flip()
+
 		// Use it so it doesn't get optimized out
 		(0..Random.nextInt(13)).forEachIndexed { i, _ ->
-			val x = spacer[0] + spacer[i]
+			val x = spacer.get(0) + spacer.get(i)
 			// Almost never true, but compiler doesn't know that
 			if (System.nanoTime() % 10000 == 0L) {
 				// Side effect that's rarely executed
@@ -32,13 +40,18 @@ class ShardedKey(
 			}
 		}
 
-		// Randomly size the storage array so it is not the exact length of an AES key
-		part2 = ByteArray(keySize + CryptographyRandom.nextInt(5, 111))
-		// Fill it with random data first
-		CryptographyRandom.nextBytes(part2)
+		// Randomly size the storage buffer so it is not the exact length of an AES key
+		val part2Size = keySize + CryptographyRandom.nextInt(5, 111)
+		part2 = ByteBuffer.allocateDirect(part2Size)
+		val part2Array = ByteArray(part2Size)
+		CryptographyRandom.nextBytes(part2Array)
+		part2.put(part2Array)
+		part2.flip()
+
 		// Now fill the data part with our XOR'd key
 		for (i in key.indices) {
-			part2[i] = (key[i].toInt() xor part1[i].toInt()).toByte()
+			val xorByte = (key[i].toInt() xor part1.get(i).toInt()).toByte()
+			part2.put(i, xorByte)
 		}
 	}
 
@@ -50,19 +63,31 @@ class ShardedKey(
 	fun reconstructKey(): ByteArray {
 		val originalKey = ByteArray(keySize)
 		for (i in originalKey.indices) {
-			originalKey[i] = (part1[i].toInt() xor part2[i].toInt()).toByte()
+			originalKey[i] = (part1.get(i).toInt() xor part2.get(i).toInt()).toByte()
 		}
 
 		return originalKey
 	}
 
 	fun evict() {
-		// Modern Android already zero's memory on allocation, but... who knows.
-		part1.let { k ->
-			k.forEachIndexed { index, _ -> k[index] = 0x00 }
-		}
-		part2.let { k ->
-			k.forEachIndexed { index, _ -> k[index] = 0x00 }
+		part1.zeroOut()
+		part2.zeroOut()
+	}
+}
+
+private fun ByteBuffer.zeroOut() {
+	if (hasArray()) {
+		val a = array()
+		val start = arrayOffset() + position()
+		val end = arrayOffset() + limit()
+		java.util.Arrays.fill(a, start, end, 0.toByte())
+	} else {
+		val pos = position()
+		val lim = limit()
+		for (i in 0 until capacity()) {
+			put(i, 0)
 		}
+		position(pos)
+		limit(lim)
 	}
-}
+}