1
package server
2
 
3
import (
4
	"bytes"
5
	"crypto/ed25519"
6
	"crypto/rand"
7
	"crypto/tls"
8
	"crypto/x509"
9
	"crypto/x509/pkix"
10
	"encoding/base64"
11
	"fmt"
12
	"io"
13
	"math/big"
14
	"strings"
15
	"testing"
16
	"time"
17
 
18
	"go.miragespace.co/specter/spec/chord"
19
	mocks "go.miragespace.co/specter/spec/mocks"
20
	"go.miragespace.co/specter/spec/pki"
21
	"go.miragespace.co/specter/spec/protocol"
22
	"go.miragespace.co/specter/spec/rpc"
23
	"go.miragespace.co/specter/spec/transport"
24
	"go.miragespace.co/specter/spec/tun"
25
 
26
	"github.com/stretchr/testify/mock"
27
	"github.com/stretchr/testify/require"
28
	"github.com/twitchtv/twirp"
29
	"go.uber.org/zap"
30
)
31
 
32
type sVNode struct {
33
	chord.VNode
34
	identity *protocol.Node
35
}
36
 
37
func (n *sVNode) ID() uint64 {
38
	return n.identity.GetId()
39
}
40
 
41
func (n *sVNode) Identity() *protocol.Node {
42
	return n.identity
43
}
44
 
45
func getVNode(n *protocol.Node) chord.VNode {
46
	return &sVNode{
47
		identity: n,
48
	}
49
}
50
 
51
func makeNodes(num int) []*protocol.Node {
52
	nodes := make([]*protocol.Node, 0)
53
	for range num {
54
		nodes = append(nodes, &protocol.Node{
55
			Address: strings.Join(generator.MustGenerate(5), "-"),
56
			Id:      chord.Random(),
57
		})
58
	}
59
	return nodes
60
}
61
 
62
func makeNodeList(num int) ([]*protocol.Node, []chord.VNode) {
63
	nodes := make([]*protocol.Node, num)
64
	list := make([]chord.VNode, num)
65
	for i := range nodes {
66
		nodes[i] = &protocol.Node{
67
			Address: strings.Join(generator.MustGenerate(5), "-"),
68
			Id:      chord.Random(),
69
		}
70
	}
71
	for i := range nodes {
72
		list[i] = getVNode(nodes[i])
73
	}
74
	return nodes, list
75
}
76
 
77
func assertNodes(got, exp []*protocol.Node) bool {
78
	for _, g := range got {
79
		for _, e := range exp {
80
			if g.GetId() == e.GetId() {
81
				return true
82
			}
83
		}
84
	}
85
	return false
86
}
87
 
88
func mustGenerateToken() []byte {
89
	b := make([]byte, 32)
90
	n, err := io.ReadFull(rand.Reader, b)
91
	if n != len(b) || err != nil {
92
		panic(fmt.Errorf("error generating token: %w", err))
93
	}
94
	return []byte(base64.StdEncoding.EncodeToString(b))
95
}
96
 
97
type extractFunc func(ed25519.PrivateKey)
98
 
99
func withExtractPrivKey(dst ed25519.PrivateKey) extractFunc {
100
	return func(pk ed25519.PrivateKey) {
101
		copy(dst, pk)
102
	}
103
}
104
 
105
func toCertificate(as *require.Assertions, logger *zap.Logger, client *protocol.Node, token *protocol.ClientToken, extra ...extractFunc) *x509.Certificate {
106
	// generate a CA
107
	caPubKey, caPrivKey, err := ed25519.GenerateKey(rand.Reader)
108
	as.NoError(err)
109
 
110
	template := x509.Certificate{
111
		SerialNumber: big.NewInt(1),
112
		Subject: pkix.Name{
113
			Organization: []string{"dev"},
114
		},
115
		NotBefore: time.Now(),
116
		NotAfter:  time.Now().Add(time.Hour * 24 * 180),
117
 
118
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
119
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
120
		BasicConstraintsValid: true,
121
	}
122
 
123
	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, caPubKey, caPrivKey)
124
	as.NoError(err)
125
 
126
	certPubKey, certPrivKey, err := ed25519.GenerateKey(rand.Reader)
127
	as.NoError(err)
128
 
