cr50: use NIST-compliant configuration of TRNG

According to NIST SP 800-90B only vetted conditioning mechanism should
be used for post-processing raw entropy.  See SP 800-90B, 3.1.5.1 Using
Vetted Conditioning Components. Use of non-vetted algorithms is governed
in 3.1.5.2, but assumes conservative coefficient 0.85 for entropy estimate,
which increase number of requests to TRNG to get desirable entropy.
More details on entropy estimate tests are in associated bug.

Entropy measurements using NIST assessment tool didn't report noticeable
change in entropy estimate. However, more changes are needed to use
DRBG instead of raw TRNG for all purposes.

TRNG changes reviewed also at https://crrev.com/c/1926384

BUG=b:138577834
TEST=test/tpm_test/nist_entropy.sh

Signed-off-by: Vadim Sukhomlinov <sukhomlinov@google.com>
Change-Id: I5a578b90b8b7a77fae6a218eec48e87e7644ab44
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/2240519
Reviewed-by: Vadim Sukhomlinov <sukhomlinov@chromium.org>
Reviewed-by: Vadim Bendebury <vbendeb@chromium.org>
Reviewed-by: Andrey Pronin <apronin@chromium.org>
Tested-by: Vadim Sukhomlinov <sukhomlinov@chromium.org>
Commit-Queue: Vadim Sukhomlinov <sukhomlinov@chromium.org>
Auto-Submit: Vadim Sukhomlinov <sukhomlinov@chromium.org>

5 files changed, 324 insertions, 113 deletions

diff --git a/board/cr50/tpm2/trng.c b/board/cr50/tpm2/trng.c
index 7cce13ff1c..ae4312be2c 100644
--- a/board/cr50/tpm2/trng.c
+++ b/board/cr50/tpm2/trng.c
@@ -9,3 +9,52 @@ CRYPT_RESULT _cpri__StirRandom(int32_t num, uint8_t *entropy)
 {
 	return CRYPT_SUCCESS;     /* NO-OP on CR50. */
 }
+
+#ifdef CRYPTO_TEST_SETUP
+#include "endian.h"
+#include "extension.h"
+#include "trng.h"
+/*
+ * This extension command is similar to TPM2_GetRandom, but made
+ * available for CRYPTO_TEST = 1 which disables TPM.
+ * Command structure, shared out of band with the test driver running
+ * on the host:
+ *
+ * field     |    size  |                  note
+ * =========================================================================
+ * text_len  |    2     | the number of random bytes to generate, big endian
+ * type      |    1     | 0 = TRNG, other values reserved for extensions
+ */
+static enum vendor_cmd_rc trng_test(enum vendor_cmd_cc code, void *buf,
+				    size_t input_size, size_t *response_size)
+{
+	uint16_t text_len;
+	uint8_t *cmd = buf;
+	uint8_t op_type = 0;
+
+	if (input_size != sizeof(text_len) + 1) {
+		*response_size = 0;
+		return VENDOR_RC_BOGUS_ARGS;
+	}
+
+	text_len = be16toh(*(uint16_t *)cmd);
+	op_type = cmd[sizeof(text_len)];
+
+	if (text_len > *response_size) {
+		*response_size = 0;
+		return VENDOR_RC_BOGUS_ARGS;
+	}
+
+	switch (op_type) {
+	case 0:
+		rand_bytes(buf, text_len);
+		break;
+	default:
+		return VENDOR_RC_BOGUS_ARGS;
+	}
+	*response_size = text_len;
+	return VENDOR_RC_SUCCESS;
+}
+
+DECLARE_VENDOR_COMMAND(VENDOR_CC_TRNG_TEST, trng_test);
+#endif /* CRYPTO_TEST_SETUP */
diff --git a/chip/g/trng.c b/chip/g/trng.c
index afd9fa86e3..b64c3c7dbe 100644
--- a/chip/g/trng.c
+++ b/chip/g/trng.c
@@ -8,6 +8,35 @@
 #include "init_chip.h"
 #include "registers.h"
 #include "trng.h"
