Refactor code to wrap gopsutils (Velocidex#2824)

The gopsutils library shells out to various programs (like ps, lsof)
under some situations. This makes it unpredictable how expensive an
operation is.

This PR wraps gopsutils carefully ensuring that no expensive operations
occur.

Fixes: Velocidex#2822

actions/client_info.go

@@ -1,12 +1,13 @@
 package actions
 
 import (
+	"context"
 	"runtime"
 
 	"github.com/Showmax/go-fqdn"
-	"github.com/shirou/gopsutil/v3/host"
 	actions_proto "www.velocidex.com/golang/velociraptor/actions/proto"
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
+	"www.velocidex.com/golang/velociraptor/vql/psutils"
 )
 
 // Return essential information about the client used for indexing
@@ -15,10 +16,11 @@ import (
 // server periodically to avoid having to issue Generic.Client.Info
 // hunts all the time.
 func GetClientInfo(
+	ctx context.Context,
 	config_obj *config_proto.Config) *actions_proto.ClientInfo {
 	result := &actions_proto.ClientInfo{}
 
-	info, err := host.Info()
+	info, err := psutils.InfoWithContext(ctx)
 	if err == nil {
 		result = &actions_proto.ClientInfo{
 			Hostname:     info.Hostname,

actions/throttler.go

@@ -8,16 +8,17 @@ package actions
 import (
 	"context"
 	"os"
+	"runtime"
 	"sync"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
-	"github.com/shirou/gopsutil/v3/cpu"
-	"github.com/shirou/gopsutil/v3/process"
 	"www.velocidex.com/golang/velociraptor/utils"
 	"www.velocidex.com/golang/vfilter"
 	"www.velocidex.com/golang/vfilter/types"
+
+	"www.velocidex.com/golang/velociraptor/vql/psutils"
 )
 
 var (
@@ -56,7 +57,8 @@ type statsCollector struct {
 	cond *sync.Cond
 	id   uint64
 
-	proc *process.Process
+	// Our own pid so we can check stats etc
+	pid int32
 
 	samples [2]sample
 
@@ -68,22 +70,12 @@ type statsCollector struct {
 	number_of_cores     float64
 }
 
-func newStatsCollector() (*statsCollector, error) {
-	proc, err := process.NewProcess(int32(os.Getpid()))
-	if err != nil || proc == nil {
-		return nil, err
-	}
-
-	number_of_cores, err := cpu.Counts(true)
-	if err != nil || number_of_cores <= 0 {
-		return nil, err
-	}
-
+func newStatsCollector(ctx context.Context) (*statsCollector, error) {
 	result := &statsCollector{
 		check_duration_msec: 300,
 		id:                  utils.GetId(),
-		proc:                proc,
-		number_of_cores:     float64(number_of_cores),
+		pid:                 int32(os.Getpid()),
+		number_of_cores:     float64(runtime.NumCPU()),
 	}
 	result.cond = sync.NewCond(&result.mu)
 
@@ -136,15 +128,15 @@ func (self *statsCollector) GetAverageIOPS() float64 {
 // process. This is called not that frequently in order to minimize
 // the overheads of making a system call.
 func (self *statsCollector) getCpuTime(ctx context.Context) float64 {
-	cpu_time, err := self.proc.TimesWithContext(ctx)
+	cpu_time, err := psutils.TimesWithContext(ctx, self.pid)
 	if err != nil {
 		return 0
 	}
-	return cpu_time.Total()
+	return cpu_time.User + cpu_time.System
 }
 
 func (self *statsCollector) getIops(ctx context.Context) float64 {
-	counters, err := self.proc.IOCountersWithContext(ctx)
+	counters, err := psutils.IOCountersWithContext(ctx, self.pid)
 	if err != nil {
 		return 0
 	}
@@ -259,7 +251,7 @@ func NewThrottler(
 
 	if stats == nil {
 		var err error
-		stats, err = newStatsCollector()
+		stats, err = newStatsCollector(ctx)
 		if err != nil {
 			return nil
 		}
@@ -298,12 +290,9 @@ func init() {
 				return stats.GetAverageIOPS()
 			}
 
-			proc, err := process.NewProcess(int32(os.Getpid()))
-			if err != nil || proc == nil {
-				return 0
-			}
-
-			counters, err := proc.IOCounters()
+			ctx := context.Background()
+			pid := int32(os.Getpid())
+			counters, err := psutils.IOCountersWithContext(ctx, pid)
 			if err != nil {
 				return 0
 			}
@@ -313,26 +302,19 @@ func init() {
 	_ = prometheus.Register(promauto.NewGaugeFunc(
 		prometheus.GaugeOpts{
 			Name: "process_cpu_used",
-			Help: "Current CPU utilization by this process",
+			Help: "Total CPU utilization by this process",
 		}, func() float64 {
 			if stats != nil {
 				return stats.GetAverageCPULoad()
 			}
 
-			proc, err := process.NewProcess(int32(os.Getpid()))
-			if err != nil || proc == nil {
-				return 0
-			}
-
-			number_of_cores, err := cpu.Counts(true)
-			if err != nil || number_of_cores <= 0 {
-				return 0
-			}
-
-			cpu_time, err := proc.CPUPercent()
+			ctx := context.Background()
+			pid := int32(os.Getpid())
+			number_of_cores := runtime.NumCPU()
+			cpu_time, err := psutils.TimesWithContext(ctx, pid)
 			if err != nil {
 				return 0
 			}
-			return cpu_time / float64(number_of_cores)
+			return float64(cpu_time.User+cpu_time.System) / float64(number_of_cores)
 		}))
 }

bin/golden.go

@@ -35,7 +35,6 @@ import (
 	"github.com/Velocidex/yaml/v2"
 	errors "github.com/go-errors/errors"
 	"github.com/sergi/go-diff/diffmatchpatch"
-	"github.com/shirou/gopsutil/v3/process"
 	"www.velocidex.com/golang/velociraptor/actions"
 	actions_proto "www.velocidex.com/golang/velociraptor/actions/proto"
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
@@ -49,6 +48,7 @@ import (
 	"www.velocidex.com/golang/velociraptor/utils"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
 	"www.velocidex.com/golang/velociraptor/vql/acl_managers"
+	"www.velocidex.com/golang/velociraptor/vql/psutils"
 	"www.velocidex.com/golang/velociraptor/vql/remapping"
 	vfilter "www.velocidex.com/golang/vfilter"
 	"www.velocidex.com/golang/vfilter/arg_parser"
@@ -128,9 +128,9 @@ func makeCtxWithTimeout(
 				// the goroutines and mutex and hard exit.
 			case <-time.After(time.Second):
 				if time.Now().Before(deadline) {
-					proc, _ := process.NewProcess(int32(os.Getpid()))
-					total_time, _ := proc.Percent(0)
-					memory, _ := proc.MemoryInfo()
+					pid := int32(os.Getpid())
+					total_time, _ := psutils.TimesWithContext(ctx, pid)
+					memory, _ := psutils.MemoryInfoWithContext(ctx, pid)
 
 					fmt.Printf("Not time to fire yet %v %v %v\n",
 						time.Now(), total_time, memory)

executor/executor.go

@@ -77,7 +77,7 @@ type ClientExecutor struct {
 }
 
 func (self *ClientExecutor) GetClientInfo() *actions_proto.ClientInfo {
-	return actions.GetClientInfo(self.config_obj)
+	return actions.GetClientInfo(self.ctx, self.config_obj)
 }
 
 func (self *ClientExecutor) FlowManager() *responder.FlowManager {

uploads/file_based.go

@@ -80,8 +80,7 @@ func (self *FileBasedUploader) Upload(
 	atime time.Time,
 	ctime time.Time,
 	btime time.Time,
-	reader io.Reader) (
-	*UploadResponse, error) {
+	reader io.Reader) (*UploadResponse, error) {
 
 	if self.UploadDir == "" {
 		scope.Log("UploadDir is not set")
@@ -152,7 +151,10 @@ func (self *FileBasedUploader) Upload(
 	}
 
 	// It is not an error if we cant set the timestamps - best effort.
-	_ = setFileTimestamps(file_path, mtime, atime, ctime)
+	err = setFileTimestamps(file_path, mtime, atime, ctime)
+	if err != nil {
+		scope.Log("FileBasedUploader: %v", err)
+	}
 
 	scope.Log("Uploaded %v (%v bytes)", file_path, offset)
 	result = &UploadResponse{
@@ -171,8 +173,7 @@ func (self *FileBasedUploader) maybeCollectSparseFile(
 	ctx context.Context,
 	reader io.Reader,
 	store_as_name *accessors.OSPath,
-	sanitized_name string) (
-	*UploadResponse, error) {
+	sanitized_name string) (*UploadResponse, error) {
 
 	// Can the reader produce ranges?
 	range_reader, ok := reader.(RangeReader)

vql/filesystem/filesystem.go

@@ -26,13 +26,13 @@ import (
 	"github.com/Velocidex/ordereddict"
 	"github.com/go-errors/errors"
 
-	"github.com/shirou/gopsutil/v3/disk"
 	"www.velocidex.com/golang/velociraptor/accessors"
 	"www.velocidex.com/golang/velociraptor/acls"
 	config_proto "www.velocidex.com/golang/velociraptor/config/proto"
 	"www.velocidex.com/golang/velociraptor/glob"
 	"www.velocidex.com/golang/velociraptor/vql"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
+	"www.velocidex.com/golang/velociraptor/vql/psutils"
 	"www.velocidex.com/golang/vfilter"
 	"www.velocidex.com/golang/vfilter/arg_parser"
 )
@@ -448,7 +448,7 @@ func init() {
 				scope vfilter.Scope,
 				args *ordereddict.Dict) []vfilter.Row {
 				var result []vfilter.Row
-				partitions, err := disk.Partitions(true)
+				partitions, err := psutils.PartitionsWithContext(ctx)
 				if err == nil {
 					for _, item := range partitions {
 						result = append(result, item)

vql/filesystem/filesystems.go

@@ -21,34 +21,31 @@ import (
 	"context"
 
 	"github.com/Velocidex/ordereddict"
-	"github.com/shirou/gopsutil/v3/disk"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
+	"www.velocidex.com/golang/velociraptor/vql/psutils"
 	"www.velocidex.com/golang/vfilter"
 	"www.velocidex.com/golang/vfilter/arg_parser"
 )
 
 type ExtendedFileSystemInfo struct {
-	Partition disk.PartitionStat
+	Partition psutils.PartitionStat
 }
 
-func (self ExtendedFileSystemInfo) Usage() *disk.UsageStat {
-	usage, err := disk.Usage(self.Partition.Mountpoint)
+func (self ExtendedFileSystemInfo) Usage() *psutils.UsageStat {
+	usage, err := psutils.Usage(self.Partition.Mountpoint)
 	if err != nil {
 		return nil
 	}
 
 	return usage
 }
 
-/*
 func (self ExtendedFileSystemInfo) SerialNumber() string {
-	return disk.GetDiskSerialNumber(self.Partition.Device)
+	res, _ := psutils.SerialNumber(self.Partition.Device)
+	return res
 }
-*/
 
-type PartitionsArgs struct {
-	All bool `vfilter:"optional,field=all,doc=If specified list all Partitions"`
-}
+type PartitionsArgs struct{}
 
 func init() {
 	vql_subsystem.RegisterPlugin(
@@ -67,7 +64,8 @@ func init() {
 					return result
 				}
 
-				if partitions, err := disk.Partitions(arg.All); err == nil {
+				partitions, err := psutils.PartitionsWithContext(ctx)
+				if err == nil {
 					for _, item := range partitions {
 						extended_info := ExtendedFileSystemInfo{item}
 						result = append(result, extended_info)

vql/functions/pskill.go

@@ -21,10 +21,10 @@ import (
 	"context"
 
 	"github.com/Velocidex/ordereddict"
-	"github.com/shirou/gopsutil/v3/process"
 	"www.velocidex.com/golang/velociraptor/acls"
 	"www.velocidex.com/golang/velociraptor/vql"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
+	"www.velocidex.com/golang/velociraptor/vql/psutils"
 	"www.velocidex.com/golang/vfilter"
 	"www.velocidex.com/golang/vfilter/arg_parser"
 )
@@ -54,13 +54,12 @@ func (self *PsKillFunction) Call(ctx context.Context,
 		return vfilter.Null{}
 	}
 
-	process_obj, err := process.NewProcess(int32(arg.Pid))
+	err = psutils.Kill(int32(arg.Pid))
 	if err != nil {
 		scope.Log("pskill: %v", err)
 		return vfilter.Null{}
 	}
-
-	return process_obj.Kill()
+	return arg.Pid
 }
 
 func (self PsKillFunction) Info(scope vfilter.Scope, type_map *vfilter.TypeMap) *vfilter.FunctionInfo {

vql/info.go

@@ -25,9 +25,9 @@ import (
 
 	fqdn "github.com/Showmax/go-fqdn"
 	"github.com/Velocidex/ordereddict"
-	"github.com/shirou/gopsutil/v3/host"
 
 	"www.velocidex.com/golang/velociraptor/acls"
+	"www.velocidex.com/golang/velociraptor/vql/psutils"
 	"www.velocidex.com/golang/vfilter"
 	"www.velocidex.com/golang/vfilter/arg_parser"
 )
@@ -36,7 +36,7 @@ var (
 	start_time = time.Now()
 )
 
-func GetInfo(host *host.InfoStat) *ordereddict.Dict {
+func GetInfo(host *psutils.InfoStat) *ordereddict.Dict {
 	me, _ := os.Executable()
 	return ordereddict.NewDict().
 		Set("Hostname", host.Hostname).
@@ -84,9 +84,9 @@ func init() {
 				// It turns out that host.Info() is
 				// actually rather slow so we cache it
 				// in the scope cache.
-				info, ok := CacheGet(scope, "__info").(*host.InfoStat)
+				info, ok := CacheGet(scope, "__info").(*psutils.InfoStat)
 				if !ok {
-					info, err = host.Info()
+					info, err = psutils.InfoWithContext(ctx)
 					if err != nil {
 						scope.Log("info: %s", err)
 						return result

vql/linux/connections.go

@@ -23,14 +23,14 @@ import (
 	"syscall"
 
 	"github.com/Velocidex/ordereddict"
-	"github.com/shirou/gopsutil/v3/net"
 	"www.velocidex.com/golang/velociraptor/acls"
 	"www.velocidex.com/golang/velociraptor/vql"
 	vql_subsystem "www.velocidex.com/golang/velociraptor/vql"
+	"www.velocidex.com/golang/velociraptor/vql/psutils"
 	"www.velocidex.com/golang/vfilter"
 )
 
-func makeDict(in net.ConnectionStat) *ordereddict.Dict {
+func makeDict(in psutils.ConnectionStat) *ordereddict.Dict {
 	var family, conn_type string
 
 	switch in.Family {
@@ -90,7 +90,8 @@ func init() {
 					return result
 				}
 
-				if cons, err := net.Connections("all"); err == nil {
+				cons, err := psutils.ConnectionsWithContext(ctx, "all")
+				if err == nil {
 					for _, item := range cons {
 						result = append(result, makeDict(item))
 					}