diff --git a/internal/command/process.go b/internal/command/process.go index 31594c1..987998c 100644 --- a/internal/command/process.go +++ b/internal/command/process.go @@ -160,6 +160,7 @@ func ResolveProcessTargets(ctx context.Context, plat platform.Platform, p Parsed } var targets []model.Proc seen := map[int32]bool{} + protection := kill.NewProtection(ps, protected) for _, proc := range ps { name := processName(proc.Exe) match := slices.Contains(s.pids, proc.PID) @@ -186,7 +187,7 @@ func ResolveProcessTargets(ctx context.Context, plat platform.Platform, p Parsed if !match || seen[proc.PID] { continue } - if reason := kill.ProtectedReason(proc, ps, protected); reason != "" { + if reason := protection.Reason(proc); reason != "" { return nil, fmt.Errorf("PID %d: %s", proc.PID, reason) } if proc.Zombie { diff --git a/internal/kill/protected_test.go b/internal/kill/protected_test.go new file mode 100644 index 0000000..0b11b0f --- /dev/null +++ b/internal/kill/protected_test.go @@ -0,0 +1,80 @@ +package kill + +import ( + "os" + "os/user" + "runtime" + "testing" + "time" + + "github.com/labstack/driga/internal/model" +) + +// PIDs far above any the test runner's own processes have. +const fakePID = 1 << 30 + +func ownedProc(t *testing.T, pid, ppid int32, exe string, start time.Time) model.Proc { + t.Helper() + p := model.Proc{PID: pid, PPID: ppid, Exe: exe, Cmdline: exe, Argv: []string{exe}, StartTime: start, OwnerKnown: true, OwnerUID: uint32(os.Getuid())} + if runtime.GOOS == "windows" { + u, err := user.Current() + if err != nil { + t.Fatal(err) + } + p.OwnerUID, p.OwnerSID = 0, u.Uid + } + return p +} + +// Windows keeps the PID of an exited parent, so a Driga process can name a +// parent PID that a newer, unrelated process now has. On Windows that process +// is neither Driga's ancestor nor its descendant; elsewhere a PPID always +// names the live parent. +func TestProtectionReusedParentPID(t *testing.T) { + now := time.Now() + driga := ownedProc(t, fakePID+1, fakePID+2, "/opt/driga", now.Add(-time.Hour)) + newer := ownedProc(t, fakePID+2, fakePID+3, "/bin/sleep", now) + shell := ownedProc(t, fakePID+3, 1, "/bin/sh", now.Add(-2*time.Hour)) + live := []model.Proc{driga, newer, shell} + if reason := newProtection(live, nil, true).Reason(newer); reason != "" { + t.Fatalf("Windows: process with a reused PID protected: %s", reason) + } + if newProtection(live, nil, true).Reason(shell) != "" { + t.Fatal("Windows: the reused PID's parent protected") + } + if newProtection(live, nil, false).Reason(newer) == "" { + t.Fatal("Unix: Driga's parent not protected") + } + // A real parent, started before Driga, and Driga's children stay + // protected everywhere, as does a child older than its root. + older := ownedProc(t, fakePID+2, fakePID+3, "/bin/sleep", now.Add(-2*time.Hour)) + child := ownedProc(t, fakePID+4, fakePID+1, "/bin/sleep", now.Add(-3*time.Hour)) + for _, windows := range []bool{false, true} { + pr := newProtection([]model.Proc{driga, older, shell, child}, nil, windows) + if pr.Reason(older) == "" || pr.Reason(shell) == "" || pr.Reason(child) == "" { + t.Fatalf("windows=%v: Driga's family not protected", windows) + } + } +} + +// A broker root protects its children even after it exited, and stops +// protecting a PID only when a clearly newer process holds it. +func TestProtectionBrokerRoots(t *testing.T) { + now := time.Now() + root := ownedProc(t, fakePID+5, 1, "/bin/agent", now.Add(-time.Hour)) + child := ownedProc(t, fakePID+6, fakePID+5, "/bin/sleep", now) + for _, windows := range []bool{false, true} { + if newProtection([]model.Proc{child}, []model.Proc{root}, windows).Reason(child) == "" { + t.Fatalf("windows=%v: exited root's child not protected", windows) + } + } + rounded := root + rounded.StartTime = root.StartTime.Add(500 * time.Millisecond) + if newProtection([]model.Proc{rounded}, []model.Proc{root}, true).Reason(rounded) == "" { + t.Fatal("root with a slightly different start time not protected") + } + reused := ownedProc(t, fakePID+5, 1, "/bin/sleep", now) + if reason := newProtection([]model.Proc{reused}, []model.Proc{root}, true).Reason(reused); reason != "" { + t.Fatalf("newer process with a root's PID protected: %s", reason) + } +} diff --git a/internal/kill/snapshot.go b/internal/kill/snapshot.go index 6a804a2..dc5f99a 100644 --- a/internal/kill/snapshot.go +++ b/internal/kill/snapshot.go @@ -42,58 +42,125 @@ func Owned(p model.Proc) bool { return p.OwnerUID == uint32(os.Getuid()) && p.OwnerSID == "" } -// ProtectedReason uses live graph membership on each call, including new Chat -// descendants and ancestors. The broker supplies its complete protected roots. +// ProtectedReason is why p may not be signalled given one live snapshot and +// the broker's protected roots, or "" when it may. func ProtectedReason(p model.Proc, live, protected []model.Proc) string { - if p.PID <= 1 { - return "system PID or process group is refused" + if reason := identityReason(p); reason != "" { + return reason } - if !Owned(p) { - return "process owner is unknown or differs from your user" + return NewProtection(live, protected).Reason(p) +} + +// Protection is the set of processes one live snapshot protects: Driga, Chat +// and the broker's protected roots, all their descendants, and the ancestors +// of all of them. It errs toward protecting: descendants follow every parent +// link, and an ancestor walk stops only at a parent that provably reused an +// exited parent's PID on Windows, where a process keeps its creator's PID. +type Protection struct{ blocked map[int32]bool } + +// NewProtection computes the protected set once, so callers checking many +// targets against one snapshot do not rebuild it per target. +func NewProtection(live, protected []model.Proc) Protection { + return newProtection(live, protected, runtime.GOOS == "windows") +} + +func newProtection(live, protected []model.Proc, windows bool) Protection { + byPID := make(map[int32]model.Proc, len(live)) + children := map[int32][]int32{} + for _, r := range live { + byPID[r.PID] = r + children[r.PPID] = append(children[r.PPID], r.PID) } - if p.Partial || p.StartTime.IsZero() || p.Exe == "" || p.Cmdline == "" || len(p.Argv) == 0 { - return "process identity is incomplete" + blocked := map[int32]bool{} + ancestors := func(pid int32) { + for seen := map[int32]bool{pid: true}; ; { + cur, ok := byPID[pid] + if !ok || cur.PPID <= 0 || seen[cur.PPID] { + return + } + parent, ok := byPID[cur.PPID] + if ok && windows && !parent.StartTime.IsZero() && !cur.StartTime.IsZero() && parent.StartTime.After(cur.StartTime) { + return + } + pid = cur.PPID + seen[pid] = true + blocked[pid] = true + } + } + rooted := map[int32]bool{} + root := func(pid int32) { + if rooted[pid] { + return + } + blocked[pid] = true + ancestors(pid) + for pending := []int32{pid}; len(pending) > 0; { + parent := pending[len(pending)-1] + pending = pending[:len(pending)-1] + rooted[parent] = true + for _, child := range children[parent] { + if !rooted[child] { + blocked[child] = true + pending = append(pending, child) + } + } + } + } + self := int32(os.Getpid()) + blocked[self] = true + if _, ok := byPID[self]; ok { + ancestors(self) + } else { + // Without itself in the snapshot, Driga still protects its parent and + // whatever of the parent's ancestry the snapshot has. + parent := int32(os.Getppid()) + blocked[parent] = true + ancestors(parent) } - roots := map[int32]bool{} - blocked := map[int32]bool{int32(os.Getpid()): true, int32(os.Getppid()): true} for _, r := range protected { - roots[r.PID] = true - blocked[r.PID] = true + // A newer process holding the PID of a broker root that exited is + // not that root. Only a clearly later start proves it; any other + // mismatch keeps the protection. + if now, ok := byPID[r.PID]; ok && !r.StartTime.IsZero() && !now.StartTime.IsZero() && now.StartTime.After(r.StartTime.Add(time.Second)) { + continue + } + root(r.PID) } for _, r := range live { name := strings.ToLower(filepath.Base(strings.ReplaceAll(r.Exe, "\\", "/"))) if name == "driga" || name == "driga.exe" || strings.HasPrefix(name, "driga-") || r.Role == "chat" || strings.Contains(r.Cmdline, model.ChatPreamble) { - roots[r.PID] = true - blocked[r.PID] = true - } - } - // Ancestors are protected, but their unrelated descendants are not. Only - // explicit Chat/Driga roots propagate protection downwards. - for changed := true; changed; { - changed = false - for _, r := range live { - if roots[r.PPID] && !roots[r.PID] { - roots[r.PID] = true - blocked[r.PID] = true - changed = true - } + root(r.PID) } } - for changed := true; changed; { - changed = false - for _, r := range live { - if blocked[r.PID] && r.PPID > 0 && !blocked[r.PPID] { - blocked[r.PPID] = true - changed = true - } - } + return Protection{blocked: blocked} +} + +// Reason is why p may not be signalled, or "" when it may. +func (pr Protection) Reason(p model.Proc) string { + if reason := identityReason(p); reason != "" { + return reason } - if blocked[p.PID] { + if pr.blocked[p.PID] { return "Driga, Chat, or an ancestor/descendant is protected" } return "" } +// identityReason refuses a target whose own identity rules out a signal, +// before any process graph is built. +func identityReason(p model.Proc) string { + if p.PID <= 1 { + return "system PID or process group is refused" + } + if !Owned(p) { + return "process owner is unknown or differs from your user" + } + if p.Partial || p.StartTime.IsZero() || p.Exe == "" || p.Cmdline == "" || len(p.Argv) == 0 { + return "process identity is incomplete" + } + return "" +} + func Snapshot(ctx context.Context, plat platform.Platform, targets, protected []model.Proc, force bool, wait time.Duration) (SnapshotResult, error) { return snapshot(ctx, plat, targets, protected, force, wait, signalSnapshotIdentity) }