2 files changed, 58 insertions, 21 deletions

diff --git a/lib/efi_loader/efi_boottime.c b/lib/efi_loader/efi_boottime.c
index 54fff85e64..1ccf54c386 100644
--- a/lib/efi_loader/efi_boottime.c
+++ b/lib/efi_loader/efi_boottime.c
@@ -921,6 +921,14 @@ static efi_status_t EFIAPI efi_close_event(struct efi_event *event)
 	list_for_each_entry_safe(item, next, &efi_register_notify_events,
 				 link) {
 		if (event == item->event) {
+			struct efi_protocol_notification *hitem, *hnext;
+
+			/* Remove signaled handles */
+			list_for_each_entry_safe(hitem, hnext, &item->handles,
+						 link) {
+				list_del(&hitem->link);
+				free(hitem);
+			}
 			list_del(&item->link);
 			free(item);
 		}
@@ -1049,8 +1057,19 @@ efi_status_t efi_add_protocol(const efi_handle_t handle,
 
 	/* Notify registered events */
 	list_for_each_entry(event, &efi_register_notify_events, link) {
-		if (!guidcmp(protocol, &event->protocol))
+		if (!guidcmp(protocol, &event->protocol)) {
+			struct efi_protocol_notification *notif;
+
+			notif = calloc(1, sizeof(*notif));
+			if (!notif) {
+				list_del(&handler->link);
+				free(handler);
+				return EFI_OUT_OF_RESOURCES;
+			}
+			notif->handle = handle;
+			list_add_tail(&notif->link, &event->handles);
 			efi_signal_event(event->event, true);
+		}
 	}
 
 	if (!guidcmp(&efi_guid_device_path, protocol))
@@ -1332,6 +1351,7 @@ static efi_status_t EFIAPI efi_register_protocol_notify(
 
 	item->event = event;
 	memcpy(&item->protocol, protocol, sizeof(efi_guid_t));
+	INIT_LIST_HEAD(&item->handles);
 
 	list_add_tail(&item->link, &efi_register_notify_events);
 
@@ -1359,7 +1379,6 @@ static int efi_search(enum efi_locate_search_type search_type,
 	switch (search_type) {
 	case ALL_HANDLES:
 		return 0;
-	case BY_REGISTER_NOTIFY:
 	case BY_PROTOCOL:
 		ret = efi_search_protocol(handle, protocol, NULL);
 		return (ret != EFI_SUCCESS);
@@ -1391,6 +1410,7 @@ static efi_status_t efi_locate_handle(
 	struct efi_object *efiobj;
 	efi_uintn_t size = 0;
 	struct efi_register_notify_event *item, *event = NULL;
+	struct efi_protocol_notification *handle = NULL;
 
 	/* Check parameters */
 	switch (search_type) {
@@ -1409,8 +1429,6 @@ static efi_status_t efi_locate_handle(
 		}
 		if (!event)
 			return EFI_INVALID_PARAMETER;
-
-		protocol = &event->protocol;
 		break;
 	case BY_PROTOCOL:
 		if (!protocol)
@@ -1421,14 +1439,23 @@ static efi_status_t efi_locate_handle(
 	}
 
 	/* Count how much space we need */
-	list_for_each_entry(efiobj, &efi_obj_list, link) {
-		if (!efi_search(search_type, protocol, efiobj))
-			size += sizeof(void *);
+	if (search_type == BY_REGISTER_NOTIFY) {
+		if (list_empty(&event->handles))
+			return EFI_NOT_FOUND;
+		handle = list_first_entry(&event->handles,
+					  struct efi_protocol_notification,
+					  link);
+		efiobj = handle->handle;
+		size += sizeof(void *);
+	} else {
+		list_for_each_entry(efiobj, &efi_obj_list, link) {
+			if (!efi_search(search_type, protocol, efiobj))
+				size += sizeof(void *);
+		}
+		if (size == 0)
+			return EFI_NOT_FOUND;
 	}
 
-	if (size == 0)
-		return EFI_NOT_FOUND;
-
 	if (!buffer_size)
 		return EFI_INVALID_PARAMETER;
 
@@ -1444,9 +1471,14 @@ static efi_status_t efi_locate_handle(
 		return EFI_INVALID_PARAMETER;
 
 	/* Then fill the array */
-	list_for_each_entry(efiobj, &efi_obj_list, link) {
-		if (!efi_search(search_type, protocol, efiobj))
-			*buffer++ = efiobj;
+	if (search_type == BY_REGISTER_NOTIFY) {
+		*buffer = efiobj;
+		list_del(&handle->link);
+	} else {
+		list_for_each_entry(efiobj, &efi_obj_list, link) {
+			if (!efi_search(search_type, protocol, efiobj))
+				*buffer++ = efiobj;
+		}
 	}
 
 	return EFI_SUCCESS;
diff --git a/lib/efi_selftest/efi_selftest_register_notify.c b/lib/efi_selftest/efi_selftest_register_notify.c
index ee0ef395de..ad763dd6cb 100644
--- a/lib/efi_selftest/efi_selftest_register_notify.c
+++ b/lib/efi_selftest/efi_selftest_register_notify.c
@@ -47,15 +47,20 @@ static void EFIAPI notify(struct efi_event *event, void *context)
 {
 	struct context *cp = context;
 	efi_status_t ret;
+	efi_uintn_t handle_count;
+	efi_handle_t *handles;
 
 	cp->notify_count++;
 
-	ret = boottime->locate_handle_buffer(BY_REGISTER_NOTIFY, NULL,
-					     cp->registration_key,
-					     &cp->handle_count,
-					     &cp->handles);
-	if (ret != EFI_SUCCESS)
-		cp->handle_count = 0;
+	for (;;) {
+		ret = boottime->locate_handle_buffer(BY_REGISTER_NOTIFY, NULL,
+						     cp->registration_key,
+						     &handle_count, &handles);
+		if (ret != EFI_SUCCESS)
+			break;
+		cp->handle_count += handle_count;
+		cp->handles = handles;
+	}
 }
 
 /*
@@ -170,7 +175,7 @@ static int execute(void)
 		efi_st_error("reinstall was notified too often\n");
 		return EFI_ST_FAILURE;
 	}
-	if (context.handle_count != 1) {
+	if (context.handle_count != 2) {
 		efi_st_error("LocateHandle failed\n");
 		return EFI_ST_FAILURE;
 	}
@@ -195,7 +200,7 @@ static int execute(void)
 		efi_st_error("install was notified too often\n");
 		return EFI_ST_FAILURE;
 	}
-	if (context.handle_count != 2) {
+	if (context.handle_count != 3) {
 		efi_st_error("LocateHandle failed\n");
 		return EFI_ST_FAILURE;
 	}