From 30da8d86b8bb339fe6b083ba5d64e99f10e1791a Mon Sep 17 00:00:00 2001
From: Jacky Zhao <j.zhao2k19@gmail.com>
Date: Mon, 26 Aug 2024 15:21:15 -0700
Subject: [PATCH] 3.2.0 add `cgroupPath` option (#41)

---
 index.d.ts          |  1 +
 package-lock.json   |  4 +--
 package.json        |  2 +-
 src/lib.rs          | 20 +++++++++++++++
 tests/index.test.ts | 60 ++++++++++++++++++++++++++++++++++++++++++++-
 5 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/index.d.ts b/index.d.ts
index f75b4c0..544fd76 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -10,6 +10,7 @@ export interface PtyOptions {
   envs?: Record<string, string>
   dir?: string
   size?: Size
+  cgroupPath?: string
   interactive?: boolean
   onExit: (err: null | Error, exitCode: number) => void
 }
diff --git a/package-lock.json b/package-lock.json
index 69ed1aa..23726ef 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "@replit/ruspty",
-  "version": "3.1.3",
+  "version": "3.2.0",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@replit/ruspty",
-      "version": "3.1.3",
+      "version": "3.2.0",
       "license": "MIT",
       "devDependencies": {
         "@napi-rs/cli": "^2.18.2",
diff --git a/package.json b/package.json
index 1886cbb..87bc796 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@replit/ruspty",
-  "version": "3.1.3",
+  "version": "3.2.0",
   "main": "dist/wrapper.js",
   "types": "dist/wrapper.d.ts",
   "author": "Szymon Kaliski <hi@szymonkaliski.com>",
diff --git a/src/lib.rs b/src/lib.rs
index 97a0a27..beb2ab2 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,6 +1,8 @@
 use std::collections::HashMap;
+use std::fs::File;
 use std::io::Error;
 use std::io::ErrorKind;
+use std::io::Write;
 use std::os::fd::{AsRawFd, OwnedFd};
 use std::os::fd::{BorrowedFd, FromRawFd, IntoRawFd, RawFd};
 use std::os::unix::process::CommandExt;
@@ -40,6 +42,7 @@ struct PtyOptions {
   pub envs: Option<HashMap<String, String>>,
   pub dir: Option<String>,
   pub size: Option<Size>,
+  pub cgroup_path: Option<String>,
   pub interactive: Option<bool>,
   #[napi(ts_type = "(err: null | Error, exitCode: number) => void")]
   pub on_exit: JsFunction,
@@ -105,6 +108,14 @@ impl Pty {
   #[napi(constructor)]
   #[allow(dead_code)]
   pub fn new(_env: Env, opts: PtyOptions) -> Result<Self, napi::Error> {
+    #[cfg(not(target_os = "linux"))]
+    if opts.cgroup_path.is_some() {
+      return Err(napi::Error::new(
+        napi::Status::GenericFailure,
+        "cgroup_path is only supported on Linux",
+      ));
+    }
+
     let size = opts.size.unwrap_or(Size { cols: 80, rows: 24 });
     let window_size = Winsize {
       ws_col: size.cols,
@@ -160,6 +171,15 @@ impl Pty {
           return Err(Error::new(ErrorKind::Other, "setsid"));
         }
 
+        // set the cgroup if specified
+        #[cfg(target_os = "linux")]
+        if let Some(cgroup_path) = &opts.cgroup_path {
+          let pid = libc::getpid();
+          let cgroup_path = format!("{}/cgroup.procs", cgroup_path);
+          let mut cgroup_file = File::create(cgroup_path)?;
+          cgroup_file.write_all(format!("{}", pid).as_bytes())?;
+        }
+
         // become the controlling tty for the program
         let err = libc::ioctl(raw_user_fd, libc::TIOCSCTTY.into(), 0);
         if err == -1 {
diff --git a/tests/index.test.ts b/tests/index.test.ts
index d0bf0c0..ae2bda0 100644
--- a/tests/index.test.ts
+++ b/tests/index.test.ts
@@ -1,13 +1,17 @@
 import { Pty, getCloseOnExec, setCloseOnExec } from '../wrapper';
 import { type Writable } from 'stream';
 import { readdirSync, readlinkSync } from 'fs';
-import { describe, test, expect } from 'vitest';
+import { describe, test, expect, beforeEach, afterEach } from 'vitest';
+import { exec as execAsync } from 'child_process';
+import { promisify } from 'util';
+const exec = promisify(execAsync);
 
 const EOT = '\x04';
 const procSelfFd = '/proc/self/fd/';
 const IS_DARWIN = process.platform === 'darwin';
 
 const testSkipOnDarwin = IS_DARWIN ? test.skip : test;
+const testOnlyOnDarwin = IS_DARWIN ? test : test.skip;
 
 type FdRecord = Record<string, string>;
 function getOpenFds(): FdRecord {
@@ -414,6 +418,60 @@ describe(
   { repeats: 50 },
 );
 
+describe('cgroup opts', () => {
+  beforeEach(async () => {
+    if (!IS_DARWIN) {
+      // create a new cgroup with the right permissions
+      await exec("sudo cgcreate -g 'cpu:/test.slice'")
+      await exec("sudo chown -R $(id -u):$(id -g) /sys/fs/cgroup/cpu/test.slice")
+    }
+  });
+
+  afterEach(async () => {
+    if (!IS_DARWIN) {
+      // remove the cgroup
+      await exec("sudo cgdelete cpu:/test.slice")
+    }
+  });
+
+  testSkipOnDarwin('basic cgroup', () => new Promise<void>((done) => {
+    const oldFds = getOpenFds();
+    let buffer = '';
+    const pty = new Pty({
+      command: '/bin/cat',
+      args: ['/proc/self/cgroup'],
+      cgroupPath: '/sys/fs/cgroup/cpu/test.slice',
+      onExit: (err, exitCode) => {
+        expect(err).toBeNull();
+        expect(exitCode).toBe(0);
+        expect(buffer).toContain('/test.slice');
+        expect(getOpenFds()).toStrictEqual(oldFds);
+        done();
+      },
+    });
+
+    const readStream = pty.read;
+    readStream.on('data', (data) => {
+      buffer = data.toString();
+    });
+  })
+  );
+
+  testOnlyOnDarwin('cgroup is not supported on darwin', () => {
+    expect(() => {
+      new Pty({
+        command: '/bin/cat',
+        args: ['/proc/self/cgroup'],
+        cgroupPath: '/sys/fs/cgroup/cpu/test.slice',
+        onExit: (err, exitCode) => {
+          expect(err).toBeNull();
+          expect(exitCode).toBe(0);
+        },
+      })
+    }).toThrowError();
+  });
+});
+
 describe('setCloseOnExec', () => {
   test('setCloseOnExec', () => {
     // stdio typically never has the close-on-exec flag since it's always expected to be