fips_rand: FIPS-compliant way to generate randoms

Add proper TRNG health tests and CR50-wide DRBG with reseeding

BUG=b:138578157
TEST=tpmtest.py -t1 fails after cr50 reboot.
rand_perf in console (kick-off FIPS TRNG test) and then
tpmtest.py -t1 and tpmtest.py -t2 should succeed.

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

6 files changed, 583 insertions, 3 deletions

diff --git a/board/cr50/build.mk b/board/cr50/build.mk
index 0396ad6c17..a268f47602 100644
--- a/board/cr50/build.mk
+++ b/board/cr50/build.mk
@@ -48,6 +48,7 @@ board-y += power_button.o
 board-y += servo_state.o
 board-y += ap_uart_state.o
 board-y += factory_mode.o
+board-y += fips_rand.o
 board-${CONFIG_RDD} += rdd.o
 board-${CONFIG_USB_SPI} += usb_spi.o
 board-${CONFIG_USB_I2C} += usb_i2c.o
diff --git a/board/cr50/fips_rand.c b/board/cr50/fips_rand.c
new file mode 100644
index 0000000000..1111e009de
--- /dev/null
+++ b/board/cr50/fips_rand.c
@@ -0,0 +1,453 @@
+/* Copyright 2020 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#include "console.h"
+#include "cryptoc/util.h"
+#include "fips_rand.h"
+#include "flash_log.h"
+#include "init_chip.h"
+#include "registers.h"
+#include "task.h"
+#include "timer.h"
+#include "trng.h"
+#include "util.h"
+
+/**
+ * FIPS-compliant CR50-wide DRBG
+ * Since raw TRNG input shouldn't be used as random number generator,
+ * all FIPS-compliant code use DRBG, seeded from TRNG
+ */
+static struct drbg_ctx fips_drbg;
+
+#define ENTROPY_SIZE_BITS  512
+#define ENTROPY_SIZE_WORDS (BITS_TO_WORDS(ENTROPY_SIZE_BITS))
+
+/**
+ * buffer for entropy condensing. initialized during
+ * fips_trng_startup(), but also used in KAT tests,
+ * thus size is enough to accommodate needs
+ */
+static uint32_t entropy_fifo[ENTROPY_SIZE_WORDS];
+
+/**
+ * NIST FIPS TRNG health tests (NIST SP 800-90B 4.3)
+ * If any of the approved continuous health tests are used by the entropy
+ * source, the false positive probability for these tests shall be set to
+ * at least 2^-50
+ */
+
+/* dummy function, will be removed with FIPS framework */
+static int fips_crypto_allowed(void)
+{
+	return 1;
+}
+
+/**
+ * rand() should be able to return error code if reading from TRNG failed
+ * return as struct with 2 params is more efficient as data is passed in
+ * registers
+ */
+struct rand_result {
+	uint32_t random_value;
+	bool valid;
+};
+
+
+/* state data for TRNG health test */
+static struct {
+	uint32_t count; /* number of 1-bits in APT test window */
+	int last_clz; /* running count of leading 0 bits */
+	int last_clo; /* running count of leading 1 bits */
+	uint8_t pops[APT_WINDOW_SIZE_NWORDS];
+	uint8_t rct_count; /* current windows size for RCT */
+	uint8_t oldest; /* position in APT window */
+	bool apt_initialized; /* flag APT window is filled with data */
+	bool drbg_initialized; /* flag DRBG is initialized */
+} rand_state;
+
+/**
+ * NIST SP 800-90B 4.4.1
+ * The repetition count test detects abnormal runs of 0s or 1s.
+ * RCT_CUTOFF_BITS must be >= 32.
+ * If a single value appears more than 100/H times in a row,
+ * the tests must detect this with high probability.
+ *
+ * This implementation assumes TRNG is configured to produce 1-bit
+ * readings, packed into 32-bit words.
+ * @return false if test failed
+ */
+static bool repetition_count_test(uint32_t rnd)
+{
+	uint32_t clz, ctz, clo, cto;
+	/* count repeating 0 and 1 bits from each side */
+	clz = __builtin_clz(rnd);  /* # of leading 0s */
+	ctz = __builtin_ctz(rnd);  /* # of trailing 0s */
+	clo = __builtin_clz(~rnd); /* # of leading 1s */
+	cto = __builtin_ctz(~rnd); /* # of trailing 1s */
+
+	/**
+	 * check that number of trailing 0/1 in current sample added to
+	 * leading 0/1 of previous sample is less than cut off value, so
+	 * we don't have long repetitive series of 0s or 1s
+	 */
+	if ((ctz + rand_state.last_clz >= RCT_CUTOFF_SAMPLES) ||
+	    (cto + rand_state.last_clo >= RCT_CUTOFF_SAMPLES))
+		return false;
+
+	/**
+	 * merge series of repetitive values - update running counters in
+	 * such way that if current sample is all 0s then add previous
+	 * counter of zeros to current number which will be 32,
+	 * otherwise (we had 1s) - just use current value. Same for 1s
+	 */
+	if (rnd == 0) /* if all 32 samples are 0s */
+		clz += rand_state.last_clz;
+
+	if (rnd == ~0) /* if all 32 samples are 1s */
+		clo += rand_state.last_clo;
+	rand_state.last_clz = clz;
+	rand_state.last_clo = clo;
+
+	/* check we collected enough bits for statistics */
+	if (rand_state.rct_count < RCT_CUTOFF_WORDS)
+		++rand_state.rct_count;
+	return true;
+}
+
+static int misbalanced(uint32_t count)
+{
+	return count > APT_CUTOFF_SAMPLES ||
+	       count < APT_WINDOW_SIZE_BITS - APT_CUTOFF_SAMPLES;
+}
+
+/* Returns number of set bits (1s) in 32-bit word */
+static int popcount(uint32_t x)
+{
+	x = x - ((x >> 1) & 0x55555555);
+	x = (x & 0x33333333) + ((x >> 2) & 0x33333333);
+	x = (x + (x >> 4)) & 0x0F0F0F0F;
+	x = x + (x >> 8);
+	x = x + (x >> 16);
+	return x & 0x0000003F;
+}
+
+/**
+ * NIST SP 800-90B 4.4.2 Adaptive Proportion Test.
+ * Implementation for 1-bit alphabet.
+ * Instead of storing actual samples we can store pop counts
+ * of each 32bit reading, which would fit in 8-bit.
+ */
+bool adaptive_proportion_test(uint32_t rnd)
+{
+	/* update rolling count */
+	rand_state.count -= rand_state.pops[rand_state.oldest];
+	/**
+	 * Since we have 1-bit samples, in order to have running window to
+	 * count ratio of 1s and 0s in it we can just store number of 1s in
+	 * each 32-bit sample which requires 1 byte vs. 4 bytes.
+	 * number of zeros  =  32  - number of 1s
+	 */
+	rand_state.pops[rand_state.oldest] = popcount(rnd);
+
+	/* update rolling count with current sample statistics */
+	rand_state.count += rand_state.pops[rand_state.oldest];
+	if (++rand_state.oldest >= APT_WINDOW_SIZE_NWORDS) {
+		rand_state.apt_initialized = true; /* saw full window */
+		rand_state.oldest = 0;
+	}
+	/* check when initialized */
+	if (rand_state.apt_initialized && misbalanced(rand_state.count))
+		return false;
+	return true;
+}
+
+static bool fips_powerup_passed(void)
+{
+	return rand_state.rct_count >= RCT_CUTOFF_WORDS &&
+	       rand_state.apt_initialized;
+}
+
+/**
+ * Attempts to read TRNG_EMPTY before reporting a stall.
+ * Practically data should be available in less than 777
+ * cycles under normal conditions. Give 4 attempts to
+ * reset before making decision TRNG is broken
+ */
+#define TRNG_EMPTY_COUNT 777
+#define TRNG_RESET_COUNT 4
+
+/**
+ * replica of rand() with interface which returns errors properly
+ */
+static struct rand_result read_rand(void)
+{
+	uint32_t empty_count = 0;
+	uint32_t reset_count = 0;
+	/**
+	 * make sure we never hang in the loop - try at max 1
+	 * reset attempt, then return error
+	 */
+	while (GREAD(TRNG, EMPTY) && (reset_count < TRNG_RESET_COUNT)) {
+		if (GREAD_FIELD(TRNG, FSM_STATE, FSM_IDLE) ||
+		    empty_count > TRNG_EMPTY_COUNT) {
+			/* TRNG timed out, restart */
+			GWRITE(TRNG, STOP_WORK, 1);
+			flash_log_add_event(FE_LOG_TRNG_STALL, 0, NULL);
+			GWRITE(TRNG, GO_EVENT, 1);
+			empty_count = 0;
+			reset_count++;
+		}
+		empty_count++;
+	}
+	return (struct rand_result){ .random_value = GREAD(TRNG, READ_DATA),
+				     .valid =
+					     (reset_count < TRNG_RESET_COUNT) };
+}
+
+/**
+ * get random from TRNG and run continuous health tests.
+ * it is also can simulate stuck-bit error
+ * @param power_up if non-zero indicates warm-up mode
+ * @return random value from TRNG
+ */
+static struct rand_result fips_trng32(int power_up)
+{
+	struct rand_result r;
+
+	/* Continuous health tests should have been initialized by now */
+	if (!fips_crypto_allowed() || (!power_up && !fips_powerup_passed()))
+		return (struct rand_result){ .random_value = 0,
+					     .valid = false };
+
+	/* get noise */
+	r = read_rand();
+
+	if (r.valid)
+		/* test #1: Repetition Count Test (a.k.a Stuck-bit) */
+		if (!repetition_count_test(r.random_value) ||
+		    /* warm-up test #2: Adaptive Proportion Test */
+		    !adaptive_proportion_test(r.random_value))
+			r.valid = false; /* TODO: add FIPS status update */
+	return r;
+}
+
+bool fips_trng_bytes(void *buffer, size_t len)
+{
+	uint8_t *buf = (uint8_t *)buffer;
+	uint32_t random_togo = 0;
+	struct rand_result r;
+	/**
+	 * Retrieve random numbers in 4 byte quantities and pack as many bytes
+	 * as needed into 'buffer'. If len is not divisible by 4, the
+	 * remaining random bytes get dropped.
+	 */
+	while (len--) {
+		if (!random_togo) {
+			r = fips_trng32(0);
+			if (!r.valid)
+				return false;
+			random_togo = sizeof(r.random_value);
+		}
+		*buf++ = (uint8_t)r.random_value;
+		random_togo--;
+		r.random_value >>= 8;
+	}
+	return true;
+}
+
+/* FIPS TRNG power-up tests */
+bool fips_trng_startup(int stage)
+{
+	if (!stage) {
+		/**
+		 * To hide TRNG latency, split it into 2 stages.
+		 * at stage 0, initialize variables
+		 */
+		rand_state.rct_count = 0;
+		rand_state.apt_initialized = 0;
+	}
+	/* Startup tests per NIST SP800-90B, Section 4 */
+	/* 4096 1-bit samples, in 2 steps, 2048 bit in each */
+	for (uint32_t i = 0; i < (TRNG_INIT_WORDS) / 2; i++) {
+		struct rand_result r = fips_trng32(1);
+
+		if (!r.valid)
+			return false;
+		/* store entropy for further use */
+		entropy_fifo[i % ARRAY_SIZE(entropy_fifo)] = r.random_value;
+	}
+	return true;
+}
+
+bool fips_drbg_init(void)
+{
+	struct rand_result nonce;
+
+	if (!fips_crypto_allowed())
+		return EC_ERROR_INVALID_CONFIG;
+
+	/**
+	 * initialize DRBG with 440 bits of entropy as required
+	 * by NIST SP 800-90A 10.1. Includes entropy and nonce,
+	 * both received from entropy source.
+	 * entropy_fifo contains 512 bits of noise with H>= 0.85
+	 * this is roughly equal to 435 bits of full entropy.
+	 * Add 32 * 0.85 = 27 bits from nonce.
+	 */
+	nonce = fips_trng32(0);
+	if (!nonce.valid)
+		return false;
+
+	/* read another 512 bits of noise */
+	if (!fips_trng_bytes(&entropy_fifo, sizeof(entropy_fifo)))
+		return false;
+
+	hmac_drbg_init(&fips_drbg, &entropy_fifo, sizeof(entropy_fifo),
+		       &nonce.random_value, sizeof(nonce.random_value), NULL,
+		       0);
+
+	rand_state.drbg_initialized = 1;
+	return true;
+}
+
+/* zeroize DRBG state */
+void fips_drbg_clear(void)
+{
+	drbg_exit(&fips_drbg);
+	rand_state.drbg_initialized = 0;
+}
+
+enum hmac_result fips_hmac_drbg_generate_reseed(struct drbg_ctx *ctx, void *out,
+						size_t out_len,
+						const void *input,
+						size_t input_len)
+{
+	enum hmac_result err =
+		hmac_drbg_generate(ctx, out, out_len, input, input_len);
+
+	while (err == HMAC_DRBG_RESEED_REQUIRED) {
+		/* read another 512 bits of noise */
+		if (!fips_trng_bytes(&entropy_fifo, sizeof(entropy_fifo))) {
+			/* TODO: Report FIPS error properly */
+			ccprintf("FIPS trng bytes failed\n");
+			return HMAC_DRBG_RESEED_REQUIRED;
+		}
+
+		hmac_drbg_reseed(ctx, entropy_fifo, sizeof(entropy_fifo), NULL,
+				 0, NULL, 0);
+		err = hmac_drbg_generate(ctx, out, out_len, input, input_len);
+	}
+	if (err != HMAC_DRBG_SUCCESS)
+		ccprintf("DRBG error %d\n", err);
+	return err;
+}
+
+bool fips_rand_bytes(void *buffer, size_t len)
+{
+	if (!fips_crypto_allowed())
+		return false;
+	/**
+	 * make sure cr50 DRBG is initialized after power-on or resume,
+	 * but do it on first use to minimize latency of board_init()
+	 */
+	if (!rand_state.drbg_initialized && !fips_drbg_init())
+		return false;
+
+	/* HMAC_DRBG can only return up to 7500 bits in a single request */
+	while (len) {
+		size_t request = (len > (7500 / 8)) ? (7500 / 8) : len;
+
+		if (fips_hmac_drbg_generate_reseed(&fips_drbg, buffer, request,
+						   NULL,
+						   0) != HMAC_DRBG_SUCCESS)
+			return false;
+		len -= request;
+		buffer += request;
+	}
+	return true;
+}
+
+/* return codes match dcrypto_p256_ecdsa_sign */
+int fips_p256_ecdsa_sign(const p256_int *key, const p256_int *message,
+			 p256_int *r, p256_int *s)
+{
+	if (!fips_crypto_allowed())
+		return 0;
+	if (!rand_state.drbg_initialized) {
+		int err;
+
+		err = fips_drbg_init();
+		if (err)
+			return 0;
+	}
+	return dcrypto_p256_ecdsa_sign(&fips_drbg, key, message, r, s);
+}
+
+#ifdef CRYPTO_TEST_SETUP
+#include "endian.h"
+#include "extension.h"
+#include "trng.h"
+#include "watchdog.h"
+
+static int cmd_rand_perf(int argc, char **argv)
+{
+	uint64_t starttime;
+	static uint32_t buf[SHA256_DIGEST_WORDS];
+	int j, k;
+
+	starttime = get_time().val;
+
+	/* run power-up tests to measure raw performance */
+	if (fips_trng_startup(0) && fips_trng_startup(1)) {
+		starttime = get_time().val - starttime;
+		ccprintf("time for fips_trng_startup = %llu\n", starttime);
+	} else
+		ccprintf("TRNG power up test failed\n");
+
+	cflush();
+
+	starttime = get_time().val;
+	fips_drbg_init();
+	starttime = get_time().val - starttime;
+	ccprintf("time for drbg_init = %llu\n", starttime);
+	cflush();
+	starttime = get_time().val;
+	for (k = 0; k < 10; k++) {
+		for (j = 0; j < 100; j++)
+			fips_rand_bytes(buf, sizeof(buf));
+		watchdog_reload();
+		cflush();
+	}
+	starttime = get_time().val - starttime;
+	ccprintf("time for 1000 drbg reads = %llu\n", starttime);
+	cflush();
+
+	starttime = get_time().val;
+	for (k = 0; k < 10; k++) {
+		for (j = 0; j < 100; j++)
+			rand_bytes(&buf, sizeof(buf));
+		watchdog_reload();
+	}
+	starttime = get_time().val - starttime;
+	ccprintf("time for 1000 rand_byte() = %llu\n", starttime);
+	cflush();
+
+	starttime = get_time().val;
+	for (k = 0; k < 10; k++) {
+		for (j = 0; j < 100; j++)
+			if (!fips_trng_bytes(&buf, sizeof(buf)))
+				ccprintf("FIPS TRNG error\n");
+		watchdog_reload();
+	}
+	starttime = get_time().val - starttime;
+	ccprintf("time for 1000 fips_trng_byte() = %llu\n", starttime);
+	cflush();
+
+	return 0;
+}
+
+DECLARE_SAFE_CONSOLE_COMMAND(rand_perf, cmd_rand_perf, NULL, NULL);
+
+#endif /* CRYPTO_TEST_SETUP */
diff --git a/board/cr50/fips_rand.h b/board/cr50/fips_rand.h
new file mode 100644
index 0000000000..14eba3fe32
--- /dev/null
+++ b/board/cr50/fips_rand.h
@@ -0,0 +1,111 @@
+/* Copyright 2020 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+#ifndef __EC_BOARD_CR50_FIPS_RAND_H
+#define __EC_BOARD_CR50_FIPS_RAND_H
+
+#include <stddef.h>
+#include <string.h>
+
+#include "common.h"
+#include "dcrypto.h"
+#include "util.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#define TRNG_SAMPLE_BITS 1
+
+/**
+ * TRNG Health Tests
+ *
+ * If any of the approved continuous health tests are used by the entropy
+ * source, the false positive probability for these tests shall be set to
+ * at least 2^(-50)  (NIST SP 800-90B 4.3).
+ * Reason for 2^(-50) vs 2^(-40) is to minimize impact to user experience
+ * due to false positives.
+ *
+ * For H1 minimal assessed entropy H >=0.85 for 1-bit samples
+ * using NIST Entropy Assessment tool.
+ */
+
+/**
+ * The entropy source's startup tests shall run the continuous health tests
+ * over at least 4096 consecutive samples. We use 1-bit samples
+ */
+#define TRNG_INIT_BITS	(4096 * TRNG_SAMPLE_BITS)
+#define TRNG_INIT_WORDS (BITS_TO_WORDS(TRNG_INIT_BITS))
+
+/**
+ * (1) Repetition Count Test (RCT) NIST SP 800-90B 4.4.1
+ * Cut off value is computed as:
+ * c = ceil(1 + (-log2 alpha)/H);
+ * alpha = 2^-50, H = 0.85; RCT_CUTOFF = CEIL(1+(50/0.85))
+ */
+#define RCT_CUTOFF_SAMPLES 60
+
+/**
+ * Number of 32-bit words containing RCT_CUTOFF_SAMPLES samples
+ */
+#define RCT_CUTOFF_WORDS (BITS_TO_WORDS(RCT_CUTOFF_SAMPLES))
+
+/**
+ * (2) Adaptive Proportion Test (APT), NIST SP 800-90B 4.4.2, Table 2
+ */
+#if TRNG_SAMPLE_BITS == 1
+/* APT Windows size W = 1024 for 1 bit samples */
+#define APT_WINDOW_SIZE_SAMPLES 1024
+#else
+/* or 512 samples if more than 1 bit per sample */
+#define APT_WINDOW_SIZE_SAMPLES 512
+#endif
+#define APT_WINDOW_SIZE_BITS   (APT_WINDOW_SIZE_SAMPLES * TRNG_SAMPLE_BITS)
+#define APT_WINDOW_SIZE_NWORDS (BITS_TO_WORDS(APT_WINDOW_SIZE_BITS))
+/**
+ * Cut off value = CRITBINOM(W, power(2,(-H)),1-α).
+ * 692 = CRITBINOM(1024, power(2,(-0.85)), 1 - 2^(-50))
+ */
+#define APT_CUTOFF_SAMPLES 692
+
+/**
+ * FIPS-compliant TRNG startup.
+ * The entropy source's startup tests shall run the continuous health tests
+ * over at least 4096 consecutive samples.
+ * Note: This function can throw FIPS_FATAL_TRNG error
+ *
+ * To hide latenccy of reading TRNG data, this test is executed in 2 stages
+ * @param stage is 0 or 1, choosing the stage. On each stage 2048
+ * samples are processed. Assuming that some other tasks can be executed
+ * between stages, when TRNG FIFO if filled with samples.
+ *
+ * Some number of samples will be available in entropy_fifo
+ */
+bool fips_trng_startup(int stage);
+
+bool fips_trng_bytes(void *buffer, size_t len);
+
+/* initialize cr50-wide DRBG replacing rand */
+bool fips_drbg_init(void);
+/* mark cr50-wide DRBG as not initialized */
+void fips_drbg_init_clear(void);
+
+/* random bytes using FIPS-compliant HMAC_DRBG */
+bool fips_rand_bytes(void *buffer, size_t len);
+
+/* wrapper around dcrypto_p256_ecdsa_sign using FIPS-compliant HMAC_DRBG */
+int fips_p256_ecdsa_sign(const p256_int *key, const p256_int *message,
+			 p256_int *r, p256_int *s);
+/**
+ * wrapper around hmac_drbg_generate to automatically reseed drbg
+ * when needed.
+ */
+enum hmac_result fips_hmac_drbg_generate_reseed(struct drbg_ctx *ctx, void *out,
+						size_t out_len,
+						const void *input,
+						size_t input_len);
+#ifdef __cplusplus
+}
+#endif
+#endif /* ! __EC_BOARD_CR50_FIPS_RAND_H */
diff --git a/board/cr50/tpm2/trng.c b/board/cr50/tpm2/trng.c
index ae4312be2c..87519b0e85 100644
--- a/board/cr50/tpm2/trng.c
+++ b/board/cr50/tpm2/trng.c
@@ -13,6 +13,7 @@ CRYPT_RESULT _cpri__StirRandom(int32_t num, uint8_t *entropy)
 #ifdef CRYPTO_TEST_SETUP
 #include "endian.h"
 #include "extension.h"
