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