Memory hotplug in secure environments requires the unaccepted memory
bitmap to grow as new memory is added. Currently, the bitmap is
implemented as a flexible array member at the end of struct
efi_unaccepted_memory, which is reserved by memblock at boot and cannot
be resized without reallocating the entire structure.

Replace the flexible array member with a pointer. This allows the bitmap
to be allocated and managed independently from the unaccepted memory
table, enabling dynamic growth to support memory hotplug.

Signed-off-by: Pratik R. Sampat <prsampat@amd.com>
---
 arch/x86/boot/compressed/efi.h                |  2 +-
 arch/x86/include/asm/unaccepted_memory.h      |  9 +++++++++
 .../firmware/efi/libstub/unaccepted_memory.c  | 11 ++++++++++-
 drivers/firmware/efi/unaccepted_memory.c      | 19 ++++++++++++++-----
 include/linux/efi.h                           |  2 +-
 5 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/arch/x86/boot/compressed/efi.h b/arch/x86/boot/compressed/efi.h
index b22300970f97..4f7027f33def 100644
--- a/arch/x86/boot/compressed/efi.h
+++ b/arch/x86/boot/compressed/efi.h
@@ -102,7 +102,7 @@ struct efi_unaccepted_memory {
 	u32 unit_size;
 	u64 phys_base;
 	u64 size;
-	unsigned long bitmap[];
+	unsigned long *bitmap;
 };
 
 static inline int efi_guidcmp (efi_guid_t left, efi_guid_t right)
diff --git a/arch/x86/include/asm/unaccepted_memory.h b/arch/x86/include/asm/unaccepted_memory.h
index f5937e9866ac..5da80e68d718 100644
--- a/arch/x86/include/asm/unaccepted_memory.h
+++ b/arch/x86/include/asm/unaccepted_memory.h
@@ -24,4 +24,13 @@ static inline struct efi_unaccepted_memory *efi_get_unaccepted_table(void)
 		return NULL;
 	return __va(efi.unaccepted);
 }
+
+static inline unsigned long *efi_get_unaccepted_bitmap(void)
+{
+	struct efi_unaccepted_memory *unaccepted = efi_get_unaccepted_table();
+
+	if (!unaccepted)
+		return NULL;
+	return __va(unaccepted->bitmap);
+}
 #endif
diff --git a/drivers/firmware/efi/libstub/unaccepted_memory.c b/drivers/firmware/efi/libstub/unaccepted_memory.c
index 757dbe734a47..c1370fc14555 100644
--- a/drivers/firmware/efi/libstub/unaccepted_memory.c
+++ b/drivers/firmware/efi/libstub/unaccepted_memory.c
@@ -63,13 +63,22 @@ efi_status_t allocate_unaccepted_bitmap(__u32 nr_desc,
 				   EFI_UNACCEPTED_UNIT_SIZE * BITS_PER_BYTE);
 
 	status = efi_bs_call(allocate_pool, EFI_ACPI_RECLAIM_MEMORY,
-			     sizeof(*unaccepted_table) + bitmap_size,
+			     sizeof(*unaccepted_table),
 			     (void **)&unaccepted_table);
 	if (status != EFI_SUCCESS) {
 		efi_err("Failed to allocate unaccepted memory config table\n");
 		return status;
 	}
 
+	status = efi_bs_call(allocate_pool, EFI_ACPI_RECLAIM_MEMORY,
+			     bitmap_size,
+			     (void **)&unaccepted_table->bitmap);
+	if (status != EFI_SUCCESS) {
+		efi_bs_call(free_pool, unaccepted_table);
+		efi_err("Failed to allocate unaccepted memory bitmap\n");
+		return status;
+	}
+
 	unaccepted_table->version = 1;
 	unaccepted_table->unit_size = EFI_UNACCEPTED_UNIT_SIZE;
 	unaccepted_table->phys_base = unaccepted_start;
diff --git a/drivers/firmware/efi/unaccepted_memory.c b/drivers/firmware/efi/unaccepted_memory.c
index c2c067eff634..4479aad258f8 100644
--- a/drivers/firmware/efi/unaccepted_memory.c
+++ b/drivers/firmware/efi/unaccepted_memory.c
@@ -36,7 +36,7 @@ void accept_memory(phys_addr_t start, unsigned long size)
 	unsigned long range_start, range_end;
 	struct accept_range range, *entry;
 	phys_addr_t end = start + size;
-	unsigned long flags;
+	unsigned long flags, *bitmap;
 	u64 unit_size;
 
 	unaccepted = efi_get_unaccepted_table();
