Change command line to string slice.

jviney · May 6, 2019 · 963546f · 963546f
1 parent a0f7169
commit 963546f
Show file tree

Hide file tree

Showing 5 changed files with 103 additions and 113 deletions.
diff --git a/README.md b/README.md
@@ -15,7 +15,7 @@ Inspired by [sys-proctable](http://github.com/djberg96/sys-proctable).
   if process := proc.GetProcess(cmd.Process.Pid); process != nil {
     process.Pid # <pid>
     process.Command # "sleep"
-    process.ComandLine # "sleep 5"
+    process.ComandLine # []string{"sleep", "5"}
   }
 ```
 

diff --git a/proc.go b/proc.go
@@ -1,20 +1,21 @@
 package proc
 
 type ProcessInfo struct {
-  Pid int
-  Command string
-  CommandLine string
+	Pid         int
+	Command     string
+	CommandLine []string
 }
 
 func GetProcessInfo(pid int) *ProcessInfo {
-  processes := ps(pid)
+	processes := ps(pid)
 
-  if len(processes) == 1 {
-    return processes[0]
-  }
-  return nil
+	if len(processes) == 1 {
+		return processes[0]
+	}
+
+	return nil
 }
 
 func GetAllProcessesInfo() []*ProcessInfo {
-  return ps(-1)
+	return ps(-1)
 }
diff --git a/proc_darwin.go b/proc_darwin.go
@@ -55,7 +55,7 @@ func ps(pid int) []*ProcessInfo {
 
     commandParts := strings.Split(command, "/")
     process.Command = strings.TrimSpace(commandParts[len(commandParts) - 1])
-    process.CommandLine += strings.Join(argv, " ")
+    process.CommandLine = argv
 
     processes = append(processes, &process)
 

diff --git a/proc_linux.go b/proc_linux.go
@@ -3,58 +3,64 @@
 package proc
 
 import (
-  "io/ioutil"
-  "regexp"
-  "strings"
-  "strconv"
+	"bytes"
+	"io/ioutil"
+	"regexp"
+	"strconv"
+	"strings"
 )
 
 func ps(pid int) []*ProcessInfo {
-  processes := []*ProcessInfo{}
-  files, _ := ioutil.ReadDir("/proc")
+	processes := []*ProcessInfo{}
+	files, _ := ioutil.ReadDir("/proc")
 
-  for _, file := range files {
-    procPid, err := strconv.Atoi(file.Name())
+	for _, file := range files {
+		procPid, err := strconv.Atoi(file.Name())
 
-    // Ignore non-numeric entries
-    if err != nil {
-      continue
-    }
+		// Ignore non-numeric entries
+		if err != nil {
+			continue
+		}
 
-    // Ignore a non-matching pid
-    if pid >= 0 && procPid != pid {
-      continue
-    }
+		// Ignore a non-matching pid
+		if pid >= 0 && procPid != pid {
+			continue
+		}
 
-    process := ProcessInfo{Pid: procPid}
+		process := ProcessInfo{Pid: procPid, CommandLine: []string{}}
 
-    if commandLine, err := ioutil.ReadFile("/proc/" + file.Name() + "/cmdline"); err != nil {
-      continue // Process terminated
-    } else {
-      process.CommandLine = strings.TrimSpace(strings.Replace(string(commandLine), "\000", " ", -1))
-    }
+		if commandLine, err := ioutil.ReadFile("/proc/" + file.Name() + "/cmdline"); err != nil {
+			continue // Process terminated
+		} else {
+			args := bytes.Split(commandLine, []byte{'\x00'})
+			for _, arg := range args {
+				strArg := strings.TrimSpace(string(arg))
+				if len(strArg) == 0 {
+					continue
+				}
 
-    if stat, err := ioutil.ReadFile("/proc/" + file.Name() + "/stat"); err != nil {
-      continue // Process terminated
-    } else {
-      statRegex := regexp.MustCompile("\\(([^\\)]+)\\)")
-      parts := statRegex.FindStringSubmatch(string(stat))
+				process.CommandLine = append(process.CommandLine, strArg)
+			}
+		}
 
-      if len(parts) == 2 {
-        process.Command = strings.TrimSpace(parts[1])
-      }
-    }
+		if stat, err := ioutil.ReadFile("/proc/" + file.Name() + "/stat"); err != nil {
+			continue // Process terminated
+		} else {
+			statRegex := regexp.MustCompile("\\(([^\\)]+)\\)")
+			parts := statRegex.FindStringSubmatch(string(stat))
 
-    if process.CommandLine == "" {
-      process.CommandLine = process.Command
-    }
+			if len(parts) == 2 {
+				process.Command = strings.TrimSpace(parts[1])
+			}
+		}
 
-    processes = append(processes, &process)
+		processes = append(processes, &process)
 
-    if process.Pid == pid {
-      break
-    }
-  }
+		// Break if this is the one process we were looking for
+		if process.Pid == pid {
+			break
+		}
+	}
 
-  return processes
+	return processes
 }
diff --git a/proc_test.go b/proc_test.go
@@ -1,91 +1,74 @@
 package proc
 
-import(
-  "os/exec"
-  "sync"
-  "testing"
+import (
+	"os/exec"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestGetProcessNegativePid(t *testing.T) {
-  if GetProcessInfo(-1) != nil {
-    t.Errorf("expected nil")
-  }
+	assert.Nil(t, GetProcessInfo(-1))
 }
 
 func TestGetProcessZero(t *testing.T) {
-  GetProcessInfo(0)
+	assert.Nil(t, GetProcessInfo(0))
 }
 
 func TestGetProcessMissingPid(t *testing.T) {
-  if GetProcessInfo(99999999) != nil {
-    t.Errorf("expected nil")
-  }
+	assert.Nil(t, GetProcessInfo(99999999))
 }
 
 func TestProcessFields(t *testing.T) {
-  cmd := exec.Command("sleep", "5")
-  cmd.Start()
-
-  if process := GetProcessInfo(cmd.Process.Pid); process == nil {
-    t.Errorf("failed to find process")
-  } else {
-    if process.Command != "sleep" {
-      t.Errorf("expected %s got %s", "sleep", process.Command)
-    }
-
-    if process.CommandLine != "sleep 5" {
-      t.Errorf("expected '%s' got '%s'", "sleep 5", process.CommandLine)
-    }
-  }
-
-  cmd.Process.Kill()
-  cmd.Wait()
+	cmd := exec.Command("sleep", "5")
+	cmd.Start()
+
+	process := GetProcessInfo(cmd.Process.Pid)
+	assert.NotNil(t, process)
+
+	assert.Equal(t, "sleep", process.Command)
+	assert.Equal(t, []string{"sleep", "5"}, process.CommandLine)
+
+	cmd.Process.Kill()
+	cmd.Wait()
 }
 
 func TestLongCommandLine(t *testing.T) {
-  cmd := exec.Command("echo", "1", "2", "3")
-  cmd.Start()
-
-  if process := GetProcessInfo(cmd.Process.Pid); process == nil {
-    t.Errorf("failed to find process")
-  } else {
-    if process.Command != "echo" {
-      t.Fatal("expected echo")
-    }
-
-    if process.CommandLine != "echo 1 2 3" {
-      t.Errorf("expected %s got %s", "echo 1 2 3", process.CommandLine)
-    }
-  }
-
-  cmd.Wait()
+	cmd := exec.Command("dd", "if=/dev/zero", "of=/dev/null")
+	cmd.Start()
+
+	process := GetProcessInfo(cmd.Process.Pid)
+	assert.NotNil(t, process)
+
+	assert.Equal(t, "dd", process.Command)
+	assert.Equal(t, []string{"dd", "if=/dev/zero", "of=/dev/null"}, process.CommandLine)
+
+	cmd.Process.Kill()
+	cmd.Wait()
 }
 
 func TestGetProcessGoRoutines(t *testing.T) {
-  cmd := exec.Command("sleep", "5")
-  cmd.Start()
+	cmd := exec.Command("sleep", "5")
+	cmd.Start()
 
-  count := 100
+	count := 100
 
-  var wg sync.WaitGroup
-  wg.Add(count)
+	var wg sync.WaitGroup
+	wg.Add(count)
 
-  for i := 0; i < count; i++ {
-    go func() {
-      if GetProcessInfo(cmd.Process.Pid) == nil {
-        t.Errorf("failed to get process from goroutine")
-      }
-      wg.Done()
-    }()
-  }
+	for i := 0; i < count; i++ {
+		go func() {
+			assert.NotNil(t, GetProcessInfo(cmd.Process.Pid))
+			wg.Done()
+		}()
+	}
 
-  wg.Wait()
+	wg.Wait()
 }
 
 func TestGetAllProcesses(t *testing.T) {
-  processes := GetAllProcessesInfo()
+	processes := GetAllProcessesInfo()
 
-  if len(processes) < 10 {
-    t.Errorf("expected at least 10 processes")
-  }
+	assert.Greater(t, len(processes), 10)
 }