tree_test.go raw
1 package mls
2
3 import (
4 "bytes"
5 "fmt"
6 "reflect"
7 "testing"
8 "time"
9 )
10
11 type treeValidationTest struct {
12 CipherSuite CipherSuite `json:"cipher_suite"`
13
14 Tree testBytes `json:"tree"`
15 GroupID testBytes `json:"group_id"`
16
17 Resolutions [][]nodeIndex `json:"resolutions"`
18 TreeHashes []testBytes `json:"tree_hashes"`
19 }
20
21 func testTreeValidation(t *testing.T, tc *treeValidationTest) {
22 var tree ratchetTree
23 if err := unmarshal([]byte(tc.Tree), &tree); err != nil {
24 t.Fatalf("unmarshal(tree) = %v", err)
25 }
26
27 for i, want := range tc.Resolutions {
28 x := nodeIndex(i)
29 res := tree.resolve(x)
30 if len(res) == 0 {
31 res = make([]nodeIndex, 0)
32 }
33 if !reflect.DeepEqual(res, want) {
34 t.Errorf("resolve(%v) = %v, want %v", x, res, want)
35 }
36 }
37
38 for i, want := range tc.TreeHashes {
39 x := nodeIndex(i)
40 if h, err := tree.computeTreeHash(tc.CipherSuite, x, nil); err != nil {
41 t.Errorf("computeTreeHash(%v) = %v", x, err)
42 } else if !bytes.Equal(h, []byte(want)) {
43 t.Errorf("computeTreeHash(%v) = %v, want %v", x, h, want)
44 }
45 }
46
47 if !tree.verifyParentHashes(tc.CipherSuite) {
48 t.Errorf("verifyParentHashes() failed")
49 }
50
51 groupID := GroupID(tc.GroupID)
52 for i, node := range tree {
53 if node == nil || node.nodeType != nodeTypeLeaf {
54 continue
55 }
56 li, ok := nodeIndex(i).leafIndex()
57 if !ok {
58 t.Errorf("leafIndex(%v) = false", i)
59 continue
60 }
61 if !node.leafNode.verifySignature(tc.CipherSuite, groupID, li) {
62 t.Errorf("verify(%v) = false", li)
63 }
64 }
65 }
66
67 func TestTreeValidation(t *testing.T) {
68 var tests []treeValidationTest
69 loadTestVector(t, "testdata/tree-validation.json", &tests)
70
71 for i, tc := range tests {
72 t.Run(fmt.Sprintf("[%v]", i), func(t *testing.T) {
73 testTreeValidation(t, &tc)
74 })
75 }
76 }
77
78 type treeKEMTest struct {
79 CipherSuite CipherSuite `json:"cipher_suite"`
80
81 GroupID testBytes `json:"group_id"`
82 Epoch uint64 `json:"epoch"`
83 ConfirmedTranscriptHash testBytes `json:"confirmed_transcript_hash"`
84
85 RatchetTree testBytes `json:"ratchet_tree"`
86
87 LeavesPrivate []struct {
88 Index leafIndex `json:"index"`
89 EncryptionPriv testBytes `json:"encryption_priv"`
90 SignaturePriv testBytes `json:"signature_priv"`
91 PathSecrets []struct {
92 Node nodeIndex `json:"node"`
93 PathSecret testBytes `json:"path_secret"`
94 } `json:"path_secrets"`
95 } `json:"leaves_private"`
96
97 UpdatePaths []struct {
98 Sender leafIndex `json:"sender"`
99 UpdatePath testBytes `json:"update_path"`
100 PathSecrets []testBytes `json:"path_secrets"`
101 CommitSecret testBytes `json:"commit_secret"`
102 TreeHashAfter testBytes `json:"tree_hash_after"`
103 } `json:"update_paths"`
104 }
105
106 func testTreeKEM(t *testing.T, tc *treeKEMTest) {
107 type privNode struct {
108 encryptionPriv []byte
109 signaturePriv []byte
110 }
111
112 if !tc.CipherSuite.Supported() {
113 t.Skipf("unsupported cipher suite %v", tc.CipherSuite)
114 }
115
116 kem, _, _ := tc.CipherSuite.hpke().Params()
117
118 for _, leafPrivate := range tc.LeavesPrivate {
119 var tree ratchetTree
120 if err := unmarshal([]byte(tc.RatchetTree), &tree); err != nil {
121 t.Fatalf("unmarshal(ratchetTree) = %v", err)
122 }
123
124 privTree := make([]privNode, len(tree))
125 privTree[int(leafPrivate.Index.nodeIndex())] = privNode{
126 encryptionPriv: leafPrivate.EncryptionPriv,
127 signaturePriv: leafPrivate.SignaturePriv,
128 }
129
130 for _, ps := range leafPrivate.PathSecrets {
131 priv, err := nodePrivFromPathSecret(tc.CipherSuite, ps.PathSecret, tree.get(ps.Node).encryptionKey())
132 if err != nil {
133 t.Fatalf("failed to derive node %v private key from path secret: %v", ps.Node, err)
134 }
135
136 privTree[int(ps.Node)] = privNode{
137 encryptionPriv: priv,
138 }
139 }
140
141 for i, privNode := range privTree {
142 if privNode.encryptionPriv == nil {
143 continue
144 }
145
146 priv, err := kem.Scheme().UnmarshalBinaryPrivateKey(privNode.encryptionPriv)
147 if err != nil {
148 t.Fatalf("UnmarshalBinaryPrivateKey() = %v", err)
149 }
150
151 pub, err := kem.Scheme().UnmarshalBinaryPublicKey(tree[i].encryptionKey())
152 if err != nil {
153 t.Fatalf("UnmarshalBinaryPublicKey() = %v", err)
154 }
155
156 if !priv.Public().Equal(pub) {
157 t.Errorf("key pair mismatch for node %v", i)
158 }
159
160 // TODO: check signature key
161 }
162 }
163
164 for _, updatePathTest := range tc.UpdatePaths {
165 var tree ratchetTree
166 if err := unmarshal([]byte(tc.RatchetTree), &tree); err != nil {
167 t.Fatalf("unmarshal(ratchetTree) = %v", err)
168 }
169
170 var up updatePath
171 if err := unmarshal([]byte(updatePathTest.UpdatePath), &up); err != nil {
172 t.Fatalf("unmarshal(updatePath) = %v", err)
173 }
174
175 // TODO: verify that UpdatePath is parent-hash valid relative to ratchet tree
176 // TODO: process UpdatePath using private leaves
177
178 if err := tree.mergeUpdatePath(tc.CipherSuite, updatePathTest.Sender, &up); err != nil {
179 t.Fatalf("ratchetTree.mergeUpdatePath() = %v", err)
180 }
181
182 treeHash, err := tree.computeRootTreeHash(tc.CipherSuite)
183 if err != nil {
184 t.Errorf("ratchetTree.computeRootTreeHash() = %v", err)
185 } else if !bytes.Equal(treeHash, []byte(updatePathTest.TreeHashAfter)) {
186 t.Errorf("ratchetTree.computeRootTreeHash() = %v, want %v", treeHash, updatePathTest.TreeHashAfter)
187 }
188
189 // TODO: create and verify new update path
190 }
191 }
192
193 func TestTreeKEM(t *testing.T) {
194 var tests []treeKEMTest
195 loadTestVector(t, "testdata/treekem.json", &tests)
196
197 for i, tc := range tests {
198 t.Run(fmt.Sprintf("[%v]", i), func(t *testing.T) {
199 testTreeKEM(t, &tc)
200 })
201 }
202 }
203
204 type treeOperationsTest struct {
205 CipherSuite CipherSuite `json:"cipher_suite"`
206
207 TreeBefore testBytes `json:"tree_before"`
208 Proposal testBytes `json:"proposal"`
209 ProposalSender leafIndex `json:"proposal_sender"`
210
211 TreeHashBefore testBytes `json:"tree_hash_before"`
212 TreeAfter testBytes `json:"tree_after"`
213 TreeHashAfter testBytes `json:"tree_hash_after"`
214 }
215
216 func testTreeOperations(t *testing.T, tc *treeOperationsTest) {
217 var tree ratchetTree
218 if err := unmarshal([]byte(tc.TreeBefore), &tree); err != nil {
219 t.Fatalf("unmarshal(tree) = %v", err)
220 }
221
222 treeHash, err := tree.computeRootTreeHash(tc.CipherSuite)
223 if err != nil {
224 t.Errorf("ratchetTree.computeRootTreeHash() = %v", err)
225 } else if !bytes.Equal(treeHash, []byte(tc.TreeHashBefore)) {
226 t.Errorf("ratchetTree.computeRootTreeHash() = %v, want %v", treeHash, tc.TreeHashBefore)
227 }
228
229 var prop proposal
230 if err := unmarshal([]byte(tc.Proposal), &prop); err != nil {
231 t.Fatalf("unmarshal(proposal) = %v", err)
232 }
233
234 switch prop.proposalType {
235 case proposalTypeAdd:
236 ctx := groupContext{
237 version: prop.add.keyPackage.version,
238 cipherSuite: prop.add.keyPackage.cipherSuite,
239 }
240 if err := prop.add.keyPackage.verify(&ctx); err != nil {
241 t.Errorf("keyPackage.verify() = %v", err)
242 }
243 tree.add(&prop.add.keyPackage.leafNode)
244 case proposalTypeUpdate:
245 signatureKeys, encryptionKeys := tree.keys()
246 err := prop.update.leafNode.verify(&leafNodeVerifyOptions{
247 cipherSuite: tc.CipherSuite,
248 groupID: nil,
249 leafIndex: tc.ProposalSender,
250 supportedCreds: tree.supportedCreds(),
251 signatureKeys: signatureKeys,
252 encryptionKeys: encryptionKeys,
253 now: func() time.Time { return time.Time{} },
254 })
255 if err != nil {
256 t.Errorf("leafNode.verify() = %v", err)
257 }
258 tree.update(tc.ProposalSender, &prop.update.leafNode)
259 case proposalTypeRemove:
260 if tree.getLeaf(prop.remove.removed) == nil {
261 t.Errorf("leaf node %v is blank", prop.remove.removed)
262 }
263 tree.remove(prop.remove.removed)
264 default:
265 panic("unreachable")
266 }
267
268 rawTree, err := marshal(&tree)
269 if err != nil {
270 t.Fatalf("marshal(tree) = %v", err)
271 } else if !bytes.Equal(rawTree, []byte(tc.TreeAfter)) {
272 t.Errorf("marshal(tree) = %v, want %v", rawTree, tc.TreeAfter)
273 }
274
275 treeHash, err = tree.computeRootTreeHash(tc.CipherSuite)
276 if err != nil {
277 t.Errorf("ratchetTree.computeRootTreeHash() = %v", err)
278 } else if !bytes.Equal(treeHash, []byte(tc.TreeHashAfter)) {
279 t.Errorf("ratchetTree.computeRootTreeHash() = %v, want %v", treeHash, tc.TreeHashAfter)
280 }
281 }
282
283 func TestTreeOperations(t *testing.T) {
284 var tests []treeOperationsTest
285 loadTestVector(t, "testdata/tree-operations.json", &tests)
286
287 for i, tc := range tests {
288 t.Run(fmt.Sprintf("[%v]", i), func(t *testing.T) {
289 testTreeOperations(t, &tc)
290 })
291 }
292 }
293