From f89c8c58bdd5348f54ac22d0d58cf797c35bdc2b Mon Sep 17 00:00:00 2001
From: Lencerf <changyuan.lv@Gmail.com>
Date: Mon, 13 Nov 2023 00:26:43 -0800
Subject: [PATCH] Add AMD Memory Encrypt detection (#140)

This patch detects additional information of the AMD Memory Encryption
feature.

[0]:
https://www.amd.com/content/dam/amd/en/documents/processor-tech-docs/programmer-references/24594.pdf
table E.4.17

Signed-off-by: Changyuan Lyu <changyuanl@google.com>
---
 cmd/cpuid/main.go |  3 +++
 cpuid.go          | 35 +++++++++++++++++++++++++++++++----
 cpuid_test.go     | 10 ++++++++++
 detect_x86.go     |  1 +
 4 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/cmd/cpuid/main.go b/cmd/cpuid/main.go
index 74a8e5b..904e427 100644
--- a/cmd/cpuid/main.go
+++ b/cmd/cpuid/main.go
@@ -82,4 +82,7 @@ func main() {
 	if cpuid.CPU.SGX.Available {
 		fmt.Printf("SGX: %+v\n", cpuid.CPU.SGX)
 	}
+	if cpuid.CPU.AMDMemEncryption.Available {
+		fmt.Printf("AMD Memory Encryption: %+v\n", cpuid.CPU.AMDMemEncryption)
+	}
 }
diff --git a/cpuid.go b/cpuid.go
index 15b7603..64b4cab 100644
--- a/cpuid.go
+++ b/cpuid.go
@@ -309,10 +309,11 @@ type CPUInfo struct {
 		L2  int // L2 Cache (per core or shared). Will be -1 if undetected
 		L3  int // L3 Cache (per core, per ccx or shared). Will be -1 if undetected
 	}
-	SGX        SGXSupport
-	AVX10Level uint8
-	maxFunc    uint32
-	maxExFunc  uint32
+	SGX              SGXSupport
+	AMDMemEncryption AMDMemEncryptionSupport
+	AVX10Level       uint8
+	maxFunc          uint32
+	maxExFunc        uint32
 }
 
 var cpuid func(op uint32) (eax, ebx, ecx, edx uint32)
@@ -1079,6 +1080,32 @@ func hasSGX(available, lc bool) (rval SGXSupport) {
 	return
 }
 
+type AMDMemEncryptionSupport struct {
+	Available          bool
+	CBitPossition      uint32
+	NumVMPL            uint32
+	PhysAddrReduction  uint32
+	NumEntryptedGuests uint32
+	MinSevNoEsAsid     uint32
+}
+
+func hasAMDMemEncryption(available bool) (rval AMDMemEncryptionSupport) {
+	rval.Available = available
+	if !available {
+		return
+	}
+
+	_, b, c, d := cpuidex(0x8000001f, 0)
+
+	rval.CBitPossition = b & 0x3f
+	rval.PhysAddrReduction = (b >> 6) & 0x3F
+	rval.NumVMPL = (b >> 12) & 0xf
+	rval.NumEntryptedGuests = c
+	rval.MinSevNoEsAsid = d
+
+	return
+}
+
 func support() flagSet {
 	var fs flagSet
 	mfi := maxFunctionID()
diff --git a/cpuid_test.go b/cpuid_test.go
index 7c1b02e..3131e8c 100644
--- a/cpuid_test.go
+++ b/cpuid_test.go
@@ -136,6 +136,16 @@ func TestSGX(t *testing.T) {
 	}
 }
 
+// TestAMDMemEncryption tests AMDMemEncryption detection
+func TestAMDMemEncryption(t *testing.T) {
+	got := CPU.AMDMemEncryption.Available
+	expected := CPU.featureSet.inSet(SME) || CPU.featureSet.inSet(SEV)
+	if got != expected {
+		t.Fatalf("AMDMemEncryption: expected %v, got %v", expected, got)
+	}
+	t.Log("AMDMemEncryption Support:", got)
+}
+
 func TestHas(t *testing.T) {
 	Detect()
 	defer Detect()
diff --git a/detect_x86.go b/detect_x86.go
index c7dfa12..799b400 100644
--- a/detect_x86.go
+++ b/detect_x86.go
@@ -27,6 +27,7 @@ func addInfo(c *CPUInfo, safe bool) {
 	c.Family, c.Model, c.Stepping = familyModel()
 	c.featureSet = support()
 	c.SGX = hasSGX(c.featureSet.inSet(SGX), c.featureSet.inSet(SGXLC))
+	c.AMDMemEncryption = hasAMDMemEncryption(c.featureSet.inSet(SME) || c.featureSet.inSet(SEV))
 	c.ThreadsPerCore = threadsPerCore()
 	c.LogicalCores = logicalCores()
 	c.PhysicalCores = physicalCores()