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)