fix(rerank): handle closed Darwin identity sockets after exec

Author: XiaoConstantine

pkg/rerank/supervisor_peer_darwin.go

@@ -18,16 +18,17 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-func rerankerProcessExecutableIdentity(pid int) (uint64, uint64, error) {
-	type uniqueIdentifierInfo struct {
-		UUID           [16]byte
-		UniqueID       uint64
-		ParentUniqueID uint64
-		IDVersion      int32
-		Reserved2      uint32
-		Reserved3      uint64
-		Reserved4      uint64
-	}
+type uniqueIdentifierInfo struct {
+	UUID           [16]byte
+	UniqueID       uint64
+	ParentUniqueID uint64
+	IDVersion      int32
+	Reserved2      uint32
+	Reserved3      uint64
+	Reserved4      uint64
+}
+
+func darwinProcessIdentity(pid int) (uniqueIdentifierInfo, error) {
 	var identity uniqueIdentifierInfo
 	const (
 		procInfoCallPIDInfo         = 0x2
@@ -44,10 +45,18 @@ func rerankerProcessExecutableIdentity(pid int) (uint64, uint64, error) {
 		unsafe.Sizeof(identity),
 	)
 	if errno != 0 {
-		return 0, 0, errno
+		return identity, errno
 	}
 	if result != unsafe.Sizeof(identity) {
-		return 0, 0, fmt.Errorf("unexpected Darwin executable identity size %d", result)
+		return identity, fmt.Errorf("unexpected Darwin executable identity size %d", result)
+	}
+	return identity, nil
+}
+
+func rerankerProcessExecutableIdentity(pid int) (uint64, uint64, error) {
+	identity, err := darwinProcessIdentity(pid)
+	if err != nil {
+		return 0, 0, err
 	}
 	first := binary.LittleEndian.Uint64(identity.UUID[:8])
 	second := binary.LittleEndian.Uint64(identity.UUID[8:])
@@ -177,10 +186,20 @@ type rerankerProcessHandle struct {
 
 func openRerankerProcessHandle(pid int, started string, identity *os.File) (*rerankerProcessHandle, error) {
 	token, err := peerAuditToken(identity)
-	if err != nil {
+	if errors.Is(err, unix.ENOTCONN) {
+		// The executed program may close inherited descriptors before we
+		// observe its image. Darwin's proc_find_audit_token (used by signal)
+		// checks only PID and pidversion; obtain both from kernel process
+		// information rather than falling back to a racy kill(pid).
+		info, infoErr := darwinProcessIdentity(pid)
+		if infoErr != nil {
+			return nil, infoErr
+		}
+		token = [8]uint32{5: uint32(pid), 7: uint32(info.IDVersion)}
+	} else if err != nil {
 		return nil, err
 	}
-	if int(token[5]) != pid || rerankerProcessStartIdentity(pid) != started {
+	if int(token[5]) != pid || started == "" || rerankerProcessStartIdentity(pid) != started {
 		return nil, fmt.Errorf("reranker process generation changed while acquiring audit token")
 	}
 	return &rerankerProcessHandle{token: token}, nil

pkg/rerank/supervisor_peer_darwin_test.go

@@ -0,0 +1,62 @@
+//go:build darwin
+
+package rerank
+
+import (
+	"errors"
+	"os"
+	"os/exec"
+	"syscall"
+	"testing"
+
+	"golang.org/x/sys/unix"
+)
+
+func TestRerankerProcessHandleAfterIdentitySocketCloses(t *testing.T) {
+	cmd := exec.Command("/bin/sleep", "30")
+	if err := cmd.Start(); err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		_ = cmd.Process.Kill()
+		_ = cmd.Wait()
+	})
+	fds, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	identity := os.NewFile(uintptr(fds[0]), "closed-peer")
+	defer func() { _ = identity.Close() }()
+	if err := unix.Close(fds[1]); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := peerAuditToken(identity); !errors.Is(err, unix.ENOTCONN) {
+		t.Fatalf("peer audit token error = %v, want ENOTCONN", err)
+	}
+	pid := cmd.Process.Pid
+	started := rerankerProcessStartIdentity(pid)
+	if started == "" {
+		t.Fatal("missing process start identity")
+	}
+	for _, wrong := range []string{"", started + "-stale"} {
+		if _, err := openRerankerProcessHandle(pid, wrong, identity); err == nil {
+			t.Fatal("accepted incorrect process generation")
+		}
+	}
+	handle, err := openRerankerProcessHandle(pid, started, identity)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = handle.Close() }()
+	stale := *handle
+	stale.token[7]++
+	if err := stale.Signal(syscall.SIGKILL); !errors.Is(err, os.ErrProcessDone) {
+		t.Fatalf("stale pidversion signal = %v, want ErrProcessDone", err)
+	}
+	if err := syscall.Kill(pid, 0); err != nil {
+		t.Fatalf("stale handle killed the live process: %v", err)
+	}
+	if err := handle.Signal(syscall.SIGKILL); err != nil {
+		t.Fatalf("terminate through handle: %v", err)
+	}
+}