staros/process_unix.go

//go:build linux
// +build linux

package staros

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"golang.org/x/sys/unix"
)

var loadCurrentKeepCaps = currentKeepCaps

// FindProcessByName 通过进程名来查询应用信息
func FindProcessByName(name string) (datas []Process, err error) {
	return FindProcess(func(in Process) bool {
		if name == in.Name {
			return true
		}
		return false
	})
}

// FindProcess 通过进程信息来查询应用信息
func FindProcess(compare func(Process) bool) (datas []Process, err error) {
	var name, main string
	var mainb []byte
	netSnapshot := loadNetSnapshot(false)
	paths, err := ioutil.ReadDir("/proc")
	if err != nil {
		return
	}
	for _, v := range paths {
		if v.IsDir() && Exists("/proc/"+v.Name()+"/comm") {
			name, err = readAsString("/proc/" + v.Name() + "/comm")
			if err != nil {
				continue
			}
			var tmp Process
			tmp.LocalPath, err = os.Readlink("/proc/" + v.Name() + "/exe")
			tmp.Path = tmp.LocalPath
			tmp.LocalPath = filepath.Dir(tmp.LocalPath)
			tmp.ExecPath, err = os.Readlink("/proc/" + v.Name() + "/cwd")
			tmp.Name = strings.TrimSpace(name)
			main, err = readAsString("/proc/" + v.Name() + "/status")
			if err != nil {
				tmp.Err = err
				if compare(tmp) {
					netSnapshot.appendTo(&tmp)
					datas = append(datas, tmp)
					continue
				}
			} else {
				data := splitBy(main, ":")
				tmp.Pid, _ = strconv.ParseInt(data["Pid"], 10, 64)
				tmp.PPid, _ = strconv.ParseInt(data["PPid"], 10, 64)
				tmp.TPid, _ = strconv.ParseInt(data["TracerPid"], 10, 64)
				uids := splitBySpace(data["Uid"])
				gids := splitBySpace(data["Gid"])
				tmp.RUID, _ = atoiField(uids, 0)
				tmp.EUID, _ = atoiField(uids, 1)
				tmp.RGID, _ = atoiField(gids, 0)
				tmp.EGID, _ = atoiField(gids, 1)
				tmp.VmPeak = parseProcStatusKB(data["VmPeak"])
				tmp.VmSize = parseProcStatusKB(data["VmSize"])
				tmp.VmHWM = parseProcStatusKB(data["VmHWM"])
				tmp.VmRSS = parseProcStatusKB(data["VmRSS"])
				tmp.VmLck = parseProcStatusKB(data["VmLck"])
				tmp.VmData = parseProcStatusKB(data["VmData"])
			}
			mainb, err = ioutil.ReadFile("/proc/" + v.Name() + "/cmdline")
			if err != nil {
				tmp.Err = err
				if compare(tmp) {
					netSnapshot.appendTo(&tmp)
					datas = append(datas, tmp)
					continue
				}
			} else {
				args := bytes.Split(mainb, []byte{0})
				for _, v := range args {
					tmp.Args = append(tmp.Args, string(v))
				}
			}
			mainb, err = ioutil.ReadFile("/proc/" + v.Name() + "/environ")
			if err != nil {
				tmp.Err = err
				if compare(tmp) {
					netSnapshot.appendTo(&tmp)
					datas = append(datas, tmp)
					continue
				}
			} else {
				args := bytes.Split(mainb, []byte{0})
				for _, v := range args {
					tmp.Env = append(tmp.Env, string(v))
				}
			}

			main, err = readAsString("/proc/" + v.Name() + "/stat")
			if err != nil {
				tmp.Err = err
				if compare(tmp) {
					netSnapshot.appendTo(&tmp)
					datas = append(datas, tmp)
					continue
				}
			} else if uptime, ok := procStartTimeFromStat([]byte(main)); ok {
				tmp.Uptime = uptime
			}
			if compare(tmp) {
				netSnapshot.appendTo(&tmp)
				datas = append(datas, tmp)
			}
		}
	}
	return
}

