From 988e8b5e2b56c3e0782197ddefc0d05b13177b1b Mon Sep 17 00:00:00 2001
From: amaslennikov <amaslennikov@nvidia.com>
Date: Fri, 15 Dec 2023 15:25:13 +0300
Subject: [PATCH] Add GetPkeyFromPci function

The function reads the from the PCI device's
infiniband dir and returns the index0 PKey

Signed-off-by: amaslennikov <amaslennikov@nvidia.com>
---
 sriovnet.go      | 29 +++++++++++++++++++++++++++++
 sriovnet_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/sriovnet.go b/sriovnet.go
index 5e9ff38..0ab4e78 100644
--- a/sriovnet.go
+++ b/sriovnet.go
@@ -520,3 +520,32 @@ func GetPciFromNetDevice(name string) (string, error) {
 	}
 	return base, nil
 }
+
+// GetPKeyFromPci returns the index0 PKey for the IB PCI device
+func GetPKeyFromPci(pciAddress string) (string, error) {
+	pciDir := filepath.Join(PciSysDir, pciAddress, "infiniband")
+	dirEntries, err := utilfs.Fs.ReadDir(pciDir)
+	if err != nil {
+		return "", err
+	}
+	if len(dirEntries) == 0 {
+		return "", fmt.Errorf("infiniband directory is empty for device: %s", pciAddress)
+	}
+
+	portsDir := filepath.Join(pciDir, dirEntries[0].Name(), "ports")
+	dirEntries, err = utilfs.Fs.ReadDir(portsDir)
+	if err != nil {
+		return "", err
+	}
+	if len(dirEntries) == 0 {
+		return "", fmt.Errorf("infiniband ports directory is empty for device: %s", pciAddress)
+	}
+
+	index0FilePath := filepath.Join(portsDir, dirEntries[0].Name(), "pkeys", "0")
+	pKeyBytes, err := utilfs.Fs.ReadFile(index0FilePath)
+	if err != nil {
+		return "", err
+	}
+
+	return string(pKeyBytes), nil
+}
diff --git a/sriovnet_test.go b/sriovnet_test.go
index 8ed9c4e..a3f96b0 100644
--- a/sriovnet_test.go
+++ b/sriovnet_test.go
@@ -246,3 +246,33 @@ func TestGetPciFromNetDeviceNotPCI(t *testing.T) {
 	assert.Error(t, err)
 	assert.Contains(t, err.Error(), "is not a PCI device")
 }
+
+func TestGetPKeyFromPci(t *testing.T) {
+	teardown := setupFakeFs(t)
+	defer teardown()
+
+	devices := map[string]struct {
+		path string
+		pkey string
+	}{
+		"0000:03:00.2": {"/sys/bus/pci/devices/0000:03:00.2/infiniband/mlx5_2/ports/1/pkeys/", "0x66"},
+		"0000:03:00.3": {"/sys/bus/pci/devices/0000:03:00.3/infiniband/mlx5_3/ports/2/pkeys/", "0x424"},
+	}
+
+	for _, v := range devices {
+		err := utilfs.Fs.MkdirAll(v.path, os.FileMode(0755))
+		assert.NoError(t, err)
+		file, err := utilfs.Fs.Create(filepath.Join(v.path, "0"))
+		assert.NoError(t, err)
+		_, err = file.Write([]byte(v.pkey))
+		assert.NoError(t, err)
+		err = file.Close()
+		assert.NoError(t, err)
+	}
+
+	for k, v := range devices {
+		pKey, err := GetPKeyFromPci(k)
+		assert.NoError(t, err)
+		assert.Equal(t, v.pkey, pKey)
+	}
+}