1package dsig
2
3import (
4	"sort"
5
6	"github.com/beevik/etree"
7	"github.com/russellhaering/goxmldsig/etreeutils"
8)
9
10// Canonicalizer is an implementation of a canonicalization algorithm.
11type Canonicalizer interface {
12	Canonicalize(el *etree.Element) ([]byte, error)
13	Algorithm() AlgorithmID
14}
15
16type c14N10ExclusiveCanonicalizer struct {
17	prefixList string
18}
19
20// MakeC14N10ExclusiveCanonicalizerWithPrefixList constructs an exclusive Canonicalizer
21// from a PrefixList in NMTOKENS format (a white space separated list).
22func MakeC14N10ExclusiveCanonicalizerWithPrefixList(prefixList string) Canonicalizer {
23	return &c14N10ExclusiveCanonicalizer{
24		prefixList: prefixList,
25	}
26}
27
28// Canonicalize transforms the input Element into a serialized XML document in canonical form.
29func (c *c14N10ExclusiveCanonicalizer) Canonicalize(el *etree.Element) ([]byte, error) {
30	err := etreeutils.TransformExcC14n(el, c.prefixList)
31	if err != nil {
32		return nil, err
33	}
34
35	return canonicalSerialize(el)
36}
37
38func (c *c14N10ExclusiveCanonicalizer) Algorithm() AlgorithmID {
39	return CanonicalXML10ExclusiveAlgorithmId
40}
41
42type c14N11Canonicalizer struct{}
43
44// MakeC14N11Canonicalizer constructs an inclusive canonicalizer.
45func MakeC14N11Canonicalizer() Canonicalizer {
46	return &c14N11Canonicalizer{}
47}
48
49// Canonicalize transforms the input Element into a serialized XML document in canonical form.
50func (c *c14N11Canonicalizer) Canonicalize(el *etree.Element) ([]byte, error) {
51	scope := make(map[string]struct{})
52	return canonicalSerialize(canonicalPrep(el, scope))
53}
54
55func (c *c14N11Canonicalizer) Algorithm() AlgorithmID {
56	return CanonicalXML11AlgorithmId
57}
58
59func composeAttr(space, key string) string {
60	if space != "" {
61		return space + ":" + key
62	}
63
64	return key
65}
66
67type c14nSpace struct {
68	a    etree.Attr
69	used bool
70}
71
72const nsSpace = "xmlns"
73
74// canonicalPrep accepts an *etree.Element and transforms it into one which is ready
75// for serialization into inclusive canonical form. Specifically this
76// entails:
77//
78// 1. Stripping re-declarations of namespaces
79// 2. Sorting attributes into canonical order
80//
81// Inclusive canonicalization does not strip unused namespaces.
82//
83// TODO(russell_h): This is very similar to excCanonicalPrep - perhaps they should
84// be unified into one parameterized function?
85func canonicalPrep(el *etree.Element, seenSoFar map[string]struct{}) *etree.Element {
86	_seenSoFar := make(map[string]struct{})
87	for k, v := range seenSoFar {
88		_seenSoFar[k] = v
89	}
90
91	ne := el.Copy()
92	sort.Sort(etreeutils.SortedAttrs(ne.Attr))
93	if len(ne.Attr) != 0 {
94		for _, attr := range ne.Attr {
95			if attr.Space != nsSpace {
96				continue
97			}
98			key := attr.Space + ":" + attr.Key
99			if _, seen := _seenSoFar[key]; seen {
100				ne.RemoveAttr(attr.Space + ":" + attr.Key)
101			} else {
102				_seenSoFar[key] = struct{}{}
103			}
104		}
105	}
106
107	for i, token := range ne.Child {
108		childElement, ok := token.(*etree.Element)
109		if ok {
110			ne.Child[i] = canonicalPrep(childElement, _seenSoFar)
111		}
112	}
113
114	return ne
115}
116
117func canonicalSerialize(el *etree.Element) ([]byte, error) {
118	doc := etree.NewDocument()
119	doc.SetRoot(el.Copy())
120
121	doc.WriteSettings = etree.WriteSettings{
122		CanonicalAttrVal: true,
123		CanonicalEndTags: true,
124		CanonicalText:    true,
125	}
126
127	return doc.WriteToBytes()
128}
129