129
	der, err := pki.GenerateCertificate(logger, tls.Certificate{
130
		Certificate: [][]byte{derBytes},
131
		PrivateKey:  caPrivKey,
132
	}, pki.IdentityRequest{
133
		Subject:   pki.MakeSubjectV1(client.GetId(), string(token.GetToken())),
134
		PublicKey: certPubKey,
135
	})
136
	as.NoError(err)
137
 
138
	cert, err := x509.ParseCertificate(der)
139
	as.NoError(err)
140
 
141
	for _, e := range extra {
142
		e(certPrivKey)
143
	}
144
 
145
	return cert
146
}
147
 
148
func TestRPCRegisterClientNewCertificateOK(t *testing.T) {
149
	as := require.New(t)
150
 
151
	logger, node, clientT, chordT, serv := getFixture(t, as)
152
	cli, _, _ := getIdentities()
153
 
154
	ctx := t.Context()
155
 
156
	testToken := &protocol.ClientToken{
157
		Token: mustGenerateToken(),
158
	}
159
 
160
	node.On("Put", mock.Anything,
161
		mock.MatchedBy(func(key []byte) bool {
162
			return bytes.Equal(key, []byte(tun.ClientTokenKey(testToken)))
163
		}),
164
		mock.Anything).Return(nil).Once()
165
	clientT.On("SendDatagram",
166
		mock.MatchedBy(func(node *protocol.Node) bool {
167
			return node.GetId() == cli.GetId()
168
		}),
169
		mock.MatchedBy(func(b []byte) bool {
170
			return bytes.Equal(b, []byte(testDatagramData))
171
		}),
172
	).Return(nil)
173
 
174
	tp := mocks.SelfTransport()
175
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
176
	go streamRouter.Accept(ctx)
177
 
178
	serv.AttachRouter(ctx, streamRouter)
179
 
180
	cRPC := rpc.DynamicTunnelClient(ctx, tp)
181
	tp.WithCertificate(toCertificate(as, logger, cli, testToken))
182
	resp, err := cRPC.RegisterIdentity(rpc.WithNode(ctx, cli), &protocol.RegisterIdentityRequest{})
183
 
184
	as.NoError(err)
185
	as.NotNil(resp.GetApex())
186
 
187
	node.AssertExpectations(t)
188
	clientT.AssertExpectations(t)
189
	chordT.AssertExpectations(t)
190
}
191
 
192
func TestRPCRegisterClientFailed(t *testing.T) {
193
	as := require.New(t)
194
 
195
	logger, node, clientT, chordT, serv := getFixture(t, as)
196
	cli, _, _ := getIdentities()
197
 
198
	ctx := t.Context()
199
 
200
	testToken := &protocol.ClientToken{
201
		Token: mustGenerateToken(),
202
	}
203
 
204
	node.On("Put", mock.Anything,
205
		mock.MatchedBy(func(key []byte) bool {
206
			return bytes.Equal(key, []byte(tun.ClientTokenKey(testToken)))
207
		}),
208
		mock.Anything).Return(fmt.Errorf("failed")).Once()
209
	clientT.On("SendDatagram", mock.Anything, mock.Anything).Return(fmt.Errorf("failed")).Once()
210
	clientT.On("SendDatagram", mock.Anything, mock.Anything).Return(nil).Once()
211
 
212
	tp := mocks.SelfTransport()
213
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
214
	go streamRouter.Accept(ctx)
215
 
216
	serv.AttachRouter(ctx, streamRouter)
217
 
218
	cRPC := rpc.DynamicTunnelClient(rpc.DisablePooling(ctx), tp)
219
 
220
	requests := []struct {
221
		Client *protocol.Node
222
	}{
223
 
224
		// verified Client but not connected
225
		{
226
			Client: cli,
227
		},
228
 
229
		// verified Client and connected but KV failed
230
		{
231
			Client: cli,
232
		},
233
	}
234
 
235
	for _, req := range requests {
236
		tp.WithCertificate(toCertificate(as, logger, req.Client, testToken))
237
		_, err := cRPC.RegisterIdentity(rpc.WithNode(ctx, cli), &protocol.RegisterIdentityRequest{})
238
		as.Error(err)
239
	}
240
 
241
	node.AssertExpectations(t)
242
	clientT.AssertExpectations(t)
243
	chordT.AssertExpectations(t)
244
}
245
 