+#include "watchdog.h"
+#include "console.h"
+
+/**
+ * The H1 TRNG uses the collapse time of a ring oscillator (RO) that is
+ * initialized in a 3x mode (three enable pulses) and eventually collapses
+ * to a stable 1x mode as a result of accumulated jitter (thermal noise).
+ * A Phase-Frequency Detector (PFD) compares the 3x RO to a reference
+ * RO (1.5x) and captures the state of a counter that is incremented
+ * from the reference RO. The resulting reference-cycles-to-collapse
+ * distribution is log-normal, and truncation of the counter bits results in
+ * a distribution that approaches uniform.
+ *
+ * TRNG_SAMPLE_BITS defines how many bits to use from the 16 bit counter
+ * output coming from the analog unit. Entropy is highest in least significant
+ * bits of counter. For FIPS-certified code use just Bit 0 - it provides
+ * highest entropy, allows better security settings for TRNG and simplifies
+ * implementation of continuous health tests.
+ */
+#ifndef TRNG_SAMPLE_BITS
+#define TRNG_SAMPLE_BITS 1
+#endif
+
+/**
+ * Attempts to read TRNG_EMPTY before reporting a stall.
+ * Practically data should be available in less than 400
+ * cycles under normal conditions.
+ */
+#define TRNG_EMPTY_COUNT 400
 
 void init_trng(void)
 {
@@ -21,12 +50,50 @@ void init_trng(void)
 	if (!runlevel_is_high())
 		return;
 #endif
+	/**
+	 * According to NIST SP 800-90B only vetted conditioning mechanism
+	 * should be used for post-processing raw entropy.
+	 * See SP 800-90B, 3.1.5.1 Using Vetted Conditioning Components.
+	 * Use of non-vetted algorithms is governed in 3.1.5.2, but
+	 * assumes conservative coefficient 0.85 for entropy estimate,
+	 * which increase number of requests to TRNG to get desirable
+	 * entropy and prevents from getting full entropy.
+	 */
+	GWRITE(TRNG, POST_PROCESSING_CTRL, 0);
+
+	/**
+	 * TRNG can return up to 16 bits at a time, but highest bits
+	 * have lower entropy. Practically on Cr50 only 13 bits can be
+	 * used - setting to higher value makes TRNG_EMPTY always set.
+	 * Entropy assessed to be reasonable (one bit H > 0.85)
+	 * for up to 8 bits [7..0].
+	 * Time to get 32bit random is roughly 160/TRNG_SAMPLE_BITS us.
+	 */
+	GWRITE(TRNG, SLICE_MAX_UPPER_LIMIT, TRNG_SAMPLE_BITS - 1);
 
-	GWRITE(TRNG, POST_PROCESSING_CTRL,
-		GC_TRNG_POST_PROCESSING_CTRL_SHUFFLE_BITS_MASK |
-		GC_TRNG_POST_PROCESSING_CTRL_CHURN_MODE_MASK);
-	GWRITE(TRNG, SLICE_MAX_UPPER_LIMIT, 1);
+	/* lowest bit have highest entropy, so always start from it */
 	GWRITE(TRNG, SLICE_MIN_LOWER_LIMIT, 0);
+
+	/**
+	 * Analog logic cannot create a value < 8 under normal operating
+	 * conditions, but there's a chance that an attacker could coax
+	 * them out.
+	 * Bit 0 - Enable rejection for values outside of range specified
+	 * by TRNG_ALLOWED_VALUES register
+	 */
+	GWRITE(TRNG, SECURE_POST_PROCESSING_CTRL, 0x1);
+
+	/**
+	 * Since for FIPS settings we use TRNG_SAMPLE_BITS = 1,
+	 * and take only bit 0 from internal 16 bit reading, no bias is
+	 * created for bit 0 if allowed_min is set to 6, which
+	 * actually means min accepted value is 8 (RTL adds +2).
+	 * TRNG_ALLOWED_VALUES_MAX=0x04 (accept all values up to 2^16-1).
+	 * So, range will be [8..65535], with probability for bit 0 and 1
+	 * remaining 50/50.
+	 */
+	GWRITE(TRNG, ALLOWED_VALUES_MIN, 0x26);
+
 	GWRITE(TRNG, TIMEOUT_COUNTER, 0x7ff);
 	GWRITE(TRNG, TIMEOUT_MAX_TRY_NUM, 4);
 	GWRITE(TRNG, POWER_DOWN_B, 1);
@@ -34,16 +101,20 @@ void init_trng(void)
 }
 
 uint32_t rand(void)
