1package token
2
3import (
4	"testing"
5
6	"github.com/stretchr/testify/require"
7)
8
9func TestStore_RegularTokens(t *testing.T) {
10	type tokens struct {
11		userSource   TokenSource
12		user         string
13		agent        string
14		agentSource  TokenSource
15		master       string
16		masterSource TokenSource
17		repl         string
18		replSource   TokenSource
19	}
20
21	tests := []struct {
22		name      string
23		set       tokens
24		raw       tokens
25		effective tokens
26	}{
27		{
28			name:      "set user - config",
29			set:       tokens{user: "U", userSource: TokenSourceConfig},
30			raw:       tokens{user: "U", userSource: TokenSourceConfig},
31			effective: tokens{user: "U", agent: "U"},
32		},
33		{
34			name:      "set user - api",
35			set:       tokens{user: "U", userSource: TokenSourceAPI},
36			raw:       tokens{user: "U", userSource: TokenSourceAPI},
37			effective: tokens{user: "U", agent: "U"},
38		},
39		{
40			name:      "set agent - config",
41			set:       tokens{agent: "A", agentSource: TokenSourceConfig},
42			raw:       tokens{agent: "A", agentSource: TokenSourceConfig},
43			effective: tokens{agent: "A"},
44		},
45		{
46			name:      "set agent - api",
47			set:       tokens{agent: "A", agentSource: TokenSourceAPI},
48			raw:       tokens{agent: "A", agentSource: TokenSourceAPI},
49			effective: tokens{agent: "A"},
50		},
51		{
52			name:      "set user and agent",
53			set:       tokens{agent: "A", user: "U"},
54			raw:       tokens{agent: "A", user: "U"},
55			effective: tokens{agent: "A", user: "U"},
56		},
57		{
58			name:      "set repl - config",
59			set:       tokens{repl: "R", replSource: TokenSourceConfig},
60			raw:       tokens{repl: "R", replSource: TokenSourceConfig},
61			effective: tokens{repl: "R"},
62		},
63		{
64			name:      "set repl - api",
65			set:       tokens{repl: "R", replSource: TokenSourceAPI},
66			raw:       tokens{repl: "R", replSource: TokenSourceAPI},
67			effective: tokens{repl: "R"},
68		},
69		{
70			name:      "set master - config",
71			set:       tokens{master: "M", masterSource: TokenSourceConfig},
72			raw:       tokens{master: "M", masterSource: TokenSourceConfig},
73			effective: tokens{master: "M"},
74		},
75		{
76			name:      "set master - api",
77			set:       tokens{master: "M", masterSource: TokenSourceAPI},
78			raw:       tokens{master: "M", masterSource: TokenSourceAPI},
79			effective: tokens{master: "M"},
80		},
81		{
82			name:      "set all",
83			set:       tokens{user: "U", agent: "A", repl: "R", master: "M"},
84			raw:       tokens{user: "U", agent: "A", repl: "R", master: "M"},
85			effective: tokens{user: "U", agent: "A", repl: "R", master: "M"},
86		},
87	}
88	for _, tt := range tests {
89		t.Run(tt.name, func(t *testing.T) {
90			s := new(Store)
91			if tt.set.user != "" {
92				require.True(t, s.UpdateUserToken(tt.set.user, tt.set.userSource))
93			}
94
95			if tt.set.agent != "" {
96				require.True(t, s.UpdateAgentToken(tt.set.agent, tt.set.agentSource))
97			}
98
99			if tt.set.repl != "" {
100				require.True(t, s.UpdateReplicationToken(tt.set.repl, tt.set.replSource))
101			}
102
103			if tt.set.master != "" {
104				require.True(t, s.UpdateAgentMasterToken(tt.set.master, tt.set.masterSource))
105			}
106
107			// If they don't change then they return false.
108			require.False(t, s.UpdateUserToken(tt.set.user, tt.set.userSource))
109			require.False(t, s.UpdateAgentToken(tt.set.agent, tt.set.agentSource))
110			require.False(t, s.UpdateReplicationToken(tt.set.repl, tt.set.replSource))
111			require.False(t, s.UpdateAgentMasterToken(tt.set.master, tt.set.masterSource))
112
113			require.Equal(t, tt.effective.user, s.UserToken())
114			require.Equal(t, tt.effective.agent, s.AgentToken())
115			require.Equal(t, tt.effective.master, s.AgentMasterToken())
116			require.Equal(t, tt.effective.repl, s.ReplicationToken())
117
118			tok, src := s.UserTokenAndSource()
119			require.Equal(t, tt.raw.user, tok)
120			require.Equal(t, tt.raw.userSource, src)
121
122			tok, src = s.AgentTokenAndSource()
123			require.Equal(t, tt.raw.agent, tok)
124			require.Equal(t, tt.raw.agentSource, src)
125
126			tok, src = s.AgentMasterTokenAndSource()
127			require.Equal(t, tt.raw.master, tok)
128			require.Equal(t, tt.raw.masterSource, src)
129
130			tok, src = s.ReplicationTokenAndSource()
131			require.Equal(t, tt.raw.repl, tok)
132			require.Equal(t, tt.raw.replSource, src)
133		})
134	}
135}
136
137func TestStore_AgentMasterToken(t *testing.T) {
138	s := new(Store)
139
140	verify := func(want bool, toks ...string) {
141		for _, tok := range toks {
142			require.Equal(t, want, s.IsAgentMasterToken(tok))
143		}
144	}
145
146	verify(false, "", "nope")
147
148	s.UpdateAgentMasterToken("master", TokenSourceConfig)
149	verify(true, "master")
150	verify(false, "", "nope")
151
152	s.UpdateAgentMasterToken("another", TokenSourceConfig)
153	verify(true, "another")
154	verify(false, "", "nope", "master")
155
156	s.UpdateAgentMasterToken("", TokenSourceConfig)
157	verify(false, "", "nope", "master", "another")
158}
159
160func TestStore_Notify(t *testing.T) {
161	s := new(Store)
162
163	newNotification := func(t *testing.T, s *Store, kind TokenKind) Notifier {
164		n := s.Notify(kind)
165		require.NotNil(t, n.Ch)
166		return n
167	}
168
169	requireNotNotified := func(t *testing.T, ch <-chan struct{}) {
170		require.Empty(t, ch)
171	}
172
173	requireNotifiedOnce := func(t *testing.T, ch <-chan struct{}) {
174		require.Len(t, ch, 1)
175		// drain the channel
176		<-ch
177		// just to be safe
178		require.Empty(t, ch)
179	}
180
181	agentNotifier := newNotification(t, s, TokenKindAgent)
182	userNotifier := newNotification(t, s, TokenKindUser)
183	agentMasterNotifier := newNotification(t, s, TokenKindAgentMaster)
184	replicationNotifier := newNotification(t, s, TokenKindReplication)
185	replicationNotifier2 := newNotification(t, s, TokenKindReplication)
186
187	// perform an update of the user token
188	require.True(t, s.UpdateUserToken("edcae2a2-3b51-4864-b412-c7a568f49cb1", TokenSourceConfig))
189	// do it again to ensure it doesn't block even though nothing has read from the 1 buffered chan yet
190	require.True(t, s.UpdateUserToken("47788919-f944-476a-bda5-446d64be1df8", TokenSourceAPI))
191
192	// ensure notifications were sent to the user and all notifiers
193	requireNotNotified(t, agentNotifier.Ch)
194	requireNotifiedOnce(t, userNotifier.Ch)
195	requireNotNotified(t, replicationNotifier.Ch)
196	requireNotNotified(t, agentMasterNotifier.Ch)
197	requireNotNotified(t, replicationNotifier2.Ch)
198
199	// now update the agent token which should send notificaitons to the agent and all notifier
200	require.True(t, s.UpdateAgentToken("5d748ec2-d536-461f-8e2a-1f7eae98d559", TokenSourceAPI))
201
202	requireNotifiedOnce(t, agentNotifier.Ch)
203	requireNotNotified(t, userNotifier.Ch)
204	requireNotNotified(t, replicationNotifier.Ch)
205	requireNotNotified(t, agentMasterNotifier.Ch)
206	requireNotNotified(t, replicationNotifier2.Ch)
207
208	// now update the agent master token which should send notificaitons to the agent master and all notifier
209	require.True(t, s.UpdateAgentMasterToken("789badc8-f850-43e1-8742-9b9f484957cc", TokenSourceAPI))
210
211	requireNotNotified(t, agentNotifier.Ch)
212	requireNotNotified(t, userNotifier.Ch)
213	requireNotNotified(t, replicationNotifier.Ch)
214	requireNotifiedOnce(t, agentMasterNotifier.Ch)
215	requireNotNotified(t, replicationNotifier2.Ch)
216
217	// now update the replication token which should send notificaitons to the replication and all notifier
218	require.True(t, s.UpdateReplicationToken("789badc8-f850-43e1-8742-9b9f484957cc", TokenSourceAPI))
219
220	requireNotNotified(t, agentNotifier.Ch)
221	requireNotNotified(t, userNotifier.Ch)
222	requireNotifiedOnce(t, replicationNotifier.Ch)
223	requireNotNotified(t, agentMasterNotifier.Ch)
224	requireNotifiedOnce(t, replicationNotifier2.Ch)
225
226	s.StopNotify(replicationNotifier2)
227
228	// now update the replication token which should send notificaitons to the replication and all notifier
229	require.True(t, s.UpdateReplicationToken("eb0b56b9-fa65-4ae1-902a-c64457c62ac6", TokenSourceAPI))
230
231	requireNotNotified(t, agentNotifier.Ch)
232	requireNotNotified(t, userNotifier.Ch)
233	requireNotifiedOnce(t, replicationNotifier.Ch)
234	requireNotNotified(t, agentMasterNotifier.Ch)
235	requireNotNotified(t, replicationNotifier2.Ch)
236
237	// request updates but that are not changes
238	require.False(t, s.UpdateAgentToken("5d748ec2-d536-461f-8e2a-1f7eae98d559", TokenSourceAPI))
239	require.False(t, s.UpdateAgentMasterToken("789badc8-f850-43e1-8742-9b9f484957cc", TokenSourceAPI))
240	require.False(t, s.UpdateUserToken("47788919-f944-476a-bda5-446d64be1df8", TokenSourceAPI))
241	require.False(t, s.UpdateReplicationToken("eb0b56b9-fa65-4ae1-902a-c64457c62ac6", TokenSourceAPI))
242
243	// ensure that notifications were not sent
244	requireNotNotified(t, agentNotifier.Ch)
245	requireNotNotified(t, userNotifier.Ch)
246	requireNotNotified(t, replicationNotifier.Ch)
247	requireNotNotified(t, agentMasterNotifier.Ch)
248}
249