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