pr comments

coder · sreya · Sep 17, 2024 · Sep 13, 2024 · Sep 13, 2024 · Sep 13, 2024
commit 86782ed823441c1ec5319cb33e814a4b4ed541b8
diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -17,6 +17,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/sqlc-dev/pqtype"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
@@ -897,18 +898,22 @@ func CustomRole(t testing.TB, db database.Store, seed database.CustomRole) datab
 func CryptoKey(t testing.TB, db database.Store, seed database.CryptoKey) database.CryptoKey {
 	t.Helper()
 
-	b := make([]byte, 96)
-	_, err := rand.Read(b)
-	require.NoError(t, err, "generate secret")
+	seed.Feature = takeFirst(seed.Feature, database.CryptoKeyFeatureWorkspaceApps)
 
-	key, err := db.InsertCryptoKey(genCtx, database.InsertCryptoKeyParams{
-		Sequence: takeFirst(seed.Sequence, 123),
-		Secret: takeFirst(seed.Secret, sql.NullString{
-			String: hex.EncodeToString(b),
+	if !seed.Secret.Valid {
+		secret, err := newCryptoKeySecret(seed.Feature)
+		require.NoError(t, err, "generate secret")
+		seed.Secret = sql.NullString{
+			String: secret,
 			Valid:  true,
-		}),
+		}
+	}
+
+	key, err := db.InsertCryptoKey(genCtx, database.InsertCryptoKeyParams{
+		Sequence:    takeFirst(seed.Sequence, 123),
+		Secret:      seed.Secret,
 		SecretKeyID: takeFirst(seed.SecretKeyID, sql.NullString{}),
-		Feature:     takeFirst(seed.Feature, database.CryptoKeyFeatureWorkspaceApps),
+		Feature:     seed.Feature,
 		StartsAt:    takeFirst(seed.StartsAt, time.Now()),
 	})
 	require.NoError(t, err, "insert crypto key")
@@ -967,3 +972,24 @@ func takeFirst[Value comparable](values ...Value) Value {
 		return v != empty
 	})
 }
+
+func newCryptoKeySecret(feature database.CryptoKeyFeature) (string, error) {
+	switch feature {
+	case database.CryptoKeyFeatureWorkspaceApps:
+		return generateCryptoKey(96)
+	case database.CryptoKeyFeatureOidcConvert:
+		return generateCryptoKey(32)
+	case database.CryptoKeyFeatureTailnetResume:
+		return generateCryptoKey(64)
+	}
+	return "", xerrors.Errorf("unknown feature: %s", feature)
+}
+
+func generateCryptoKey(length int) (string, error) {
+	b := make([]byte, length)
+	_, err := rand.Read(b)
+	if err != nil {
+		return "", xerrors.Errorf("rand read: %w", err)
+	}
+	return hex.EncodeToString(b), nil
+}
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -1448,6 +1448,8 @@ func (q *FakeQuerier) DeleteCryptoKey(_ context.Context, arg database.DeleteCryp
 		if key.Feature == arg.Feature && key.Sequence == arg.Sequence {
 			q.cryptoKeys[i].Secret.String = ""
 			q.cryptoKeys[i].Secret.Valid = false
+			q.cryptoKeys[i].SecretKeyID.String = ""
+			q.cryptoKeys[i].SecretKeyID.Valid = false
 			return q.cryptoKeys[i], nil
 		}
 	}
@@ -2871,11 +2873,10 @@ func (q *FakeQuerier) GetLatestCryptoKeyByFeature(_ context.Context, feature dat
 			latestKey = key
 		}
 	}
-
-	if latestKey.Secret.Valid {
-		return latestKey, nil
+	if latestKey.StartsAt.IsZero() {
+		return database.CryptoKey{}, sql.ErrNoRows
 	}
-	return database.CryptoKey{}, sql.ErrNoRows
+	return latestKey, nil
 }
 
 func (q *FakeQuerier) GetLatestWorkspaceBuildByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) (database.WorkspaceBuild, error) {

diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000250_crypto_keys.up.sql b/coderd/database/migrations/000250_crypto_keys.up.sql
@@ -1,7 +1,7 @@
 CREATE TYPE crypto_key_feature AS ENUM (
     'workspace_apps',
     'oidc_convert',
-    'peer_reconnect'
+    'tailnet_resume'
 );
 
 CREATE TABLE crypto_keys (

diff --git a/coderd/database/models.go b/coderd/database/models.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/crypto_keys.sql b/coderd/database/queries/crypto_keys.sql
@@ -10,7 +10,6 @@ WHERE feature = $1
 ORDER BY sequence DESC
 LIMIT 1;
 
-
 -- name: GetCryptoKeyByFeatureAndSequence :one
 SELECT *
 FROM crypto_keys
@@ -29,13 +28,13 @@ INSERT INTO crypto_keys (
     sequence,
     secret,
     starts_at,
-	secret_key_id
+    secret_key_id
 ) VALUES (
     $1,
     $2,
     $3,
     $4,
-	$5
+    $5
 ) RETURNING *;
 
 -- name: UpdateCryptoKeyDeletesAt :one

diff --git a/enterprise/dbcrypt/dbcrypt_internal_test.go b/enterprise/dbcrypt/dbcrypt_internal_test.go
@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"io"
 	"testing"
+	"time"
 
 	"github.com/lib/pq"
 	"github.com/stretchr/testify/require"
@@ -353,8 +354,6 @@ func TestCryptoKeys(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
-	// We don't write a GetCryptoKeyByFeatureAndSequence test
-	// because it's basically the same as InsertCryptoKey.
 	t.Run("InsertCryptoKey", func(t *testing.T) {
 		t.Parallel()
 
@@ -432,6 +431,25 @@ func TestCryptoKeys(t *testing.T) {
 		require.Equal(t, ciphers[0].HexDigest(), key.SecretKeyID.String)
 	})
 
+	t.Run("UpdateCryptoKeyDeletesAt", func(t *testing.T) {
+		t.Parallel()
+		_, crypt, ciphers := setup(t)
+		key := dbgen.CryptoKey(t, crypt, database.CryptoKey{
+			Secret: sql.NullString{String: "test", Valid: true},
+		})
+		key, err := crypt.UpdateCryptoKeyDeletesAt(ctx, database.UpdateCryptoKeyDeletesAtParams{
+			Feature:  key.Feature,
+			Sequence: key.Sequence,
+			DeletesAt: sql.NullTime{
+				Time:  time.Now().Add(time.Hour),
+				Valid: true,
+			},
+		})
+		require.NoError(t, err)
+		require.Equal(t, "test", key.Secret.String)
+		require.Equal(t, ciphers[0].HexDigest(), key.SecretKeyID.String)
+	})
+
 	t.Run("DecryptErr", func(t *testing.T) {
 		t.Parallel()
 		db, crypt, ciphers := setup(t)