246
func TestRPCPingOK(t *testing.T) {
247
	as := require.New(t)
248
 
249
	logger, node, clientT, chordT, serv := getFixture(t, as)
250
	cli, _, tn := getIdentities()
251
 
252
	ctx := t.Context()
253
 
254
	clientT.On("Identity").Return(tn)
255
 
256
	tp := mocks.SelfTransport()
257
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
258
	go streamRouter.Accept(ctx)
259
 
260
	serv.AttachRouter(ctx, streamRouter)
261
 
262
	cRPC := rpc.DynamicTunnelClient(ctx, tp)
263
 
264
	tp.WithCertificate(toCertificate(as, logger, cli, &protocol.ClientToken{}))
265
	resp, err := cRPC.Ping(rpc.WithNode(ctx, cli), &protocol.ClientPingRequest{})
266
 
267
	as.NoError(err)
268
	as.Equal(testRootDomain, resp.GetApex())
269
	as.Equal(tn.GetId(), resp.GetNode().GetId())
270
 
271
	node.AssertExpectations(t)
272
	clientT.AssertExpectations(t)
273
	chordT.AssertExpectations(t)
274
}
275
 
276
func TestRPCGetNodesUnique(t *testing.T) {
277
	as := require.New(t)
278
 
279
	logger, node, clientT, chordT, serv := getFixture(t, as)
280
	cli, cht, tn := getIdentities()
281
 
282
	ctx := t.Context()
283
 
284
	token := &protocol.ClientToken{
285
		Token: mustGenerateToken(),
286
	}
287
 
288
	pair := &protocol.TunnelDestination{
289
		Chord:  cht,
290
		Tunnel: tn,
291
	}
292
	pairBuf, err := pair.MarshalVT()
293
	as.Nil(err)
294
 
295
	node.On("Identity").Return(cht)
296
	node.On("GetSuccessors").Return([]chord.VNode{getVNode(cht)}, nil)
297
	node.On("Get", mock.Anything, []byte(tun.DestinationByChordKey(cht))).Return(pairBuf, nil)
298
 
299
	tp := mocks.SelfTransport()
300
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
301
	go streamRouter.Accept(ctx)
302
 
303
	serv.AttachRouter(ctx, streamRouter)
304
 
305
	cRPC := rpc.DynamicTunnelClient(rpc.DisablePooling(ctx), tp)
306
 
307
	_, err = cRPC.GetNodes(rpc.WithNode(ctx, cli), &protocol.GetNodesRequest{})
308
	as.Equal(twirp.Unauthenticated, err.(twirp.Error).Code())
309
 
310
	tp.WithCertificate(toCertificate(as, logger, cli, token))
311
	resp, err := cRPC.GetNodes(rpc.WithNode(ctx, cli), &protocol.GetNodesRequest{})
312
	as.NoError(err)
313
 
314
	// should only have ourself
315
	as.Len(resp.GetNodes(), 1)
316
	as.True(assertNodes(resp.GetNodes(), []*protocol.Node{tn}))
317
	node.AssertNotCalled(t, "Get", mock.Anything, []byte(tun.ClientTokenKey(token)))
318
	node.AssertNotCalled(t, "Put", mock.Anything, mock.Anything, mock.Anything)
319
 
320
	node.AssertExpectations(t)
321
	clientT.AssertExpectations(t)
322
	chordT.AssertExpectations(t)
323
}
324
 
