From ddaa38853d386e5b9f9fa1c3813048872c8ad687 Mon Sep 17 00:00:00 2001
From: niuyongwen <niuyongwen@hygon.cn>
Date: Sun, 29 Sep 2024 09:45:15 +0800
Subject: [PATCH] hw/misc/psp: Pin the hugepage memory specified by mem2 during
 use for psp

Signed-off-by: niuyongwen <niuyongwen@hygon.cn>
---
 hw/misc/psp.c | 138 +++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 121 insertions(+), 17 deletions(-)

diff --git a/hw/misc/psp.c b/hw/misc/psp.c
index 4eb5ca0e0b..03e8663027 100644
--- a/hw/misc/psp.c
+++ b/hw/misc/psp.c
@@ -17,6 +17,7 @@
 #include "sysemu/runstate.h"
 #include "exec/memory.h"
 #include "exec/address-spaces.h"
+#include "exec/ramblock.h"
 #include "hw/i386/e820_memory_layout.h"
 #include <sys/ioctl.h>
 
@@ -38,6 +39,8 @@ struct PSPDevState {
      * the TKM module uses different key spaces based on different vids.
     */
     uint32_t vid;
+    /* pinned hugepage numbers */
+    int hp_num;
 };
 
 #define PSP_DEV_PATH "/dev/hygon_psp_config"
