jq.go

  1package shell
  2
  3import (
  4	"context"
  5	"encoding/json"
  6	"fmt"
  7	"io"
  8	"os"
  9	"strings"
 10
 11	"github.com/itchyny/gojq"
 12	"mvdan.cc/sh/v3/interp"
 13)
 14
 15const jqUsage = `jq - Go implementation of jq (gojq 0.12.19 builtin)
 16
 17Synopsis:
 18  %% echo '{"foo": 128}' | jq '.foo'
 19
 20Usage:
 21  jq [OPTIONS] [FILTER] [FILE...]
 22
 23Options:
 24  -r, --raw-output              output raw strings
 25  -j, --join-output             implies -r with no newline delimiter
 26  -c, --compact-output          output without pretty-printing
 27  -s, --slurp                   read all inputs into an array
 28  -n, --null-input              use null as input value
 29  -e, --exit-status             exit 1 when the last value is false or null
 30  -R, --raw-input               read input as raw strings
 31      --arg name value          set a string value to a variable
 32      --argjson name value      set a JSON value to a variable
 33  -h, --help                    display this help
 34`
 35
 36// handleJQ implements the jq builtin using gojq. It supports a subset of jq
 37// flags: -r (raw output), -c (compact output), -s (slurp), -n (null input),
 38// -e (exit status), -R (raw input), and --arg name value.
 39//
 40// ctx is polled at each iteration of the output loop and at each reader in
 41// [readInputs] so that hook timeouts or other cancellations can interrupt
 42// long-running queries. A cancelled context surfaces as ctx.Err(), not an
 43// [interp.ExitStatus], so callers (e.g. the hook runner) can distinguish
 44// "filter exited non-zero" from "we ran out of time".
 45//
 46// Note that this is somewhat of a reimplmentation of the CLI of the glorious
 47// github.com/itchyny/gojq, and we'd ideally get the CLI exposed upstream to
 48// avoid this falling out of sync.
 49func handleJQ(ctx context.Context, args []string, stdin io.Reader, stdout, stderr io.Writer) error {
 50	var (
 51		rawOutput  bool
 52		compact    bool
 53		slurp      bool
 54		nullInput  bool
 55		exitStatus bool
 56		rawInput   bool
 57		joinOutput bool
 58		argNames   []string
 59		argValues  []any
 60	)
 61
 62	// Parse flags and extract the query.
 63	var queryStr string
 64	var fileArgs []string
 65	i := 1 // skip "jq"
 66	for i < len(args) {
 67		arg := args[i]
 68		switch {
 69		case arg == "-h" || arg == "--help":
 70			fmt.Fprint(stdout, jqUsage)
 71			return nil
 72		case arg == "-r" || arg == "--raw-output":
 73			rawOutput = true
 74		case arg == "-j" || arg == "--join-output":
 75			joinOutput = true
 76			rawOutput = true
 77		case arg == "-c" || arg == "--compact-output":
 78			compact = true
 79		case arg == "-s" || arg == "--slurp":
 80			slurp = true
 81		case arg == "-n" || arg == "--null-input":
 82			nullInput = true
 83		case arg == "-e" || arg == "--exit-status":
 84			exitStatus = true
 85		case arg == "-R" || arg == "--raw-input":
 86			rawInput = true
 87		case arg == "--arg":
 88			if i+2 >= len(args) {
 89				fmt.Fprintf(stderr, "jq: --arg requires name and value\n")
 90				return interp.ExitStatus(2)
 91			}
 92			argNames = append(argNames, "$"+args[i+1])
 93			argValues = append(argValues, args[i+2])
 94			i += 2
 95		case arg == "--argjson":
 96			if i+2 >= len(args) {
 97				fmt.Fprintf(stderr, "jq: --argjson requires name and value\n")
 98				return interp.ExitStatus(2)
 99			}
100			var val any
101			if err := json.Unmarshal([]byte(args[i+2]), &val); err != nil {
102				fmt.Fprintf(stderr, "jq: invalid JSON for --argjson %s: %s\n", args[i+1], err)
103				return interp.ExitStatus(2)
104			}
105			argNames = append(argNames, "$"+args[i+1])
106			argValues = append(argValues, val)
107			i += 2
108		case arg == "--":
109			i++
110			// Remaining args are file arguments.
111			for i < len(args) {
112				fileArgs = append(fileArgs, args[i])
113				i++
114			}
115			continue
116		case strings.HasPrefix(arg, "-") && queryStr != "":
117			fmt.Fprintf(stderr, "jq: unknown option: %s\n", arg)
118			return interp.ExitStatus(2)
119		default:
120			if queryStr == "" {
121				queryStr = arg
122			} else {
123				fileArgs = append(fileArgs, arg)
124			}
125		}
126		i++
127	}
128
129	if queryStr == "" {
130		queryStr = "."
131	}
132
133	query, err := gojq.Parse(queryStr)
134	if err != nil {
135		fmt.Fprintf(stderr, "jq: %s\n", err)
136		return interp.ExitStatus(3)
137	}
138
139	opts := []gojq.CompilerOption{
140		gojq.WithEnvironLoader(os.Environ),
141	}
142	if len(argNames) > 0 {
143		opts = append(opts, gojq.WithVariables(argNames))
144	}
145
146	code, err := gojq.Compile(query, opts...)
147	if err != nil {
148		fmt.Fprintf(stderr, "jq: %s\n", err)
149		return interp.ExitStatus(3)
150	}
151
152	// Build input values.
153	inputs, err := readInputs(ctx, stdin, fileArgs, nullInput, rawInput, slurp)
154	if err != nil {
155		// Prefer surfacing ctx cancellation verbatim so timeouts are
156		// distinguishable from user input errors.
157		if ctxErr := ctx.Err(); ctxErr != nil {
158			return ctxErr
159		}
160		fmt.Fprintf(stderr, "jq: %s\n", err)
161		return interp.ExitStatus(2)
162	}
163
164	var lastFalsy bool
165	for _, input := range inputs {
166		iter := code.Run(input, argValues...)
167		for {
168			// Poll ctx on every value so a long-running filter (e.g. a
169			// generator over a slurped array) can be interrupted by hook
170			// timeouts without waiting for iter.Next to yield.
171			if err := ctx.Err(); err != nil {
172				return err
173			}
174			v, ok := iter.Next()
175			if !ok {
176				break
177			}
178			if err, ok := v.(error); ok {
179				fmt.Fprintf(stderr, "jq: %s\n", err)
180				return interp.ExitStatus(5)
181			}
182			if exitStatus {
183				lastFalsy = v == nil || v == false
184			}
185			if err := writeValue(stdout, v, rawOutput, compact, joinOutput); err != nil {
186				return err
187			}
188		}
189	}
190
191	if exitStatus && lastFalsy {
192		return interp.ExitStatus(1)
193	}
194	return nil
195}
196
197// readInputs reads JSON (or raw) input values from stdin or files.
198//
199// ctx is polled in three places so that a cancellation observed mid-read
200// short-circuits promptly:
201//   - between readers (before opening the next file / consuming stdin);
202//   - on every io.Read call via ctxReader, so io.ReadAll on a large but
203//     non-blocking source (e.g. the bytes.NewReader payload the hook
204//     runner supplies) returns ctx.Err() on the next chunk boundary;
205//   - inside the post-read value accumulation loops (raw-input line
206//     split and JSON stream decode), which are otherwise unbounded in
207//     the size of the input.
208//
209// A reader that blocks forever in Read (e.g. an unterminated pipe) can
210// still outlast ctx; the outer abandon-goroutine path in the hook
211// runner (internal/hooks/runner.go) is the authoritative enforcer for
212// that case.
213func readInputs(ctx context.Context, stdin io.Reader, files []string, nullInput, rawInput, slurp bool) ([]any, error) {
214	if nullInput {
215		return []any{nil}, nil
216	}
217
218	var readers []io.Reader
219	if len(files) > 0 {
220		for _, f := range files {
221			file, err := os.Open(f)
222			if err != nil {
223				return nil, err
224			}
225			defer file.Close()
226			readers = append(readers, file)
227		}
228	} else {
229		readers = []io.Reader{stdin}
230	}
231
232	var vals []any
233	for _, r := range readers {
234		if err := ctx.Err(); err != nil {
235			return nil, err
236		}
237		data, err := io.ReadAll(ctxReader{ctx: ctx, r: r})
238		if err != nil {
239			// ctxReader surfaces ctx.Err() verbatim; preserve it so the
240			// caller can distinguish cancellation from a parse error.
241			if ctxErr := ctx.Err(); ctxErr != nil {
242				return nil, ctxErr
243			}
244			return nil, err
245		}
246
247		if rawInput {
248			lines := strings.Split(string(data), "\n")
249			if slurp {
250				vals = append(vals, strings.Join(lines, "\n"))
251			} else {
252				for _, line := range lines {
253					if err := ctx.Err(); err != nil {
254						return nil, err
255					}
256					if line != "" || !slurp {
257						vals = append(vals, line)
258					}
259				}
260			}
261			continue
262		}
263
264		// Decode potentially multiple JSON values from the stream.
265		dec := json.NewDecoder(strings.NewReader(string(data)))
266		var streamVals []any
267		for {
268			if err := ctx.Err(); err != nil {
269				return nil, err
270			}
271			var v any
272			if err := dec.Decode(&v); err != nil {
273				if err == io.EOF {
274					break
275				}
276				return nil, fmt.Errorf("parse error: %w", err)
277			}
278			streamVals = append(streamVals, v)
279		}
280
281		if slurp {
282			vals = append(vals, streamVals)
283		} else {
284			vals = append(vals, streamVals...)
285		}
286	}
287
288	if len(vals) == 0 {
289		return []any{nil}, nil
290	}
291	return vals, nil
292}
293
294// ctxReader wraps an io.Reader so that each Read call checks ctx first.
295// This makes io.ReadAll over a large but non-blocking source (e.g. a
296// bytes.Reader of the hook stdin payload) cancellable on the next chunk
297// boundary. A reader that itself blocks in Read will still outlast ctx —
298// the hook runner's abandon-goroutine path is the enforcer of last resort
299// for that case.
300type ctxReader struct {
301	ctx context.Context
302	r   io.Reader
303}
304
305func (cr ctxReader) Read(p []byte) (int, error) {
306	if err := cr.ctx.Err(); err != nil {
307		return 0, err
308	}
309	return cr.r.Read(p)
310}
311
312// writeValue writes a single jq output value.
313func writeValue(w io.Writer, v any, raw, compact, join bool) error {
314	if raw {
315		if s, ok := v.(string); ok {
316			if _, err := fmt.Fprint(w, s); err != nil {
317				return err
318			}
319			if !join {
320				_, err := fmt.Fprint(w, "\n")
321				return err
322			}
323			return nil
324		}
325	}
326
327	var bs []byte
328	var err error
329	if compact {
330		bs, err = gojq.Marshal(v)
331	} else {
332		bs, err = json.MarshalIndent(v, "", "  ")
333	}
334	if err != nil {
335		return err
336	}
337	if _, writeErr := w.Write(bs); writeErr != nil {
338		return writeErr
339	}
340	_, err = fmt.Fprint(w, "\n")
341	return err
342}