325
func TestRPCGetNodes(t *testing.T) {
326
	as := require.New(t)
327
 
328
	logger, node, clientT, chordT, serv := getFixture(t, as)
329
	cli, cht, tn := getIdentities()
330
 
331
	ctx := t.Context()
332
 
333
	token := &protocol.ClientToken{
334
		Token: mustGenerateToken(),
335
	}
336
 
337
	node.On("Identity").Return(cht)
338
 
339
	nodes, vlist := makeNodeList(chord.ExtendedSuccessorEntries)
340
 
341
	pair := &protocol.TunnelDestination{
342
		Chord:  cht,
343
		Tunnel: tn,
344
	}
345
	pairBuf, err := pair.MarshalVT()
346
	as.Nil(err)
347
 
348
	node.On("GetSuccessors").Return(vlist, nil)
349
	node.On("Get", mock.Anything, mock.MatchedBy(func(k []byte) bool {
350
		exp := make([][]byte, len(nodes)+1)
351
		exp[0] = []byte(tun.DestinationByChordKey(cht))
352
		for i := 1; i < len(exp); i++ {
353
			exp[i] = []byte(tun.DestinationByChordKey(nodes[i-1]))
354
		}
355
		return assertBytes(k, exp...)
356
	})).Return(pairBuf, nil)
357
 
358
	tp := mocks.SelfTransport()
359
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
360
	go streamRouter.Accept(ctx)
361
 
362
	serv.AttachRouter(ctx, streamRouter)
363
 
364
	cRPC := rpc.DynamicTunnelClient(ctx, tp)
365
 
366
	tp.WithCertificate(toCertificate(as, logger, cli, token))
367
	resp, err := cRPC.GetNodes(rpc.WithNode(ctx, cli), &protocol.GetNodesRequest{})
368
 
369
	as.NoError(err)
370
	as.Len(resp.GetNodes(), tun.NumRedundantLinks)
371
	as.True(assertNodes(resp.GetNodes(), []*protocol.Node{tn, nodes[0], nodes[1]}))
372
	as.False(assertNodes(resp.GetNodes(), []*protocol.Node{nodes[2]}))
373
 
374
	node.AssertExpectations(t)
375
	clientT.AssertExpectations(t)
376
	chordT.AssertExpectations(t)
377
}
378
 
379
func TestRPCRequestHostnameOK(t *testing.T) {
380
	as := require.New(t)
381
 
382
	logger, node, clientT, chordT, serv := getFixture(t, as)
383
	cli, _, _ := getIdentities()
384
 
385
	ctx := t.Context()
386
 
387
	token := &protocol.ClientToken{
388
		Token: mustGenerateToken(),
389
	}
390
 
391
	clientBuf, err := cli.MarshalVT()
392
	as.NoError(err)
393
 
394
	node.On("Get",
395
		mock.Anything,
396
		mock.MatchedBy(func(key []byte) bool {
397
			return bytes.Equal(key, []byte(tun.ClientTokenKey(token)))
398
		}),
399
	).Return(clientBuf, nil)
400
	node.On("PrefixAppend",
401
		mock.Anything,
402
		mock.MatchedBy(func(prefix []byte) bool {
403
			return bytes.Equal(prefix, []byte(tun.ClientHostnamesPrefix(token)))
404
		}),
405
		mock.Anything,
406
	).Return(nil)
407
 
408
	tp := mocks.SelfTransport()
409
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
410
	go streamRouter.Accept(ctx)
411
 
412
	serv.AttachRouter(ctx, streamRouter)
413
 
414
	cRPC := rpc.DynamicTunnelClient(ctx, tp)
415
 
416
	tp.WithCertificate(toCertificate(as, logger, cli, token))
417
	resp, err := cRPC.GenerateHostname(rpc.WithNode(ctx, cli), &protocol.GenerateHostnameRequest{})
418
	as.NoError(err)
419
 
420
	as.NotEmpty(resp.GetHostname())
421
 
422
	node.AssertExpectations(t)
423
	clientT.AssertExpectations(t)
424
	chordT.AssertExpectations(t)
425
}
426
 
427
func TestRPCRegisteredHostnames(t *testing.T) {
428
	as := require.New(t)
429
 
430
	logger, node, clientT, chordT, serv := getFixture(t, as)
431
	cli, _, _ := getIdentities()
432
 
433
	ctx := t.Context()
434
 
435
	token := &protocol.ClientToken{
436
		Token: mustGenerateToken(),
437
	}
438
 
439
	clientBuf, err := cli.MarshalVT()
440
	as.NoError(err)
441
 
442
	hostnames := []string{"hostname-A", "hostname-B"}
443
	hostnameBytes := make([][]byte, len(hostnames))
444
	for i, h := range hostnames {
445
		hostnameBytes[i] = []byte(h)
446
	}
447
 
448
	node.On("Get",
449
		mock.Anything,
450
		mock.MatchedBy(func(key []byte) bool {
451
			return bytes.Equal(key, []byte(tun.ClientTokenKey(token)))
452
		}),
453
	).Return(clientBuf, nil)
454
	node.On("PrefixList",
455
		mock.Anything,
456
		mock.MatchedBy(func(prefix []byte) bool {
457
			return bytes.Equal(prefix, []byte(tun.ClientHostnamesPrefix(token)))
458
		}),
459
	).Return(hostnameBytes, nil)
460
 
461
	tp := mocks.SelfTransport()
462
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
463
	go streamRouter.Accept(ctx)
464
 
465
	serv.AttachRouter(ctx, streamRouter)
466
 
467
	cRPC := rpc.DynamicTunnelClient(ctx, tp)
468
 
469
	tp.WithCertificate(toCertificate(as, logger, cli, token))
470
	resp, err := cRPC.RegisteredHostnames(rpc.WithNode(ctx, cli), &protocol.RegisteredHostnamesRequest{})
471
	as.NoError(err)
472
 
473
	as.Len(resp.GetHostnames(), len(hostnames))
474
	as.EqualValues(hostnames, resp.GetHostnames())
475
 
476
	node.AssertExpectations(t)
477
	clientT.AssertExpectations(t)
478
	chordT.AssertExpectations(t)
479
}
480
 
