44 "bytes"
55 "context"
66 "database/sql"
7+ "encoding/binary"
78 "errors"
89 "fmt"
910 "testing"
3536 testEntryValue = []byte {1 , 2 , 3 }
3637)
3738
39+ // rootKeyMockStore is a mock implementation of a macaroon service store that
40+ // can be used to generate mock root keys for testing.
41+ type rootKeyMockStore struct {
42+ // rootKeys is a slice of all root keys that have been added to the
43+ // store.
44+ rootKeys [][]byte
45+ }
46+
47+ // addRootKeyFromIDSuffix adds a new root key to the store, using the passed
48+ // 4 byte suffix. The function generates a root key that ends with the 4 byte
49+ // suffix, prefixed by 4 random bytes.
50+ func (r * rootKeyMockStore ) addRootKeyFromIDSuffix (suffix [4 ]byte ) uint64 {
51+ // As a real root key is 8 bytes, we need to generate a random 4 byte
52+ // prefix to prepend to the passed 4 byte suffix.
53+ rootKey := append (randomBytes (4 ), suffix [:]... )
54+ r .rootKeys = append (r .rootKeys , rootKey )
55+
56+ return binary .BigEndian .Uint64 (rootKey [:])
57+ }
58+
59+ // addRootKeyFromAcctID adds a new root key to the store, using the first 4
60+ // bytes of the passed account ID as the suffix for the root key, prefixed by 4
61+ // random bytes.
62+ func (r * rootKeyMockStore ) addRootKeyFromAcctID (id accounts.AccountID ) uint64 {
63+ var acctPrefix [4 ]byte
64+ copy (acctPrefix [:], id [:4 ])
65+
66+ return r .addRootKeyFromIDSuffix (acctPrefix )
67+ }
68+
69+ // addRandomRootKey adds a new random root key to the store, and returns the
70+ // root key ID as an uint64.
71+ func (r * rootKeyMockStore ) addRandomRootKey () uint64 {
72+ rootKey := randomBytes (8 )
73+ r .rootKeys = append (r .rootKeys , rootKey )
74+
75+ return binary .BigEndian .Uint64 (rootKey [:])
76+ }
77+
78+ // getAllRootKeys returns all root keys that have been added to the store.
79+ func (r * rootKeyMockStore ) getAllRootKeys () [][]byte {
80+ return r .rootKeys
81+ }
82+
3883// expectedResult represents the expected result of a migration test.
3984type expectedResult struct {
4085 kvEntries []* kvEntry
@@ -294,13 +339,16 @@ func TestFirewallDBMigration(t *testing.T) {
294339 tests := []struct {
295340 name string
296341 populateDB func (t * testing.T , ctx context.Context ,
297- boltDB * BoltDB , sessionStore session.Store ) * expectedResult
342+ boltDB * BoltDB , sessionStore session.Store ,
343+ accountsStore accounts.Store ,
344+ rKeyStore * rootKeyMockStore ) * expectedResult
298345 }{
299346 {
300347 name : "empty" ,
301348 populateDB : func (t * testing.T , ctx context.Context ,
302- boltDB * BoltDB ,
303- sessionStore session.Store ) * expectedResult {
349+ boltDB * BoltDB , sessionStore session.Store ,
350+ accountsStore accounts.Store ,
351+ rKeyStore * rootKeyMockStore ) * expectedResult {
304352
305353 // Don't populate the DB, and return empty kv
306354 // records and privacy pairs.
@@ -384,9 +432,12 @@ func TestFirewallDBMigration(t *testing.T) {
384432 require .NoError (t , firewallStore .Close ())
385433 })
386434
435+ rootKeyStore := & rootKeyMockStore {}
436+
387437 // Populate the kv store.
388438 entries := test .populateDB (
389439 t , ctx , firewallStore , sessionsStore ,
440+ accountStore , rootKeyStore ,
390441 )
391442
392443 // Create the SQL store that we will migrate the data
@@ -412,7 +463,8 @@ func TestFirewallDBMigration(t *testing.T) {
412463// globalEntries populates the kv store with one global entry for the temp
413464// store, and one for the perm store.
414465func globalEntries (t * testing.T , ctx context.Context , boltDB * BoltDB ,
415- _ session.Store ) * expectedResult {
466+ _ session.Store , _ accounts.Store ,
467+ _ * rootKeyMockStore ) * expectedResult {
416468
417469 return insertTempAndPermEntry (
418470 t , ctx , boltDB , testRuleName , fn .None [[]byte ](),
@@ -424,7 +476,8 @@ func globalEntries(t *testing.T, ctx context.Context, boltDB *BoltDB,
424476// entry for the local temp store, and one session specific entry for the perm
425477// local store.
426478func sessionSpecificEntries (t * testing.T , ctx context.Context , boltDB * BoltDB ,
427- sessionStore session.Store ) * expectedResult {
479+ sessionStore session.Store , _ accounts.Store ,
480+ _ * rootKeyMockStore ) * expectedResult {
428481
429482 groupAlias := getNewSessionAlias (t , ctx , sessionStore )
430483
@@ -438,7 +491,8 @@ func sessionSpecificEntries(t *testing.T, ctx context.Context, boltDB *BoltDB,
438491// entry for the local temp store, and one feature specific entry for the perm
439492// local store.
440493func featureSpecificEntries (t * testing.T , ctx context.Context , boltDB * BoltDB ,
441- sessionStore session.Store ) * expectedResult {
494+ sessionStore session.Store , _ accounts.Store ,
495+ _ * rootKeyMockStore ) * expectedResult {
442496
443497 groupAlias := getNewSessionAlias (t , ctx , sessionStore )
444498
@@ -456,7 +510,8 @@ func featureSpecificEntries(t *testing.T, ctx context.Context, boltDB *BoltDB,
456510// any entries when the entry set is more complex than just a single entry at
457511// each level.
458512func allEntryCombinations (t * testing.T , ctx context.Context , boltDB * BoltDB ,
459- sessionStore session.Store ) * expectedResult {
513+ sessionStore session.Store , acctStore accounts.Store ,
514+ rStore * rootKeyMockStore ) * expectedResult {
460515
461516 var result []* kvEntry
462517 add := func (entry * expectedResult ) {
@@ -465,9 +520,13 @@ func allEntryCombinations(t *testing.T, ctx context.Context, boltDB *BoltDB,
465520
466521 // First lets create standard entries at all levels, which represents
467522 // 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 ))
523+ add (globalEntries (t , ctx , boltDB , sessionStore , acctStore , rStore ))
524+ add (sessionSpecificEntries (
525+ t , ctx , boltDB , sessionStore , acctStore , rStore ,
526+ ))
527+ add (featureSpecificEntries (
528+ t , ctx , boltDB , sessionStore , acctStore , rStore ,
529+ ))
471530
472531 groupAlias := getNewSessionAlias (t , ctx , sessionStore )
473532
@@ -647,7 +706,8 @@ func insertKvEntry(t *testing.T, ctx context.Context,
647706// across all possible combinations of different levels of entries in the kv
648707// store. All values and different bucket names are randomly generated.
649708func randomKVEntries (t * testing.T , ctx context.Context ,
650- boltDB * BoltDB , sessionStore session.Store ) * expectedResult {
709+ boltDB * BoltDB , sessionStore session.Store , _ accounts.Store ,
710+ _ * rootKeyMockStore ) * expectedResult {
651711
652712 var (
653713 // We set the number of entries to insert to 1000, as that
@@ -769,23 +829,26 @@ func randomKVEntries(t *testing.T, ctx context.Context,
769829// oneSessionAndPrivPair inserts 1 session with 1 privacy pair into the
770830// boltDB.
771831func oneSessionAndPrivPair (t * testing.T , ctx context.Context ,
772- boltDB * BoltDB , sessionStore session.Store ) * expectedResult {
832+ boltDB * BoltDB , sessionStore session.Store , _ accounts.Store ,
833+ _ * rootKeyMockStore ) * expectedResult {
773834
774835 return createPrivacyPairs (t , ctx , boltDB , sessionStore , 1 , 1 )
775836}
776837
777838// oneSessionsMultiplePrivPairs inserts 1 session with 10 privacy pairs into the
778839// boltDB.
779840func oneSessionsMultiplePrivPairs (t * testing.T , ctx context.Context ,
780- boltDB * BoltDB , sessionStore session.Store ) * expectedResult {
841+ boltDB * BoltDB , sessionStore session.Store , _ accounts.Store ,
842+ _ * rootKeyMockStore ) * expectedResult {
781843
782844 return createPrivacyPairs (t , ctx , boltDB , sessionStore , 1 , 10 )
783845}
784846
785847// multipleSessionsAndPrivacyPairs inserts 5 sessions with 10 privacy pairs
786848// per session into the boltDB.
787849func multipleSessionsAndPrivacyPairs (t * testing.T , ctx context.Context ,
788- boltDB * BoltDB , sessionStore session.Store ) * expectedResult {
850+ boltDB * BoltDB , sessionStore session.Store , _ accounts.Store ,
851+ _ * rootKeyMockStore ) * expectedResult {
789852
790853 return createPrivacyPairs (t , ctx , boltDB , sessionStore , 5 , 10 )
791854}
@@ -847,7 +910,8 @@ func createPrivacyPairs(t *testing.T, ctx context.Context,
847910
848911// randomPrivacyPairs creates a random number of privacy pairs to 10 sessions.
849912func randomPrivacyPairs (t * testing.T , ctx context.Context ,
850- boltDB * BoltDB , sessionStore session.Store ) * expectedResult {
913+ boltDB * BoltDB , sessionStore session.Store , _ accounts.Store ,
914+ _ * rootKeyMockStore ) * expectedResult {
851915
852916 numSessions := 10
853917 maxPairsPerSession := 20
@@ -905,10 +969,15 @@ func randomPrivacyPairs(t *testing.T, ctx context.Context,
905969// TODO(viktor): Extend this function to also populate it with random action
906970// entries, once the actions migration has been implemented.
907971func randomFirewallDBEntries (t * testing.T , ctx context.Context ,
908- boltDB * BoltDB , sessionStore session.Store ) * expectedResult {
972+ boltDB * BoltDB , sessionStore session.Store , acctStore accounts.Store ,
973+ rStore * rootKeyMockStore ) * expectedResult {
909974
910- kvEntries := randomKVEntries (t , ctx , boltDB , sessionStore )
911- privPairs := randomPrivacyPairs (t , ctx , boltDB , sessionStore )
975+ kvEntries := randomKVEntries (
976+ t , ctx , boltDB , sessionStore , acctStore , rStore ,
977+ )
978+ privPairs := randomPrivacyPairs (
979+ t , ctx , boltDB , sessionStore , acctStore , rStore ,
980+ )
912981
913982 return & expectedResult {
914983 kvEntries : kvEntries .kvEntries ,
@@ -927,3 +996,12 @@ func randomString(n int) string {
927996 }
928997 return string (b )
929998}
999+
1000+ // randomBytes generates a random byte array of the passed length n.
1001+ func randomBytes (n int ) []byte {
1002+ b := make ([]byte , n )
1003+ for i := range b {
1004+ b [i ] = byte (rand .Intn (256 )) // Random int between 0-255, then cast to byte
1005+ }
1006+ return b
1007+ }
0 commit comments