@@ -124,8 +124,12 @@ void accept_memory(phys_addr_t start, unsigned long size)
 	list_add(&range.list, &accepting_list);
 
 	range_start = range.start;
-	for_each_set_bitrange_from(range_start, range_end, unaccepted->bitmap,
-				   range.end) {
+
+	bitmap = efi_get_unaccepted_bitmap();
+	if (!bitmap)
+		return;
+
+	for_each_set_bitrange_from(range_start, range_end, bitmap, range.end) {
 		unsigned long phys_start, phys_end;
 		unsigned long len = range_end - range_start;
 
@@ -147,7 +151,7 @@ void accept_memory(phys_addr_t start, unsigned long size)
 		arch_accept_memory(phys_start, phys_end);
 
 		spin_lock(&unaccepted_memory_lock);
-		bitmap_clear(unaccepted->bitmap, range_start, len);
+		bitmap_clear(bitmap, range_start, len);
 	}
 
 	list_del(&range.list);
@@ -197,7 +201,12 @@ bool range_contains_unaccepted_memory(phys_addr_t start, unsigned long size)
 
 	spin_lock_irqsave(&unaccepted_memory_lock, flags);
 	while (start < end) {
-		if (test_bit(start / unit_size, unaccepted->bitmap)) {
+		unsigned long *bitmap = efi_get_unaccepted_bitmap();
+
+		if (!bitmap)
+			break;
+
+		if (test_bit(start / unit_size, bitmap)) {
 			ret = true;
 			break;
 		}
diff --git a/include/linux/efi.h b/include/linux/efi.h
index a98cc39e7aaa..a74b393c54d8 100644
--- a/include/linux/efi.h
+++ b/include/linux/efi.h
@@ -545,7 +545,7 @@ struct efi_unaccepted_memory {
 	u32 unit_size;
 	u64 phys_base;
 	u64 size;
-	unsigned long bitmap[];
+	unsigned long *bitmap;
 };
 
 /*
-- 
2.51.1


The unaccepted memory structure currently only supports accepting memory
present at boot time. The unaccepted table uses a fixed-size bitmap
reserved in memblock based on the initial memory layout, preventing
dynamic addition of memory ranges after boot. This causes guest
termination when memory is hot-added in a secure virtual machine due to
accessing pages that have not transitioned to private before use.

Extend the unaccepted memory framework to handle hotplugged memory by
dynamically managing the unaccepted bitmap. Allocate a new bitmap when
hotplugged ranges exceed the reserved bitmap capacity and switch to
kernel-managed allocation.

Hotplugged memory also follows the same acceptance policy using the
accept_memory=[eager|lazy] kernel parameter to accept memory either
up-front when added or before first use.

Signed-off-by: Pratik R. Sampat <prsampat@amd.com>
---
 arch/x86/boot/compressed/efi.h                |  1 +
 .../firmware/efi/libstub/unaccepted_memory.c  |  1 +
 drivers/firmware/efi/unaccepted_memory.c      | 83 +++++++++++++++++++
 include/linux/efi.h                           |  1 +
 include/linux/mm.h                            | 11 +++
 mm/memory_hotplug.c                           |  7 ++
 mm/page_alloc.c                               |  2 +
 7 files changed, 106 insertions(+)

diff --git a/arch/x86/boot/compressed/efi.h b/arch/x86/boot/compressed/efi.h
index 4f7027f33def..a220a1966cae 100644
--- a/arch/x86/boot/compressed/efi.h
+++ b/arch/x86/boot/compressed/efi.h
@@ -102,6 +102,7 @@ struct efi_unaccepted_memory {
 	u32 unit_size;
 	u64 phys_base;
 	u64 size;
+	bool mem_reserved;
 	unsigned long *bitmap;
 };
 
diff --git a/drivers/firmware/efi/libstub/unaccepted_memory.c b/drivers/firmware/efi/libstub/unaccepted_memory.c
index c1370fc14555..b16bd61c12bf 100644
--- a/drivers/firmware/efi/libstub/unaccepted_memory.c
+++ b/drivers/firmware/efi/libstub/unaccepted_memory.c
@@ -83,6 +83,7 @@ efi_status_t allocate_unaccepted_bitmap(__u32 nr_desc,
 	unaccepted_table->unit_size = EFI_UNACCEPTED_UNIT_SIZE;
 	unaccepted_table->phys_base = unaccepted_start;
 	unaccepted_table->size = bitmap_size;
+	unaccepted_table->mem_reserved = true;
 	memset(unaccepted_table->bitmap, 0, bitmap_size);
 
 	status = efi_bs_call(install_configuration_table,
diff --git a/drivers/firmware/efi/unaccepted_memory.c b/drivers/firmware/efi/unaccepted_memory.c
index 4479aad258f8..8537812346e2 100644
--- a/drivers/firmware/efi/unaccepted_memory.c
+++ b/drivers/firmware/efi/unaccepted_memory.c
@@ -218,6 +218,89 @@ bool range_contains_unaccepted_memory(phys_addr_t start, unsigned long size)
 	return ret;
 }
 
+static int extend_unaccepted_bitmap(phys_addr_t mem_range_start,
+				    unsigned long mem_range_size)
+{
+	struct efi_unaccepted_memory *unacc_tbl;
+	unsigned long *old_bitmap, *new_bitmap;
+	phys_addr_t start, end, mem_range_end;
+	u64 phys_base, size, unit_size;
+	unsigned long flags;
+
+	unacc_tbl = efi_get_unaccepted_table();
+	if (!unacc_tbl || !unacc_tbl->unit_size)
+		return -EIO;
+
+	unit_size = unacc_tbl->unit_size;
+	phys_base = unacc_tbl->phys_base;
+
+	mem_range_end = round_up(mem_range_start + mem_range_size, unit_size);
+	size = DIV_ROUND_UP(mem_range_end - phys_base, unit_size * BITS_PER_BYTE);
+
+	/* Translate to offsets from the beginning of the bitmap */
+	start = mem_range_start - phys_base;
+	end = mem_range_end - phys_base;
+
+	old_bitmap = efi_get_unaccepted_bitmap();
+	if (!old_bitmap)
+		return -EIO;
+
+	/* If the bitmap is already large enough, just set the bits */
+	if (unacc_tbl->size >= size) {
+		spin_lock_irqsave(&unaccepted_memory_lock, flags);
+		bitmap_set(old_bitmap, start / unit_size, (end - start) / unit_size);
+		spin_unlock_irqrestore(&unaccepted_memory_lock, flags);
+
+		return 0;
+	}
+
+	/* Reserved memblocks cannot be extended so allocate a new bitmap */
+	if (unacc_tbl->mem_reserved) {
+		new_bitmap = kzalloc(size, GFP_KERNEL);
+		if (!new_bitmap)
+			return -ENOMEM;
+
+		spin_lock_irqsave(&unaccepted_memory_lock, flags);
+		memcpy(new_bitmap, old_bitmap, unacc_tbl->size);
+		unacc_tbl->mem_reserved = false;
+		free_reserved_area(old_bitmap, old_bitmap + unacc_tbl->size, -1, NULL);
+		spin_unlock_irqrestore(&unaccepted_memory_lock, flags);
+	} else {
+		new_bitmap = krealloc(old_bitmap, size, GFP_KERNEL);
+		if (!new_bitmap)
+			return -ENOMEM;
+
+		/* Zero the bitmap from the range it was extended from */
+		memset(new_bitmap + unacc_tbl->size, 0, size - unacc_tbl->size);
+	}
+
+	bitmap_set(new_bitmap, start / unit_size, (end - start) / unit_size);
+
+	spin_lock_irqsave(&unaccepted_memory_lock, flags);
+	unacc_tbl->size = size;
+	unacc_tbl->bitmap = (unsigned long *)__pa(new_bitmap);
+	spin_unlock_irqrestore(&unaccepted_memory_lock, flags);
+
+	return 0;
+}
+
+int accept_hotplug_memory(phys_addr_t mem_range_start, unsigned long mem_range_size)
+{
+	int ret;
+
+	if (!IS_ENABLED(CONFIG_UNACCEPTED_MEMORY))
+		return 0;
+
+	ret = extend_unaccepted_bitmap(mem_range_start, mem_range_size);
+	if (ret)
+		return ret;
+
+	if (!mm_lazy_accept_enabled())
+		accept_memory(mem_range_start, mem_range_size);
+
+	return 0;
+}
+
 #ifdef CONFIG_PROC_VMCORE
 static bool unaccepted_memory_vmcore_pfn_is_ram(struct vmcore_cb *cb,
 						unsigned long pfn)
diff --git a/include/linux/efi.h b/include/linux/efi.h
index a74b393c54d8..1021eb78388f 100644
--- a/include/linux/efi.h
+++ b/include/linux/efi.h
@@ -545,6 +545,7 @@ struct efi_unaccepted_memory {
 	u32 unit_size;
 	u64 phys_base;
 	u64 size;
+	bool mem_reserved;
 	unsigned long *bitmap;
 };
 
diff --git a/include/linux/mm.h b/include/linux/mm.h
index 1ae97a0b8ec7..bb43876e6c47 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -4077,6 +4077,9 @@ int set_anon_vma_name(unsigned long addr, unsigned long size,
 
 bool range_contains_unaccepted_memory(phys_addr_t start, unsigned long size);
 void accept_memory(phys_addr_t start, unsigned long size);
+int accept_hotplug_memory(phys_addr_t mem_range_start,
+			  unsigned long mem_range_size);
+bool mm_lazy_accept_enabled(void);
 
 #else
 
@@ -4090,6 +4093,14 @@ static inline void accept_memory(phys_addr_t start, unsigned long size)
 {
 }
 
+static inline int accept_hotplug_memory(phys_addr_t mem_range_start,
+					unsigned long mem_range_size)
+{
+	return 0;
+}
+
+static inline bool mm_lazy_accept_enabled(void) { return false; }
+
 #endif
 
 static inline bool pfn_is_unaccepted_memory(unsigned long pfn)
diff --git a/mm/memory_hotplug.c b/mm/memory_hotplug.c
index 74318c787715..bf8086682b66 100644
--- a/mm/memory_hotplug.c
+++ b/mm/memory_hotplug.c
@@ -1581,6 +1581,13 @@ int add_memory_resource(int nid, struct resource *res, mhp_t mhp_flags)
 	if (!strcmp(res->name, "System RAM"))
 		firmware_map_add_hotplug(start, start + size, "System RAM");
 
+	ret = accept_hotplug_memory(start, size);
+	if (ret) {
+		remove_memory_block_devices(start, size);
+		arch_remove_memory(start, size, params.altmap);
+		goto error;
+	}
+
 	/* device_online() will take the lock when calling online_pages() */
 	mem_hotplug_done();
 
diff --git a/mm/page_alloc.c b/mm/page_alloc.c
index d1d037f97c5f..d0c298dcaf9d 100644
--- a/mm/page_alloc.c
+++ b/mm/page_alloc.c
@@ -7331,6 +7331,8 @@ bool has_managed_dma(void)
 
 static bool lazy_accept = true;
 
+bool mm_lazy_accept_enabled(void) { return lazy_accept; }
+
 static int __init accept_memory_parse(char *p)
 {
 	if (!strcmp(p, "lazy")) {
-- 
2.51.1


When hot-removing memory in a SEV-SNP environment, pages must be set to
shared state so they can be reused by the hypervisor. This also applies
when memory is intended to be hotplugged back in later, as those pages
will need to be re-accepted after crossing the trust boundary.

However, memory can already be set to shared state externally. In such
cases, the pvalidate rescind operation will not change the validated bit
in the RMP table, setting the carry flag and causing the guest to
terminate.

Since memory hotplug is arguably unique, introduce a guest-maintained
memory state tracking structure that maintains a bitmap to track the
state (private vs shared) of all hotplugged memory supplemented with a
flag to indicate intent. This allows for memory that is already marked
as shared in the hotplug bitmap to avoid performing the pvalidate
rescind operation. Additionally, tracking page state changes from the
guest's perspective, enables the detection of inconsistencies if the
hypervisor changes states unexpectedly. For example, if the guest bitmap
reports memory as private but the hypervisor has already changed the RMP
state to shared, the guest detects this inconsistency when attempting to
share the memory and terminate rather than skipping over the pvalidate
rescind operation.

Signed-off-by: Pratik R. Sampat <prsampat@amd.com>
---
 arch/x86/coco/sev/core.c                 | 104 +++++++++++++++++++++--
 arch/x86/include/asm/sev.h               |  32 +++++++
 arch/x86/include/asm/unaccepted_memory.h |  13 +++
 drivers/firmware/efi/unaccepted_memory.c |   2 +-
 4 files changed, 143 insertions(+), 8 deletions(-)

diff --git a/arch/x86/coco/sev/core.c b/arch/x86/coco/sev/core.c
index 14ef5908fb27..a5c9615a6e0c 100644
--- a/arch/x86/coco/sev/core.c
+++ b/arch/x86/coco/sev/core.c
@@ -46,6 +46,8 @@
 #include <asm/cmdline.h>
 #include <asm/msr.h>
 
+struct snp_hotplug_memory *snp_hp_mem;
+
 /* AP INIT values as documented in the APM2  section "Processor Initialization State" */
 #define AP_INIT_CS_LIMIT		0xffff
 #define AP_INIT_DS_LIMIT		0xffff
@@ -453,9 +455,54 @@ static int vmgexit_psc(struct ghcb *ghcb, struct snp_psc_desc *desc)
 	return ret;
 }
 
+static bool snp_hotplug_state_shared(unsigned long vaddr)
+{
+	phys_addr_t paddr = __pa(vaddr);
+	u64 hotplug_bit;
+
+	if (!snp_is_hotplug_memory(paddr))
+		return false;
+
+	hotplug_bit = (paddr - snp_hp_mem->phys_base) / snp_hp_mem->unit_size;
+
+	return !test_bit(hotplug_bit, snp_hp_mem->bitmap);
+}
+
+static void snp_set_hotplug_bit(unsigned long vaddr, bool private)
+{
+	phys_addr_t paddr = __pa(vaddr);
+	u64 hotplug_bit;
+
+	if (!snp_is_hotplug_memory(paddr))
+		return;
+
+	hotplug_bit = (paddr - snp_hp_mem->phys_base) / snp_hp_mem->unit_size;
+	if (private)
+		set_bit(hotplug_bit, snp_hp_mem->bitmap);
+	else
+		clear_bit(hotplug_bit, snp_hp_mem->bitmap);
+}
+
+static void set_hotplug_pages_state(struct snp_psc_desc *desc)
+{
+	struct psc_entry *e;
+	unsigned long vaddr;
+	bool op;
+	int i;
+
+	for (i = 0; i <= desc->hdr.end_entry; i++) {
+		e = &desc->entries[i];
+		vaddr = (unsigned long)pfn_to_kaddr(e->gfn);
+		op = e->operation == SNP_PAGE_STATE_PRIVATE;
+
+		snp_set_hotplug_bit(vaddr, op);
+	}
+}
+
 static unsigned long __set_pages_state(struct snp_psc_desc *data, unsigned long vaddr,
-				       unsigned long vaddr_end, int op)
+				       unsigned long vaddr_end, int op, u8 psc_flags)
 {
+	unsigned long vaddr_base;
 	struct ghcb_state state;
 	bool use_large_entry;
 	struct psc_hdr *hdr;
@@ -465,6 +512,7 @@ static unsigned long __set_pages_state(struct snp_psc_desc *data, unsigned long
 	struct ghcb *ghcb;
 	int i;
 
+	vaddr_base = vaddr;
 	hdr = &data->hdr;
 	e = data->entries;
 
@@ -499,7 +547,8 @@ static unsigned long __set_pages_state(struct snp_psc_desc *data, unsigned long
 	}
 
 	/* Page validation must be rescinded before changing to shared */
-	if (op == SNP_PAGE_STATE_SHARED)
+	if (op == SNP_PAGE_STATE_SHARED &&
+	    !(snp_hotplug_state_shared(vaddr_base) && (psc_flags & SNP_PSC_SHARED_TO_SHARED)))
 		pvalidate_pages(data);
 
 	local_irq_save(flags);
@@ -522,10 +571,12 @@ static unsigned long __set_pages_state(struct snp_psc_desc *data, unsigned long
 	if (op == SNP_PAGE_STATE_PRIVATE)
 		pvalidate_pages(data);
 
+	set_hotplug_pages_state(data);
+
 	return vaddr;
 }
 
-static void set_pages_state(unsigned long vaddr, unsigned long npages, int op)
+static void set_pages_state(unsigned long vaddr, unsigned long npages, int op, u8 psc_flags)
 {
 	struct snp_psc_desc desc;
 	unsigned long vaddr_end;
@@ -538,7 +589,7 @@ static void set_pages_state(unsigned long vaddr, unsigned long npages, int op)
 	vaddr_end = vaddr + (npages << PAGE_SHIFT);
 
 	while (vaddr < vaddr_end)
-		vaddr = __set_pages_state(&desc, vaddr, vaddr_end, op);
+		vaddr = __set_pages_state(&desc, vaddr, vaddr_end, op, psc_flags);
 }
 
 void snp_set_memory_shared(unsigned long vaddr, unsigned long npages)
@@ -546,7 +597,7 @@ void snp_set_memory_shared(unsigned long vaddr, unsigned long npages)
 	if (!cc_platform_has(CC_ATTR_GUEST_SEV_SNP))
 		return;
 
-	set_pages_state(vaddr, npages, SNP_PAGE_STATE_SHARED);
+	set_pages_state(vaddr, npages, SNP_PAGE_STATE_SHARED, 0);
 }
 
 void snp_set_memory_private(unsigned long vaddr, unsigned long npages)
@@ -554,7 +605,7 @@ void snp_set_memory_private(unsigned long vaddr, unsigned long npages)
 	if (!cc_platform_has(CC_ATTR_GUEST_SEV_SNP))
 		return;
 
-	set_pages_state(vaddr, npages, SNP_PAGE_STATE_PRIVATE);
+	set_pages_state(vaddr, npages, SNP_PAGE_STATE_PRIVATE, 0);
 }
 
 void snp_accept_memory(phys_addr_t start, phys_addr_t end)
@@ -567,7 +618,46 @@ void snp_accept_memory(phys_addr_t start, phys_addr_t end)
 	vaddr = (unsigned long)__va(start);
 	npages = (end - start) >> PAGE_SHIFT;
 
-	set_pages_state(vaddr, npages, SNP_PAGE_STATE_PRIVATE);
+	set_pages_state(vaddr, npages, SNP_PAGE_STATE_PRIVATE, 0);
+}
+
+int snp_extend_hotplug_memory_state_bitmap(phys_addr_t start,
+					   unsigned long size,
+					   uint64_t unit_size)
+{
+	u64 hp_mem_size = DIV_ROUND_UP(size, unit_size * BITS_PER_BYTE);
+
+	if (snp_hp_mem) {
+		u64 old_size = snp_hp_mem->size;
+		unsigned long *bitmap;
+
+		bitmap = krealloc(snp_hp_mem->bitmap, hp_mem_size, GFP_KERNEL);
+		if (!bitmap)
+			return -ENOMEM;
+
+		memset(bitmap + old_size, 0, hp_mem_size - old_size);
+		snp_hp_mem->size = hp_mem_size;
+		snp_hp_mem->bitmap = bitmap;
+
+		return 0;
+	}
+
+	snp_hp_mem = kzalloc(sizeof(*snp_hp_mem), GFP_KERNEL);
+	if (!snp_hp_mem)
+		return -ENOMEM;
+
+	snp_hp_mem->bitmap = kzalloc(hp_mem_size, GFP_KERNEL);
+	if (!snp_hp_mem->bitmap) {
+		kfree(snp_hp_mem);
+		return -ENOMEM;
+	}
+
+	snp_hp_mem->phys_base = start;
+	snp_hp_mem->phys_end = start + hp_mem_size;
+	snp_hp_mem->size = hp_mem_size;
+	snp_hp_mem->unit_size = unit_size;
+
+	return 0;
 }
 
 static int vmgexit_ap_control(u64 event, struct sev_es_save_area *vmsa, u32 apic_id)
diff --git a/arch/x86/include/asm/sev.h b/arch/x86/include/asm/sev.h
index 465b19fd1a2d..eb605892645c 100644
--- a/arch/x86/include/asm/sev.h
+++ b/arch/x86/include/asm/sev.h
@@ -464,6 +464,38 @@ static __always_inline void sev_es_nmi_complete(void)
 extern int __init sev_es_efi_map_ghcbs_cas(pgd_t *pgd);
 extern void sev_enable(struct boot_params *bp);
 
+#define SNP_PSC_SHARED_TO_SHARED	0x1
+
+struct snp_hotplug_memory {
+	u64 phys_base;
+	u64 phys_end;
+	u32 unit_size;
+	u64 size;
+	/* bitmap bit unset: shared, set: private */
+	unsigned long *bitmap;
+};
+
+extern struct snp_hotplug_memory *snp_hp_mem;
+
+#ifdef CONFIG_UNACCEPTED_MEMORY
+int snp_extend_hotplug_memory_state_bitmap(phys_addr_t start,
+					   unsigned long size,
+					   uint64_t unit_size);
+static inline bool snp_is_hotplug_memory(phys_addr_t paddr)
+{
+	return snp_hp_mem && paddr >= snp_hp_mem->phys_base && paddr < snp_hp_mem->phys_end;
+}
+#else /* !CONFIG_UNACCEPTED_MEMORY */
+static inline int snp_extend_hotplug_memory_state_bitmap(phys_addr_t start,
+							 unsigned long size,
+							 uint64_t unit_size)
+{
+	return 0;
+}
+
+static inline bool snp_is_hotplug_memory(phys_addr_t paddr) { return false; }
+#endif
+
 /*
  * RMPADJUST modifies the RMP permissions of a page of a lesser-
  * privileged (numerically higher) VMPL.
diff --git a/arch/x86/include/asm/unaccepted_memory.h b/arch/x86/include/asm/unaccepted_memory.h
index 5da80e68d718..abdf5472de9e 100644
--- a/arch/x86/include/asm/unaccepted_memory.h
+++ b/arch/x86/include/asm/unaccepted_memory.h
@@ -33,4 +33,17 @@ static inline unsigned long *efi_get_unaccepted_bitmap(void)
 		return NULL;
 	return __va(unaccepted->bitmap);
 }
+
+static inline int arch_set_unaccepted_mem_state(phys_addr_t start, unsigned long size)
+{
+	struct efi_unaccepted_memory *unaccepted = efi_get_unaccepted_table();
+
+	if (!unaccepted)
+		return -EIO;
+
+	if (cc_platform_has(CC_ATTR_GUEST_SEV_SNP))
+		return snp_extend_hotplug_memory_state_bitmap(start, size, unaccepted->unit_size);
+
+	return 0;
+}
 #endif
diff --git a/drivers/firmware/efi/unaccepted_memory.c b/drivers/firmware/efi/unaccepted_memory.c
index 8537812346e2..6796042a64aa 100644
--- a/drivers/firmware/efi/unaccepted_memory.c
+++ b/drivers/firmware/efi/unaccepted_memory.c
@@ -281,7 +281,7 @@ static int extend_unaccepted_bitmap(phys_addr_t mem_range_start,
 	unacc_tbl->bitmap = (unsigned long *)__pa(new_bitmap);
 	spin_unlock_irqrestore(&unaccepted_memory_lock, flags);
 
-	return 0;
+	return arch_set_unaccepted_mem_state(mem_range_start, mem_range_size);
 }
 
 int accept_hotplug_memory(phys_addr_t mem_range_start, unsigned long mem_range_size)
-- 
2.51.1


Transition memory to shared during a hot-remove operation so that it can
be re-used by the hypervisor. During lazy acceptance, only memory that
was used has been accepted, therefore during hot-remove only mark pages
as shared that were previously accepted / made private.

Signed-off-by: Pratik R. Sampat <prsampat@amd.com>
---
 arch/x86/coco/sev/core.c                 | 23 +++++++++++++++
 arch/x86/include/asm/sev.h               |  2 ++
 arch/x86/include/asm/unaccepted_memory.h |  9 ++++++
 drivers/firmware/efi/unaccepted_memory.c | 37 ++++++++++++++++++++++++
 include/linux/mm.h                       |  7 +++++
 mm/memory_hotplug.c                      |  2 ++
 6 files changed, 80 insertions(+)

diff --git a/arch/x86/coco/sev/core.c b/arch/x86/coco/sev/core.c
index a5c9615a6e0c..c05fc91d10a1 100644
--- a/arch/x86/coco/sev/core.c
+++ b/arch/x86/coco/sev/core.c
@@ -621,6 +621,29 @@ void snp_accept_memory(phys_addr_t start, phys_addr_t end)
 	set_pages_state(vaddr, npages, SNP_PAGE_STATE_PRIVATE, 0);
 }
 
+void snp_unaccept_memory(phys_addr_t start, phys_addr_t end)
+{
+	unsigned long vaddr, npages;
+
+	if (!cc_platform_has(CC_ATTR_GUEST_SEV_SNP))
+		return;
+
+	vaddr = (unsigned long)__va(start);
+	npages = (end - start) >> PAGE_SHIFT;
+
+	/*
+	 * Hotplugged memory can be set to shared externally. Attempting to
+	 * re-share the memory (during hot-remove) will cause the pvalidate
+	 * operation to not make any changes to the RMP table triggering the
+	 * PVALIDATE_FAIL_NOUPDATE condition
+	 *
+	 * Since the memory hotplug case is unique, specify this intent so that
+	 * if the page is part of hotplugged memory a pvalidate rescind
+	 * operation is not performed
+	 */
+	set_pages_state(vaddr, npages, SNP_PAGE_STATE_SHARED, SNP_PSC_SHARED_TO_SHARED);
+}
+
 int snp_extend_hotplug_memory_state_bitmap(phys_addr_t start,
 					   unsigned long size,
 					   uint64_t unit_size)
diff --git a/arch/x86/include/asm/sev.h b/arch/x86/include/asm/sev.h
index eb605892645c..8f3c5b878fd7 100644
--- a/arch/x86/include/asm/sev.h
+++ b/arch/x86/include/asm/sev.h
@@ -547,6 +547,7 @@ void __noreturn snp_abort(void);
 void snp_dmi_setup(void);
 int snp_issue_svsm_attest_req(u64 call_id, struct svsm_call *call, struct svsm_attest_call *input);
 void snp_accept_memory(phys_addr_t start, phys_addr_t end);
+void snp_unaccept_memory(phys_addr_t start, phys_addr_t end);
 u64 snp_get_unsupported_features(u64 status);
 u64 sev_get_status(void);
 void sev_show_status(void);
@@ -639,6 +640,7 @@ static inline int snp_issue_svsm_attest_req(u64 call_id, struct svsm_call *call,
 	return -ENOTTY;
 }
 static inline void snp_accept_memory(phys_addr_t start, phys_addr_t end) { }
+static inline void snp_unaccept_memory(phys_addr_t start, phys_addr_t end) { }
 static inline u64 snp_get_unsupported_features(u64 status) { return 0; }
 static inline u64 sev_get_status(void) { return 0; }
 static inline void sev_show_status(void) { }
diff --git a/arch/x86/include/asm/unaccepted_memory.h b/arch/x86/include/asm/unaccepted_memory.h
index abdf5472de9e..ad392294b71b 100644
--- a/arch/x86/include/asm/unaccepted_memory.h
+++ b/arch/x86/include/asm/unaccepted_memory.h
@@ -18,6 +18,15 @@ static inline void arch_accept_memory(phys_addr_t start, phys_addr_t end)
 	}
 }
 