481
func TestRPCOtherFailed(t *testing.T) {
482
	as := require.New(t)
483
 
484
	logger, node, clientT, chordT, serv := getFixture(t, as)
485
	cli, _, _ := getIdentities()
486
 
487
	ctx := t.Context()
488
 
489
	node.On("Get", mock.Anything, mock.Anything).Return(nil, nil)
490
 
491
	tp := mocks.SelfTransport()
492
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
493
	go streamRouter.Accept(ctx)
494
 
495
	serv.AttachRouter(ctx, streamRouter)
496
 
497
	cRPC := rpc.DynamicTunnelClient(rpc.DisablePooling(ctx), tp)
498
 
499
	requests := []*protocol.ClientToken{
500
		// non-existent token
501
		{
502
			Token: []byte("nah"),
503
		},
504
	}
505
 
506
	for _, req := range requests {
507
		tp.WithCertificate(toCertificate(as, logger, cli, req))
508
		resp, err := cRPC.GenerateHostname(rpc.WithNode(ctx, cli), &protocol.GenerateHostnameRequest{})
509
		as.Error(err)
510
		as.Nil(resp)
511
	}
512
 
513
	node.AssertExpectations(t)
514
	clientT.AssertExpectations(t)
515
	chordT.AssertExpectations(t)
516
 
517
}
518
 
519
func TestRPCPublishTunnelOK(t *testing.T) {
520
	as := require.New(t)
521
 
522
	logger, node, _, _, serv := getFixture(t, as)
523
	cli, cht, tn := getIdentities()
524
 
525
	ctx := t.Context()
526
 
527
	nodes, _ := makeNodeList(tun.NumRedundantLinks)
528
 
529
	pair := &protocol.TunnelDestination{
530
		Chord:  cht,
531
		Tunnel: tn,
532
	}
533
	pairBuf, err := pair.MarshalVT()
534
	as.Nil(err)
535
 
536
	hostname := "test-1234"
537
	token := &protocol.ClientToken{
538
		Token: mustGenerateToken(),
539
	}
540
 
541
	clientBuf, err := cli.MarshalVT()
542
	as.NoError(err)
543
 
544
	node.On("Get",
545
		mock.Anything,
546
		mock.MatchedBy(func(key []byte) bool {
547
			return bytes.Equal(key, []byte(tun.ClientTokenKey(token)))
548
		}),
549
	).Return(clientBuf, nil).Once()
550
 
551
	node.On("Get",
552
		mock.Anything,
553
		mock.MatchedBy(func(k []byte) bool {
554
			exp := make([][]byte, len(nodes)+1)
555
			exp[0] = []byte(tun.DestinationByTunnelKey(cht))
556
			for i := 1; i < len(exp); i++ {
557
				exp[i] = []byte(tun.DestinationByTunnelKey(nodes[i-1]))
558
			}
559
			return assertBytes(k, exp...)
560
		}),
561
	).Return(pairBuf, nil)
562
 
563
	fakeLease := uint64(1234)
564
	node.On("Acquire",
565
		mock.Anything,
566
		mock.MatchedBy(func(k []byte) bool {
567
			return bytes.Equal(k, []byte(tun.ClientLeaseKey(token)))
568
		}),
569
		mock.Anything,
570
	).Return(fakeLease, nil)
571
 
572
	node.On("Release",
573
		mock.Anything,
574
		mock.MatchedBy(func(k []byte) bool {
575
			return bytes.Equal(k, []byte(tun.ClientLeaseKey(token)))
576
		}),
577
		fakeLease,
578
	).Return(nil)
579
 
580
	node.On("Put",
581
		mock.Anything,
582
		mock.MatchedBy(func(k []byte) bool {
583
			return true
584
		}), mock.MatchedBy(func(v []byte) bool {
585
			return true
586
		}),
587
	).Return(nil)
588
 
589
	node.On("PrefixContains",
590
		mock.Anything,
591
		mock.MatchedBy(func(prefix []byte) bool {
592
			return bytes.Equal(prefix, []byte(tun.ClientHostnamesPrefix(token)))
593
		}),
594
		mock.MatchedBy(func(child []byte) bool {
595
			return bytes.Equal(child, []byte(hostname))
596
		}),
597
	).Return(true, nil)
598
 
599
	tp := mocks.SelfTransport()
600
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
601
	go streamRouter.Accept(ctx)
602
 
603
	serv.AttachRouter(ctx, streamRouter)
604
 
605
	cRPC := rpc.DynamicTunnelClient(ctx, tp)
606
 
607
	tp.WithCertificate(toCertificate(as, logger, cli, token))
608
	resp, err := cRPC.PublishTunnel(rpc.WithNode(ctx, cli), &protocol.PublishTunnelRequest{
609
		Hostname: hostname,
610
		Servers:  nodes,
611
	})
612
 
613
	as.NoError(err)
614
	as.NotNil(resp)
615
	as.Len(resp.GetPublished(), len(nodes))
616
 
617
	node.AssertExpectations(t)
618
}
619
 
