From 3fc631926644deb38eb97360a8e3e81b8d7881e0 Mon Sep 17 00:00:00 2001
From: Piotr Resztak <piotr.resztak@gmail.com>
Date: Sat, 8 Oct 2022 21:50:08 +0200
Subject: [PATCH] Add uid,huid descriptors

Signed-off-by: Piotr Resztak <piotr.resztak@gmail.com>
---
 README.md        |  2 ++
 psgo.go          | 33 +++++++++++++++++++++++++++++++++
 test/format.bats | 18 +++++++++++++++++-
 test/list.bats   |  2 +-
 4 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 684c80a..e54423c 100644
--- a/README.md
+++ b/README.md
@@ -83,6 +83,8 @@ The ps library is compatible with all AIX format descriptors of the ps command-l
   - The corresponding host PID of a container process.
 - **huser**
   - The corresponding effective user of a container process on the host.
+- **huid**
+  - The corresponding host UID of a container process.
 - **label**
   - Current security attributes of the process.
 - **seccomp**
diff --git a/psgo.go b/psgo.go
index d6cfcef..5c2c4d0 100644
--- a/psgo.go
+++ b/psgo.go
@@ -181,6 +181,11 @@ var (
 			header: "USER",
 			procFn: processUSER,
 		},
+		{
+			normal: "uid",
+			header: "UID",
+			procFn: processUID,
+		},
 		{
 			code:   "%a",
 			normal: "args",
@@ -294,6 +299,12 @@ var (
 			onHost: true,
 			procFn: processHUSER,
 		},
+		{
+			normal: "huid",
+			header: "HUID",
+			onHost: true,
+			procFn: processHUID,
+		},
 		{
 			normal: "hgroup",
 			header: "HGROUP",
@@ -648,6 +659,11 @@ func processUSER(p *process.Process, ctx *psContext) (string, error) {
 	return process.LookupUID(p.Status.Uids[1])
 }
 
+// processUID returns the effective UID of the process as the decimal representation.
+func processUID(p *process.Process, ctx *psContext) (string, error) {
+	return p.Status.Uids[1], nil
+}
+
 // processRUSER returns the effective user name of the process.  This will be
 // the textual user ID, if it can be obtained, or a decimal representation
 // otherwise.
@@ -857,6 +873,23 @@ func processHUSER(p *process.Process, ctx *psContext) (string, error) {
 	return "?", nil
 }
 
+// processHUID returns the effective UID of the corresponding host process
+// of the (container) as the decimal representation or "?" if no corresponding
+// process could be found.
+func processHUID(p *process.Process, ctx *psContext) (string, error) {
+	if hp := findHostProcess(p, ctx); hp != nil {
+		if ctx.opts != nil && len(ctx.opts.UIDMap) > 0 {
+			// Return uid without searching its textual representation.
+			lookupFunc := func(uid string) (string, error) {
+				return uid, nil
+			}
+			return findID(hp.Status.Uids[1], ctx.opts.UIDMap, lookupFunc, "/proc/sys/fs/overflowuid")
+		}
+		return hp.Status.Uids[1], nil
+	}
+	return "?", nil
+}
+
 // processHGROUP returns the effective group ID of the corresponding host
 // process of the (container) or "?" if no corresponding process could be
 // found.
diff --git a/test/format.bats b/test/format.bats
index e0979e2..80a2e17 100644
--- a/test/format.bats
+++ b/test/format.bats
@@ -59,6 +59,12 @@
 	[[ ${lines[0]} =~ "USER" ]]
 }
 
+@test "UID header" {
+	run ./bin/psgo -format "uid"
+	[ "$status" -eq 0 ]
+	[[ ${lines[0]} =~ "UID" ]]
+}
+
 @test "COMMAND (args) header" {
 	run ./bin/psgo -format "%a"
 	[ "$status" -eq 0 ]
@@ -211,6 +217,14 @@
 	[[ ${lines[1]} =~ "?" ]]
 }
 
+@test "HUID header" {
+	run ./bin/psgo -format "huid"
+	[ "$status" -eq 0 ]
+	[[ ${lines[0]} =~ "HUID" ]]
+	# host UIDs are only extracted with `-pid`
+	[[ ${lines[1]} =~ "?" ]]
+}
+
 @test "HGROUP header" {
 	run ./bin/psgo -format "hgroup"
 	[ "$status" -eq 0 ]
@@ -265,7 +279,7 @@ function is_labeling_enabled() {
 }
 
 @test "ALL header" {
-	run ./bin/psgo -format "pcpu, group, groups, ppid, user, args, comm, rgroup, nice, pid, pgid, etime, ruser, time, tty, vsz, capamb, capinh, capprm, capeff, capbnd, seccomp, hpid, huser, hgroup, hgroups, rss, state"
+	run ./bin/psgo -format "pcpu, group, groups, ppid, user, uid, args, comm, rgroup, nice, pid, pgid, etime, ruser, time, tty, vsz, capamb, capinh, capprm, capeff, capbnd, seccomp, hpid, huser, huid, hgroup, hgroups, rss, state"
 	[ "$status" -eq 0 ]
 
 	[[ ${lines[0]} =~ "%CPU" ]]
@@ -273,6 +287,7 @@ function is_labeling_enabled() {
 	[[ ${lines[0]} =~ "GROUPS" ]]
 	[[ ${lines[0]} =~ "PPID" ]]
 	[[ ${lines[0]} =~ "USER" ]]
+	[[ ${lines[0]} =~ "UID" ]]
 	[[ ${lines[0]} =~ "COMMAND" ]]
 	[[ ${lines[0]} =~ "COMMAND" ]]
 	[[ ${lines[0]} =~ "RGROUP" ]]
@@ -291,6 +306,7 @@ function is_labeling_enabled() {
 	[[ ${lines[0]} =~ "SECCOMP" ]]
 	[[ ${lines[0]} =~ "HPID" ]]
 	[[ ${lines[0]} =~ "HUSER" ]]
+	[[ ${lines[0]} =~ "HUID" ]]
 	[[ ${lines[0]} =~ "HGROUP" ]]
 	[[ ${lines[0]} =~ "HGROUPS" ]]
 	[[ ${lines[0]} =~ "RSS" ]]
diff --git a/test/list.bats b/test/list.bats
index 4edfb1f..f31d00a 100644
--- a/test/list.bats
+++ b/test/list.bats
@@ -3,5 +3,5 @@
 @test "List descriptors" {
 	run ./bin/psgo -list
 	[ "$status" -eq 0 ]
-	[[ ${lines[0]} =~ "args, capamb, capbnd, capeff, capinh, capprm, comm, etime, group, groups, hgroup, hgroups, hpid, huser, label, nice, pcpu, pgid, pid, ppid, rgroup, rss, ruser, seccomp, state, stime, time, tty, user, vsz" ]]
+	[[ ${lines[0]} =~ "args, capamb, capbnd, capeff, capinh, capprm, comm, etime, group, groups, hgroup, hgroups, hpid, huid, huser, label, nice, pcpu, pgid, pid, ppid, rgroup, rss, ruser, seccomp, state, stime, time, tty, uid, user, vsz" ]]
 }