+static inline void arch_unaccept_memory(phys_addr_t start, phys_addr_t end)
+{
+	if (cc_platform_has(CC_ATTR_GUEST_SEV_SNP)) {
+		snp_unaccept_memory(start, end);
+	} else {
+		panic("Cannot accept memory: unknown platform\n");
+	}
+}
+
 static inline struct efi_unaccepted_memory *efi_get_unaccepted_table(void)
 {
 	if (efi.unaccepted == EFI_INVALID_TABLE_ADDR)
diff --git a/drivers/firmware/efi/unaccepted_memory.c b/drivers/firmware/efi/unaccepted_memory.c
index 6796042a64aa..662cf0d6715f 100644
--- a/drivers/firmware/efi/unaccepted_memory.c
+++ b/drivers/firmware/efi/unaccepted_memory.c
@@ -301,6 +301,43 @@ int accept_hotplug_memory(phys_addr_t mem_range_start, unsigned long mem_range_s
 	return 0;
 }
 
+void unaccept_hotplug_memory(phys_addr_t mem_range_start, unsigned long mem_range_size)
+{
+	u64 unit_size, phys_base, bit_start, bit_end, addr;
+	struct efi_unaccepted_memory *unacc_tbl;
+	unsigned long flags, *bitmap;
+	phys_addr_t start, end;
+	int i;
+
+	unacc_tbl = efi_get_unaccepted_table();
+	if (!unacc_tbl)
+		return;
+
+	phys_base = unacc_tbl->phys_base;
+	unit_size = unacc_tbl->unit_size;
+
+	start = mem_range_start - phys_base;
+	end = (mem_range_start + mem_range_size) - phys_base;
+
+	bit_start = start / unit_size;
+	bit_end = end / unit_size;
+
+	/* Only unaccept memory that was previously accepted in the range */
+	for (i = bit_start; i < bit_end; i++) {
+		spin_lock_irqsave(&unaccepted_memory_lock, flags);
+		bitmap = efi_get_unaccepted_bitmap();
+		if (!bitmap || test_bit(i, bitmap)) {
+			spin_unlock_irqrestore(&unaccepted_memory_lock, flags);
+			continue;
+		}
+
+		addr = phys_base + i * unit_size;
+
+		arch_unaccept_memory(addr, addr + unit_size);
+		spin_unlock_irqrestore(&unaccepted_memory_lock, flags);
+	}
+}
+
 #ifdef CONFIG_PROC_VMCORE
 static bool unaccepted_memory_vmcore_pfn_is_ram(struct vmcore_cb *cb,
 						unsigned long pfn)
diff --git a/include/linux/mm.h b/include/linux/mm.h
index bb43876e6c47..34d48693dc86 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -4079,6 +4079,8 @@ bool range_contains_unaccepted_memory(phys_addr_t start, unsigned long size);
 void accept_memory(phys_addr_t start, unsigned long size);
 int accept_hotplug_memory(phys_addr_t mem_range_start,
 			  unsigned long mem_range_size);
+void unaccept_hotplug_memory(phys_addr_t mem_range_start,
+			     unsigned long mem_range_size);
 bool mm_lazy_accept_enabled(void);
 
 #else
@@ -4099,6 +4101,11 @@ static inline int accept_hotplug_memory(phys_addr_t mem_range_start,
 	return 0;
 }
 
+static inline void unaccept_hotplug_memory(phys_addr_t mem_range_start,
+					   unsigned long mem_range_size)
+{
+}
+
 static inline bool mm_lazy_accept_enabled(void) { return false; }
 
 #endif
diff --git a/mm/memory_hotplug.c b/mm/memory_hotplug.c
index bf8086682b66..0b14b14e53fe 100644
--- a/mm/memory_hotplug.c
+++ b/mm/memory_hotplug.c
@@ -2254,6 +2254,8 @@ static int try_remove_memory(u64 start, u64 size)
 
 	mem_hotplug_begin();
 
+	unaccept_hotplug_memory(start, size);
+
 	rc = memory_blocks_have_altmaps(start, size);
 	if (rc < 0) {
 		mem_hotplug_done();
-- 
2.51.1