620
func TestRPCPublishTunnelFailed(t *testing.T) {
621
	as := require.New(t)
622
 
623
	logger, node, _, _, serv := getFixture(t, as)
624
	cli, _, _ := getIdentities()
625
 
626
	ctx := t.Context()
627
 
628
	token := &protocol.ClientToken{
629
		Token: mustGenerateToken(),
630
	}
631
 
632
	clientBuf, err := cli.MarshalVT()
633
	as.NoError(err)
634
 
635
	node.On("Get",
636
		mock.Anything,
637
		mock.MatchedBy(func(key []byte) bool {
638
			return bytes.Equal(key, []byte(tun.ClientTokenKey(token)))
639
		}),
640
	).Return(clientBuf, nil)
641
 
642
	node.On("PrefixContains", mock.Anything, mock.Anything, mock.Anything).Return(false, nil).Once()
643
 
644
	fakeLease := uint64(1234)
645
	node.On("Acquire",
646
		mock.Anything,
647
		mock.MatchedBy(func(k []byte) bool {
648
			return bytes.Equal(k, []byte(tun.ClientLeaseKey(token)))
649
		}),
650
		mock.Anything,
651
	).Return(fakeLease, nil).Once()
652
 
653
	node.On("Release",
654
		mock.Anything,
655
		mock.MatchedBy(func(k []byte) bool {
656
			return bytes.Equal(k, []byte(tun.ClientLeaseKey(token)))
657
		}),
658
		fakeLease,
659
	).Return(nil).Once()
660
 
661
	tp := mocks.SelfTransport()
662
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
663
	go streamRouter.Accept(ctx)
664
 
665
	serv.AttachRouter(ctx, streamRouter)
666
 
667
	cRPC := rpc.DynamicTunnelClient(rpc.DisablePooling(ctx), tp)
668
 
669
	requests := []*protocol.PublishTunnelRequest{
670
		// has token but hostname not requested
671
		{
672
			Hostname: "nil",
673
			Servers: []*protocol.Node{
674
				{
675
					Id: chord.Random(),
676
				},
677
			},
678
		},
679
		// has token but not enough servers
680
		{
681
			Servers: nil,
682
		},
683
 
684
		// has token but too many servers
685
		{
686
			Servers: makeNodes(tun.NumRedundantLinks * 2),
687
		},
688
	}
689
 
690
	for _, req := range requests {
691
		tp.WithCertificate(toCertificate(as, logger, cli, token))
692
		resp, err := cRPC.PublishTunnel(rpc.WithNode(ctx, cli), req)
693
		as.Error(err)
694
		as.Nil(resp)
695
	}
696
 
697
	node.AssertExpectations(t)
698
}
699
 
