1package systests 2 3import ( 4 "testing" 5 "time" 6 7 "github.com/keybase/client/go/libkb" 8 "github.com/keybase/client/go/protocol/gregor1" 9 keybase1 "github.com/keybase/client/go/protocol/keybase1" 10 "github.com/keybase/client/go/teambot" 11 "github.com/keybase/client/go/teams" 12 "github.com/keybase/clockwork" 13 "github.com/stretchr/testify/require" 14) 15 16func checkNewTeambotKeyNotifications(tc *libkb.TestContext, notifications *teamNotifyHandler, 17 expectedArgs []keybase1.NewTeambotKeyArg) { 18 matches := map[keybase1.NewTeambotKeyArg]struct{}{} 19 numFound := 0 20 for { 21 select { 22 case arg := <-notifications.newTeambotKeyCh: 23 for _, expectedArg := range expectedArgs { 24 if expectedArg == arg { 25 matches[arg] = struct{}{} 26 break 27 } 28 } 29 // make don't have any unexpected notifications 30 if len(matches) <= numFound { 31 require.Fail(tc.T, "unexpected newTeamKeyNeeded notification", arg) 32 } 33 if len(matches) == len(expectedArgs) { 34 return 35 } 36 numFound++ 37 case <-time.After(5 * time.Second * libkb.CITimeMultiplier(tc.G)): 38 require.Fail(tc.T, "no notification on newTeambotKey") 39 } 40 } 41} 42 43func checkTeambotKeyNeededNotifications(tc *libkb.TestContext, notifications *teamNotifyHandler, 44 expectedArg keybase1.TeambotKeyNeededArg) { 45 select { 46 case arg := <-notifications.teambotKeyNeededCh: 47 require.Equal(tc.T, expectedArg, arg) 48 return 49 case <-time.After(5 * time.Second * libkb.CITimeMultiplier(tc.G)): 50 require.Fail(tc.T, "no notification on teambotKeyNeeded") 51 } 52} 53 54func noNewTeambotKeyNotification(tc *libkb.TestContext, notifications *teamNotifyHandler) { 55 select { 56 case arg := <-notifications.newTeambotKeyCh: 57 require.Fail(tc.T, "unexpected newTeambotKey notification", arg) 58 default: 59 } 60} 61 62func noTeambotKeyNeeded(tc *libkb.TestContext, notifications *teamNotifyHandler) { 63 select { 64 case arg := <-notifications.teambotKeyNeededCh: 65 require.Fail(tc.T, "unexpected teambotKeyNeeded notification", arg) 66 default: 67 } 68} 69 70func TestTeambotKey(t *testing.T) { 71 tt := newTeamTester(t) 72 defer tt.cleanup() 73 74 fc := clockwork.NewFakeClockAt(time.Now()) 75 76 user1 := tt.addUser("one") 77 user2 := tt.addUserWithPaper("two") 78 botua := tt.addUser("botua") 79 botuaUID := gregor1.UID(botua.uid.ToBytes()) 80 mctx1 := libkb.NewMetaContextForTest(*user1.tc) 81 mctx2 := libkb.NewMetaContextForTest(*user2.tc) 82 mctx3 := libkb.NewMetaContextForTest(*botua.tc) 83 memberKeyer1 := mctx1.G().GetTeambotMemberKeyer() 84 memberKeyer2 := mctx2.G().GetTeambotMemberKeyer() 85 botKeyer := mctx3.G().GetTeambotBotKeyer().(*teambot.BotKeyer) 86 botKeyer.SetClock(fc) 87 88 teamID, teamName := user1.createTeam2() 89 user1.addTeamMember(teamName.String(), user2.username, keybase1.TeamRole_WRITER) 90 user1.addRestrictedBotTeamMember(teamName.String(), botua.username, keybase1.TeamBotSettings{}) 91 92 // bot gets a key on addition to the team 93 newKeyArgs := []keybase1.NewTeambotKeyArg{ 94 { 95 Id: teamID, 96 Generation: 1, 97 Application: keybase1.TeamApplication_CHAT, 98 }, 99 { 100 Id: teamID, 101 Generation: 1, 102 Application: keybase1.TeamApplication_KVSTORE, 103 }, 104 } 105 checkNewTeambotKeyNotifications(botua.tc, botua.notifications, newKeyArgs) 106 107 // grab the latest chat application key and make sure the generation lines 108 // up with the teambotKey 109 team, err := teams.Load(mctx1.Ctx(), mctx1.G(), keybase1.LoadTeamArg{ 110 ID: teamID, 111 ForceRepoll: true, 112 }) 113 require.NoError(t, err) 114 appKey1, err := team.ChatKey(mctx1.Ctx()) 115 require.NoError(t, err) 116 117 // now created = false since we published on member addition 118 teambotKey, created, err := memberKeyer1.GetOrCreateTeambotKey(mctx1, teamID, botuaUID, appKey1) 119 require.NoError(t, err) 120 require.False(t, created) 121 require.Equal(t, appKey1.Generation(), teambotKey.Generation()) 122 123 teambotKey2, err := botKeyer.GetLatestTeambotKey(mctx3, teamID, keybase1.TeamApplication_CHAT) 124 require.NoError(t, err) 125 require.Equal(t, teambotKey, teambotKey2) 126 127 // delete the initial key to check regeneration flows 128 err = teambot.DeleteTeambotKeyForTest(mctx3, teamID, keybase1.TeamApplication_CHAT, 129 teambotKey.Metadata.Generation) 130 131 require.NoError(t, err) 132 133 // initial get, bot has no key to access 134 _, err = botKeyer.GetLatestTeambotKey(mctx3, teamID, keybase1.TeamApplication_CHAT) 135 require.IsType(t, teambot.TeambotTransientKeyError{}, err) 136 137 // cry for help has been issued. 138 keyNeededArg := keybase1.TeambotKeyNeededArg{ 139 Id: teamID, 140 Uid: botua.uid, 141 Generation: 1, 142 Application: keybase1.TeamApplication_CHAT, 143 } 144 checkTeambotKeyNeededNotifications(user1.tc, user1.notifications, keyNeededArg) 145 checkTeambotKeyNeededNotifications(user2.tc, user2.notifications, keyNeededArg) 146 147 // and answered. 148 newKeyArgs = []keybase1.NewTeambotKeyArg{{ 149 Id: teamID, 150 Generation: 1, 151 Application: keybase1.TeamApplication_CHAT, 152 }} 153 checkNewTeambotKeyNotifications(botua.tc, botua.notifications, newKeyArgs) 154 155 // bot can access the key 156 teambotKey2, err = botKeyer.GetLatestTeambotKey(mctx3, teamID, keybase1.TeamApplication_CHAT) 157 require.NoError(t, err) 158 require.Equal(t, teambotKey, teambotKey2) 159 noTeambotKeyNeeded(user1.tc, user1.notifications) 160 noTeambotKeyNeeded(user2.tc, user2.notifications) 161 noNewTeambotKeyNotification(botua.tc, botua.notifications) 162 163 // check for wrong application 164 _, err = botKeyer.GetLatestTeambotKey(mctx3, teamID, keybase1.TeamApplication_KBFS) 165 require.IsType(t, teambot.TeambotTransientKeyError{}, err) 166 167 // cry for help has been issued. 168 keyNeededArg = keybase1.TeambotKeyNeededArg{ 169 Id: teamID, 170 Uid: botua.uid, 171 Generation: 1, 172 Application: keybase1.TeamApplication_KBFS, 173 } 174 checkTeambotKeyNeededNotifications(user1.tc, user1.notifications, keyNeededArg) 175 checkTeambotKeyNeededNotifications(user2.tc, user2.notifications, keyNeededArg) 176 177 // and answered. 178 newKeyArgs = []keybase1.NewTeambotKeyArg{{ 179 Id: teamID, 180 Generation: 1, 181 Application: keybase1.TeamApplication_KBFS, 182 }} 183 checkNewTeambotKeyNotifications(botua.tc, botua.notifications, newKeyArgs) 184 185 _, err = botKeyer.GetLatestTeambotKey(mctx3, teamID, keybase1.TeamApplication_KBFS) 186 require.NoError(t, err) 187 noTeambotKeyNeeded(user1.tc, user1.notifications) 188 noTeambotKeyNeeded(user2.tc, user2.notifications) 189 noNewTeambotKeyNotification(botua.tc, botua.notifications) 190 191 // Test the AtGeneration flow 192 teambotKey2b, err := botKeyer.GetTeambotKeyAtGeneration(mctx3, teamID, 193 keybase1.TeamApplication_CHAT, teambotKey2.Metadata.Generation) 194 require.NoError(t, err) 195 require.Equal(t, teambotKey2, teambotKey2b) 196 noTeambotKeyNeeded(user1.tc, user1.notifications) 197 noTeambotKeyNeeded(user2.tc, user2.notifications) 198 noNewTeambotKeyNotification(botua.tc, botua.notifications) 199 200 // force a PTK rotation 201 user2.revokePaperKey() 202 user1.waitForRotateByID(teamID, keybase1.Seqno(4)) 203 204 // bot gets a new key on rotation 205 newKeyArgs = []keybase1.NewTeambotKeyArg{ 206 { 207 Id: teamID, 208 Generation: 2, 209 Application: keybase1.TeamApplication_CHAT, 210 }, 211 { 212 Id: teamID, 213 Generation: 2, 214 Application: keybase1.TeamApplication_KVSTORE, 215 }, 216 } 217 checkNewTeambotKeyNotifications(botua.tc, botua.notifications, newKeyArgs) 218 219 // delete to check regeneration flow 220 err = teambot.DeleteTeambotKeyForTest(mctx3, teamID, keybase1.TeamApplication_CHAT, 2) 221 require.NoError(t, err) 222 223 // Force a wrongKID error on the bot user by expiring the wrongKID cache 224 key := teambot.TeambotKeyWrongKIDCacheKey(teamID, botua.uid, teambotKey2.Metadata.Generation, 225 keybase1.TeamApplication_CHAT) 226 expired := keybase1.ToTime(fc.Now()) 227 err = mctx3.G().GetKVStore().PutObj(key, nil, expired) 228 require.NoError(t, err) 229 permitted, ctime, err := teambot.TeambotKeyWrongKIDPermitted(mctx3, teamID, botua.uid, 230 keybase1.TeamApplication_CHAT, teambotKey2.Metadata.Generation, keybase1.ToTime(fc.Now())) 231 require.NoError(t, err) 232 require.True(t, permitted) 233 require.Equal(t, expired, ctime) 234 235 fc.Advance(teambot.MaxTeambotKeyWrongKIDPermitted) // expire wrong KID cache 236 permitted, ctime, err = teambot.TeambotKeyWrongKIDPermitted(mctx3, teamID, botua.uid, 237 keybase1.TeamApplication_CHAT, teambotKey2.Metadata.Generation, keybase1.ToTime(fc.Now())) 238 require.NoError(t, err) 239 require.False(t, permitted) 240 require.Equal(t, expired, ctime) 241 242 _, err = botKeyer.GetLatestTeambotKey(mctx3, teamID, keybase1.TeamApplication_CHAT) 243 require.IsType(t, teambot.TeambotPermanentKeyError{}, err) 244 require.False(t, created) 245 keyNeededArg = keybase1.TeambotKeyNeededArg{ 246 Id: teamID, 247 Uid: botua.uid, 248 Generation: teambotKey2.Metadata.Generation + 1, 249 Application: keybase1.TeamApplication_CHAT, 250 } 251 checkTeambotKeyNeededNotifications(user1.tc, user1.notifications, keyNeededArg) 252 checkTeambotKeyNeededNotifications(user2.tc, user2.notifications, keyNeededArg) 253 newKeyArgs = []keybase1.NewTeambotKeyArg{{ 254 Id: teamID, 255 Generation: teambotKey2.Metadata.Generation + 1, 256 Application: keybase1.TeamApplication_CHAT, 257 }} 258 checkNewTeambotKeyNotifications(botua.tc, botua.notifications, newKeyArgs) 259 260 teambotKey3, err := botKeyer.GetLatestTeambotKey(mctx3, teamID, keybase1.TeamApplication_CHAT) 261 require.NoError(t, err) 262 require.Equal(t, teambotKey3.Metadata.Generation, teambotKey2.Metadata.Generation+1) 263 require.Equal(t, keybase1.TeamApplication_CHAT, teambotKey3.Metadata.Application) 264 265 // another PTK rotation happens, this time the bot can proceed with a key 266 // signed by the old PTK since the wrongKID cache did not expire 267 user1.removeTeamMember(teamName.String(), user2.username) 268 user1.addTeamMember(teamName.String(), user2.username, keybase1.TeamRole_WRITER) 269 user2.waitForNewlyAddedToTeamByID(teamID) 270 botua.waitForNewlyAddedToTeamByID(teamID) 271 272 newKeyArgs = []keybase1.NewTeambotKeyArg{ 273 { 274 Id: teamID, 275 Generation: 3, 276 Application: keybase1.TeamApplication_CHAT, 277 }, 278 { 279 Id: teamID, 280 Generation: 3, 281 Application: keybase1.TeamApplication_KVSTORE, 282 }, 283 } 284 checkNewTeambotKeyNotifications(botua.tc, botua.notifications, newKeyArgs) 285 286 err = teambot.DeleteTeambotKeyForTest(mctx3, teamID, keybase1.TeamApplication_CHAT, 3) 287 require.NoError(t, err) 288 289 // bot can access the old teambotKey, but asks for a new one to 290 // be created since it was signed by the old PTK 291 teambotKey4, err := botKeyer.GetLatestTeambotKey(mctx3, teamID, keybase1.TeamApplication_CHAT) 292 require.NoError(t, err) 293 require.Equal(t, teambotKey3, teambotKey4) 294 keyNeededArg = keybase1.TeambotKeyNeededArg{ 295 Id: teamID, 296 Uid: botua.uid, 297 Generation: teambotKey4.Metadata.Generation + 1, 298 Application: keybase1.TeamApplication_CHAT, 299 } 300 checkTeambotKeyNeededNotifications(user1.tc, user1.notifications, keyNeededArg) 301 checkTeambotKeyNeededNotifications(user2.tc, user2.notifications, keyNeededArg) 302 303 newKeyArgs = []keybase1.NewTeambotKeyArg{{ 304 Id: teamID, 305 Generation: teambotKey4.Metadata.Generation + 1, 306 Application: keybase1.TeamApplication_CHAT, 307 }} 308 checkNewTeambotKeyNotifications(botua.tc, botua.notifications, newKeyArgs) 309 310 team, err = teams.Load(mctx1.Ctx(), mctx1.G(), keybase1.LoadTeamArg{ 311 ID: teamID, 312 ForceRepoll: true, 313 }) 314 require.NoError(t, err) 315 appKey2, err := team.ApplicationKey(mctx1.Ctx(), keybase1.TeamApplication_CHAT) 316 require.NoError(t, err) 317 teambotKey, _, err = memberKeyer1.GetOrCreateTeambotKey(mctx1, teamID, botuaUID, appKey2) 318 require.NoError(t, err) 319 require.Equal(t, appKey1.Generation()+2, teambotKey.Generation()) 320 321 teambotKey2, err = botKeyer.GetLatestTeambotKey(mctx3, teamID, keybase1.TeamApplication_CHAT) 322 require.NoError(t, err) 323 require.Equal(t, teambotKey, teambotKey2) 324 noTeambotKeyNeeded(user1.tc, user1.notifications) 325 noTeambotKeyNeeded(user2.tc, user2.notifications) 326 noNewTeambotKeyNotification(botua.tc, botua.notifications) 327 328 // kill the cache and make sure we don't republish 329 memberKeyer1.PurgeCache(mctx1) 330 teambotKeyNoCache, created, err := memberKeyer1.GetOrCreateTeambotKey(mctx1, teamID, botuaUID, appKey2) 331 require.NoError(t, err) 332 // created is True since we attempt to publish but the generation remains 333 require.True(t, created) 334 require.Equal(t, teambotKey.Metadata.Generation, teambotKeyNoCache.Metadata.Generation) 335 require.Equal(t, keybase1.TeamApplication_CHAT, teambotKey.Metadata.Application) 336 337 // Make sure we can access the teambotKey at various generations 338 for i := keybase1.TeambotKeyGeneration(1); i < teambotKey.Metadata.Generation; i++ { 339 teambotKeyBot, err := botKeyer.GetTeambotKeyAtGeneration(mctx3, teamID, keybase1.TeamApplication_CHAT, i) 340 require.NoError(t, err) 341 noTeambotKeyNeeded(user1.tc, user1.notifications) 342 noTeambotKeyNeeded(user2.tc, user2.notifications) 343 noNewTeambotKeyNotification(botua.tc, botua.notifications) 344 345 appKey, err := team.ApplicationKeyAtGeneration(mctx1.Ctx(), keybase1.TeamApplication_CHAT, keybase1.PerTeamKeyGeneration(i)) 346 require.NoError(t, err) 347 348 teambotKeyNonBot1, _, err := memberKeyer1.GetOrCreateTeambotKey(mctx1, teamID, botuaUID, appKey) 349 require.NoError(t, err) 350 require.Equal(t, teambotKeyBot, teambotKeyNonBot1) 351 352 teambotKeyNonBot2, _, err := memberKeyer2.GetOrCreateTeambotKey(mctx2, teamID, botuaUID, appKey) 353 require.NoError(t, err) 354 require.Equal(t, teambotKeyBot, teambotKeyNonBot2) 355 } 356 357 // bot asks for a non-existent generation, no new key is created. 358 badGen := teambotKey.Metadata.Generation + 50 359 _, err = botKeyer.GetTeambotKeyAtGeneration(mctx3, teamID, keybase1.TeamApplication_CHAT, badGen) 360 require.IsType(t, teambot.TeambotTransientKeyError{}, err) 361 keyNeededArg = keybase1.TeambotKeyNeededArg{ 362 Id: teamID, 363 Uid: botua.uid, 364 Generation: badGen, 365 Application: keybase1.TeamApplication_CHAT, 366 } 367 checkTeambotKeyNeededNotifications(user1.tc, user1.notifications, keyNeededArg) 368 checkTeambotKeyNeededNotifications(user2.tc, user2.notifications, keyNeededArg) 369 noNewTeambotKeyNotification(botua.tc, botua.notifications) 370} 371 372func TestTeambotKeyRemovedMember(t *testing.T) { 373 tt := newTeamTester(t) 374 defer tt.cleanup() 375 376 user1 := tt.addUser("one") 377 botua := tt.addUser("botua") 378 botuaUID := gregor1.UID(botua.uid.ToBytes()) 379 mctx1 := libkb.NewMetaContextForTest(*user1.tc) 380 ekLib1 := mctx1.G().GetEKLib() 381 memberKeyer1 := mctx1.G().GetTeambotMemberKeyer() 382 383 teamID, teamName := user1.createTeam2() 384 user1.addRestrictedBotTeamMember(teamName.String(), botua.username, keybase1.TeamBotSettings{}) 385 newKeyArgs := []keybase1.NewTeambotKeyArg{ 386 { 387 Id: teamID, 388 Generation: 1, 389 Application: keybase1.TeamApplication_CHAT, 390 }, 391 { 392 Id: teamID, 393 Generation: 1, 394 Application: keybase1.TeamApplication_KVSTORE, 395 }, 396 } 397 checkNewTeambotKeyNotifications(botua.tc, botua.notifications, newKeyArgs) 398 user1.removeTeamMember(teamName.String(), botua.username) 399 400 team, err := teams.Load(mctx1.Ctx(), mctx1.G(), keybase1.LoadTeamArg{ 401 ID: teamID, 402 ForceRepoll: true, 403 }) 404 require.NoError(t, err) 405 appKey, err := team.ChatKey(mctx1.Ctx()) 406 require.NoError(t, err) 407 408 _, created, err := memberKeyer1.GetOrCreateTeambotKey(mctx1, teamID, botuaUID, appKey) 409 require.False(t, created) 410 require.NoError(t, err) 411 noNewTeambotKeyNotification(botua.tc, botua.notifications) 412 413 err = ekLib1.KeygenIfNeeded(mctx1) 414 require.NoError(t, err) 415 _, created, err = ekLib1.GetOrCreateLatestTeambotEK(mctx1, teamID, botuaUID) 416 require.False(t, created) 417 require.NoError(t, err) 418 noNewTeambotEKNotification(botua.tc, botua.notifications) 419} 420