feat: display process tree with environment variables

checkpointctl.go

@@ -20,6 +20,7 @@ var (
 	pID       uint32
 	psTree    bool
 	psTreeCmd bool
+	psTreeEnv bool
 	files     bool
 	showAll   bool
 )
@@ -110,6 +111,12 @@ func setupInspect() *cobra.Command {
 		false,
 		"Display an overview of processes in the container checkpoint with full command line arguments",
 	)
+	flags.BoolVar(
+		&psTreeEnv,
+		"ps-tree-env",
+		false,
+		"Display an overview of processes in the container checkpoint with their environment variables",
+	)
 	flags.BoolVar(
 		&files,
 		"files",
@@ -168,8 +175,8 @@ func inspect(cmd *cobra.Command, args []string) error {
 		)
 	}
 
-	if psTreeCmd {
-		// Enable displaying process tree when using --ps-tree-cmd.
+	if psTreeCmd || psTreeEnv {
+		// Enable displaying process tree when using --ps-tree-cmd or --ps-tree-env.
 		psTree = true
 		requiredFiles = append(
 			requiredFiles,

container.go

@@ -328,3 +328,23 @@ func getCmdline(checkpointOutputDir string, pid uint32) (cmdline string, err err
 	cmdline = strings.Join(strings.Split(buffer.String(), "\x00"), " ")
 	return
 }
+
+func getPsEnvVars(checkpointOutputDir string, pid uint32) (envVars []string, err error) {
+	mr, err := crit.NewMemoryReader(filepath.Join(checkpointOutputDir, metadata.CheckpointDirectory), pid, pageSize)
+	if err != nil {
+		return
+	}
+
+	buffer, err := mr.GetPsEnvVars()
+	if err != nil {
+		return
+	}
+
+	for _, envVar := range strings.Split(buffer.String(), "\x00") {
+		if envVar != "" {
+			envVars = append(envVars, envVar)
+		}
+	}
+
+	return
+}

test/checkpointctl.bats

@@ -315,6 +315,37 @@ function teardown() {
 	[[ ${lines[0]} == *"failed to process command line arguments"* ]]
 }
 
+@test "Run checkpointctl inspect with tar file and --ps-tree-env" {
+	cp data/config.dump \
+		data/spec.dump "$TEST_TMP_DIR1"
+	mkdir "$TEST_TMP_DIR1"/checkpoint
+	cp test-imgs/pstree.img \
+		test-imgs/core-*.img \
+		test-imgs/pagemap-*.img \
+		test-imgs/pages-*.img \
+		test-imgs/mm-*.img "$TEST_TMP_DIR1"/checkpoint
+	( cd "$TEST_TMP_DIR1" && tar cf "$TEST_TMP_DIR2"/test.tar . )
+	checkpointctl inspect "$TEST_TMP_DIR2"/test.tar --ps-tree-env
+	[ "$status" -eq 0 ]
+	[[ ${lines[8]} == *"Process tree"* ]]
+	[[ ${lines[9]} == *"piggie"* ]]
+	[[ ${lines[10]} == *"="* ]]
+}
+
+@test "Run checkpointctl inspect with tar file and --ps-tree-env and missing pages-*.img {
+	cp data/config.dump \
+		data/spec.dump "$TEST_TMP_DIR1"
+	mkdir "$TEST_TMP_DIR1"/checkpoint
+	cp test-imgs/pstree.img \
+		test-imgs/core-*.img \
+		test-imgs/pagemap-*.img \
+		test-imgs/mm-*.img "$TEST_TMP_DIR1"/checkpoint
+	( cd "$TEST_TMP_DIR1" && tar cf "$TEST_TMP_DIR2"/test.tar . )
+	checkpointctl inspect "$TEST_TMP_DIR2"/test.tar --ps-tree-env
+	[ "$status" -eq 1 ]
+	[[ ${lines[0]} == *"no such file or directory"* ]]
+}
+
 @test "Run checkpointctl inspect with tar file and --files" {
 	cp data/config.dump \
 		data/spec.dump "$TEST_TMP_DIR1"

tree.go

@@ -43,7 +43,7 @@ func renderTreeView(tasks []task) error {
 				}
 			}
 
-			if err = addPsTreeToTree(tree, psTree); err != nil {
+			if err = addPsTreeToTree(tree, psTree, task.outputDir); err != nil {
 				return fmt.Errorf("failed to get process tree: %w", err)
 			}
 		}
@@ -118,7 +118,7 @@ func addDumpStatsToTree(tree treeprint.Tree, dumpStats *stats_pb.DumpStatsEntry)
 	statsTree.AddBranch(fmt.Sprintf("Pages written: %d", dumpStats.GetPagesWritten()))
 }
 
-func addPsTreeToTree(tree treeprint.Tree, psTree *crit.PsTree) error {
+func addPsTreeToTree(tree treeprint.Tree, psTree *crit.PsTree, checkpointOutputDir string) error {
 	psRoot := psTree
 	if pID != 0 {
 		// dfs performs a short-circuiting depth-first search.
@@ -145,17 +145,30 @@ func addPsTreeToTree(tree treeprint.Tree, psTree *crit.PsTree) error {
 	// processNodes is a recursive function to create
 	// a new branch for each process and add its child
 	// processes as child nodes of the branch.
-	var processNodes func(treeprint.Tree, *crit.PsTree)
-	processNodes = func(tree treeprint.Tree, root *crit.PsTree) {
+	var processNodes func(treeprint.Tree, *crit.PsTree) error
+	processNodes = func(tree treeprint.Tree, root *crit.PsTree) error {
 		node := tree.AddMetaBranch(root.PID, root.Comm)
+		// attach environment variables to process
+		if psTreeEnv {
+			envVars, err := getPsEnvVars(checkpointOutputDir, root.PID)
+			if err != nil {
+				return err
+			}
+
+			for _, env := range envVars {
+				node.AddBranch(env)
+			}
+		}
 		for _, child := range root.Children {
-			processNodes(node, child)
+			if err := processNodes(node, child); err != nil {
+				return err
+			}
 		}
+		return nil
 	}
 	psTreeNode := tree.AddBranch("Process tree")
-	processNodes(psTreeNode, psRoot)
 
-	return nil
+	return processNodes(psTreeNode, psRoot)
 }
 
 func addFdsToTree(tree treeprint.Tree, fds []*crit.Fd) {