3535	testEntryValue  =  []byte {1 , 2 , 3 }
3636)
3737
38+ // rootKeyMockStore is a mock implementation of a macaroon service store that 
39+ // can be used to generate mock root keys for testing. 
40+ type  rootKeyMockStore  struct  {
41+ 	// rootKeys is a slice of all root keys that have been added to the 
42+ 	// store. 
43+ 	rootKeys  [][]byte 
44+ }
45+ 
46+ // addRootKeyFromIDSuffix adds a new root key to the store, using the passed 
47+ // 4 byte suffix. The function generates a root key that ends with the 4 byte 
48+ // suffix, prefixed by 4 random bytes. 
49+ func  (r  * rootKeyMockStore ) addRootKeyFromIDSuffix (suffix  [4 ]byte ) [4 ]byte  {
50+ 	// As a real root key is 8 bytes, we need to generate a random 4 byte 
51+ 	// prefix to append to the passed 4 byte suffix. 
52+ 	rootKey  :=  append (randomBytes (4 ), suffix [:]... )
53+ 	r .rootKeys  =  append (r .rootKeys , rootKey )
54+ 
55+ 	return  suffix 
56+ }
57+ 
58+ // addRootKeyFromAcctID adds a new root key to the store, using the first 4 
59+ // bytes of the passed account ID as the suffix for the root key, prefixed by 4 
60+ // random bytes. 
61+ func  (r  * rootKeyMockStore ) addRootKeyFromAcctID (id  accounts.AccountID ) [4 ]byte  {
62+ 	var  acctPrefix  [4 ]byte 
63+ 	copy (acctPrefix [:], id [:4 ])
64+ 
65+ 	return  r .addRootKeyFromIDSuffix (acctPrefix )
66+ }
67+ 
68+ // addRandomRootKey adds a new random root key to the store, and returns the 
69+ // first 4 bytes of the root key as the root key ID. 
70+ func  (r  * rootKeyMockStore ) addRandomRootKey () [4 ]byte  {
71+ 	rootKey  :=  randomBytes (8 )
72+ 	r .rootKeys  =  append (r .rootKeys , rootKey )
73+ 
74+ 	// As we only return the first 4 bytes as the root key ID, we copy 
75+ 	// those into a fixed size array. 
76+ 	var  shortID  [4 ]byte 
77+ 	copy (shortID [:], rootKey [len (rootKey )- 4 :])
78+ 
79+ 	return  shortID 
80+ }
81+ 
82+ // getAllRootKeys returns all root keys that have been added to the store. 
83+ func  (r  * rootKeyMockStore ) getAllRootKeys () [][]byte  {
84+ 	return  r .rootKeys 
85+ }
86+ 
3887// expectedResult represents the expected result of a migration test. 
3988type  expectedResult  struct  {
4089	kvEntries  []* kvEntry 
@@ -294,13 +343,16 @@ func TestFirewallDBMigration(t *testing.T) {
294343	tests  :=  []struct  {
295344		name        string 
296345		populateDB  func (t  * testing.T , ctx  context.Context ,
297- 			boltDB  * BoltDB , sessionStore  session.Store ) * expectedResult 
346+ 			boltDB  * BoltDB , sessionStore  session.Store ,
347+ 			accountsStore  accounts.Store ,
348+ 			rKeyStore  * rootKeyMockStore ) * expectedResult 
298349	}{
299350		{
300351			name : "empty" ,
301352			populateDB : func (t  * testing.T , ctx  context.Context ,
302- 				boltDB  * BoltDB ,
303- 				sessionStore  session.Store ) * expectedResult  {
353+ 				boltDB  * BoltDB , sessionStore  session.Store ,
354+ 				accountsStore  accounts.Store ,
355+ 				rKeyStore  * rootKeyMockStore ) * expectedResult  {
304356
305357				// Don't populate the DB, and return empty kv 
306358				// records and privacy pairs. 
@@ -384,9 +436,12 @@ func TestFirewallDBMigration(t *testing.T) {
384436				require .NoError (t , firewallStore .Close ())
385437			})
386438
439+ 			rootKeyStore  :=  & rootKeyMockStore {}
440+ 
387441			// Populate the kv store. 
388442			entries  :=  test .populateDB (
389443				t , ctx , firewallStore , sessionsStore ,
444+ 				accountStore , rootKeyStore ,
390445			)
391446
392447			// Create the SQL store that we will migrate the data 
@@ -412,7 +467,8 @@ func TestFirewallDBMigration(t *testing.T) {
412467// globalEntries populates the kv store with one global entry for the temp 
413468// store, and one for the perm store. 
414469func  globalEntries (t  * testing.T , ctx  context.Context , boltDB  * BoltDB ,
415- 	_  session.Store ) * expectedResult  {
470+ 	_  session.Store , _  accounts.Store ,
471+ 	_  * rootKeyMockStore ) * expectedResult  {
416472
417473	return  insertTempAndPermEntry (
418474		t , ctx , boltDB , testRuleName , fn .None [[]byte ](),
@@ -424,7 +480,8 @@ func globalEntries(t *testing.T, ctx context.Context, boltDB *BoltDB,
424480// entry for the local temp store, and one session specific entry for the perm 
425481// local store. 
426482func  sessionSpecificEntries (t  * testing.T , ctx  context.Context , boltDB  * BoltDB ,
427- 	sessionStore  session.Store ) * expectedResult  {
483+ 	sessionStore  session.Store , _  accounts.Store ,
484+ 	_  * rootKeyMockStore ) * expectedResult  {
428485
429486	groupAlias  :=  getNewSessionAlias (t , ctx , sessionStore )
430487
@@ -438,7 +495,8 @@ func sessionSpecificEntries(t *testing.T, ctx context.Context, boltDB *BoltDB,
438495// entry for the local temp store, and one feature specific entry for the perm 
439496// local store. 
440497func  featureSpecificEntries (t  * testing.T , ctx  context.Context , boltDB  * BoltDB ,
441- 	sessionStore  session.Store ) * expectedResult  {
498+ 	sessionStore  session.Store , _  accounts.Store ,
499+ 	_  * rootKeyMockStore ) * expectedResult  {
442500
443501	groupAlias  :=  getNewSessionAlias (t , ctx , sessionStore )
444502
@@ -456,7 +514,8 @@ func featureSpecificEntries(t *testing.T, ctx context.Context, boltDB *BoltDB,
456514// any entries when the entry set is more complex than just a single entry at 
457515// each level. 
458516func  allEntryCombinations (t  * testing.T , ctx  context.Context , boltDB  * BoltDB ,
459- 	sessionStore  session.Store ) * expectedResult  {
517+ 	sessionStore  session.Store , acctStore  accounts.Store ,
518+ 	rStore  * rootKeyMockStore ) * expectedResult  {
460519
461520	var  result  []* kvEntry 
462521	add  :=  func (entry  * expectedResult ) {
@@ -465,9 +524,13 @@ func allEntryCombinations(t *testing.T, ctx context.Context, boltDB *BoltDB,
465524
466525	// First lets create standard entries at all levels, which represents 
467526	// the entries added by other tests. 
468- 	add (globalEntries (t , ctx , boltDB , sessionStore ))
469- 	add (sessionSpecificEntries (t , ctx , boltDB , sessionStore ))
470- 	add (featureSpecificEntries (t , ctx , boltDB , sessionStore ))
527+ 	add (globalEntries (t , ctx , boltDB , sessionStore , acctStore , rStore ))
528+ 	add (sessionSpecificEntries (
529+ 		t , ctx , boltDB , sessionStore , acctStore , rStore ,
530+ 	))
531+ 	add (featureSpecificEntries (
532+ 		t , ctx , boltDB , sessionStore , acctStore , rStore ,
533+ 	))
471534
472535	groupAlias  :=  getNewSessionAlias (t , ctx , sessionStore )
473536
@@ -647,7 +710,8 @@ func insertKvEntry(t *testing.T, ctx context.Context,
647710// across all possible combinations of different levels of entries in the kv 
648711// store. All values and different bucket names are randomly generated. 
649712func  randomKVEntries (t  * testing.T , ctx  context.Context ,
650- 	boltDB  * BoltDB , sessionStore  session.Store ) * expectedResult  {
713+ 	boltDB  * BoltDB , sessionStore  session.Store , _  accounts.Store ,
714+ 	_  * rootKeyMockStore ) * expectedResult  {
651715
652716	var  (
653717		// We set the number of entries to insert to 1000, as that 
@@ -769,23 +833,26 @@ func randomKVEntries(t *testing.T, ctx context.Context,
769833// oneSessionAndPrivPair inserts 1 session with 1 privacy pair into the 
770834// boltDB. 
771835func  oneSessionAndPrivPair (t  * testing.T , ctx  context.Context ,
772- 	boltDB  * BoltDB , sessionStore  session.Store ) * expectedResult  {
836+ 	boltDB  * BoltDB , sessionStore  session.Store , _  accounts.Store ,
837+ 	_  * rootKeyMockStore ) * expectedResult  {
773838
774839	return  createPrivacyPairs (t , ctx , boltDB , sessionStore , 1 , 1 )
775840}
776841
777842// oneSessionsMultiplePrivPairs inserts 1 session with 10 privacy pairs into the 
778843// boltDB. 
779844func  oneSessionsMultiplePrivPairs (t  * testing.T , ctx  context.Context ,
780- 	boltDB  * BoltDB , sessionStore  session.Store ) * expectedResult  {
845+ 	boltDB  * BoltDB , sessionStore  session.Store , _  accounts.Store ,
846+ 	_  * rootKeyMockStore ) * expectedResult  {
781847
782848	return  createPrivacyPairs (t , ctx , boltDB , sessionStore , 1 , 10 )
783849}
784850
785851// multipleSessionsAndPrivacyPairs inserts 5 sessions with 10 privacy pairs 
786852// per session into the boltDB. 
787853func  multipleSessionsAndPrivacyPairs (t  * testing.T , ctx  context.Context ,
788- 	boltDB  * BoltDB , sessionStore  session.Store ) * expectedResult  {
854+ 	boltDB  * BoltDB , sessionStore  session.Store , _  accounts.Store ,
855+ 	_  * rootKeyMockStore ) * expectedResult  {
789856
790857	return  createPrivacyPairs (t , ctx , boltDB , sessionStore , 5 , 10 )
791858}
@@ -847,7 +914,8 @@ func createPrivacyPairs(t *testing.T, ctx context.Context,
847914
848915// randomPrivacyPairs creates a random number of privacy pairs to 10 sessions. 
849916func  randomPrivacyPairs (t  * testing.T , ctx  context.Context ,
850- 	boltDB  * BoltDB , sessionStore  session.Store ) * expectedResult  {
917+ 	boltDB  * BoltDB , sessionStore  session.Store , _  accounts.Store ,
918+ 	_  * rootKeyMockStore ) * expectedResult  {
851919
852920	numSessions  :=  10 
853921	maxPairsPerSession  :=  20 
@@ -905,10 +973,15 @@ func randomPrivacyPairs(t *testing.T, ctx context.Context,
905973// TODO(viktor): Extend this function to also populate it with random action 
906974// entries, once the actions migration has been implemented. 
907975func  randomFirewallDBEntries (t  * testing.T , ctx  context.Context ,
908- 	boltDB  * BoltDB , sessionStore  session.Store ) * expectedResult  {
976+ 	boltDB  * BoltDB , sessionStore  session.Store , acctStore  accounts.Store ,
977+ 	rStore  * rootKeyMockStore ) * expectedResult  {
909978
910- 	kvEntries  :=  randomKVEntries (t , ctx , boltDB , sessionStore )
911- 	privPairs  :=  randomPrivacyPairs (t , ctx , boltDB , sessionStore )
979+ 	kvEntries  :=  randomKVEntries (
980+ 		t , ctx , boltDB , sessionStore , acctStore , rStore ,
981+ 	)
982+ 	privPairs  :=  randomPrivacyPairs (
983+ 		t , ctx , boltDB , sessionStore , acctStore , rStore ,
984+ 	)
912985
913986	return  & expectedResult {
914987		kvEntries : kvEntries .kvEntries ,
@@ -927,3 +1000,12 @@ func randomString(n int) string {
9271000	}
9281001	return  string (b )
9291002}
1003+ 
1004+ // randomBytes generates a random byte array of the passed length n. 
1005+ func  randomBytes (n  int ) []byte  {
1006+ 	b  :=  make ([]byte , n )
1007+ 	for  i  :=  range  b  {
1008+ 		b [i ] =  byte (rand .Intn (256 )) // Random int between 0-255, then cast to byte 
1009+ 	}
1010+ 	return  b 
1011+ }
0 commit comments