Add support reference

decode.go

@@ -22,20 +22,23 @@ import (
 	"io"
 	"math"
 	"reflect"
+	"sort"
 	"strconv"
+	"strings"
 	"time"
 )
 
 // ----------------------------------------------------------------------------
 // Parser, produces a node tree out of a libyaml event stream.
 
 type parser struct {
-	parser   yaml_parser_t
-	event    yaml_event_t
-	doc      *Node
-	anchors  map[string]*Node
-	doneInit bool
-	textless bool
+	parser     yaml_parser_t
+	event      yaml_event_t
+	doc        *Node
+	anchors    map[string]*Node
+	references map[string][]*Node
+	doneInit   bool
+	textless   bool
 }
 
 func newParser(b []byte) *parser {
@@ -64,6 +67,7 @@ func (p *parser) init() {
 		return
 	}
 	p.anchors = make(map[string]*Node)
+	p.references = make(map[string][]*Node)
 	p.expect(yaml_STREAM_START_EVENT)
 	p.doneInit = true
 }
@@ -144,6 +148,28 @@ func (p *parser) anchor(n *Node, anchor []byte) {
 	}
 }
 
+func (p *parser) reference(n *Node, ref *Reference) {
+	if n.References == nil {
+		n.References = make(map[string]*References)
+	}
+	if n.References[ref.Name] == nil {
+		n.References[ref.Name] = &References{
+			Name: ref.Name,
+		}
+		f := false
+		for _, m := range p.references[ref.Name] {
+			if m == n {
+				f = true
+				break
+			}
+		}
+		if !f {
+			p.references[ref.Name] = append(p.references[ref.Name], n)
+		}
+	}
+	n.References[ref.Name].References = append(n.References[ref.Name].References, ref)
+}
+
 func (p *parser) parse() *Node {
 	p.init()
 	switch p.peek() {
@@ -244,9 +270,17 @@ func (p *parser) scalar() *Node {
 	} else {
 		defaultTag = strTag
 	}
-	n := p.node(ScalarNode, defaultTag, nodeTag, nodeValue)
+	kind := ScalarNode
+	refs := parseReferences(nodeValue)
+	if len(refs) > 0 {
+		kind = ScalarReferenceNode
+	}
+	n := p.node(kind, defaultTag, nodeTag, nodeValue)
 	n.Style |= nodeStyle
 	p.anchor(n, p.event.anchor)
+	for _, ref := range refs {
+		p.reference(n, ref)
+	}
 	p.expect(yaml_SCALAR_EVENT)
 	return n
 }
@@ -318,11 +352,13 @@ type decoder struct {
 	stringMapType  reflect.Type
 	generalMapType reflect.Type
 
-	knownFields bool
-	uniqueKeys  bool
-	decodeCount int
-	aliasCount  int
-	aliasDepth  int
+	knownFields     bool
+	allowReferences bool
+	uniqueKeys      bool
+	decodeCount     int
+	aliasCount      int
+	aliasDepth      int
+	referenceDepth  int
 
 	mergedFields map[interface{}]bool
 }
@@ -504,6 +540,8 @@ func (d *decoder) unmarshal(n *Node, out reflect.Value) (good bool) {
 		return good
 	}
 	switch n.Kind {
+	case ScalarReferenceNode:
+		good = d.scalarReference(n, out)
 	case ScalarNode:
 		good = d.scalar(n, out)
 	case MappingNode:
@@ -562,6 +600,61 @@ func (d *decoder) null(out reflect.Value) bool {
 	return false
 }
 
+func (d *decoder) scalarReference(n *Node, out reflect.Value) bool {
+	if d.allowReferences {
+		if d.referenceDepth > 100 {
+			failf("document contains excessive references")
+		}
+		var ranges []*ReferenceReverse
+		for name, refs := range n.References {
+			for _, ref := range refs.References {
+				r := &ReferenceReverse{
+					Range:      &ref.Range,
+					References: refs,
+				}
+				ranges = append(ranges, r)
+			}
+			if refs.Target == nil {
+				p := d.doc
+				for _, k := range strings.Split(name, ".") {
+					var t map[string]Node
+					if err := p.Decode(&t); err != nil {
+						failf(err.Error())
+					}
+					q := t[k]
+					if q.Kind == 0 {
+						failf("parsing %s reference failed", name)
+					}
+					p = &q
+				}
+				refs.Target = p
+			}
+		}
+		sort.SliceStable(ranges, func(i, j int) bool {
+			return ranges[i].Range.End <= ranges[j].Range.Begin
+		})
+		var result string
+		i := 0
+		for _, r := range ranges {
+			var t string
+			d.referenceDepth++
+			if good := d.unmarshal(r.References.Target, reflect.ValueOf(&t).Elem()); !good {
+				failf("unmarshal failed for %s", r.References.Name)
+			}
+			d.referenceDepth--
+			result += n.Value[i:r.Range.Begin]
+			result += t
+			i = r.Range.End
+		}
+		result += n.Value[i:]
+		m := Node(*n)
+		m.Kind = ScalarNode
+		m.Value = result
+		return d.scalar(&m, out)
+	}
+	return d.scalar(n, out)
+}
+
 func (d *decoder) scalar(n *Node, out reflect.Value) bool {
 	var tag string
 	var resolved interface{}

decode_test.go

@@ -1703,6 +1703,91 @@ func (s *S) TestUnmarshalKnownFields(c *C) {
 	}
 }
 
+var unmarshalReferencesTests = []struct {
+	refer bool
+	data  string
+	value interface{}
+	error string
+}{{
+	refer: false,
+	data:  "a: x\nb: ${a}\n",
+	value: struct{ A, B string }{A: "x", B: "${a}"},
+}, {
+	refer: true,
+	data:  "a: x\nb: ${a}\n",
+	value: struct{ A, B string }{A: "x", B: "x"},
+}, {
+	refer: true,
+	data:  "a: x\nb: ${c}\n",
+	value: struct{ A, B string }{A: "x"},
+	error: "yaml: parsing c reference failed",
+}, {
+	refer: true,
+	data:  "a: x\nb: ${b}\n",
+	value: struct{ A, B string }{A: "x"},
+	error: "yaml: document contains excessive references",
+}, {
+	refer: true,
+	data:  "a: x\nb: ${d}\nc: ${b}\nd: ${c}\n",
+	value: struct{ A, B, C, D string }{A: "x"},
+	error: "yaml: document contains excessive references",
+}, {
+	refer: true,
+	data:  "a: x\nb: y\nc: ${a}${b}\n",
+	value: struct{ A, B, C string }{A: "x", B: "y", C: "xy"},
+}, {
+	refer: true,
+	data:  "a: x\nb: y\nc: ${a} ${b}\n",
+	value: struct{ A, B, C string }{A: "x", B: "y", C: "x y"},
+}, {
+	refer: true,
+	data:  "a: x\nb: y\nc: ${a} ${b}\nd: ${c} z\n",
+	value: struct{ A, B, C, D string }{A: "x", B: "y", C: "x y", D: "x y z"},
+}, {
+	refer: true,
+	data:  "a: x\nb: y\nc: ${a} ${b}\nd: ${c} z\n",
+	value: struct{ A, B, D string }{A: "x", B: "y", D: "x y z"},
+}, {
+	refer: true,
+	data:  "x:\n  a: x\nc: ${x.a}\n",
+	value: struct {
+		X struct{ A string }
+		C string
+	}{X: struct{ A string }{A: "x"}, C: "x"},
+}, {
+	refer: true,
+	data:  "x:\n  a: x\nc: ${x.b}\n",
+	value: struct {
+		X struct{ A string }
+		C string
+	}{X: struct{ A string }{A: "x"}},
+	error: "yaml: parsing x.b reference failed",
+}, {
+	refer: true,
+	data:  "x:\n  a: x\n  b: y\nc: ${x.a} ${x.b}\n",
+	value: struct {
+		X struct{ A, B string }
+		C string
+	}{X: struct{ A, B string }{A: "x", B: "y"}, C: "x y"},
+}}
+
+func (s *S) TestUnmarshalAllowReferences(c *C) {
+	for i, item := range unmarshalReferencesTests {
+		c.Logf("test %d: %q", i, item.data)
+		t := reflect.ValueOf(item.value).Type()
+		value := reflect.New(t)
+		dec := yaml.NewDecoder(bytes.NewBuffer([]byte(item.data)))
+		dec.AllowReferences(item.refer)
+		err := dec.Decode(value.Interface())
+		c.Assert(value.Elem().Interface(), DeepEquals, item.value)
+		if item.error == "" {
+			c.Assert(err, IsNil)
+		} else {
+			c.Assert(err, ErrorMatches, item.error)
+		}
+	}
+}
+
 type textUnmarshaler struct {
 	S string
 }

reference.go

@@ -0,0 +1,74 @@
+//
+// Copyright (c) 2011-2019 Canonical Ltd
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package yaml implements YAML support for the Go language.
+//
+// Source code and other details for the project are available at GitHub:
+//
+//	https://github.com/go-yaml/yaml
+package yaml
+
+import (
+	"regexp"
+	"strings"
+)
+
+// Range is a programming construct used to represent a contiguous range of values.
+type Range struct {
+	Begin int
+	End   int
+}
+
+// Reference is an object that consists of a Name string and an array of
+// Range objects. It allows for organizing and referencing multiple
+// contiguous value ranges associated with a specific name or identifier.
+type Reference struct {
+	Name  string
+	Range Range
+}
+
+// References to represent a group of References with the same Name and Value,
+// pointing to the same Node
+type References struct {
+	Name       string
+	Target     *Node
+	References []*Reference
+}
+
+// ReferenceReverse is an object containing Range and References, which helps
+// to identify and replace a range with the corresponding references
+type ReferenceReverse struct {
+	Range      *Range
+	References *References
+}
+
+var yamlRegexReference = regexp.MustCompile(`\${(.+?)}`)
+
+func parseReferences(value string) []*Reference {
+	items := make([]*Reference, 0)
+	if yamlRegexReference.MatchString(value) {
+		for _, x := range yamlRegexReference.FindAllStringSubmatchIndex(value, -1) {
+			item := &Reference{
+				Name: strings.TrimSpace(value[x[2]:x[3]]),
+				Range: Range{
+					Begin: x[0],
+					End:   x[1],
+				},
+			}
+			items = append(items, item)
+		}
+	}
+	return items
+}

yaml.go

@@ -91,8 +91,9 @@ func Unmarshal(in []byte, out interface{}) (err error) {
 
 // A Decoder reads and decodes YAML values from an input stream.
 type Decoder struct {
-	parser      *parser
-	knownFields bool
+	parser          *parser
+	knownFields     bool
+	allowReferences bool
 }
 
 // NewDecoder returns a new decoder that reads from r.
@@ -111,6 +112,12 @@ func (dec *Decoder) KnownFields(enable bool) {
 	dec.knownFields = enable
 }
 
+// AllowReferences allows the use of ${<field_name>} to refer to
+// other fields in the same struct being decoded into.
+func (dec *Decoder) AllowReferences(enable bool) {
+	dec.allowReferences = enable
+}
+
 // Decode reads the next YAML-encoded value from its input
 // and stores it in the value pointed to by v.
 //
@@ -119,6 +126,7 @@ func (dec *Decoder) KnownFields(enable bool) {
 func (dec *Decoder) Decode(v interface{}) (err error) {
 	d := newDecoder()
 	d.knownFields = dec.knownFields
+	d.allowReferences = dec.allowReferences
 	defer handleErr(&err)
 	node := dec.parser.parse()
 	if node == nil {
@@ -328,6 +336,7 @@ const (
 	MappingNode
 	ScalarNode
 	AliasNode
+	ScalarReferenceNode
 )
 
 type Style uint32
@@ -396,6 +405,9 @@ type Node struct {
 	// Alias holds the node that this alias points to. Only valid when Kind is AliasNode.
 	Alias *Node
 
+	// References holds the references for this node
+	References map[string]*References
+
 	// Content holds contained nodes for documents, mappings, and sequences.
 	Content []*Node