// FindProcessByPid 通过Pid来查询应用信息
func FindProcessByPid(pid int64) (datas Process, err error) {
	var name, main string
	var mainb []byte
	if !Exists("/proc/" + fmt.Sprint(pid) + "/comm") {
		err = errors.New("Not Found")
		return
	}
	name, err = readAsString("/proc/" + fmt.Sprint(pid) + "/comm")
	if err != nil {
		return
	}
	main, err = readAsString("/proc/" + fmt.Sprint(pid) + "/status")
	if err != nil {
		return
	}
	data := splitBy(main, ":")
	datas.Name = strings.TrimSpace(name)
	datas.Pid, _ = strconv.ParseInt(data["Pid"], 10, 64)
	datas.PPid, _ = strconv.ParseInt(data["PPid"], 10, 64)
	datas.TPid, _ = strconv.ParseInt(data["TracerPid"], 10, 64)
	uids := splitBySpace(data["Uid"])
	gids := splitBySpace(data["Gid"])
	datas.RUID, _ = atoiField(uids, 0)
	datas.EUID, _ = atoiField(uids, 1)
	datas.RGID, _ = atoiField(gids, 0)
	datas.EGID, _ = atoiField(gids, 1)
	datas.VmPeak = parseProcStatusKB(data["VmPeak"])
	datas.VmSize = parseProcStatusKB(data["VmSize"])
	datas.VmHWM = parseProcStatusKB(data["VmHWM"])
	datas.VmRSS = parseProcStatusKB(data["VmRSS"])
	datas.VmLck = parseProcStatusKB(data["VmLck"])
	datas.VmData = parseProcStatusKB(data["VmData"])
	loadNetSnapshot(false).appendTo(&datas)
	mainb, err = ioutil.ReadFile("/proc/" + fmt.Sprint(pid) + "/cmdline")
	if err != nil {
		datas.Err = err
		err = nil
	} else {
		args := bytes.Split(mainb, []byte{0})
		for _, v := range args {
			datas.Args = append(datas.Args, string(v))
		}
	}

	mainb, err = ioutil.ReadFile("/proc/" + fmt.Sprint(pid) + "/environ")
	if err != nil {
		datas.Err = err
		err = nil
	} else {
		args := bytes.Split(mainb, []byte{0})
		for _, v := range args {
			datas.Env = append(datas.Env, string(v))
		}
	}

	datas.LocalPath, err = os.Readlink("/proc/" + fmt.Sprint(pid) + "/exe")
	datas.Path = datas.LocalPath
	datas.LocalPath = filepath.Dir(datas.LocalPath)
	datas.ExecPath, err = os.Readlink("/proc/" + fmt.Sprint(pid) + "/cwd")
	main, err = readAsString("/proc/" + fmt.Sprint(pid) + "/stat")
	if err != nil {
		return
	}
	if uptime, ok := procStartTimeFromStat([]byte(main)); ok {
		datas.Uptime = uptime
	}
	return
}

func procStartTimeFromStat(content []byte) (time.Time, bool) {
	fields := splitProcStat(content)
	if len(fields) <= 22 {
		return time.Time{}, false
	}
	startTicks, err := strconv.ParseInt(strings.TrimSpace(fields[22]), 10, 64)
	if err != nil {
		return time.Time{}, false
	}
	ticks := int64(clockTicks())
	seconds := startTicks / ticks
	nanos := (startTicks % ticks) * int64(time.Second) / ticks
	return time.Unix(StartTime().Unix()+seconds, nanos), true
}

func atoiField(fields []string, index int) (int, error) {
	if index < 0 || index >= len(fields) {
		return 0, errors.New("field index out of range")
	}
	return strconv.Atoi(fields[index])
}

func parseProcStatusKB(value string) int64 {
	fields := splitBySpace(value)
	if len(fields) == 0 || fields[0] == "" {
		return 0
	}
	size, err := strconv.ParseInt(fields[0], 10, 64)
	if err != nil {
		return 0
	}
	return size * 1024
}

type netSnapshot struct {
	conns []NetConn
	err   error
}