+#include "fips_rand.h"
 #include "trng.h"
 /*
  * This extension command is similar to TPM2_GetRandom, but made
@@ -23,7 +24,8 @@ CRYPT_RESULT _cpri__StirRandom(int32_t num, uint8_t *entropy)
  * field     |    size  |                  note
  * =========================================================================
  * text_len  |    2     | the number of random bytes to generate, big endian
- * type      |    1     | 0 = TRNG, other values reserved for extensions
+ * type      |    1     | 0 = TRNG, 1 = FIPS TRNG, 2 = FIPS DRBG
+ *           |          | 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)
@@ -49,6 +51,15 @@ static enum vendor_cmd_rc trng_test(enum vendor_cmd_cc code, void *buf,
 	case 0:
 		rand_bytes(buf, text_len);
 		break;
+	case 1:
+		if (!fips_trng_bytes(buf, text_len))
+			return VENDOR_RC_INTERNAL_ERROR;
+		break;
+	case 2:
+		if (!fips_rand_bytes(buf, text_len))
+			return VENDOR_RC_INTERNAL_ERROR;
+		break;
+
 	default:
 		return VENDOR_RC_BOGUS_ARGS;
 	}
diff --git a/include/util.h b/include/util.h
index a6de67e3d8..43d9a71992 100644
--- a/include/util.h
+++ b/include/util.h
@@ -66,6 +66,9 @@ extern "C" {
 #define DIV_ROUND_UP(x, y) (((x) + ((y) - 1)) / (y))
 #define DIV_ROUND_NEAREST(x, y) (((x) + ((y) / 2)) / (y))
 
+/* Convert number of bits to number of 32-bit words */
+#define BITS_TO_WORDS(x) (((x) + 31) / 32)
+
 /*
  * Swap two variables (requires c99)
  *
diff --git a/test/tpm_test/trng_test.py b/test/tpm_test/trng_test.py
index 60faa32e99..bdf0477a32 100644
--- a/test/tpm_test/trng_test.py
+++ b/test/tpm_test/trng_test.py
@@ -20,7 +20,8 @@ TRNG_SAMPLE_COUNT = 1000000 # NIST require at least 1000000 of 8-bit samples
 # field     |    size  |                  note
 # ===================================================================
 # text_len  |    2     | number of bytes to read, big endian
-# type      |    1     | 0 = TRNG, other values reserved for extensions
+# type      |    1     | 0 = TRNG, 1 = FIPS TRNG, 2 = FIPS DRBG
+#           |          | other values reserved for extensions
 def get_random_command(size, trng_op):
     """Encode get_random command"""
     return struct.pack(TRNG_TEST_FMT, size, trng_op)
@@ -68,7 +69,7 @@ def trng_test(tpm, trng_output, trng_mode, tsb=1):
         subcmd.TpmTestError: on unexpected target responses
     """
 
-    if trng_mode not in [0]:
+    if trng_mode not in [0, 1, 2]:
         raise subcmd.TpmTestError('Unknown random source: %d' % trng_mode)
 
     # minimal recommended by NIST is 1000 samples per block