From 7565a96525edb518129c7892843130bf0948eeb6 Mon Sep 17 00:00:00 2001
From: Robin Jarry <robin@jarry.cc>
Date: Tue, 1 Nov 2022 13:16:21 +0100
Subject: [PATCH] address-book-cmd: ignore completion above 100 items

Avoid aerc from consuming all memory on the system if an address book
command returns 12 million addresses.

Read at most the first 100 lines and kill the command if it has not
finished. Display a warning in the logs for good measure.

The command is now assigned an different PGID (equal to its PID) to
allow killing it *and* all of its children. When the address book
command is a shell script that forks a process which never exits, it
will avoid killing the shell process and leaving its children without
parents.

Signed-off-by: Robin Jarry <robin@jarry.cc>
Tested-by: Bence Ferdinandy <bence@ferdinandy.com>
---
 completer/completer.go | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/completer/completer.go b/completer/completer.go
index bbeb5f1..f3f7b9f 100644
--- a/completer/completer.go
+++ b/completer/completer.go
@@ -10,7 +10,9 @@ import (
 	"os/exec"
 	"regexp"
 	"strings"
+	"syscall"
 
+	"git.sr.ht/~rjarry/aerc/logging"
 	"github.com/google/shlex"
 )
 
@@ -71,6 +73,10 @@ func isAddressHeader(h string) bool {
 	return false
 }
 
+const maxCompletionLines = 100
+
+var tooManyLines = fmt.Errorf("returned more than %d lines", maxCompletionLines)
+
 // completeAddress uses the configured address book completion command to fetch
 // completions for the specified string, returning a slice of completions and
 // a prefix to be prepended to the selected completion, or an error.
@@ -88,6 +94,8 @@ func (c *Completer) completeAddress(s string) ([]string, string, error) {
 	if err != nil {
 		return nil, "", fmt.Errorf("stderr: %w", err)
 	}
+	// reset the process group id to allow killing all its children
+	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
 	if err := cmd.Start(); err != nil {
 		return nil, "", fmt.Errorf("cmd start: %w", err)
 	}
@@ -99,6 +107,12 @@ func (c *Completer) completeAddress(s string) ([]string, string, error) {
 
 	completions, err := readCompletions(stdout)
 	if err != nil {
+		// make sure to kill the process *and* all its children
+		//nolint:errcheck // who cares?
+		syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
+		logging.Warnf("command %s killed: %s", cmd, err)
+	}
+	if err != nil && !errors.Is(err, tooManyLines) {
 		buf, _ := io.ReadAll(stderr)
 		msg := strings.TrimSpace(string(buf))
 		if msg != "" {
@@ -168,6 +182,9 @@ func readCompletions(r io.Reader) ([]string, error) {
 			return nil, fmt.Errorf("could not decode MIME string: %w", err)
 		}
 		completions = append(completions, decoded)
+		if len(completions) >= maxCompletionLines {
+			return completions, tooManyLines
+		}
 	}
 }