Fix another test

fleetdm · Jan 22, 2025 · f29d49d · f29d49d
1 parent d4d2a7f
commit f29d49d
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 9 deletions.
diff --git a/server/datastore/mysql/hosts_test.go b/server/datastore/mysql/hosts_test.go
@@ -8033,6 +8033,14 @@ func checkEncryptionKeyStatus(t *testing.T, ds *Datastore, hostID uint, expected
 	require.NoError(t, err)
 	require.Equal(t, expectedKey, got.Base64Encrypted)
 	require.Equal(t, expectedDecryptable, got.Decryptable)
+	if expectedKey != "" {
+		var archiveKey string
+		ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
+			return sqlx.GetContext(context.Background(), q, &archiveKey,
+				`SELECT base64_encrypted FROM host_disk_encryption_keys_archive WHERE host_id = ? ORDER BY created_at DESC LIMIT 1`, hostID)
+		})
+		assert.Equal(t, expectedKey, archiveKey)
+	}
 }
 
 func testLUKSDatastoreFunctions(t *testing.T, ds *Datastore) {
@@ -8117,17 +8125,34 @@ func testLUKSDatastoreFunctions(t *testing.T, ds *Datastore) {
 	require.NoError(t, err)
 	require.NoError(t, ds.AssertHasNoEncryptionKeyStored(ctx, host1.ID))
 	require.Error(t, ds.AssertHasNoEncryptionKeyStored(ctx, host2.ID))
-	key, err := ds.GetHostDiskEncryptionKey(ctx, host2.ID)
-	require.NoError(t, err)
-	require.Equal(t, "bazqux", key.Base64Encrypted)
+	checkLUKSEncryptionKey(t, ds, host2.ID, "bazqux", "fuzzmuffin")
 
 	// persists when host hasn't had anything queued
 	err = ds.SaveLUKSData(ctx, host3, "newstuff", "fuzzball", 1)
 	require.NoError(t, err)
 	require.Error(t, ds.AssertHasNoEncryptionKeyStored(ctx, host3.ID))
-	key, err = ds.GetHostDiskEncryptionKey(ctx, host3.ID)
+	checkLUKSEncryptionKey(t, ds, host3.ID, "newstuff", "fuzzball")
+}
+
+func checkLUKSEncryptionKey(t *testing.T, ds *Datastore, hostID uint, expectedKey string, expectedSalt string) {
+	got, err := ds.GetHostDiskEncryptionKey(context.Background(), hostID)
 	require.NoError(t, err)
-	require.Equal(t, "newstuff", key.Base64Encrypted)
+	require.Equal(t, expectedKey, got.Base64Encrypted)
+	if expectedKey != "" {
+		var archiveKey string
+		ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
+			return sqlx.GetContext(context.Background(), q, &archiveKey,
+				`SELECT base64_encrypted FROM host_disk_encryption_keys_archive WHERE host_id = ? ORDER BY created_at DESC LIMIT 1`, hostID)
+		})
+		assert.Equal(t, expectedKey, archiveKey)
+		var archiveSalt string
+		ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
+			return sqlx.GetContext(context.Background(), q, &archiveSalt,
+				`SELECT base64_encrypted_salt FROM host_disk_encryption_keys_archive WHERE host_id = ? ORDER BY created_at DESC LIMIT 1`,
+				hostID)
+		})
+		assert.Equal(t, expectedSalt, archiveSalt)
+	}
 }
 
 func testHostsSetOrUpdateHostDisksEncryptionKey(t *testing.T, ds *Datastore) {

diff --git a/server/datastore/mysql/microsoft_mdm.go b/server/datastore/mysql/microsoft_mdm.go
@@ -512,7 +512,7 @@ func (ds *Datastore) whereBitLockerStatus(status fleet.DiskEncryptionStatus) str
 		whereEncrypted        = `(hd.encrypted IS NOT NULL AND hd.encrypted = 1)`
 		whereHostDisksUpdated = `(hd.updated_at IS NOT NULL AND hdek.updated_at IS NOT NULL AND hd.updated_at >= hdek.updated_at)`
 		whereClientError      = `(hdek.client_error IS NOT NULL AND hdek.client_error != '')`
-		withinGracePeriod     = `(hdek.updated_at IS NOT NULL AND hdek.updated_at >= DATE_SUB(NOW(), INTERVAL 1 HOUR))`
+		withinGracePeriod     = `(hdek.updated_at IS NOT NULL AND hdek.updated_at >= DATE_SUB(NOW(6), INTERVAL 1 HOUR))`
 	)
 
 	// TODO: what if windows sends us a key for an already encrypted volumne? could it get stuck

diff --git a/server/datastore/mysql/migrations/tables/20250121094045_AddHostDiskEncryptionKeysArchive.go b/server/datastore/mysql/migrations/tables/20250121094045_AddHostDiskEncryptionKeysArchive.go
@@ -10,7 +10,14 @@ func init() {
 }
 
 func Up_20250121094045(tx *sql.Tx) error {
-	_, err := tx.Exec(`ALTER TABLE host_disk_encryption_keys
+	_, err := tx.Exec(`ALTER TABLE host_disks
+		MODIFY COLUMN created_at TIMESTAMP(6) NOT NULL DEFAULT NOW(6),
+		MODIFY COLUMN updated_at TIMESTAMP(6) NULL DEFAULT NOW(6) ON UPDATE NOW(6)`)
+	if err != nil {
+		return fmt.Errorf("failed to alter host_disks table: %w", err)
+	}
+
+	_, err = tx.Exec(`ALTER TABLE host_disk_encryption_keys
 		MODIFY COLUMN created_at TIMESTAMP(6) NOT NULL DEFAULT NOW(6),
 		MODIFY COLUMN updated_at TIMESTAMP(6) NULL DEFAULT NOW(6) ON UPDATE NOW(6)`)
 	if err != nil {

diff --git a/server/datastore/mysql/schema.sql b/server/datastore/mysql/schema.sql
@@ -334,8 +334,8 @@ CREATE TABLE `host_disks` (
   `host_id` int unsigned NOT NULL,
   `gigs_disk_space_available` decimal(10,2) NOT NULL DEFAULT '0.00',
   `percent_disk_space_available` decimal(10,2) NOT NULL DEFAULT '0.00',
-  `created_at` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,
-  `updated_at` timestamp NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
+  `created_at` timestamp(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
+  `updated_at` timestamp(6) NULL DEFAULT CURRENT_TIMESTAMP(6) ON UPDATE CURRENT_TIMESTAMP(6),
   `encrypted` tinyint(1) DEFAULT NULL,
   `gigs_total_disk_space` decimal(10,2) NOT NULL DEFAULT '0.00',
   PRIMARY KEY (`host_id`),