if (!keys)
                return -ENOMEM;
 
+       down_read(¤t->mm->mmap_sem);
        for (i = 0; i < args->count; i++) {
                hva = gfn_to_hva(kvm, args->start_gfn + i);
                if (kvm_is_error_hva(hva)) {
                        r = -EFAULT;
-                       goto out;
+                       break;
                }
 
                curkey = get_guest_storage_key(current->mm, hva);
                if (IS_ERR_VALUE(curkey)) {
                        r = curkey;
-                       goto out;
+                       break;
                }
                keys[i] = curkey;
        }
+       up_read(¤t->mm->mmap_sem);
+
+       if (!r) {
+               r = copy_to_user((uint8_t __user *)args->skeydata_addr, keys,
+                                sizeof(uint8_t) * args->count);
+               if (r)
+                       r = -EFAULT;
+       }
 
-       r = copy_to_user((uint8_t __user *)args->skeydata_addr, keys,
-                        sizeof(uint8_t) * args->count);
-       if (r)
-               r = -EFAULT;
-out:
        kvfree(keys);
        return r;
 }
        if (r)
                goto out;
 
+       down_read(¤t->mm->mmap_sem);
        for (i = 0; i < args->count; i++) {
                hva = gfn_to_hva(kvm, args->start_gfn + i);
                if (kvm_is_error_hva(hva)) {
                        r = -EFAULT;
-                       goto out;
+                       break;
                }
 
                /* Lowest order bit is reserved */
                if (keys[i] & 0x01) {
                        r = -EINVAL;
-                       goto out;
+                       break;
                }
 
                r = set_guest_storage_key(current->mm, hva,
                                          (unsigned long)keys[i], 0);
                if (r)
-                       goto out;
+                       break;
        }
+       up_read(¤t->mm->mmap_sem);
 out:
        kvfree(keys);
        return r;
 
 
                        if (rc)
                                return rc;
-                       if (set_guest_storage_key(current->mm, useraddr,
+                       down_read(¤t->mm->mmap_sem);
+                       rc = set_guest_storage_key(current->mm, useraddr,
                                        vcpu->run->s.regs.gprs[reg1] & PFMF_KEY,
-                                       vcpu->run->s.regs.gprs[reg1] & PFMF_NQ))
+                                       vcpu->run->s.regs.gprs[reg1] & PFMF_NQ);
+                       up_read(¤t->mm->mmap_sem);
+                       if (rc)
                                return kvm_s390_inject_program_int(vcpu, PGM_ADDRESSING);
                }
 
 
        pgste_t old, new;
        pte_t *ptep;
 
-       down_read(&mm->mmap_sem);
        ptep = get_locked_pte(mm, addr, &ptl);
-       if (unlikely(!ptep)) {
-               up_read(&mm->mmap_sem);
+       if (unlikely(!ptep))
                return -EFAULT;
-       }
 
        new = old = pgste_get_lock(ptep);
        pgste_val(new) &= ~(PGSTE_GR_BIT | PGSTE_GC_BIT |
 
        pgste_set_unlock(ptep, new);
        pte_unmap_unlock(ptep, ptl);
-       up_read(&mm->mmap_sem);
        return 0;
 }
 EXPORT_SYMBOL(set_guest_storage_key);
        pgste_t pgste;
        pte_t *ptep;
 
-       down_read(&mm->mmap_sem);
        ptep = get_locked_pte(mm, addr, &ptl);
-       if (unlikely(!ptep)) {
-               up_read(&mm->mmap_sem);
+       if (unlikely(!ptep))
                return -EFAULT;
-       }
-       pgste = pgste_get_lock(ptep);
 
+       pgste = pgste_get_lock(ptep);
        if (pte_val(*ptep) & _PAGE_INVALID) {
                key  = (pgste_val(pgste) & PGSTE_ACC_BITS) >> 56;
                key |= (pgste_val(pgste) & PGSTE_FP_BIT) >> 56;
                if (pgste_val(pgste) & PGSTE_GC_BIT)
                        key |= _PAGE_CHANGED;
        }
-
        pgste_set_unlock(ptep, pgste);
        pte_unmap_unlock(ptep, ptl);
-       up_read(&mm->mmap_sem);
        return key;
 }
 EXPORT_SYMBOL(get_guest_storage_key);