From 80fb95e06ddc5f311e0c9dc20b19d00058ca39af Mon Sep 17 00:00:00 2001
From: Brian Chen <brianc118@fb.com>
Date: Tue, 31 Aug 2021 14:59:57 -0700
Subject: [PATCH] Fix procfs reading cgroups with colon in name

Summary:
When reading cgroup membership, we split by ':' and expect there to
be 3 parts. This is not always the case since cgroup name can
contain ':'.

Reviewed By: lnyng

Differential Revision: D30681159

fbshipit-source-id: 895e1b26965ae33454a19c2ec1bc5478a5f95635
---
 below/procfs/src/lib.rs  |  2 +-
 below/procfs/src/test.rs | 34 ++++++++++++++++++++++++++++++----
 2 files changed, 31 insertions(+), 5 deletions(-)

--- a/below/procfs/src/lib.rs
+++ b/below/procfs/src/lib.rs
@@ -565,7 +565,7 @@ impl ProcReader {
             // A line starting with "0::" would be an entry for cgroup v2.
             // Otherwise, the line containing "pids" controller is what we want
             // for cgroup v1.
-            let parts: Vec<_> = line.split(':').collect();
+            let parts: Vec<_> = line.splitn(3, ':').collect();
             if parts.len() == 3 {
                 if parts[0] == "0" && parts[1] == "" {
                     cgroup_path = Some(parts[2].to_owned());
--- a/below/procfs/src/test.rs
+++ b/below/procfs/src/test.rs
@@ -804,7 +804,7 @@ cancelled_write_bytes: 5431947264

 #[test]
 fn test_pid_cgroupv2() {
-    let cgroup = b"0::/user.slice/user-119756.slice/session-3.scope
+    let cgroup = b"0::/user.slice/user:with:colon.slice/session-3.scope
 ";

     let procfs = TestProcfs::new();
@@ -814,12 +814,12 @@ fn test_pid_cgroupv2() {
         .read_pid_cgroup(1024)
         .expect("Failed to read pid cgroup file");

-    assert_eq!(cgroup, "/user.slice/user-119756.slice/session-3.scope");
+    assert_eq!(cgroup, "/user.slice/user:with:colon.slice/session-3.scope");
 }

 #[test]
 fn test_pid_cgroupv1() {
-    let cgroup = b"11:pids:/init.scope
+    let cgroup = b"11:pids:/cgroup-path:colon
 10:perf_event:/
 9:hugetlb:/
 8:cpu,cpuacct:/init.scope
@@ -838,7 +838,33 @@ fn test_pid_cgroupv1() {
         .read_pid_cgroup(1024)
         .expect("Failed to read pid cgroup file");

-    assert_eq!(cgroup, "/init.scope");
+    assert_eq!(cgroup, "/cgroup-path:colon");
+}
+
+#[test]
+fn test_pid_cgroupv1and2() {
+    let cgroup = b"11:pids:/cgroup-path:colon
+10:perf_event:/
+9:hugetlb:/
+8:cpu,cpuacct:/init.scope
+7:blkio:/init.scope
+6:freezer:/
+5:cpuset:/
+4:memory:/init.scope
+3:devices:/init.scope
+2:net_cls,net_prio:/
+1:name=systemd:/init.scope
+0::/user.slice/user:with:colon.slice/session-3.scope";
+
+    let procfs = TestProcfs::new();
+    procfs.create_pid_file_with_content(1024, "cgroup", cgroup);
+    let reader = procfs.get_reader();
+    let cgroup = reader
+        .read_pid_cgroup(1024)
+        .expect("Failed to read pid cgroup file");
+
+    // When we see both cgroup v1 and v2, v2 takes precedence
+    assert_eq!(cgroup, "/user.slice/user:with:colon.slice/session-3.scope");
 }

 #[test]