1 files changed, 7 insertions, 2 deletions

diff --git a/drivers/infiniband/hw/ipath/ipath_diag.c b/drivers/infiniband/hw/ipath/ipath_diag.c
index af59bf380ca2..c9bfd827e71b 100644
--- a/drivers/infiniband/hw/ipath/ipath_diag.c
+++ b/drivers/infiniband/hw/ipath/ipath_diag.c
@@ -332,12 +332,17 @@ static ssize_t ipath_diagpkt_write(struct file *fp,
         u64 val;
         u32 l_state, lt_state; /* LinkState, LinkTrainingState */
 
-        if (count != sizeof(dp)) {
+        if (count < sizeof(odp)) {
                 ret = -EINVAL;
                 goto bail;
         }
 
-        if (copy_from_user(&dp, data, sizeof(dp))) {
+        if (count == sizeof(dp)) {
+                if (copy_from_user(&dp, data, sizeof(dp))) {
+                        ret = -EFAULT;
+                        goto bail;
+                }
+        } else if (copy_from_user(&odp, data, sizeof(odp))) {
                 ret = -EFAULT;
                 goto bail;
         }