3 files changed, 56 insertions, 17 deletions

diff --git a/sys/dev/tpm/tpm20.c b/sys/dev/tpm/tpm20.c
index 067e7ccae8f9..6c587818058d 100644
--- a/sys/dev/tpm/tpm20.c
+++ b/sys/dev/tpm/tpm20.c
@@ -42,7 +42,7 @@
 MALLOC_DEFINE(M_TPM20, "tpm_buffer", "buffer for tpm 2.0 driver");
 
 static void tpm20_discard_buffer(void *arg);
-#ifdef TPM_HARVEST
+#if defined TPM_HARVEST || defined RANDOM_ENABLE_TPM
 static void tpm20_harvest(void *arg, int unused);
 #endif
 static int  tpm20_save_state(device_t dev, bool suspend);
@@ -184,7 +184,7 @@ tpm20_ioctl(struct cdev *dev, u_long cmd, caddr_t data,
 	return (ENOTTY);
 }
 
-#ifdef TPM_HARVEST
+#if defined TPM_HARVEST || defined RANDOM_ENABLE_TPM
 static const struct random_source random_tpm = {
 	.rs_ident = "TPM",
 	.rs_source = RANDOM_PURE_TPM,
@@ -212,7 +212,7 @@ tpm20_init(struct tpm_sc *sc)
 	if (result != 0)
 		tpm20_release(sc);
 
-#ifdef TPM_HARVEST
+#if defined TPM_HARVEST || defined RANDOM_ENABLE_TPM
 	random_source_register(&random_tpm);
 	TIMEOUT_TASK_INIT(taskqueue_thread, &sc->harvest_task, 0,
 	    tpm20_harvest, sc);
@@ -227,7 +227,7 @@ void
 tpm20_release(struct tpm_sc *sc)
 {
 
-#ifdef TPM_HARVEST
+#if defined TPM_HARVEST || defined RANDOM_ENABLE_TPM
 	if (device_is_attached(sc->dev))
 		taskqueue_drain_timeout(taskqueue_thread, &sc->harvest_task);
 	random_source_deregister(&random_tpm);
@@ -254,7 +254,7 @@ tpm20_shutdown(device_t dev)
 	return (tpm20_save_state(dev, false));
 }
 
-#ifdef TPM_HARVEST
+#if defined TPM_HARVEST || defined RANDOM_ENABLE_TPM
 /*
  * Get TPM_HARVEST_SIZE random bytes and add them
  * into system entropy pool.
diff --git a/sys/dev/tpm/tpm20.h b/sys/dev/tpm/tpm20.h
index 7c2ccd30143a..b2cfcd4f25bd 100644
--- a/sys/dev/tpm/tpm20.h
+++ b/sys/dev/tpm/tpm20.h
@@ -128,7 +128,7 @@ struct tpm_sc {
 	lwpid_t		owner_tid;
 
 	struct callout 	discard_buffer_callout;
-#ifdef TPM_HARVEST
+#if defined TPM_HARVEST || defined RANDOM_ENABLE_TPM
 	struct timeout_task 	harvest_task;
 #endif
 
diff --git a/sys/dev/tpm/tpm_crb.c b/sys/dev/tpm/tpm_crb.c
index 017ebd45c7ea..28b4f21eccfb 100644
--- a/sys/dev/tpm/tpm_crb.c
+++ b/sys/dev/tpm/tpm_crb.c
@@ -301,6 +301,48 @@ tpmcrb_cancel_cmd(struct tpm_sc *sc)
 	return (true);
 }
 
+static bool
+tpmcrb_state_idle(struct tpmcrb_sc *crb_sc, bool wait)
+{
+	struct tpm_sc *sc;
+	int mask, timeout;
+
+	timeout = wait ? TPM_TIMEOUT_C : 0;
+
+	sc = &crb_sc->base;
+	OR4(sc, TPM_CRB_CTRL_REQ, TPM_CRB_CTRL_REQ_GO_IDLE);
+
+	if (timeout > 0) {
+		mask = TPM_CRB_CTRL_STS_IDLE_BIT;
+		if (!tpm_wait_for_u32(sc, TPM_CRB_CTRL_STS, mask, mask,
+		    timeout))
+			return (false);
+	}
+
+	return (true);
+}
+
+static bool
+tpmcrb_state_ready(struct tpmcrb_sc *crb_sc, bool wait)
+{
+	struct tpm_sc *sc;
+	int mask, timeout;
+
+	timeout = wait ? TPM_TIMEOUT_C : 0;
+
+	sc = &crb_sc->base;
+	OR4(sc, TPM_CRB_CTRL_REQ, TPM_CRB_CTRL_REQ_GO_READY);
+
+	if (timeout > 0) {
+		mask = TPM_CRB_CTRL_REQ_GO_READY;
+		if (!tpm_wait_for_u32(sc, TPM_CRB_CTRL_STS, mask, !mask,
+		    timeout))
+			return (false);
+	}
+
+	return (true);
+}
+
 int
 tpmcrb_transmit(device_t dev, size_t length)
 {
@@ -335,22 +377,15 @@ tpmcrb_transmit(device_t dev, size_t length)
 
 	/* Switch device to idle state if necessary */
 	if (!(TPM_READ_4(dev, TPM_CRB_CTRL_STS) & TPM_CRB_CTRL_STS_IDLE_BIT)) {
-		OR4(sc, TPM_CRB_CTRL_REQ, TPM_CRB_CTRL_REQ_GO_IDLE);
-
-		mask = TPM_CRB_CTRL_STS_IDLE_BIT;
-		if (!tpm_wait_for_u32(sc, TPM_CRB_CTRL_STS,
-			    mask, mask, TPM_TIMEOUT_C)) {
+		if (!tpmcrb_state_idle(crb_sc, true)) {
 			device_printf(dev,
 			    "Failed to transition to idle state\n");
 			return (EIO);
 		}
 	}
-	/* Switch to ready state */
-	OR4(sc, TPM_CRB_CTRL_REQ, TPM_CRB_CTRL_REQ_GO_READY);
 
-	mask = TPM_CRB_CTRL_REQ_GO_READY;
-	if (!tpm_wait_for_u32(sc, TPM_CRB_CTRL_STS,
-		    mask, !mask, TPM_TIMEOUT_C)) {
+	/* Switch to ready state */
+	if (!tpmcrb_state_ready(crb_sc, true)) {
 		device_printf(dev,
 		    "Failed to transition to ready state\n");
 		return (EIO);
@@ -394,7 +429,11 @@ tpmcrb_transmit(device_t dev, size_t length)
 	bus_read_region_stream_1(sc->mem_res, crb_sc->rsp_off + TPM_HEADER_SIZE,
 	      &sc->buf[TPM_HEADER_SIZE], bytes_available - TPM_HEADER_SIZE);
 
-	OR4(sc, TPM_CRB_CTRL_REQ, TPM_CRB_CTRL_REQ_GO_IDLE);
+	if (!tpmcrb_state_idle(crb_sc, false)) {
+		device_printf(dev,
+		    "Failed to transition to idle state post-send\n");
+		return (EIO);
+	}
 
 	tpmcrb_relinquish_locality(sc);
 	sc->pending_data_length = bytes_available;