vmsa.reserved = 0;
        vmsa.handle = to_kvm_svm(kvm)->sev_info.handle;
-       vmsa.address = __sme_pa(svm->vmsa);
+       vmsa.address = __sme_pa(svm->sev_es.vmsa);
        vmsa.len = PAGE_SIZE;
 -      return sev_issue_cmd(kvm, SEV_CMD_LAUNCH_UPDATE_VMSA, &vmsa, error);
 +      ret = sev_issue_cmd(kvm, SEV_CMD_LAUNCH_UPDATE_VMSA, &vmsa, error);
 +      if (ret)
 +        return ret;
 +
 +      vcpu->arch.guest_state_protected = true;
 +      return 0;
  }
  
  static int sev_launch_update_vmsa(struct kvm *kvm, struct kvm_sev_cmd *argp)
  
  int sev_es_string_io(struct vcpu_svm *svm, int size, unsigned int port, int in)
  {
 -      if (!setup_vmgexit_scratch(svm, in, svm->vmcb->control.exit_info_2))
 +      int count;
 +      int bytes;
 +
 +      if (svm->vmcb->control.exit_info_2 > INT_MAX)
 +              return -EINVAL;
 +
 +      count = svm->vmcb->control.exit_info_2;
 +      if (unlikely(check_mul_overflow(count, size, &bytes)))
 +              return -EINVAL;
 +
 +      if (!setup_vmgexit_scratch(svm, in, bytes))
                return -EINVAL;
  
-       return kvm_sev_es_string_io(&svm->vcpu, size, port, svm->ghcb_sa, count, in);
+       return kvm_sev_es_string_io(&svm->vcpu, size, port, svm->sev_es.ghcb_sa,
 -                                  svm->sev_es.ghcb_sa_len / size, in);
++                                  count, in);
  }
  
  void sev_es_init_vmcb(struct vcpu_svm *svm)
 
        bool initialized;
  };
  
 -      u64 ghcb_sa_len;
+ struct vcpu_sev_es_state {
+       /* SEV-ES support */
+       struct vmcb_save_area *vmsa;
+       struct ghcb *ghcb;
+       struct kvm_host_map ghcb_map;
+       bool received_first_sipi;
+ 
+       /* SEV-ES scratch area support */
+       void *ghcb_sa;
++      u32 ghcb_sa_len;
+       bool ghcb_sa_sync;
+       bool ghcb_sa_free;
+ };
+ 
  struct vcpu_svm {
        struct kvm_vcpu vcpu;
        /* vmcb always points at current_vmcb->ptr, it's purely a shorthand. */