adding tests for role based key support on server

Signed-off-by: David Lawrence <david.lawrence@docker.com> (github: endophage)

server/storage/database.go

@@ -156,16 +156,17 @@ func (db *SQLStorage) GetKey(gun, role string) (algorithm string, public []byte,
 func (db *SQLStorage) SetKey(gun, role, algorithm string, public []byte) error {
 
 	entry := Key{
-		Gun:    gun,
-		Role:   role,
-		Cipher: string(algorithm),
-		Public: public,
+		Gun:  gun,
+		Role: role,
 	}
 
 	if !db.Where(&entry).First(&Key{}).RecordNotFound() {
 		return &ErrKeyExists{gun: gun, role: role}
 	}
 
+	entry.Cipher = algorithm
+	entry.Public = public
+
 	return translateOldVersionError(
 		db.FirstOrCreate(&Key{}, &entry).Error)
 }

server/storage/database_test.go

@@ -310,6 +310,78 @@ func TestSQLSetKeyExists(t *testing.T) {
 	dbStore.DB.Close()
 }
 
+func TestSQLSetKeyMultipleRoles(t *testing.T) {
+	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
+	gormDB, dbStore := SetUpSQLite(t, tempBaseDir)
+	defer os.RemoveAll(tempBaseDir)
+
+	err = dbStore.SetKey("testGUN", data.CanonicalTimestampRole, "testCipher", []byte("1"))
+	assert.NoError(t, err, "Inserting timestamp into empty DB should succeed")
+
+	err = dbStore.SetKey("testGUN", data.CanonicalSnapshotRole, "testCipher", []byte("1"))
+	assert.NoError(t, err, "Inserting snapshot key into DB with timestamp key should succeed")
+
+	var rows []Key
+	query := gormDB.Select("ID, Gun, Role, Cipher, Public").Find(&rows)
+	assert.NoError(t, query.Error)
+
+	expectedTS := Key{Gun: "testGUN", Role: "timestamp", Cipher: "testCipher",
+		Public: []byte("1")}
+	expectedTS.Model = gorm.Model{ID: 1}
+
+	expectedSN := Key{Gun: "testGUN", Role: "snapshot", Cipher: "testCipher",
+		Public: []byte("1")}
+	expectedSN.Model = gorm.Model{ID: 2}
+
+	assert.Equal(t, []Key{expectedTS, expectedSN}, rows)
+
+	dbStore.DB.Close()
+}
+
+func TestSQLSetKeyMultipleGuns(t *testing.T) {
+	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
+	gormDB, dbStore := SetUpSQLite(t, tempBaseDir)
+	defer os.RemoveAll(tempBaseDir)
+
+	err = dbStore.SetKey("testGUN", data.CanonicalTimestampRole, "testCipher", []byte("1"))
+	assert.NoError(t, err, "Inserting timestamp into empty DB should succeed")
+
+	err = dbStore.SetKey("testAnotherGUN", data.CanonicalTimestampRole, "testCipher", []byte("1"))
+	assert.NoError(t, err, "Inserting snapshot key into DB with timestamp key should succeed")
+
+	var rows []Key
+	query := gormDB.Select("ID, Gun, Role, Cipher, Public").Find(&rows)
+	assert.NoError(t, query.Error)
+
+	expected1 := Key{Gun: "testGUN", Role: "timestamp", Cipher: "testCipher",
+		Public: []byte("1")}
+	expected1.Model = gorm.Model{ID: 1}
+
+	expected2 := Key{Gun: "testAnotherGUN", Role: "timestamp", Cipher: "testCipher",
+		Public: []byte("1")}
+	expected2.Model = gorm.Model{ID: 2}
+
+	assert.Equal(t, []Key{expected1, expected2}, rows)
+
+	dbStore.DB.Close()
+}
+
+func TestSQLSetKeySameRoleGun(t *testing.T) {
+	tempBaseDir, err := ioutil.TempDir("", "notary-test-")
+	_, dbStore := SetUpSQLite(t, tempBaseDir)
+	defer os.RemoveAll(tempBaseDir)
+
+	err = dbStore.SetKey("testGUN", data.CanonicalTimestampRole, "testCipher", []byte("1"))
+	assert.NoError(t, err, "Inserting timestamp into empty DB should succeed")
+
+	err = dbStore.SetKey("testGUN", data.CanonicalTimestampRole, "testCipher", []byte("2"))
+	assert.Error(t, err)
+	assert.IsType(t, &ErrKeyExists{}, err,
+		"Expected ErrKeyExists from SetKey")
+
+	dbStore.DB.Close()
+}
+
 // TestDBCheckHealthTableMissing asserts that the health check fails if one or
 // both the tables are missing.
 func TestDBCheckHealthTableMissing(t *testing.T) {

server/storage/memory_test.go

@@ -54,9 +54,39 @@ func TestGetTimestampKey(t *testing.T) {
 
 func TestSetKey(t *testing.T) {
 	s := NewMemStorage()
-	s.SetKey("gun", data.CanonicalTimestampRole, data.RSAKey, []byte("test"))
+	err := s.SetKey("gun", data.CanonicalTimestampRole, data.RSAKey, []byte("test"))
+	assert.NoError(t, err)
 
-	err := s.SetKey("gun", data.CanonicalTimestampRole, data.RSAKey, []byte("test2"))
+	k := s.keys["gun"][data.CanonicalTimestampRole]
+	assert.Equal(t, data.RSAKey, k.algorithm, "Expected algorithm to be rsa, received %s", k.algorithm)
+	assert.Equal(t, []byte("test"), k.public, "Public key did not match expected")
+
+}
+
+func TestSetKeyMultipleRoles(t *testing.T) {
+	s := NewMemStorage()
+	err := s.SetKey("gun", data.CanonicalTimestampRole, data.RSAKey, []byte("test"))
+	assert.NoError(t, err)
+
+	err = s.SetKey("gun", data.CanonicalSnapshotRole, data.RSAKey, []byte("test"))
+	assert.NoError(t, err)
+
+	k := s.keys["gun"][data.CanonicalTimestampRole]
+	assert.Equal(t, data.RSAKey, k.algorithm, "Expected algorithm to be rsa, received %s", k.algorithm)
+	assert.Equal(t, []byte("test"), k.public, "Public key did not match expected")
+
+	k = s.keys["gun"][data.CanonicalSnapshotRole]
+	assert.Equal(t, data.RSAKey, k.algorithm, "Expected algorithm to be rsa, received %s", k.algorithm)
+	assert.Equal(t, []byte("test"), k.public, "Public key did not match expected")
+}
+
+func TestSetKeySameRoleGun(t *testing.T) {
+	s := NewMemStorage()
+	err := s.SetKey("gun", data.CanonicalTimestampRole, data.RSAKey, []byte("test"))
+	assert.NoError(t, err)
+
+	// set diff algo and bytes so we can confirm data didn't get replaced
+	err = s.SetKey("gun", data.CanonicalTimestampRole, data.ECDSAKey, []byte("test2"))
 	assert.IsType(t, &ErrKeyExists{}, err, "Expected err to be ErrKeyExists")
 
 	k := s.keys["gun"][data.CanonicalTimestampRole]

server/storage/models.go

@@ -19,8 +19,8 @@ func (g TUFFile) TableName() string {
 // Key represents a single timestamp key in the database
 type Key struct {
 	gorm.Model
-	Gun    string `sql:"type:varchar(255);not null"`
-	Role   string `sql:"type:varchar(255);not null"`
+	Gun    string `sql:"type:varchar(255);not null;unique_index:gun_role"`
+	Role   string `sql:"type:varchar(255);not null;unique_index:gun_role"`
 	Cipher string `sql:"type:varchar(30);not null"`
 	Public []byte `sql:"type:blob;not null"`
 }