@@ -45,6 +48,8 @@ struct PSPDevState {
 #define PSP_IOC_MUTEX_ENABLE    _IOWR(HYGON_PSP_IOC_TYPE, 1, NULL)
 #define PSP_IOC_MUTEX_DISABLE   _IOWR(HYGON_PSP_IOC_TYPE, 2, NULL)
 #define PSP_IOC_VPSP_OPT        _IOWR(HYGON_PSP_IOC_TYPE, 3, NULL)
+#define PSP_IOC_PIN_USER_PAGE   _IOWR(HYGON_PSP_IOC_TYPE, 4, NULL)
+#define PSP_IOC_UNPIN_USER_PAGE _IOWR(HYGON_PSP_IOC_TYPE, 5, NULL)
 
 enum VPSP_DEV_CTRL_OPCODE {
     VPSP_OP_VID_ADD,
@@ -69,6 +74,109 @@ struct psp_dev_ctrl {
     } __attribute__ ((packed)) data;
 };
 
+static MemoryRegion *find_memory_region_by_name(MemoryRegion *root, const char *name) {
+    MemoryRegion *subregion;
+    MemoryRegion *result;
+
+    if (strcmp(root->name, name) == 0)
+        return root;
+
+    QTAILQ_FOREACH(subregion, &root->subregions, subregions_link) {
+        result = find_memory_region_by_name(subregion, name);
+        if (result) {
+            return result;
+        }
+    }
+
+    return NULL;
+}
+
+static int pin_user_hugepage(int fd, uint64_t vaddr)
+{
+    int ret;
+
+    ret = ioctl(fd, PSP_IOC_PIN_USER_PAGE, vaddr);
+    /* 22: Invalid argument, some old kernel doesn't support this ioctl command */
+    if (ret != 0 && errno == EINVAL) {
+        ret = 0;
+    }
+    return ret;
+}
+
+static int unpin_user_hugepage(int fd, uint64_t vaddr)
+{
+    int ret;
+
+    ret = ioctl(fd, PSP_IOC_UNPIN_USER_PAGE, vaddr);
+    /* 22: Invalid argument, some old kernel doesn't support this ioctl command */
+    if (ret != 0 && errno == EINVAL) {
+        ret = 0;
+    }
+    return ret;
+}
+
+static int pin_psp_user_hugepages(struct PSPDevState *state, MemoryRegion *root)
+{
+    int ret = 0;
+    char mr_name[128] = {0};
+    int i, pinned_num;
+    MemoryRegion *find_mr = NULL;
+
+    for (i = 0 ; i < state->hp_num; ++i) {
+        sprintf(mr_name, "mem2-%d", i);
+        find_mr = find_memory_region_by_name(root, mr_name);
+        if (!find_mr) {
+            error_report("fail to find memory region by name %s.", mr_name);
+            ret = -ENOMEM;
+            goto end;
+        }
+
+        ret = pin_user_hugepage(state->dev_fd, (uint64_t)find_mr->ram_block->host);
+        if (ret) {
+            error_report("fail to pin_user_hugepage, ret: %d.", ret);
+            goto end;
+        }
+    }
+end:
+    if (ret) {
+        pinned_num = i;
+        for (i = 0 ; i < pinned_num; ++i) {
+            sprintf(mr_name, "mem2-%d", i);
+            find_mr = find_memory_region_by_name(root, mr_name);
+            if (!find_mr) {
+                continue;
+            }
+            unpin_user_hugepage(state->dev_fd, (uint64_t)find_mr->ram_block->host);
+        }
+
+    }
+    return ret;
+}
+
+static int unpin_psp_user_hugepages(struct PSPDevState *state, MemoryRegion *root)
+{
+    int ret = 0;
+    char mr_name[128] = {0};
+    int i;
+    MemoryRegion *find_mr = NULL;
+
+    for (i = 0 ; i < state->hp_num; ++i) {
+        sprintf(mr_name, "mem2-%d", i);
+        find_mr = find_memory_region_by_name(root, mr_name);
+        if (!find_mr) {
+            continue;
+        }
+
+        ret = unpin_user_hugepage(state->dev_fd, (uint64_t)find_mr->ram_block->host);
+        if (ret) {
+            error_report("fail to unpin_user_hugepage, ret: %d.", ret);
+            goto end;
+        }
+    }
+end:
+    return ret;
+}
+
 static void psp_dev_destroy(PSPDevState *state)
 {
     struct psp_dev_ctrl ctrl = { 0 };
@@ -77,6 +185,11 @@ static void psp_dev_destroy(PSPDevState *state)
             ctrl.op = VPSP_OP_VID_DEL;
             if (ioctl(state->dev_fd, PSP_IOC_VPSP_OPT, &ctrl) < 0) {
                 error_report("VPSP_OP_VID_DEL: %d", -errno);
+            }
+
+            /* Unpin hugepage memory */
+            if (unpin_psp_user_hugepages(state, get_system_memory())) {
+                error_report("unpin_psp_user_hugepages failed");
             } else {
                 state->enabled = false;
             }
@@ -99,23 +212,6 @@ static void psp_dev_shutdown_notify(Notifier *notifier, void *data)
     psp_dev_destroy(state);
 }
 
-static MemoryRegion *find_memory_region_by_name(MemoryRegion *root, const char *name) {
-    MemoryRegion *subregion;
-    MemoryRegion *result;
-
-    if (strcmp(root->name, name) == 0)
-        return root;
-
-    QTAILQ_FOREACH(subregion, &root->subregions, subregions_link) {
-        result = find_memory_region_by_name(subregion, name);
-        if (result) {
-            return result;
-        }
-    }
-
-    return NULL;
-}
-
 static void psp_dev_realize(DeviceState *dev, Error **errp)
 {
     int i;
@@ -150,6 +246,8 @@ static void psp_dev_realize(DeviceState *dev, Error **errp)
         ram2_end = find_mr->addr + find_mr->size - 1;
     }
 
+    state->hp_num = i;
+
     if (ram2_start != ram2_end) {
         ctrl.op = VPSP_OP_SET_GPA;
         ctrl.data.gpa.gpa_start = ram2_start;
@@ -159,6 +257,12 @@ static void psp_dev_realize(DeviceState *dev, Error **errp)
                         ram2_start, ram2_end, -errno);
             goto del_vid;
         }
+
+        /* Pin hugepage memory */
+        if(pin_psp_user_hugepages(state, root_mr)) {
+            error_setg(errp, "pin_psp_user_hugepages failed.");
+            goto del_vid;
+        }
     }
 
     state->enabled = true;
-- 
2.41.0.windows.1