5 files changed, 47 insertions, 6 deletions

diff --git a/drivers/s390/crypto/ap_bus.c b/drivers/s390/crypto/ap_bus.c
index 4feb27215ab6..03e4d6246d87 100644
--- a/drivers/s390/crypto/ap_bus.c
+++ b/drivers/s390/crypto/ap_bus.c
@@ -468,6 +468,8 @@ int ap_recv(ap_qid_t qid, unsigned long long *psmid, void *msg, size_t length)
 {
 	struct ap_queue_status status;
 
+	if (msg == NULL)
+		return -EINVAL;
 	status = __ap_recv(qid, psmid, msg, length);
 	switch (status.response_code) {
 	case AP_RESPONSE_NORMAL:
@@ -617,6 +619,8 @@ static enum ap_wait ap_sm_read(struct ap_device *ap_dev)
 {
 	struct ap_queue_status status;
 
+	if (!ap_dev->reply)
+		return AP_WAIT_NONE;
 	status = ap_sm_recv(ap_dev);
 	switch (status.response_code) {
 	case AP_RESPONSE_NORMAL:
@@ -638,6 +642,31 @@ static enum ap_wait ap_sm_read(struct ap_device *ap_dev)
 }
 
 /**
+ * ap_sm_suspend_read(): Receive pending reply messages from an AP device
+ * without changing the device state in between. In suspend mode we don't
+ * allow sending new requests, therefore just fetch pending replies.
+ * @ap_dev: pointer to the AP device
+ *
+ * Returns AP_WAIT_NONE or AP_WAIT_AGAIN
+ */
+static enum ap_wait ap_sm_suspend_read(struct ap_device *ap_dev)
+{
+	struct ap_queue_status status;
+
+	if (!ap_dev->reply)
+		return AP_WAIT_NONE;
+	status = ap_sm_recv(ap_dev);
+	switch (status.response_code) {
+	case AP_RESPONSE_NORMAL:
+		if (ap_dev->queue_count > 0)
+			return AP_WAIT_AGAIN;
+		/* fall through */
+	default:
+		return AP_WAIT_NONE;
+	}
+}
+
+/**
  * ap_sm_write(): Send messages from the request queue to an AP device.
  * @ap_dev: pointer to the AP device
  *
@@ -738,7 +767,7 @@ static enum ap_wait ap_sm_reset_wait(struct ap_device *ap_dev)
 	struct ap_queue_status status;
 	unsigned long info;
 
-	if (ap_dev->queue_count > 0)
+	if (ap_dev->queue_count > 0 && ap_dev->reply)
 		/* Try to read a completed message and get the status */
 		status = ap_sm_recv(ap_dev);
 	else
@@ -778,7 +807,7 @@ static enum ap_wait ap_sm_setirq_wait(struct ap_device *ap_dev)
 	struct ap_queue_status status;
 	unsigned long info;
 
-	if (ap_dev->queue_count > 0)
+	if (ap_dev->queue_count > 0 && ap_dev->reply)
 		/* Try to read a completed message and get the status */
 		status = ap_sm_recv(ap_dev);
 	else
@@ -834,7 +863,7 @@ static ap_func_t *ap_jumptable[NR_AP_STATES][NR_AP_EVENTS] = {
 		[AP_EVENT_TIMEOUT] = ap_sm_reset,
 	},
 	[AP_STATE_SUSPEND_WAIT] = {
-		[AP_EVENT_POLL] = ap_sm_read,
+		[AP_EVENT_POLL] = ap_sm_suspend_read,
 		[AP_EVENT_TIMEOUT] = ap_sm_nop,
 	},
 	[AP_STATE_BORKED] = {
@@ -1335,6 +1364,17 @@ static struct bus_type ap_bus_type = {
 	.resume = ap_dev_resume,
 };
 
+void ap_device_init_reply(struct ap_device *ap_dev,
+			  struct ap_message *reply)
+{
+	ap_dev->reply = reply;
+
+	spin_lock_bh(&ap_dev->lock);
+	ap_sm_wait(ap_sm_event(ap_dev, AP_EVENT_POLL));
+	spin_unlock_bh(&ap_dev->lock);
+}
+EXPORT_SYMBOL(ap_device_init_reply);
+
 static int ap_device_probe(struct device *dev)
 {
 	struct ap_device *ap_dev = to_ap_dev(dev);
diff --git a/drivers/s390/crypto/ap_bus.h b/drivers/s390/crypto/ap_bus.h
index 6adcbdf225d1..d7fdf5c024d7 100644
--- a/drivers/s390/crypto/ap_bus.h
+++ b/drivers/s390/crypto/ap_bus.h
@@ -262,6 +262,7 @@ void ap_queue_message(struct ap_device *ap_dev, struct ap_message *ap_msg);
 void ap_cancel_message(struct ap_device *ap_dev, struct ap_message *ap_msg);
 void ap_flush_queue(struct ap_device *ap_dev);
 void ap_bus_force_rescan(void);
+void ap_device_init_reply(struct ap_device *ap_dev, struct ap_message *ap_msg);
 
 int ap_module_init(void);
 void ap_module_exit(void);
diff --git a/drivers/s390/crypto/zcrypt_cex2a.c b/drivers/s390/crypto/zcrypt_cex2a.c
index 1e849d6e1dfe..15104aaa075a 100644
--- a/drivers/s390/crypto/zcrypt_cex2a.c
+++ b/drivers/s390/crypto/zcrypt_cex2a.c
@@ -126,7 +126,7 @@ static int zcrypt_cex2a_probe(struct ap_device *ap_dev)
 					   MSGTYPE50_VARIANT_DEFAULT);
 	zdev->ap_dev = ap_dev;
 	zdev->online = 1;
-	ap_dev->reply = &zdev->reply;
+	ap_device_init_reply(ap_dev, &zdev->reply);
 	ap_dev->private = zdev;
 	rc = zcrypt_device_register(zdev);
 	if (rc) {
diff --git a/drivers/s390/crypto/zcrypt_cex4.c b/drivers/s390/crypto/zcrypt_cex4.c
index bb3908818505..ccb2e78ebf0e 100644
--- a/drivers/s390/crypto/zcrypt_cex4.c
+++ b/drivers/s390/crypto/zcrypt_cex4.c
@@ -147,7 +147,7 @@ static int zcrypt_cex4_probe(struct ap_device *ap_dev)
 		return -ENODEV;
 	zdev->ap_dev = ap_dev;
 	zdev->online = 1;
-	ap_dev->reply = &zdev->reply;
+	ap_device_init_reply(ap_dev, &zdev->reply);
 	ap_dev->private = zdev;
 	rc = zcrypt_device_register(zdev);
 	if (rc) {
diff --git a/drivers/s390/crypto/zcrypt_pcixcc.c b/drivers/s390/crypto/zcrypt_pcixcc.c
index f41852768953..df8f0c4dacb7 100644
--- a/drivers/s390/crypto/zcrypt_pcixcc.c
+++ b/drivers/s390/crypto/zcrypt_pcixcc.c
@@ -327,7 +327,7 @@ static int zcrypt_pcixcc_probe(struct ap_device *ap_dev)
 	else
 		zdev->ops = zcrypt_msgtype_request(MSGTYPE06_NAME,
 						   MSGTYPE06_VARIANT_NORNG);
-	ap_dev->reply = &zdev->reply;
+	ap_device_init_reply(ap_dev, &zdev->reply);
 	ap_dev->private = zdev;
 	rc = zcrypt_device_register(zdev);
 	if (rc)