Retrieve the number of user entries with get_user() first and return -E2BIG early if 'user_caps' is too small to fit 'caps'. Allocate memory for 'caps' only after checking the user buffer's number of entries, thus removing two gotos and the need for premature freeing. Use struct_size() instead of manually calculating the number of bytes to allocate for 'caps', including the nested flexible array. Finally, copy 'caps' to user space with a single copy_to_user() call. Signed-off-by: Thorsten Blum --- arch/x86/kvm/vmx/tdx.c | 32 ++++++++++++-------------------- 1 file changed, 12 insertions(+), 20 deletions(-) diff --git a/arch/x86/kvm/vmx/tdx.c b/arch/x86/kvm/vmx/tdx.c index 0a49c863c811..23d638b4a003 100644 --- a/arch/x86/kvm/vmx/tdx.c +++ b/arch/x86/kvm/vmx/tdx.c @@ -2282,37 +2282,29 @@ static int tdx_get_capabilities(struct kvm_tdx_cmd *cmd) if (cmd->flags) return -EINVAL; - caps = kzalloc(sizeof(*caps) + - sizeof(struct kvm_cpuid_entry2) * td_conf->num_cpuid_config, - GFP_KERNEL); - if (!caps) - return -ENOMEM; - user_caps = u64_to_user_ptr(cmd->data); - if (get_user(nr_user_entries, &user_caps->cpuid.nent)) { - ret = -EFAULT; - goto out; - } + ret = get_user(nr_user_entries, &user_caps->cpuid.nent); + if (ret) + return ret; - if (nr_user_entries < td_conf->num_cpuid_config) { - ret = -E2BIG; - goto out; - } + if (nr_user_entries < td_conf->num_cpuid_config) + return -E2BIG; + + caps = kzalloc(struct_size(caps, cpuid.entries, + td_conf->num_cpuid_config), GFP_KERNEL); + if (!caps) + return -ENOMEM; ret = init_kvm_tdx_caps(td_conf, caps); if (ret) goto out; - if (copy_to_user(user_caps, caps, sizeof(*caps))) { + if (copy_to_user(user_caps, caps, struct_size(caps, cpuid.entries, + caps->cpuid.nent))) { ret = -EFAULT; goto out; } - if (copy_to_user(user_caps->cpuid.entries, caps->cpuid.entries, - caps->cpuid.nent * - sizeof(caps->cpuid.entries[0]))) - ret = -EFAULT; - out: /* kfree() accepts NULL. */ kfree(caps); -- 2.51.1