-{
+{	uint32_t empty_count = 0;
+
 	while (GREAD(TRNG, EMPTY)) {
-		if (GREAD_FIELD(TRNG, FSM_STATE, FSM_IDLE)) {
+		if (GREAD_FIELD(TRNG, FSM_STATE, FSM_IDLE) ||
+		    empty_count > TRNG_EMPTY_COUNT) {
 			/* TRNG timed out, restart */
 			GWRITE(TRNG, STOP_WORK, 1);
 #if !defined(SECTION_IS_RO) && defined(CONFIG_FLASH_LOG)
 			flash_log_add_event(FE_LOG_TRNG_STALL, 0, NULL);
 #endif
 			GWRITE(TRNG, GO_EVENT, 1);
+			empty_count = 0;
 		}
+		empty_count++;
 	}
 	return GREAD(TRNG, READ_DATA);
 }
@@ -70,57 +141,40 @@ void rand_bytes(void *buffer, size_t len)
 	}
 }
 
-#if !defined(SECTION_IS_RO) && defined(TEST_TRNG)
+#if !defined(SECTION_IS_RO) && defined(CRYPTO_TEST_SETUP)
 #include "console.h"
 #include "watchdog.h"
 
-static uint32_t histogram[256];
-static int command_rand(int argc, char **argv)
+static void  print_rand_stat(uint32_t *histogram, size_t size)
 {
-	int count = 1000; /* Default number of cycles. */
 	struct pair {
 		uint32_t value;
 		uint32_t count;
 	};
 	struct pair min;
 	struct pair max;
-
-	if (argc == 2)
-		count = strtoi(argv[1], NULL, 10);
-
-	memset(histogram, 0, sizeof(histogram));
-	ccprintf("Retrieving %d random words.\n", count);
-	while (count-- > 0) {
-		uint32_t rvalue;
-		int size;
-
-		rvalue = rand();
-		for (size = 0; size < sizeof(rvalue); size++)
-			histogram[((uint8_t *)&rvalue)[size]]++;
-
-		if (!(count % 10000))
-			watchdog_reload();
-	}
+	size_t count;
 
 	min.count = ~0;
 	max.count = 0;
-	for (count = 0; count < ARRAY_SIZE(histogram); count++) {
+	max.value = ~0;
+	min.value = ~0;
+
+	for (count = 0; count < size; count++) {
 		if (histogram[count] > max.count) {
 			max.count = histogram[count];
 			max.value = count;
-			continue;
 		}
-		if (histogram[count] >= min.count)
-			continue;
-
-		min.count = histogram[count];
-		min.value = count;
+		if (histogram[count] < min.count) {
+			min.count = histogram[count];
+			min.value = count;
+		}
 	}
 
 	ccprintf("min %d(%d), max %d(%d)", min.count, min.value,
 		 max.count, max.value);
 
-	for (count = 0; count < ARRAY_SIZE(histogram); count++) {
+	for (count = 0; count < size; count++) {
 		if (!(count % 8)) {
 			ccprintf("\n");
 			cflush();
@@ -128,43 +182,54 @@ static int command_rand(int argc, char **argv)
 		ccprintf(" %6d", histogram[count]);
 	}
 	ccprintf("\n");
-	return EC_SUCCESS;
 }
-DECLARE_CONSOLE_COMMAND(rand, command_rand, NULL, NULL);
-#endif /* !defined(SECTION_IS_RO) && defined(TEST_TRNG) */
-
-#ifdef CRYPTO_TEST_SETUP
-#include "extension.h"
-/*
- * This extension command is similar to TPM2_GetRandom, but made
- * available for CRYPTO_TEST = 1 which disables TPM
- * Command structure, shared out of band with the test driver running
- * on the host:
- *
- * field     |    size  |                  note
- * ===================================================================
- * text_len  |    2     | size of the text to process, big endian
- */
-static enum vendor_cmd_rc trng_test(enum vendor_cmd_cc code, void *buf,
-				    size_t input_size, size_t *response_size)
+
+/* histogram at byte level */
+static uint32_t histogram[256];
+/* histogram at level of TRNG samples */
+static uint32_t histogram_trng[1 << TRNG_SAMPLE_BITS];
+
+static int command_rand(int argc, char **argv)
 {
-	uint16_t text_len;
-	uint8_t *cmd;
-	size_t response_room = *response_size;
+	int count = 1000; /* Default number of cycles. */
+	uint32_t val = 0, bits = 0;
+
+	if (argc == 2)
+		count = strtoi(argv[1], NULL, 10);
+
+	memset(histogram, 0, sizeof(histogram));
+	memset(histogram_trng, 0, sizeof(histogram_trng));
+	ccprintf("Retrieving %d 32-bit random words.\n", count);
+	while (count-- > 0) {
+		uint32_t rvalue;
+		int size;
+
+		rvalue = rand();
+		/* update byte-level histogram */
+		for (size = 0; size < sizeof(rvalue); size++)
+			histogram[((uint8_t *)&rvalue)[size]]++;
 
-	if (input_size != sizeof(text_len)) {
-		*response_size = 0;
-		return VENDOR_RC_BOGUS_ARGS;
+		/* update histogram on TRNG sample size level */
+		val = (val | (rvalue << bits)) & ((1 << TRNG_SAMPLE_BITS) - 1);
+		rvalue >>= TRNG_SAMPLE_BITS - bits;
+		bits += 32;
+		while (bits >= TRNG_SAMPLE_BITS) {
+			histogram_trng[val]++;
+			val = rvalue & ((1 << TRNG_SAMPLE_BITS) - 1);
+			rvalue >>= TRNG_SAMPLE_BITS;
+			bits -= TRNG_SAMPLE_BITS;
+		};
+
+		if (!(count % 10000))
+			watchdog_reload();
 	}
-	cmd = buf;
-	text_len = *cmd++;
-	text_len = text_len * 256 + *cmd++;
-	text_len = MIN(text_len, response_room);
-	rand_bytes(buf, text_len);
-	*response_size = text_len;
-	return VENDOR_RC_SUCCESS;
-}
+	ccprintf("Byte-level histogram:\n");
+	print_rand_stat(histogram, ARRAY_SIZE(histogram));
+	ccprintf("\nSample-level (%d bits) histogram:\n", TRNG_SAMPLE_BITS);
+	print_rand_stat(histogram_trng, ARRAY_SIZE(histogram_trng));
 
-DECLARE_VENDOR_COMMAND(VENDOR_CC_TRNG_TEST, trng_test);
+	return EC_SUCCESS;
+}
+DECLARE_SAFE_CONSOLE_COMMAND(rand, command_rand, NULL, NULL);
 
-#endif   /* CRYPTO_TEST_SETUP */
+#endif /* CRYPTO_TEST_SETUP */
diff --git a/test/tpm_test/nist_entropy.sh b/test/tpm_test/nist_entropy.sh
index f69c5652cd..5344a49890 100755
--- a/test/tpm_test/nist_entropy.sh
+++ b/test/tpm_test/nist_entropy.sh
@@ -3,19 +3,29 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-# NIST toolset needs sudo emerge dev-libs/libdivsufsort
-rm -rf /tmp/ea
-git clone --depth 1 https://github.com/usnistgov/SP800-90B_EntropyAssessment.git /tmp/ea/
-make -C /tmp/ea/cpp/ non_iid
-make -C /tmp/ea/cpp/ restart
-TRNG_OUT=/tmp/trng_output
-rm -f $TRNG_OUT
-./tpmtest.py -t
-if [ ! -f "$TRNG_OUT" ]; then
-    echo "$TRNG_OUT does not exist"
+# NIST toolset needs sudo emerge dev-libs/libdivsufsort and bz2
+set -e
+TMP_PATH="/tmp/ea"
+NIST_URL="https://github.com/usnistgov/SP800-90B_EntropyAssessment.git"
+TRNG_OUT="${TMP_PATH}/trng_output"
+EA_LOG="ea_non_iid.log"
+rm -rf "${TMP_PATH}"
+git clone --depth 1 "${NIST_URL}" "${TMP_PATH}"
+# build entropy assessment tool using mode for non independent and identically
+# distributed data, as for  H1 TRNG we can't justify the oppposite
+make -j -C "${TMP_PATH}/cpp/" non_iid restart
+rm -f "${TRNG_OUT}"
+# -t0 requests raw random samples from TRNG
+./tpmtest.py -t0 -o "${TRNG_OUT}"
+if [[ ! -f "${TRNG_OUT}" ]]; then
+    echo "${TRNG_OUT} does not exist"
     exit 1
 fi
-/tmp/ea/cpp/ea_non_iid -a $TRNG_OUT | tee ea_non_iid.log
-entropy=`grep min ea_non_iid.log | awk '{ print $5 }'`
-echo "Minimal entropy" $entropy
-/tmp/ea/cpp/ea_restart $TRNG_OUT $entropy | tee -a ea_non_iid.log
+rm -f "${EA_LOG}"
+"${TMP_PATH}/cpp/ea_non_iid" -a "${TRNG_OUT}" | tee "${EA_LOG}"
+entropy="$(awk '/min/ {print $5}' "${EA_LOG}")"
+if [[ -z "${entropy}" ]]; then
+    entropy="$(awk '/H_original/ {print $2}' "${EA_LOG}")"
+fi
+echo "Minimal entropy ${entropy}"
+"${TMP_PATH}/cpp/ea_restart" "${TRNG_OUT}" "${entropy}" | tee -a "${EA_LOG}"
diff --git a/test/tpm_test/tpmtest.py b/test/tpm_test/tpmtest.py
index 0ec33ce8f1..04567d67cb 100755
--- a/test/tpm_test/tpmtest.py
+++ b/test/tpm_test/tpmtest.py
@@ -148,34 +148,47 @@ class TPM:
 
 def usage():
     """Print usage information"""
-    print('Syntax: tpmtest.py [-d | -t | -h ]\n'
-          '     -d   -  prints additional debug information during tests\n'
-          '     -t   -  dump raw output from TRNG to /tmp/trng_output\n'
-          '     -h   -  this help\n')
+    print('Syntax: tpmtest.py [-d] | [-t source [-o file] [-s bits] ]| -h ]\n'
+          '     -d - prints additional debug information during tests\n'
+          '     -t source - only dump raw output from TRNG. source values:\n'
+          '        0 - raw TRNG'
+          '        [-o file] - set output file, default /tmp/trng_output\n'
+          '        [-s bits] - TRNG sample size in bit, default = 1\n'
+          '     -h - this help\n')
 
 def main():
     """Run TPM tests"""
     try:
-        opts, _ = getopt.getopt(sys.argv[1:], 'dth', 'help')
+        opts, _ = getopt.getopt(sys.argv[1:], 'dt:hs:o:', 'help')
     except getopt.GetoptError as err:
         print(str(err))
         usage()
         sys.exit(2)
     debug_needed = False
     trng_only = False
-    for option, _ in opts:
+    trng_output = '/tmp/trng_output'
+    trng_sample_bits = 1
+    trng_mode = 0
+
+    for option, arg in opts:
         if option == '-d':
             debug_needed = True
         elif option == '-t':
             trng_only = True
+            trng_mode = int(arg)
+        elif option == '-o':
+            trng_output = arg
+        elif option == '-s':
+            trng_sample_bits = int(arg)
         elif option in ('-h', '--help'):
             usage()
             sys.exit(0)
     try:
         tpm_object = TPM(debug_mode=debug_needed)
         if trng_only:
-            trng_test.trng_test(tpm_object)
-            sys.exit(1)
+            trng_test.trng_test(tpm_object, trng_output,
+                                trng_mode, trng_sample_bits)
+            sys.exit(0)
         crypto_test.crypto_tests(tpm_object, os.path.join(ROOT_DIR,
                                                           'crypto_test.xml'))
         drbg_test.drbg_test(tpm_object)
diff --git a/test/tpm_test/trng_test.py b/test/tpm_test/trng_test.py
index c4d9395e80..60faa32e99 100644
--- a/test/tpm_test/trng_test.py
+++ b/test/tpm_test/trng_test.py
@@ -3,21 +3,27 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 """Tests for trng."""
-from __future__ import print_function
-import struct
 
+from math import gcd
+import struct
 import subcmd
 import utils
 
-TRNG_TEST_FMT = '>H'
+TRNG_TEST_FMT = '>HB'
 TRNG_TEST_RSP_FMT = '>H2IH'
 TRNG_TEST_CC = 0x33
+
 TRNG_SAMPLE_SIZE = 1000 # minimal recommended by NIST is 1000 bytes per sample
-TRNG_SAMPLE_COUNT = 1000 # NIST require at least 1000000 of 8-bit samples
+TRNG_SAMPLE_COUNT = 1000000 # NIST require at least 1000000 of 8-bit samples
 
-def get_random_command(size):
+# Command structure, shared out of band with the test running on the target:
+# field     |    size  |                  note
+# ===================================================================
+# text_len  |    2     | number of bytes to read, big endian
+# type      |    1     | 0 = TRNG, other values reserved for extensions
+def get_random_command(size, trng_op):
     """Encode get_random command"""
-    return struct.pack(TRNG_TEST_FMT, size)
+    return struct.pack(TRNG_TEST_FMT, size, trng_op)
 
 def get_random_command_rsp(size):
     """Create expected response to get_random"""
@@ -25,29 +31,97 @@ def get_random_command_rsp(size):
                        struct.calcsize(TRNG_TEST_RSP_FMT) + size,
                        0, TRNG_TEST_CC)
 
+def to_bitstring(s, n=1):
+    """Split bytes into individual samples
 
-def trng_test(tpm):
-    """Download entropy samples from TRNG
-
-    Command structure, shared out of band with the test running on the target:
+    convert input packed byte array to n-bits in a byte representation
+    used by NIST tests. It's designed to recover sequence of samples
+    as it comes from TRNG, including the fact that rand_bytes() reverse
+    byte order in every 32-bit chunk.
+    """
+    out = b''
+    val_left = 0
+    bits_left = 0
+    while s:
+        val = (struct.unpack('>I', s[0:4].rjust(4, b'\0'))[0] << bits_left) +\
+              val_left
+        bits_left += 8 * len(s[0:4])
+        s = s[4:]
+        while bits_left >= n:
+            out += struct.pack('B', val & ((1 << n) - 1))
+            val >>= n
+            bits_left -= n
+        val_left = val
+    return out
 
-    field     |    size  |                  note
-    ===================================================================
-    text_len  |    2     | size of the text to process, big endian
+def trng_test(tpm, trng_output, trng_mode, tsb=1):
+    """Download entropy samples from TRNG
 
     Args:
         tpm: a tpm object used to communicate with the device
+        trng_output: file name
+        trng_mode: source of randomness [0 - TRNG]
+        tsb: number of bits in each TRNG sample, should be in sync with
+             TRNG configuration. Default is 1 bit per sample.
 
     Raises:
         subcmd.TpmTestError: on unexpected target responses
     """
-    with open('/tmp/trng_output', 'wb') as out_file:
-        for block in range(0, TRNG_SAMPLE_COUNT):
+
+    if trng_mode not in [0]:
+        raise subcmd.TpmTestError('Unknown random source: %d' % trng_mode)
+
+    # minimal recommended by NIST is 1000 samples per block
+    # TRNG_SAMPLE_BITS is internal setting for TRNG which is important for
+    # entropy analysis. TRNG internally gets a 16 bit sample (measurement of
+    # time to collapse ring oscillator). Then some slice of these bits
+    # [0:TRNG_SAMPLE_BITS] is added (packed) to TRNG FIFO which has internal
+    # size of 2048 bits. Reading from TRNG always return 32 bits from FIFO.
+    # To extract actual samples from packed 32-bit reading, need to reverse
+    # and split read 32-bit into individual samples, each size TRNG_SAMPLE_BITS
+    # bits. As it's only possible to read 32 bits at once and TRNG_SAMPLE_BITS
+    # can be anything from 1 to 16, including non-power of 2, need to ensure
+    # readings to result in non-fractional number of samples. At the same time
+    # NIST requires minimum 1000 samples at once, and we also want to reduce
+    # number of read requests which are adding overhead.
+    # this variable should be divisible by 4 to match 32bit reads from TRNG
+
+    if not 8 >= tsb > 0:
+        raise subcmd.TpmTestError('NIST only supports 1 to 8 bits per sample')
+
+    # compute number of bytes, which is multiple of 4 containing whole number
+    # of samples, each size tsb bits. This a reduction of
+    # tsb * 32 // gcd(tsb, 32) // 8
+    lcm = (tsb * 4) // gcd(tsb, 4)
+
+    # combine reads of small batches if possible, 2032 is max size of TPM
+    # command response, adjusted to header size
+    bytes_per_read = (2032 // lcm) * lcm
+
+    samples_per_read = 8 * bytes_per_read // tsb
+
+    if samples_per_read < 1000:
+        raise subcmd.TpmTestError("Can't meet NIST requirement of min 1000 "
+                                  'samples in batch - %d bytes only contain'
+                                  ' %d samples' % (bytes_per_read,
+                                  samples_per_read))
+    print('TRNG bits per sample = ', tsb, 'Read size (bytes) =',
+          bytes_per_read, 'Samples per read =', samples_per_read)
+
+    remaining_samples = TRNG_SAMPLE_COUNT
+    with open(trng_output, 'wb') as out_file:
+        while remaining_samples:
             response = tpm.command(tpm.wrap_ext_command(TRNG_TEST_CC,
-                                   get_random_command(TRNG_SAMPLE_SIZE)))
-            if response[:12] != get_random_command_rsp(TRNG_SAMPLE_SIZE):
-                raise subcmd.TpmTestError("Unexpected response to \'%s\': %s" %
-                                          ('trng', utils.hex_dump(response)))
-            out_file.write(response[12:])
-            print('%s %d%%\r' % (utils.cursor_back(), (block//10)), end='')
-    print('%sSUCCESS: %s' % (utils.cursor_back(), 'trng'))
+                                            get_random_command(bytes_per_read,
+                                                               trng_mode)))
+            if response[:12] != get_random_command_rsp(bytes_per_read):
+                raise subcmd.TpmTestError("Unexpected response to '%s': %s" %
+                                        ('trng', utils.hex_dump(response)))
+            bits = to_bitstring(response[12:], tsb)
+            bits = bits[:remaining_samples]
+            out_file.write(bits)
+            remaining_samples -= len(bits)
+            print('%s %d%%\r' % (utils.cursor_back(),
+                                ((TRNG_SAMPLE_COUNT - remaining_samples)*100)\
+                                // TRNG_SAMPLE_COUNT), end='')
+    print('%sSUCCESS: %s' % (utils.cursor_back(), trng_output))