func loadNetSnapshot(analysePid bool) netSnapshot {
	netInfo, err := NetConnections(analysePid, "")
	return netSnapshot{conns: netInfo, err: err}
}

func appendNetInfo(p *Process, analysePid bool) {
	loadNetSnapshot(analysePid).appendTo(p)
}

func (snapshot netSnapshot) appendTo(p *Process) {
	if snapshot.err != nil {
		p.netErr = snapshot.err
		return
	}
	fds, err := ioutil.ReadDir("/proc/" + strconv.Itoa(int(p.Pid)) + "/fd")
	if err != nil {
		if Exists("/proc/" + strconv.Itoa(int(p.Pid)) + "/fd") {
			p.netErr = err
		}
		return
	}
	for _, fd := range fds {
		socket, err := os.Readlink("/proc/" + strconv.Itoa(int(p.Pid)) + "/fd/" + fd.Name())
		if err != nil {
			continue
		}
		start := strings.Index(socket, "[")
		if start < 0 {
			continue
		}
		sid := socket[start+1 : len(socket)-1]
		for _, v := range snapshot.conns {
			if v.Inode == sid {
				v.Pid = p.Pid
				v.Process = p
				p.netConn = append(p.netConn, v)
			}
		}
	}
}

func Daemon(path string, args ...string) (int, error) {
	cmd := exec.Command(path, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid: true,
	}
	if err := cmd.Start(); err != nil {
		return -1, err
	}
	pid := cmd.Process.Pid
	err := cmd.Process.Release()
	return pid, err
}

func DaemonWithUser(uid, gid uint32, groups []uint32, path string, args ...string) (int, error) {
	cmd := exec.Command(path, args...)
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Credential: &syscall.Credential{
			Uid:    uid,
			Gid:    gid,
			Groups: groups,
		},
		Setsid: true,
	}
	if err := cmd.Start(); err != nil {
		return -1, err
	}
	pid := cmd.Process.Pid
	err := cmd.Process.Release()
	return pid, err
}

func (starcli *StarCmd) SetRunUser(uid, gid uint32, groups []uint32) {
	_ = starcli.SetRunUserE(uid, gid, groups)
}

func (starcli *StarCmd) SetRunUserE(uid, gid uint32, groups []uint32) error {
	if starcli == nil || starcli.CMD == nil {
		return errNilCommand
	}
	if atomic.LoadInt32(&starcli.started) != 0 {
		return errCommandAlreadyStarted
	}
	if starcli.CMD.SysProcAttr == nil {
		starcli.CMD.SysProcAttr = &syscall.SysProcAttr{}
	}
	if starcli.CMD.SysProcAttr.Credential == nil {
		starcli.CMD.SysProcAttr.Credential = &syscall.Credential{}
	}
	starcli.CMD.SysProcAttr.Credential.Uid = uid
	starcli.CMD.SysProcAttr.Credential.Gid = gid
	starcli.CMD.SysProcAttr.Credential.Groups = append([]uint32(nil), groups...)
	starcli.CMD.SysProcAttr.Setsid = true
	return nil
}

func (starcli *StarCmd) Release() error {
	return starcli.ReleaseE()
}

func (starcli *StarCmd) Detach() error {
	return starcli.DetachE()
}

func (starcli *StarCmd) ReleaseE() error {
	if starcli == nil || starcli.CMD == nil {
		return errNilCommand
	}
	if !atomic.CompareAndSwapInt32(&starcli.released, 0, 1) {
		return errCommandAlreadyReleased
	}
	if atomic.LoadInt32(&starcli.started) != 0 {
		if starcli.CMD.Process == nil {
			starcli.lock.Lock()
			err := starcli.runerr
			starcli.lock.Unlock()
			if err != nil {
				atomic.StoreInt32(&starcli.released, 0)
				return err
			}
			atomic.StoreInt32(&starcli.released, 0)
			return errCommandAlreadyStarted
		}
		return nil
	}
	if starcli.CMD.SysProcAttr == nil {
		starcli.CMD.SysProcAttr = &syscall.SysProcAttr{
			Setsid: true,
		}
	} else {
		if !starcli.CMD.SysProcAttr.Setsid {
			starcli.CMD.SysProcAttr.Setsid = true
		}
	}
	if err := starcli.Start(); err != nil {
		atomic.StoreInt32(&starcli.released, 0)
		return err
	}
	return nil
}

