From 27a611c3c14c3fd53684c19637880c3c4d73fd81 Mon Sep 17 00:00:00 2001
From: lfbzhm <lifubang@acmcoder.com>
Date: Sat, 28 Sep 2024 15:25:26 +0000
Subject: [PATCH] capability: add TestAmbientCapSet

Signed-off-by: lfbzhm <lifubang@acmcoder.com>
---
 capability/capability_test.go | 63 +++++++++++++++++++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/capability/capability_test.go b/capability/capability_test.go
index bb405e6..1120247 100644
--- a/capability/capability_test.go
+++ b/capability/capability_test.go
@@ -21,6 +21,16 @@ const (
 	maxLastCap = CAP_CHECKPOINT_RESTORE
 )
 
+func requirePCapSet(t *testing.T) {
+	pid, err := NewPid(0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !pid.Get(EFFECTIVE, CAP_SETPCAP) {
+		t.Skip("The test needs `CAP_SETPCAP`.")
+	}
+}
+
 func TestLastCap(t *testing.T) {
 	last, err := LastCap()
 	switch runtime.GOOS {
@@ -68,3 +78,56 @@ func TestListSupported(t *testing.T) {
 		t.Fatalf("result is too short (got %d, want %d): +%v", len(list), minLen, list)
 	}
 }
+
+func TestAmbientCapSet(t *testing.T) {
+	if runtime.GOOS != "linux" {
+		return
+	}
+	requirePCapSet(t)
+
+	capBounding := []Cap{CAP_KILL, CAP_CHOWN, CAP_SYSLOG}
+	capPermitted := []Cap{CAP_KILL, CAP_CHOWN}
+	capEffective := []Cap{CAP_KILL}
+	capInheritable := []Cap{CAP_KILL, CAP_CHOWN}
+	capAmbient := []Cap{CAP_KILL, CAP_CHOWN}
+
+	pid, err := newPid(0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pid.Set(BOUNDING, capBounding...)
+	pid.Set(PERMITTED, capPermitted...)
+	pid.Set(EFFECTIVE, capEffective...)
+	pid.Set(INHERITABLE, capInheritable...)
+	pid.Set(AMBIENT, capAmbient...)
+	if err = pid.Apply(CAPS | BOUNDING | AMBIENT); err != nil {
+		t.Fatal(err)
+	}
+
+	// Restore the cap set data from current process
+	if err = pid.Load(); err != nil {
+		t.Fatal(err)
+	}
+	for _, cap := range capAmbient {
+		if !pid.Get(AMBIENT, cap) {
+			t.Fatalf("expected ambient cap(%d) to be set but it's not", cap)
+		}
+	}
+
+	// Remove a ambient cap, to check `PR_CAP_AMBIENT_CLEAR_ALL` work.
+	pid.Clear(AMBIENT)
+	pid.Set(AMBIENT, capAmbient[0])
+	if err = pid.Apply(CAPS | BOUNDING | AMBIENT); err != nil {
+		t.Fatal(err)
+	}
+
+	if err = pid.Load(); err != nil {
+		t.Fatal(err)
+	}
+	if !pid.Get(AMBIENT, capAmbient[0]) {
+		t.Fatalf("expected ambient cap(%d) to be set but it's not", capAmbient[0])
+	}
+	if pid.Get(AMBIENT, capAmbient[1]) {
+		t.Fatalf("expected ambient cap(%d) not to be set but it has been set", capAmbient[1])
+	}
+}