700
func TestUnpublishTunnel(t *testing.T) {
701
	as := require.New(t)
702
 
703
	logger, node, _, _, serv := getFixture(t, as)
704
	cli, _, _ := getIdentities()
705
 
706
	ctx := t.Context()
707
 
708
	hostname := "test-1234"
709
	token := &protocol.ClientToken{
710
		Token: mustGenerateToken(),
711
	}
712
 
713
	clientBuf, err := cli.MarshalVT()
714
	as.NoError(err)
715
 
716
	fakeLease := uint64(1234)
717
	acquireCall := node.On("Acquire",
718
		mock.Anything,
719
		mock.MatchedBy(func(k []byte) bool {
720
			return bytes.Equal(k, []byte(tun.ClientLeaseKey(token)))
721
		}),
722
		mock.Anything,
723
	).Return(fakeLease, nil)
724
 
725
	node.On("Release",
726
		mock.Anything,
727
		mock.MatchedBy(func(k []byte) bool {
728
			return bytes.Equal(k, []byte(tun.ClientLeaseKey(token)))
729
		}),
730
		fakeLease,
731
	).Return(nil).NotBefore(acquireCall)
732
 
733
	node.On("Get",
734
		mock.Anything,
735
		mock.MatchedBy(func(key []byte) bool {
736
			return bytes.Equal(key, []byte(tun.ClientTokenKey(token)))
737
		}),
738
	).Return(clientBuf, nil).Once()
739
 
740
	node.On("PrefixContains",
741
		mock.Anything,
742
		mock.MatchedBy(func(prefix []byte) bool {
743
			return bytes.Equal(prefix, []byte(tun.ClientHostnamesPrefix(token)))
744
		}),
745
		mock.MatchedBy(func(child []byte) bool {
746
			return bytes.Equal(child, []byte(hostname))
747
		}),
748
	).Return(true, nil)
749
 
750
	for i := range tun.NumRedundantLinks {
751
		key := tun.RoutingKey(hostname, i+1)
752
		node.On("Delete", mock.Anything, []byte(key)).Return(nil)
753
	}
754
 
755
	tp := mocks.SelfTransport()
756
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
757
	go streamRouter.Accept(ctx)
758
 
759
	serv.AttachRouter(ctx, streamRouter)
760
 
761
	cRPC := rpc.DynamicTunnelClient(ctx, tp)
762
 
763
	tp.WithCertificate(toCertificate(as, logger, cli, token))
764
	resp, err := cRPC.UnpublishTunnel(rpc.WithNode(ctx, cli), &protocol.UnpublishTunnelRequest{
765
		Hostname: hostname,
766
	})
767
 
768
	as.NoError(err)
769
	as.NotNil(resp)
770
 
771
	node.AssertExpectations(t)
772
}
773
 
774
func TestReleaseTunnel(t *testing.T) {
775
	as := require.New(t)
776
 
777
	logger, node, _, _, serv := getFixture(t, as)
778
	cli, _, _ := getIdentities()
779
 
780
	ctx := t.Context()
781
 
782
	hostname := "test-1234"
783
	token := &protocol.ClientToken{
784
		Token: mustGenerateToken(),
785
	}
786
 
787
	clientBuf, err := cli.MarshalVT()
788
	as.NoError(err)
789
 
790
	fakeLease := uint64(1234)
791
	acquireCall := node.On("Acquire",
792
		mock.Anything,
793
		mock.MatchedBy(func(k []byte) bool {
794
			return bytes.Equal(k, []byte(tun.ClientLeaseKey(token)))
795
		}),
796
		mock.Anything,
797
	).Return(fakeLease, nil)
798
 
799
	node.On("Release",
800
		mock.Anything,
801
		mock.MatchedBy(func(k []byte) bool {
802
			return bytes.Equal(k, []byte(tun.ClientLeaseKey(token)))
803
		}),
804
		fakeLease,
805
	).Return(nil).NotBefore(acquireCall)
806
 
807
	node.On("Get",
808
		mock.Anything,
809
		mock.MatchedBy(func(key []byte) bool {
810
			return bytes.Equal(key, []byte(tun.ClientTokenKey(token)))
811
		}),
812
	).Return(clientBuf, nil).Once()
813
 
814
	node.On("PrefixContains",
815
		mock.Anything,
816
		mock.MatchedBy(func(prefix []byte) bool {
817
			return bytes.Equal(prefix, []byte(tun.ClientHostnamesPrefix(token)))
818
		}),
819
		mock.MatchedBy(func(child []byte) bool {
820
			return bytes.Equal(child, []byte(hostname))
821
		}),
822
	).Return(true, nil)
823
 
824
	deleteCalls := make([]*mock.Call, 0)
825
	for i := range tun.NumRedundantLinks {
826
		key := tun.RoutingKey(hostname, i+1)
827
		deleteCall := node.On("Delete", mock.Anything, []byte(key)).Return(nil)
828
		deleteCalls = append(deleteCalls, deleteCall)
829
	}
830
 
831
	prefixCall := node.On("PrefixRemove",
832
		mock.Anything,
833
		mock.MatchedBy(func(prefix []byte) bool {
834
			return bytes.Equal(prefix, []byte(tun.ClientHostnamesPrefix(token)))
835
		}),
836
		mock.MatchedBy(func(child []byte) bool {
837
			return bytes.Equal(child, []byte(hostname))
838
		}),
839
	).Return(nil).NotBefore(deleteCalls...)
840
 
841
	// This is needed for custom hostname
842
	node.On("Delete",
843
		mock.Anything,
844
		mock.MatchedBy(func(prefix []byte) bool {
845
			return bytes.Equal(prefix, []byte(tun.CustomHostnameKey(hostname)))
846
		}),
847
	).Return(nil).NotBefore(prefixCall)
848
 
849
	tp := mocks.SelfTransport()
850
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
851
	go streamRouter.Accept(ctx)
852
 
853
	serv.AttachRouter(ctx, streamRouter)
854
 
855
	cRPC := rpc.DynamicTunnelClient(ctx, tp)
856
 
857
	tp.WithCertificate(toCertificate(as, logger, cli, token))
858
	resp, err := cRPC.ReleaseTunnel(rpc.WithNode(ctx, cli), &protocol.ReleaseTunnelRequest{
859
		Hostname: hostname,
860
	})
861
 
862
	as.NoError(err)
863
	as.NotNil(resp)
864
 
865
	node.AssertExpectations(t)
866
}
867
 
