From 8b9ac394a058daf6022308c3f310e25aad53d917 Mon Sep 17 00:00:00 2001 From: unkin-agent Date: Tue, 13 Oct 2026 20:38:55 +1100 Subject: [PATCH 1/2] Query piped nodes for -R/-F with -C; count input only for bare -C --- main.go | 47 +++++++++++++++++++--------------- main_test.go | 72 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 20 deletions(-) diff --git a/main.go b/main.go index 2b22e86..af30b01 100644 --- a/main.go +++ b/main.go @@ -357,15 +357,32 @@ func run(cfg config, nodeName, factName, match string, showRole, partial, invers return nil } - if (nodeOnly || valueOnly || count || ansible) && !showRole && factName == "" { + 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 { + var lines []string + scanner := bufio.NewScanner(reader) + for scanner.Scan() { + lines = append(lines, scanner.Text()) + } + fmt.Println(strings.Join(countResults(lines), "\n")) + return nil + } + var collected []fact - var stdinLines []string doQuery := func(node string) error { query := buildQuery(node, factName, match, cfg.RoleFact, showRole, partial, inverse) @@ -377,21 +394,15 @@ func run(cfg config, nodeName, factName, match string, showRole, partial, invers return nil } - if reader, ok := stdinReader(os.Stdin); ok && nodeName == "" { + if piped { scanner := bufio.NewScanner(reader) - if count { - for scanner.Scan() { - stdinLines = append(stdinLines, scanner.Text()) + for scanner.Scan() { + fields := strings.Fields(scanner.Text()) + if len(fields) == 0 { + continue } - } else { - for scanner.Scan() { - fields := strings.Fields(scanner.Text()) - if len(fields) == 0 { - continue - } - if err := doQuery(qualifyNode(fields[0], cfg.Domain)); err != nil { - fmt.Fprintln(os.Stderr, "error:", err) - } + if err := doQuery(qualifyNode(fields[0], cfg.Domain)); err != nil { + fmt.Fprintln(os.Stderr, "error:", err) } } } else { @@ -410,11 +421,7 @@ func run(cfg config, nodeName, factName, match string, showRole, partial, invers _ = enc.Encode(factsByHost(collected, factName, cfg.RoleFact, showRole)) case count: - values := stdinLines - if len(values) == 0 { - values = returnData - } - fmt.Println(strings.Join(countResults(values), "\n")) + fmt.Println(strings.Join(countResults(returnData), "\n")) case ansible: // Attach each host's queried fact(s) as inventory host vars, e.g. diff --git a/main_test.go b/main_test.go index c65f7b5..a595bcd 100644 --- a/main_test.go +++ b/main_test.go @@ -8,6 +8,7 @@ import ( "os" "path/filepath" "strings" + "sync/atomic" "testing" "gopkg.in/yaml.v3" @@ -872,3 +873,74 @@ func TestRun_AllFacts_PrintsSortedByName(t *testing.T) { t.Fatalf("facts not sorted by name: %v", lines) } } + +// ---- run: count with piped stdin -------------------------------------------- + +func withStdin(t *testing.T, content string) { + t.Helper() + path := filepath.Join(t.TempDir(), "stdin") + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatal(err) + } + f, err := os.Open(path) + if err != nil { + t.Fatal(err) + } + old := os.Stdin + os.Stdin = f + t.Cleanup(func() { os.Stdin = old; _ = f.Close() }) +} + +const pipedLookup = "hosta.main.unkin.net puppetdbapi.service.consul\nhostb.main.unkin.net puppetdbapi.service.consul\n" + +func TestRun_CountPipedWithRoleQueriesEachNode(t *testing.T) { + withStdin(t, pipedLookup) + var queried atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + queried.Add(1) + _ = json.NewEncoder(w).Encode([]fact{{Certname: "h", Name: "enc_role", Value: rawJSON("roles::vault")}}) + })) + t.Cleanup(srv.Close) + cfg := config{PuppetDBURL: srv.URL, RoleFact: "enc_role"} + + out := strings.TrimSpace(captureStdout(t, func() { + if err := run(cfg, "", "", "", true, false, false, false, false, true, false, false, false); err != nil { + t.Fatalf("run returned error: %v", err) + } + })) + if out != "2: roles::vault" { + t.Fatalf("expected role 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_CountPipedWithoutQueryCountsInput(t *testing.T) { + withStdin(t, pipedLookup+"hostc.main.unkin.net other\n") + cfg := config{PuppetDBURL: "http://unused", RoleFact: "enc_role"} + + out := strings.TrimSpace(captureStdout(t, func() { + if err := run(cfg, "", "", "", false, false, false, false, false, true, false, false, false); err != nil { + t.Fatalf("run returned error: %v", err) + } + })) + if out != "1: other\n2: puppetdbapi.service.consul" { + t.Fatalf("expected counts of piped values, got: %q", out) + } +} + +func TestRun_CountWithoutStdinRequiresRoleOrFact(t *testing.T) { + devnull, err := os.Open(os.DevNull) + if err != nil { + t.Fatal(err) + } + old := os.Stdin + os.Stdin = devnull + t.Cleanup(func() { os.Stdin = old; _ = devnull.Close() }) + + err = run(config{PuppetDBURL: "http://unused"}, "", "", "", false, false, false, false, false, true, false, false, false) + if err == nil || !strings.Contains(err.Error(), "-R or -F") { + t.Fatalf("expected -R/-F requirement error, got: %v", err) + } +} -- 2.47.3 From fae7dad63041f2bdc38ecd71b4b20b3cbb12fc0d Mon Sep 17 00:00:00 2001 From: unkin-agent Date: Tue, 13 Oct 2026 20:41:31 +1100 Subject: [PATCH 2/2] 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") + } + }) + } +} -- 2.47.3