sjson.go

  1// Package sjson provides setting json values.
  2package sjson
  3
  4import (
  5	jsongo "encoding/json"
  6	"sort"
  7	"strconv"
  8	"unsafe"
  9
 10	"github.com/tidwall/gjson"
 11)
 12
 13type errorType struct {
 14	msg string
 15}
 16
 17func (err *errorType) Error() string {
 18	return err.msg
 19}
 20
 21// Options represents additional options for the Set and Delete functions.
 22type Options struct {
 23	// Optimistic is a hint that the value likely exists which
 24	// allows for the sjson to perform a fast-track search and replace.
 25	Optimistic bool
 26	// ReplaceInPlace is a hint to replace the input json rather than
 27	// allocate a new json byte slice. When this field is specified
 28	// the input json will not longer be valid and it should not be used
 29	// In the case when the destination slice doesn't have enough free
 30	// bytes to replace the data in place, a new bytes slice will be
 31	// created under the hood.
 32	// The Optimistic flag must be set to true and the input must be a
 33	// byte slice in order to use this field.
 34	ReplaceInPlace bool
 35}
 36
 37type pathResult struct {
 38	part  string // current key part
 39	gpart string // gjson get part
 40	path  string // remaining path
 41	force bool   // force a string key
 42	more  bool   // there is more path to parse
 43}
 44
 45func isSimpleChar(ch byte) bool {
 46	switch ch {
 47	case '|', '#', '@', '*', '?':
 48		return false
 49	default:
 50		return true
 51	}
 52}
 53
 54func parsePath(path string) (res pathResult, simple bool) {
 55	var r pathResult
 56	if len(path) > 0 && path[0] == ':' {
 57		r.force = true
 58		path = path[1:]
 59	}
 60	for i := 0; i < len(path); i++ {
 61		if path[i] == '.' {
 62			r.part = path[:i]
 63			r.gpart = path[:i]
 64			r.path = path[i+1:]
 65			r.more = true
 66			return r, true
 67		}
 68		if !isSimpleChar(path[i]) {
 69			return r, false
 70		}
 71		if path[i] == '\\' {
 72			// go into escape mode. this is a slower path that
 73			// strips off the escape character from the part.
 74			epart := []byte(path[:i])
 75			gpart := []byte(path[:i+1])
 76			i++
 77			if i < len(path) {
 78				epart = append(epart, path[i])
 79				gpart = append(gpart, path[i])
 80				i++
 81				for ; i < len(path); i++ {
 82					if path[i] == '\\' {
 83						gpart = append(gpart, '\\')
 84						i++
 85						if i < len(path) {
 86							epart = append(epart, path[i])
 87							gpart = append(gpart, path[i])
 88						}
 89						continue
 90					} else if path[i] == '.' {
 91						r.part = string(epart)
 92						r.gpart = string(gpart)
 93						r.path = path[i+1:]
 94						r.more = true
 95						return r, true
 96					} else if !isSimpleChar(path[i]) {
 97						return r, false
 98					}
 99					epart = append(epart, path[i])
100					gpart = append(gpart, path[i])
101				}
102			}
103			// append the last part
104			r.part = string(epart)
105			r.gpart = string(gpart)
106			return r, true
107		}
108	}
109	r.part = path
110	r.gpart = path
111	return r, true
112}
113
114func mustMarshalString(s string) bool {
115	for i := 0; i < len(s); i++ {
116		if s[i] < ' ' || s[i] > 0x7f || s[i] == '"' || s[i] == '\\' {
117			return true
118		}
119	}
120	return false
121}
122
123// appendStringify makes a json string and appends to buf.
124func appendStringify(buf []byte, s string) []byte {
125	if mustMarshalString(s) {
126		b, _ := jsongo.Marshal(s)
127		return append(buf, b...)
128	}
129	buf = append(buf, '"')
130	buf = append(buf, s...)
131	buf = append(buf, '"')
132	return buf
133}
134
135// appendBuild builds a json block from a json path.
136func appendBuild(buf []byte, array bool, paths []pathResult, raw string,
137	stringify bool) []byte {
138	if !array {
139		buf = appendStringify(buf, paths[0].part)
140		buf = append(buf, ':')
141	}
142	if len(paths) > 1 {
143		n, numeric := atoui(paths[1])
144		if numeric || (!paths[1].force && paths[1].part == "-1") {
145			buf = append(buf, '[')
146			buf = appendRepeat(buf, "null,", n)
147			buf = appendBuild(buf, true, paths[1:], raw, stringify)
148			buf = append(buf, ']')
149		} else {
150			buf = append(buf, '{')
151			buf = appendBuild(buf, false, paths[1:], raw, stringify)
152			buf = append(buf, '}')
153		}
154	} else {
155		if stringify {
156			buf = appendStringify(buf, raw)
157		} else {
158			buf = append(buf, raw...)
159		}
160	}
161	return buf
162}
163
164// atoui does a rip conversion of string -> unigned int.
165func atoui(r pathResult) (n int, ok bool) {
166	if r.force {
167		return 0, false
168	}
169	for i := 0; i < len(r.part); i++ {
170		if r.part[i] < '0' || r.part[i] > '9' {
171			return 0, false
172		}
173		n = n*10 + int(r.part[i]-'0')
174	}
175	return n, true
176}
177
178// appendRepeat repeats string "n" times and appends to buf.
179func appendRepeat(buf []byte, s string, n int) []byte {
180	for i := 0; i < n; i++ {
181		buf = append(buf, s...)
182	}
183	return buf
184}
185
186// trim does a rip trim
187func trim(s string) string {
188	for len(s) > 0 {
189		if s[0] <= ' ' {
190			s = s[1:]
191			continue
192		}
193		break
194	}
195	for len(s) > 0 {
196		if s[len(s)-1] <= ' ' {
197			s = s[:len(s)-1]
198			continue
199		}
200		break
201	}
202	return s
203}
204
205// deleteTailItem deletes the previous key or comma.
206func deleteTailItem(buf []byte) ([]byte, bool) {
207loop:
208	for i := len(buf) - 1; i >= 0; i-- {
209		// look for either a ',',':','['
210		switch buf[i] {
211		case '[':
212			return buf, true
213		case ',':
214			return buf[:i], false
215		case ':':
216			// delete tail string
217			i--
218			for ; i >= 0; i-- {
219				if buf[i] == '"' {
220					i--
221					for ; i >= 0; i-- {
222						if buf[i] == '"' {
223							i--
224							if i >= 0 && buf[i] == '\\' {
225								i--
226								continue
227							}
228							for ; i >= 0; i-- {
229								// look for either a ',','{'
230								switch buf[i] {
231								case '{':
232									return buf[:i+1], true
233								case ',':
234									return buf[:i], false
235								}
236							}
237						}
238					}
239					break
240				}
241			}
242			break loop
243		}
244	}
245	return buf, false
246}
247
248var errNoChange = &errorType{"no change"}
249
250func appendRawPaths(buf []byte, jstr string, paths []pathResult, raw string,
251	stringify, del bool) ([]byte, error) {
252	var err error
253	var res gjson.Result
254	var found bool
255	if del {
256		if paths[0].part == "-1" && !paths[0].force {
257			res = gjson.Get(jstr, "#")
258			if res.Int() > 0 {
259				res = gjson.Get(jstr, strconv.FormatInt(int64(res.Int()-1), 10))
260				found = true
261			}
262		}
263	}
264	if !found {
265		res = gjson.Get(jstr, paths[0].gpart)
266	}
267	if res.Index > 0 {
268		if len(paths) > 1 {
269			buf = append(buf, jstr[:res.Index]...)
270			buf, err = appendRawPaths(buf, res.Raw, paths[1:], raw,
271				stringify, del)
272			if err != nil {
273				return nil, err
274			}
275			buf = append(buf, jstr[res.Index+len(res.Raw):]...)
276			return buf, nil
277		}
278		buf = append(buf, jstr[:res.Index]...)
279		var exidx int // additional forward stripping
280		if del {
281			var delNextComma bool
282			buf, delNextComma = deleteTailItem(buf)
283			if delNextComma {
284				i, j := res.Index+len(res.Raw), 0
285				for ; i < len(jstr); i, j = i+1, j+1 {
286					if jstr[i] <= ' ' {
287						continue
288					}
289					if jstr[i] == ',' {
290						exidx = j + 1
291					}
292					break
293				}
294			}
295		} else {
296			if stringify {
297				buf = appendStringify(buf, raw)
298			} else {
299				buf = append(buf, raw...)
300			}
301		}
302		buf = append(buf, jstr[res.Index+len(res.Raw)+exidx:]...)
303		return buf, nil
304	}
305	if del {
306		return nil, errNoChange
307	}
308	n, numeric := atoui(paths[0])
309	isempty := true
310	for i := 0; i < len(jstr); i++ {
311		if jstr[i] > ' ' {
312			isempty = false
313			break
314		}
315	}
316	if isempty {
317		if numeric {
318			jstr = "[]"
319		} else {
320			jstr = "{}"
321		}
322	}
323	jsres := gjson.Parse(jstr)
324	if jsres.Type != gjson.JSON {
325		if numeric {
326			jstr = "[]"
327		} else {
328			jstr = "{}"
329		}
330		jsres = gjson.Parse(jstr)
331	}
332	var comma bool
333	for i := 1; i < len(jsres.Raw); i++ {
334		if jsres.Raw[i] <= ' ' {
335			continue
336		}
337		if jsres.Raw[i] == '}' || jsres.Raw[i] == ']' {
338			break
339		}
340		comma = true
341		break
342	}
343	switch jsres.Raw[0] {
344	default:
345		return nil, &errorType{"json must be an object or array"}
346	case '{':
347		end := len(jsres.Raw) - 1
348		for ; end > 0; end-- {
349			if jsres.Raw[end] == '}' {
350				break
351			}
352		}
353		buf = append(buf, jsres.Raw[:end]...)
354		if comma {
355			buf = append(buf, ',')
356		}
357		buf = appendBuild(buf, false, paths, raw, stringify)
358		buf = append(buf, '}')
359		return buf, nil
360	case '[':
361		var appendit bool
362		if !numeric {
363			if paths[0].part == "-1" && !paths[0].force {
364				appendit = true
365			} else {
366				return nil, &errorType{
367					"cannot set array element for non-numeric key '" +
368						paths[0].part + "'"}
369			}
370		}
371		if appendit {
372			njson := trim(jsres.Raw)
373			if njson[len(njson)-1] == ']' {
374				njson = njson[:len(njson)-1]
375			}
376			buf = append(buf, njson...)
377			if comma {
378				buf = append(buf, ',')
379			}
380
381			buf = appendBuild(buf, true, paths, raw, stringify)
382			buf = append(buf, ']')
383			return buf, nil
384		}
385		buf = append(buf, '[')
386		ress := jsres.Array()
387		for i := 0; i < len(ress); i++ {
388			if i > 0 {
389				buf = append(buf, ',')
390			}
391			buf = append(buf, ress[i].Raw...)
392		}
393		if len(ress) == 0 {
394			buf = appendRepeat(buf, "null,", n-len(ress))
395		} else {
396			buf = appendRepeat(buf, ",null", n-len(ress))
397			if comma {
398				buf = append(buf, ',')
399			}
400		}
401		buf = appendBuild(buf, true, paths, raw, stringify)
402		buf = append(buf, ']')
403		return buf, nil
404	}
405}
406
407func isOptimisticPath(path string) bool {
408	for i := 0; i < len(path); i++ {
409		if path[i] < '.' || path[i] > 'z' {
410			return false
411		}
412		if path[i] > '9' && path[i] < 'A' {
413			return false
414		}
415		if path[i] > 'z' {
416			return false
417		}
418	}
419	return true
420}
421
422// Set sets a json value for the specified path.
423// A path is in dot syntax, such as "name.last" or "age".
424// This function expects that the json is well-formed, and does not validate.
425// Invalid json will not panic, but it may return back unexpected results.
426// An error is returned if the path is not valid.
427//
428// A path is a series of keys separated by a dot.
429//
430//  {
431//    "name": {"first": "Tom", "last": "Anderson"},
432//    "age":37,
433//    "children": ["Sara","Alex","Jack"],
434//    "friends": [
435//      {"first": "James", "last": "Murphy"},
436//      {"first": "Roger", "last": "Craig"}
437//    ]
438//  }
439//  "name.last"          >> "Anderson"
440//  "age"                >> 37
441//  "children.1"         >> "Alex"
442//
443func Set(json, path string, value interface{}) (string, error) {
444	return SetOptions(json, path, value, nil)
445}
446
447// SetBytes sets a json value for the specified path.
448// If working with bytes, this method preferred over
449// Set(string(data), path, value)
450func SetBytes(json []byte, path string, value interface{}) ([]byte, error) {
451	return SetBytesOptions(json, path, value, nil)
452}
453
454// SetRaw sets a raw json value for the specified path.
455// This function works the same as Set except that the value is set as a
456// raw block of json. This allows for setting premarshalled json objects.
457func SetRaw(json, path, value string) (string, error) {
458	return SetRawOptions(json, path, value, nil)
459}
460
461// SetRawOptions sets a raw json value for the specified path with options.
462// This furnction works the same as SetOptions except that the value is set
463// as a raw block of json. This allows for setting premarshalled json objects.
464func SetRawOptions(json, path, value string, opts *Options) (string, error) {
465	var optimistic bool
466	if opts != nil {
467		optimistic = opts.Optimistic
468	}
469	res, err := set(json, path, value, false, false, optimistic, false)
470	if err == errNoChange {
471		return json, nil
472	}
473	return string(res), err
474}
475
476// SetRawBytes sets a raw json value for the specified path.
477// If working with bytes, this method preferred over
478// SetRaw(string(data), path, value)
479func SetRawBytes(json []byte, path string, value []byte) ([]byte, error) {
480	return SetRawBytesOptions(json, path, value, nil)
481}
482
483type dtype struct{}
484
485// Delete deletes a value from json for the specified path.
486func Delete(json, path string) (string, error) {
487	return Set(json, path, dtype{})
488}
489
490// DeleteBytes deletes a value from json for the specified path.
491func DeleteBytes(json []byte, path string) ([]byte, error) {
492	return SetBytes(json, path, dtype{})
493}
494
495type stringHeader struct {
496	data unsafe.Pointer
497	len  int
498}
499
500type sliceHeader struct {
501	data unsafe.Pointer
502	len  int
503	cap  int
504}
505
506func set(jstr, path, raw string,
507	stringify, del, optimistic, inplace bool) ([]byte, error) {
508	if path == "" {
509		return []byte(jstr), &errorType{"path cannot be empty"}
510	}
511	if !del && optimistic && isOptimisticPath(path) {
512		res := gjson.Get(jstr, path)
513		if res.Exists() && res.Index > 0 {
514			sz := len(jstr) - len(res.Raw) + len(raw)
515			if stringify {
516				sz += 2
517			}
518			if inplace && sz <= len(jstr) {
519				if !stringify || !mustMarshalString(raw) {
520					jsonh := *(*stringHeader)(unsafe.Pointer(&jstr))
521					jsonbh := sliceHeader{
522						data: jsonh.data, len: jsonh.len, cap: jsonh.len}
523					jbytes := *(*[]byte)(unsafe.Pointer(&jsonbh))
524					if stringify {
525						jbytes[res.Index] = '"'
526						copy(jbytes[res.Index+1:], []byte(raw))
527						jbytes[res.Index+1+len(raw)] = '"'
528						copy(jbytes[res.Index+1+len(raw)+1:],
529							jbytes[res.Index+len(res.Raw):])
530					} else {
531						copy(jbytes[res.Index:], []byte(raw))
532						copy(jbytes[res.Index+len(raw):],
533							jbytes[res.Index+len(res.Raw):])
534					}
535					return jbytes[:sz], nil
536				}
537				return []byte(jstr), nil
538			}
539			buf := make([]byte, 0, sz)
540			buf = append(buf, jstr[:res.Index]...)
541			if stringify {
542				buf = appendStringify(buf, raw)
543			} else {
544				buf = append(buf, raw...)
545			}
546			buf = append(buf, jstr[res.Index+len(res.Raw):]...)
547			return buf, nil
548		}
549	}
550	var paths []pathResult
551	r, simple := parsePath(path)
552	if simple {
553		paths = append(paths, r)
554		for r.more {
555			r, simple = parsePath(r.path)
556			if !simple {
557				break
558			}
559			paths = append(paths, r)
560		}
561	}
562	if !simple {
563		if del {
564			return []byte(jstr),
565				&errorType{"cannot delete value from a complex path"}
566		}
567		return setComplexPath(jstr, path, raw, stringify)
568	}
569	njson, err := appendRawPaths(nil, jstr, paths, raw, stringify, del)
570	if err != nil {
571		return []byte(jstr), err
572	}
573	return njson, nil
574}
575
576func setComplexPath(jstr, path, raw string, stringify bool) ([]byte, error) {
577	res := gjson.Get(jstr, path)
578	if !res.Exists() || !(res.Index != 0 || len(res.Indexes) != 0) {
579		return []byte(jstr), errNoChange
580	}
581	if res.Index != 0 {
582		njson := []byte(jstr[:res.Index])
583		if stringify {
584			njson = appendStringify(njson, raw)
585		} else {
586			njson = append(njson, raw...)
587		}
588		njson = append(njson, jstr[res.Index+len(res.Raw):]...)
589		jstr = string(njson)
590	}
591	if len(res.Indexes) > 0 {
592		type val struct {
593			index int
594			res   gjson.Result
595		}
596		vals := make([]val, 0, len(res.Indexes))
597		res.ForEach(func(_, vres gjson.Result) bool {
598			vals = append(vals, val{res: vres})
599			return true
600		})
601		if len(res.Indexes) != len(vals) {
602			return []byte(jstr), errNoChange
603		}
604		for i := 0; i < len(res.Indexes); i++ {
605			vals[i].index = res.Indexes[i]
606		}
607		sort.SliceStable(vals, func(i, j int) bool {
608			return vals[i].index > vals[j].index
609		})
610		for _, val := range vals {
611			vres := val.res
612			index := val.index
613			njson := []byte(jstr[:index])
614			if stringify {
615				njson = appendStringify(njson, raw)
616			} else {
617				njson = append(njson, raw...)
618			}
619			njson = append(njson, jstr[index+len(vres.Raw):]...)
620			jstr = string(njson)
621		}
622	}
623	return []byte(jstr), nil
624}
625
626// SetOptions sets a json value for the specified path with options.
627// A path is in dot syntax, such as "name.last" or "age".
628// This function expects that the json is well-formed, and does not validate.
629// Invalid json will not panic, but it may return back unexpected results.
630// An error is returned if the path is not valid.
631func SetOptions(json, path string, value interface{},
632	opts *Options) (string, error) {
633	if opts != nil {
634		if opts.ReplaceInPlace {
635			// it's not safe to replace bytes in-place for strings
636			// copy the Options and set options.ReplaceInPlace to false.
637			nopts := *opts
638			opts = &nopts
639			opts.ReplaceInPlace = false
640		}
641	}
642	jsonh := *(*stringHeader)(unsafe.Pointer(&json))
643	jsonbh := sliceHeader{data: jsonh.data, len: jsonh.len, cap: jsonh.len}
644	jsonb := *(*[]byte)(unsafe.Pointer(&jsonbh))
645	res, err := SetBytesOptions(jsonb, path, value, opts)
646	return string(res), err
647}
648
649// SetBytesOptions sets a json value for the specified path with options.
650// If working with bytes, this method preferred over
651// SetOptions(string(data), path, value)
652func SetBytesOptions(json []byte, path string, value interface{},
653	opts *Options) ([]byte, error) {
654	var optimistic, inplace bool
655	if opts != nil {
656		optimistic = opts.Optimistic
657		inplace = opts.ReplaceInPlace
658	}
659	jstr := *(*string)(unsafe.Pointer(&json))
660	var res []byte
661	var err error
662	switch v := value.(type) {
663	default:
664		b, merr := jsongo.Marshal(value)
665		if merr != nil {
666			return nil, merr
667		}
668		raw := *(*string)(unsafe.Pointer(&b))
669		res, err = set(jstr, path, raw, false, false, optimistic, inplace)
670	case dtype:
671		res, err = set(jstr, path, "", false, true, optimistic, inplace)
672	case string:
673		res, err = set(jstr, path, v, true, false, optimistic, inplace)
674	case []byte:
675		raw := *(*string)(unsafe.Pointer(&v))
676		res, err = set(jstr, path, raw, true, false, optimistic, inplace)
677	case bool:
678		if v {
679			res, err = set(jstr, path, "true", false, false, optimistic, inplace)
680		} else {
681			res, err = set(jstr, path, "false", false, false, optimistic, inplace)
682		}
683	case int8:
684		res, err = set(jstr, path, strconv.FormatInt(int64(v), 10),
685			false, false, optimistic, inplace)
686	case int16:
687		res, err = set(jstr, path, strconv.FormatInt(int64(v), 10),
688			false, false, optimistic, inplace)
689	case int32:
690		res, err = set(jstr, path, strconv.FormatInt(int64(v), 10),
691			false, false, optimistic, inplace)
692	case int64:
693		res, err = set(jstr, path, strconv.FormatInt(int64(v), 10),
694			false, false, optimistic, inplace)
695	case uint8:
696		res, err = set(jstr, path, strconv.FormatUint(uint64(v), 10),
697			false, false, optimistic, inplace)
698	case uint16:
699		res, err = set(jstr, path, strconv.FormatUint(uint64(v), 10),
700			false, false, optimistic, inplace)
701	case uint32:
702		res, err = set(jstr, path, strconv.FormatUint(uint64(v), 10),
703			false, false, optimistic, inplace)
704	case uint64:
705		res, err = set(jstr, path, strconv.FormatUint(uint64(v), 10),
706			false, false, optimistic, inplace)
707	case float32:
708		res, err = set(jstr, path, strconv.FormatFloat(float64(v), 'f', -1, 64),
709			false, false, optimistic, inplace)
710	case float64:
711		res, err = set(jstr, path, strconv.FormatFloat(float64(v), 'f', -1, 64),
712			false, false, optimistic, inplace)
713	}
714	if err == errNoChange {
715		return json, nil
716	}
717	return res, err
718}
719
720// SetRawBytesOptions sets a raw json value for the specified path with options.
721// If working with bytes, this method preferred over
722// SetRawOptions(string(data), path, value, opts)
723func SetRawBytesOptions(json []byte, path string, value []byte,
724	opts *Options) ([]byte, error) {
725	jstr := *(*string)(unsafe.Pointer(&json))
726	vstr := *(*string)(unsafe.Pointer(&value))
727	var optimistic, inplace bool
728	if opts != nil {
729		optimistic = opts.Optimistic
730		inplace = opts.ReplaceInPlace
731	}
732	res, err := set(jstr, path, vstr, false, false, optimistic, inplace)
733	if err == errNoChange {
734		return json, nil
735	}
736	return res, err
737}