elmprodvpn/selective-vpn-api/app/trafficmode/cgroup.go

package trafficmode

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

func CgroupCandidates(entry string, cgroupRootPath string) []string {
	v := strings.TrimSpace(entry)
	if v == "" {
		return nil
	}
	vc := filepath.Clean(v)
	vals := []string{}
	if filepath.IsAbs(vc) {
		if strings.HasPrefix(vc, cgroupRootPath) {
			vals = append(vals, vc)
		} else {
			vals = append(vals, filepath.Join(cgroupRootPath, strings.TrimPrefix(vc, "/")))
		}
	} else {
		vals = append(vals,
			filepath.Join(cgroupRootPath, strings.TrimPrefix(vc, "/")),
			filepath.Join(cgroupRootPath, "system.slice", strings.TrimPrefix(vc, "/")),
			filepath.Join(cgroupRootPath, "user.slice", strings.TrimPrefix(vc, "/")),
		)
	}
	seen := map[string]struct{}{}
	out := make([]string, 0, len(vals))
	for _, p := range vals {
		cp := filepath.Clean(p)
		if cp == "." || cp == "" {
			continue
		}
		if _, ok := seen[cp]; ok {
			continue
		}
		seen[cp] = struct{}{}
		out = append(out, cp)
	}
	return out
}

func ResolveCgroupPath(entry string, cgroupRootPath string) (string, string) {
	for _, cand := range CgroupCandidates(entry, cgroupRootPath) {
		fi, err := os.Stat(cand)
		if err != nil || !fi.IsDir() {
			continue
		}
		return cand, ""
	}
	return "", "cgroup not found: " + strings.TrimSpace(entry)
}

func CollectPIDsFromCgroup(root string) (map[int]struct{}, string) {
	const (
		maxDirs = 5000
		maxPIDs = 50000
	)

	pids := map[int]struct{}{}
	dirs := 0
	warn := ""

	_ = filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error {
		if err != nil || d == nil || !d.IsDir() {
			return nil
		}
		dirs++
		if dirs > maxDirs {
			warn = "cgroup scan truncated by directory limit"
			return filepath.SkipDir
		}
		data, err := os.ReadFile(filepath.Join(path, "cgroup.procs"))
		if err != nil {
			return nil
		}
		for _, ln := range strings.Split(string(data), "\n") {
			ln = strings.TrimSpace(ln)
			if ln == "" {
				continue
			}
			pid, err := strconv.Atoi(ln)
			if err != nil || pid <= 0 {
				continue
			}
			pids[pid] = struct{}{}
			if len(pids) > maxPIDs {
				warn = "cgroup scan truncated by pid limit"
				return filepath.SkipDir
			}
		}
		return nil
	})
	return pids, warn
}

func UIDRangeForPID(pid int) (string, bool) {
	data, err := os.ReadFile(fmt.Sprintf("/proc/%d/status", pid))
	if err != nil {
		return "", false
	}
	for _, ln := range strings.Split(string(data), "\n") {
		ln = strings.TrimSpace(ln)
		if !strings.HasPrefix(ln, "Uid:") {
			continue
		}
		fields := strings.Fields(ln)
		if len(fields) < 2 {
			return "", false
		}
		v, ok := NormalizeUIDToken(fields[1])
		return v, ok
	}
	return "", false
}

func ResolveCgroupUIDRanges(entries []string, cgroupRootPath string) ([]string, string) {
	var uids []string
	var warnings []string

	for _, entry := range NormalizeCgroupList(entries) {
		root, warn := ResolveCgroupPath(entry, cgroupRootPath)
		if root == "" {
			if warn != "" {
				warnings = append(warnings, warn)
			}
			continue
		}
		pids, scanWarn := CollectPIDsFromCgroup(root)
		if scanWarn != "" {
			warnings = append(warnings, scanWarn)
		}
		if len(pids) == 0 {
			warnings = append(warnings, "cgroup has no processes: "+entry)
			continue
		}
		for pid := range pids {
			uidRange, ok := UIDRangeForPID(pid)
			if !ok || uidRange == "" {
				continue
			}
			uids = append(uids, uidRange)
		}
	}
	seenWarn := map[string]struct{}{}
	uniqWarn := make([]string, 0, len(warnings))
	for _, w := range warnings {
		ww := strings.TrimSpace(w)
		if ww == "" {
			continue
		}
		if _, ok := seenWarn[ww]; ok {
			continue
		}
		seenWarn[ww] = struct{}{}
		uniqWarn = append(uniqWarn, ww)
	}
	return NormalizeUIDList(uids), strings.Join(uniqWarn, "; ")
}