868
func TestTokenUpgrade(t *testing.T) {
869
	as := require.New(t)
870
 
871
	logger, node, clientT, chordT, serv := getFixture(t, as)
872
	cli, _, _ := getIdentities()
873
 
874
	ctx := t.Context()
875
 
876
	token := &protocol.ClientToken{
877
		Token: mustGenerateToken(),
878
	}
879
 
880
	old := &protocol.Node{
881
		Id: cli.GetId(),
882
		// note that previous version has no Rendezvous set
883
	}
884
	clientBuf, err := old.MarshalVT()
885
	as.NoError(err)
886
 
887
	cert := toCertificate(as, logger, cli, token)
888
 
889
	node.On("Get",
890
		mock.Anything,
891
		mock.MatchedBy(func(key []byte) bool {
892
			return bytes.Equal(key, []byte(tun.ClientTokenKey(token)))
893
		}),
894
	).Return(clientBuf, nil)
895
	node.On("Put",
896
		mock.Anything,
897
		mock.MatchedBy(func(key []byte) bool {
898
			return bytes.Equal(key, []byte(tun.ClientTokenKey(token)))
899
		}),
900
		mock.MatchedBy(func(val []byte) bool {
901
			c := &protocol.Node{}
902
			err := c.UnmarshalVT(val)
903
			as.NoError(err)
904
 
905
			t.Log(c)
906
 
907
			return strings.Contains(cert.Subject.CommonName, c.GetAddress()) && c.GetRendezvous()
908
		}),
909
	).Return(nil).Once()
910
 
911
	node.On("PrefixList", mock.Anything, []byte(tun.ClientHostnamesPrefix(token))).Return([][]byte{[]byte("test-host")}, nil).Once()
912
 
913
	tp := mocks.SelfTransport()
914
	streamRouter := transport.NewStreamRouter(logger, nil, tp)
915
	go streamRouter.Accept(ctx)
916
 
917
	serv.AttachRouter(ctx, streamRouter)
918
 
919
	cRPC := rpc.DynamicTunnelClient(ctx, tp)
920
 
921
	tp.WithCertificate(cert)
922
	resp, err := cRPC.RegisteredHostnames(rpc.WithNode(ctx, cli), &protocol.RegisteredHostnamesRequest{})
923
	as.NoError(err)
924
	as.Equal([]string{"test-host"}, resp.GetHostnames())
925
 
926
	node.AssertExpectations(t)
927
	clientT.AssertExpectations(t)
928
	chordT.AssertExpectations(t)
929
}