Castor-Gemini/proctitle_linux.go

139 lines
4.1 KiB
Go

//go:build linux
package main
/*
#include <string.h>
#include <stdlib.h>
// Pointers to the original argv memory block.
// These are populated by the init() constructor function before Go's main runs.
static char **argv_start = NULL;
static int argv_len = 0;
// A C constructor function that runs before the Go runtime initializes.
// It saves the original argc and argv passed to the program.
__attribute__((constructor))
static void init(int argc, char **argv) {
if (argc == 0 || argv == NULL || *argv == NULL) {
return;
}
// Calculate the total length of the memory occupied by the arguments.
// This is the maximum length our new title can be.
for (int i = 0; i < argc; i++) {
// The space for the argument string plus the null terminator.
argv_len += strlen(argv[i]) + 1;
}
// The last argument is followed by a NULL pointer, so we back up one byte
// from that to get to the final null terminator of the last argument string.
argv_len--;
argv_start = argv;
}
// set_title performs the actual memory modification.
static void set_title(const char *title) {
// If argv_start is null, the constructor didn't run, so we can't do anything.
if (argv_start == NULL) {
return;
}
// Check if the new title is too long.
int title_len = strlen(title);
if (title_len >= argv_len) {
// If it is, we can't safely set it, so we do nothing.
return;
}
// Copy the new title into the argument memory.
strcpy(argv_start[0], title);
// Zero out the rest of the original argument memory space.
// This is important to prevent old argument data from appearing in `ps`.
memset(argv_start[0] + title_len, 0, argv_len - title_len);
// Some systems might need the pointer after the first argument to be NULL.
if (argv_start[1] != NULL) {
argv_start[1] = NULL;
}
}
*/
import "C"
import (
"bytes"
"fmt"
"os"
"time"
"unsafe"
)
// SetProcTitle sets the process title for `ps` to display.
// This is a Linux-only implementation using cgo.
// The new title cannot be longer than the original command line string.
func SetProcTitle(title string) {
// C.CString allocates memory in the C heap. We must free it.
cs := C.CString(title)
defer C.free(unsafe.Pointer(cs))
C.set_title(cs)
}
// GetProcTitle retrieves the current process title as seen by `ps`.
// This is a Linux-only implementation that reads the special /proc/self/cmdline file.
func GetProcTitle() (string, error) {
// On Linux, the kernel exposes the command line of a process in /proc/[pid]/cmdline.
// /proc/self is a symlink to the current process's directory.
// The arguments are separated by null bytes.
cmdlineBytes, err := os.ReadFile("/proc/self/cmdline")
if err != nil {
return "", fmt.Errorf("could not read /proc/self/cmdline: %w", err)
}
// The process title we set is the first null-terminated string in this file.
// We find the first null byte to isolate the title.
firstNull := bytes.IndexByte(cmdlineBytes, 0)
if firstNull == -1 {
// This would be unusual, but if there are no nulls, return the whole content.
return string(cmdlineBytes), nil
}
return string(cmdlineBytes[:firstNull]), nil
}
func main() {
pid := os.Getpid()
fmt.Printf("PID: %d\n", pid)
// Get and show the initial title
initialTitle, err := GetProcTitle()
if err != nil {
fmt.Fprintf(os.Stderr, "Could not get initial title: %v\n", err)
} else {
fmt.Printf("Initial title via GetProcTitle(): '%s'\n", initialTitle)
}
fmt.Printf("--> Check it now with: ps -f -p %d\n", pid)
fmt.Println("Waiting 15 seconds before changing the title...")
time.Sleep(15 * time.Second)
// Now, set a new title.
newTitle := "MyGoProcess (processing tasks)"
fmt.Printf("\nSetting title to: '%s'\n", newTitle)
SetProcTitle(newTitle)
// Get and show the new title to confirm it was set
readTitle, err := GetProcTitle()
if err != nil {
fmt.Fprintf(os.Stderr, "Could not get new title: %v\n", err)
} else {
fmt.Printf("Retrieved title via GetProcTitle(): '%s'\n", readTitle)
}
fmt.Printf("--> Check it again with: ps -f -p %d\n", pid)
fmt.Println("The process will exit in 15 seconds.")
time.Sleep(15 * time.Second)
fmt.Println("\nDone.")
}