From fae7dad63041f2bdc38ecd71b4b20b3cbb12fc0d Mon Sep 17 00:00:00 2001
From: unkin-agent <unkin-agent@unkin.net>
Date: Tue, 13 Oct 2026 20:41:31 +1100
Subject: [PATCH] validate flags before reading stdin and reject bare -C with
 -j/-1/-2

---
 main.go      | 26 +++++++++++-------
 main_test.go | 78 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 10 deletions(-)

diff --git a/main.go b/main.go
index af30b01..e42a55e 100644
--- a/main.go
+++ b/main.go
@@ -357,22 +357,28 @@ func run(cfg config, nodeName, factName, match string, showRole, partial, invers
 		return nil
 	}
 
+	hasQuery := showRole || factName != ""
+	// Bare -C counts already-looked-up piped lines; only it needs stdin to validate.
+	bareCount := count && !hasQuery && !nodeOnly && !valueOnly && !ansible && !jsonMode && nodeName == ""
+	countErr := fmt.Errorf("-R or -F must be used with -1, -2, -C, or -A")
+
+	if (nodeOnly || valueOnly || (count && !bareCount) || ansible) && !hasQuery {
+		return countErr
+	}
+	if (match != "" || partial || inverse) && !hasQuery {
+		return fmt.Errorf("-R or -F must be used with -m, -p, or -i")
+	}
+
 	var reader *bufio.Reader
 	piped := false
 	if nodeName == "" {
 		reader, piped = stdinReader(os.Stdin)
 	}
-	// Bare -C on piped input counts the already-looked-up lines instead of querying.
-	passthrough := count && piped && !showRole && factName == ""
 
-	if (nodeOnly || valueOnly || (count && !passthrough) || ansible) && !showRole && factName == "" {
-		return fmt.Errorf("-R or -F must be used with -1, -2, -C, or -A")
-	}
-	if (match != "" || partial || inverse) && !showRole && factName == "" {
-		return fmt.Errorf("-R or -F must be used with -m, -p, or -i")
-	}
-
-	if passthrough {
+	if bareCount {
+		if !piped {
+			return countErr
+		}
 		var lines []string
 		scanner := bufio.NewScanner(reader)
 		for scanner.Scan() {
diff --git a/main_test.go b/main_test.go
index a595bcd..4636ac7 100644
--- a/main_test.go
+++ b/main_test.go
@@ -10,6 +10,7 @@ import (
 	"strings"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"gopkg.in/yaml.v3"
 )
@@ -944,3 +945,80 @@ func TestRun_CountWithoutStdinRequiresRoleOrFact(t *testing.T) {
 		t.Fatalf("expected -R/-F requirement error, got: %v", err)
 	}
 }
+
+func TestRun_CountPipedWithFactQueriesEachNode(t *testing.T) {
+	withStdin(t, pipedLookup)
+	var queried atomic.Int32
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !strings.Contains(r.URL.Query().Get("query"), "osfamily") {
+			t.Errorf("expected osfamily query, got %q", r.URL.Query().Get("query"))
+		}
+		queried.Add(1)
+		_ = json.NewEncoder(w).Encode([]fact{{Certname: "h", Name: "osfamily", Value: rawJSON("RedHat")}})
+	}))
+	t.Cleanup(srv.Close)
+	cfg := config{PuppetDBURL: srv.URL, RoleFact: "enc_role"}
+
+	out := strings.TrimSpace(captureStdout(t, func() {
+		if err := run(cfg, "", "osfamily", "", false, false, false, false, false, true, false, false, false); err != nil {
+			t.Fatalf("run returned error: %v", err)
+		}
+	}))
+	if out != "2: RedHat" {
+		t.Fatalf("expected fact counts of piped nodes, got: %q", out)
+	}
+	if n := queried.Load(); n != 2 {
+		t.Fatalf("expected one query per piped node, got %d", n)
+	}
+}
+
+func TestRun_CountPipedWithOtherOutputFlagsRequiresRoleOrFact(t *testing.T) {
+	for name, mutate := range map[string]func(*runArgs){
+		"json":      func(a *runArgs) { a.jsonMode = true },
+		"nodeOnly":  func(a *runArgs) { a.nodeOnly = true },
+		"valueOnly": func(a *runArgs) { a.valueOnly = true },
+	} {
+		t.Run(name, func(t *testing.T) {
+			withStdin(t, pipedLookup)
+			a := runArgs{cfg: config{PuppetDBURL: "http://unused"}, count: true}
+			mutate(&a)
+			var err error
+			out := captureStdout(t, func() {
+				err = run(a.cfg, "", "", "", false, false, false, a.nodeOnly, a.valueOnly, true, false, a.jsonMode, false)
+			})
+			if err == nil || !strings.Contains(err.Error(), "-R or -F") {
+				t.Fatalf("expected -R/-F requirement error, got: %v (output %q)", err, out)
+			}
+		})
+	}
+}
+
+func TestRun_InvalidFlagsErrorWithoutReadingOpenPipe(t *testing.T) {
+	for name, a := range map[string]runArgs{
+		"match":     {match: "foo"},
+		"countJSON": {count: true, jsonMode: true},
+	} {
+		t.Run(name, func(t *testing.T) {
+			r, w, err := os.Pipe()
+			if err != nil {
+				t.Fatal(err)
+			}
+			old := os.Stdin
+			os.Stdin = r
+			t.Cleanup(func() { os.Stdin = old; _ = w.Close(); _ = r.Close() })
+
+			done := make(chan error, 1)
+			go func() {
+				done <- run(config{PuppetDBURL: "http://unused"}, "", "", a.match, false, false, false, false, false, a.count, false, a.jsonMode, false)
+			}()
+			select {
+			case err := <-done:
+				if err == nil || !strings.Contains(err.Error(), "-R or -F") {
+					t.Fatalf("expected -R/-F requirement error, got: %v", err)
+				}
+			case <-time.After(2 * time.Second):
+				t.Fatal("run blocked on an open, empty stdin pipe before validating flags")
+			}
+		})
+	}
+}