talm

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
package director_test
6

7
import (
8
	"context"
9
	"regexp"
10
	"testing"
11

12
	"github.com/siderolabs/grpc-proxy/proxy"
13
	"github.com/stretchr/testify/suite"
14
	"google.golang.org/grpc/codes"
15
	"google.golang.org/grpc/metadata"
16
	"google.golang.org/grpc/status"
17

18
	"github.com/aenix-io/talm/internal/app/apid/pkg/director"
19
)
20

21
type DirectorSuite struct {
22
	suite.Suite
23

24
	localBackend *mockBackend
25
	router       *director.Router
26
}
27

28
func (suite *DirectorSuite) SetupSuite() {
29
	suite.localBackend = &mockBackend{}
30
	suite.router = director.NewRouter(
31
		mockBackendFactory,
32
		suite.localBackend,
33
		&mockLocalAddressProvider{
34
			local: map[string]struct{}{
35
				"localhost": {},
36
			},
37
		},
38
	)
39
}
40

41
func (suite *DirectorSuite) TestStreamedDetector() {
42
	suite.Assert().False(suite.router.StreamedDetector("/service.Service/someMethod"))
43

44
	suite.router.RegisterStreamedRegex("^" + regexp.QuoteMeta("/service.Service/someMethod") + "$")
45

46
	suite.Assert().True(suite.router.StreamedDetector("/service.Service/someMethod"))
47
	suite.Assert().False(suite.router.StreamedDetector("/service.Service/someMethod2"))
48
	suite.Assert().False(suite.router.StreamedDetector("/servicexService/someMethod"))
49

50
	suite.router.RegisterStreamedRegex("Stream$")
51

52
	suite.Assert().True(suite.router.StreamedDetector("/service.Service/getStream"))
53
	suite.Assert().False(suite.router.StreamedDetector("/service.Service/getStreamItem"))
54
}
55

56
func (suite *DirectorSuite) TestDirectorAggregate() {
57
	ctx := context.Background()
58

59
	md := metadata.New(nil)
60
	md.Set("nodes", "127.0.0.1", "127.0.0.2")
61
	mode, backends, err := suite.router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
62
	suite.Assert().Equal(proxy.One2Many, mode)
63
	suite.Assert().Len(backends, 2)
64
	suite.Assert().Equal("127.0.0.1", backends[0].(*mockBackend).target)
65
	suite.Assert().Equal("127.0.0.2", backends[1].(*mockBackend).target)
66
	suite.Assert().NoError(err)
67

68
	md = metadata.New(nil)
69
	md.Set("nodes", "127.0.0.1")
70
	mode, backends, err = suite.router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
71
	suite.Assert().Equal(proxy.One2Many, mode)
72
	suite.Assert().Len(backends, 1)
73
	suite.Assert().Equal("127.0.0.1", backends[0].(*mockBackend).target)
74
	suite.Assert().NoError(err)
75
}
76

77
func (suite *DirectorSuite) TestDirectorSingleNode() {
78
	ctx := context.Background()
79

80
	md := metadata.New(nil)
81
	md.Set("node", "127.0.0.1")
82
	mode, backends, err := suite.router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
83
	suite.Assert().Equal(proxy.One2One, mode)
84
	suite.Assert().Len(backends, 1)
85
	suite.Assert().Equal("127.0.0.1", backends[0].(*mockBackend).target)
86
	suite.Assert().NoError(err)
87

88
	md = metadata.New(nil)
89
	md.Set("node", "127.0.0.1", "127.0.0.2")
90
	_, _, err = suite.router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
91
	suite.Assert().Equal(codes.InvalidArgument, status.Code(err))
92
}
93

94
func (suite *DirectorSuite) TestDirectorLocal() {
95
	ctx := context.Background()
96

97
	md := metadata.New(nil)
98
	mode, backends, err := suite.router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
99
	suite.Assert().Equal(proxy.One2One, mode)
100
	suite.Assert().Len(backends, 1)
101
	suite.Assert().Equal(suite.localBackend, backends[0])
102
	suite.Assert().NoError(err)
103
}
104

105
func (suite *DirectorSuite) TestDirectorNoRemoteBackend() {
106
	// override the router to have no remote backends
107
	router := director.NewRouter(
108
		nil,
109
		suite.localBackend,
110
		&mockLocalAddressProvider{
111
			local: map[string]struct{}{
112
				"localhost": {},
113
			},
114
		},
115
	)
116

117
	ctx := context.Background()
118

119
	// request forwarding via node/nodes is disabled
120
	md := metadata.New(nil)
121
	md.Set("node", "127.0.0.1")
122
	_, _, err := router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
123
	suite.Assert().Error(err)
124
	suite.Assert().Equal(codes.PermissionDenied, status.Code(err))
125

126
	md = metadata.New(nil)
127
	md.Set("nodes", "127.0.0.1", "127.0.0.2")
128
	_, _, err = router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
129
	suite.Assert().Error(err)
130
	suite.Assert().Equal(codes.PermissionDenied, status.Code(err))
131

132
	// no request forwarding, allowed
133
	md = metadata.New(nil)
134
	mode, backends, err := router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
135
	suite.Assert().Equal(proxy.One2One, mode)
136
	suite.Assert().Len(backends, 1)
137
	suite.Assert().Equal(suite.localBackend, backends[0])
138
	suite.Assert().NoError(err)
139

140
	// request forwarding to local address, allowed
141
	md = metadata.New(nil)
142
	md.Set("node", "localhost")
143
	mode, backends, err = router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
144
	suite.Assert().Equal(proxy.One2One, mode)
145
	suite.Assert().Len(backends, 1)
146
	suite.Assert().Equal(suite.localBackend, backends[0])
147
	suite.Assert().NoError(err)
148

149
	md = metadata.New(nil)
150
	md.Set("nodes", "localhost")
151
	mode, backends, err = router.Director(metadata.NewIncomingContext(ctx, md), "/service.Service/method")
152
	suite.Assert().Equal(proxy.One2One, mode)
153
	suite.Assert().Len(backends, 1)
154
	suite.Assert().Equal(suite.localBackend, backends[0])
155
	suite.Assert().NoError(err)
156
}
157

158
func TestDirectorSuite(t *testing.T) {
159
	suite.Run(t, new(DirectorSuite))
160
}
161