func (starcli *StarCmd) DetachE() error {
	if starcli == nil || starcli.CMD == nil {
		return errNilCommand
	}
	if !atomic.CompareAndSwapInt32(&starcli.detached, 0, 1) {
		return errCommandAlreadyDetached
	}
	if atomic.LoadInt32(&starcli.started) != 0 {
		atomic.StoreInt32(&starcli.detached, 0)
		return errCommandAlreadyStarted
	}
	cmd := exec.Command(starcli.CMD.Path, starcli.CMD.Args[1:]...)
	cmd.Dir = starcli.CMD.Dir
	cmd.Env = append([]string(nil), starcli.CMD.Env...)
	if starcli.CMD.SysProcAttr != nil {
		attr := *starcli.CMD.SysProcAttr
		cmd.SysProcAttr = &attr
	} else {
		cmd.SysProcAttr = &syscall.SysProcAttr{
			Setsid: true,
		}
	}
	if !cmd.SysProcAttr.Setsid {
		cmd.SysProcAttr.Setsid = true
	}
	devNull, err := os.OpenFile(os.DevNull, os.O_RDWR, 0)
	if err != nil {
		atomic.StoreInt32(&starcli.detached, 0)
		return err
	}
	defer devNull.Close()
	cmd.Stdin = devNull
	cmd.Stdout = devNull
	cmd.Stderr = devNull
	if err := cmd.Start(); err != nil {
		atomic.StoreInt32(&starcli.detached, 0)
		return err
	}
	starcli.CMD.Process = cmd.Process
	atomic.StoreInt32(&starcli.started, 1)
	starcli.setRunning(false)
	starcli.finish()
	if err := cmd.Process.Release(); err != nil {
		atomic.StoreInt32(&starcli.detached, 0)
		return err
	}
	return nil
}

func (starcli *StarCmd) SetKeepCaps() error {
	if err := starcli.ensureConfigurable(); err != nil {
		return err
	}
	caps, err := loadCurrentKeepCaps()
	if err != nil {
		return err
	}
	starcli.lock.Lock()
	defer starcli.lock.Unlock()
	if starcli.CMD.SysProcAttr == nil {
		starcli.CMD.SysProcAttr = &syscall.SysProcAttr{}
	}
	starcli.CMD.SysProcAttr.AmbientCaps = mergeAmbientCaps(starcli.CMD.SysProcAttr.AmbientCaps, caps)
	return nil
}

func SetKeepCaps() error {
	return unix.Prctl(unix.PR_SET_KEEPCAPS, 1, 0, 0, 0)
}

func currentKeepCaps() ([]uintptr, error) {
	hdr := unix.CapUserHeader{Version: unix.LINUX_CAPABILITY_VERSION_3}
	data := [2]unix.CapUserData{}
	if err := unix.Capget(&hdr, &data[0]); err != nil {
		return nil, err
	}
	return capsFromCapData(data), nil
}

func capsFromCapData(data [2]unix.CapUserData) []uintptr {
	var caps []uintptr
	for index, item := range data {
		mask := item.Permitted
		for bit := uint(0); bit < 32; bit++ {
			if mask&(1<<bit) == 0 {
				continue
			}
			caps = append(caps, uintptr(index*32)+uintptr(bit))
		}
	}
	return caps
}

func mergeAmbientCaps(existing, extra []uintptr) []uintptr {
	if len(existing) == 0 && len(extra) == 0 {
		return nil
	}
	merged := append(append([]uintptr(nil), existing...), extra...)
	sort.Slice(merged, func(i, j int) bool {
		return merged[i] < merged[j]
	})
	out := merged[:0]
	var last uintptr
	for idx, cap := range merged {
		if idx == 0 || cap != last {
			out = append(out, cap)
			last = cap
		}
	}
	return out
}