feat: Longer lived api keys for cli (#1935)

Emyrk · web-flow · commit 913c0f5e7fd3 · 2022-06-01T14:58:55.000-05:00
* feat: Longer lived api keys for cli
* feat: Refresh tokens based on their lifetime set in the db
* test: Add unit test for refreshing
diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -1128,9 +1128,14 @@ func (q *fakeQuerier) InsertAPIKey(_ context.Context, arg database.InsertAPIKeyP
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
 
+	if arg.LifetimeSeconds == 0 {
+		arg.LifetimeSeconds = 86400
+	}
+
 	//nolint:gosimple
 	key := database.APIKey{
 		ID:                arg.ID,
+		LifetimeSeconds:   arg.LifetimeSeconds,
 		HashedSecret:      arg.HashedSecret,
 		UserID:            arg.UserID,
 		ExpiresAt:         arg.ExpiresAt,
diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000016_api_key_lifetime.down.sql b/coderd/database/migrations/000016_api_key_lifetime.down.sql
@@ -0,0 +1 @@
+ALTER TABLE api_keys DROP COLUMN lifetime_seconds;
diff --git a/coderd/database/migrations/000016_api_key_lifetime.up.sql b/coderd/database/migrations/000016_api_key_lifetime.up.sql
@@ -0,0 +1,2 @@
+-- Default lifetime is 24hours.
+ALTER TABLE api_keys ADD COLUMN lifetime_seconds bigint default 86400 NOT NULL;
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/apikeys.sql b/coderd/database/queries/apikeys.sql
@@ -12,6 +12,7 @@ LIMIT
 INSERT INTO
 	api_keys (
 		id,
+		lifetime_seconds,
 		hashed_secret,
 		user_id,
 		last_used,
@@ -25,7 +26,13 @@ INSERT INTO
 		oauth_expiry
 	)
 VALUES
-	($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) RETURNING *;
+	(@id,
+	 -- If the lifetime is set to 0, default to 24hrs
+	 CASE @lifetime_seconds::bigint
+	     WHEN 0 THEN 86400
+		 ELSE @lifetime_seconds::bigint
+	 END
+	 , @hashed_secret, @user_id, @last_used, @expires_at, @created_at, @updated_at, @login_type, @oauth_access_token, @oauth_refresh_token, @oauth_id_token, @oauth_expiry) RETURNING *;
 
 -- name: UpdateAPIKeyByID :exec
 UPDATE
diff --git a/coderd/httpmw/apikey.go b/coderd/httpmw/apikey.go
@@ -166,7 +166,7 @@ func ExtractAPIKey(db database.Store, oauth *OAuth2Configs) func(http.Handler) h
 			}
 			// Only update the ExpiresAt once an hour to prevent database spam.
 			// We extend the ExpiresAt to reduce re-authentication.
-			apiKeyLifetime := 24 * time.Hour
+			apiKeyLifetime := time.Duration(key.LifetimeSeconds) * time.Second
 			if key.ExpiresAt.Sub(now) <= apiKeyLifetime-time.Hour {
 				key.ExpiresAt = now.Add(apiKeyLifetime)
 				changed = true
diff --git a/coderd/users.go b/coderd/users.go
@@ -660,9 +660,14 @@ func (api *API) postAPIKey(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	lifeTime := time.Hour * 24 * 7
 	sessionToken, created := api.createAPIKey(rw, r, database.InsertAPIKeyParams{
 		UserID:    user.ID,
 		LoginType: database.LoginTypePassword,
+		// All api generated keys will last 1 week. Browser login tokens have
+		// a shorter life.
+		ExpiresAt:       database.Now().Add(lifeTime),
+		LifetimeSeconds: int64(lifeTime.Seconds()),
 	})
 	if !created {
 		return
@@ -723,10 +728,21 @@ func (api *API) createAPIKey(rw http.ResponseWriter, r *http.Request, params dat
 	}
 	hashed := sha256.Sum256([]byte(keySecret))
 
+	// Default expires at to now+lifetime, or just 24hrs if not set
+	if params.ExpiresAt.IsZero() {
+		if params.LifetimeSeconds != 0 {
+			params.ExpiresAt = database.Now().Add(time.Duration(params.LifetimeSeconds) * time.Second)
+		} else {
+			params.ExpiresAt = database.Now().Add(24 * time.Hour)
+		}
+	}
+
 	_, err = api.Database.InsertAPIKey(r.Context(), database.InsertAPIKeyParams{
-		ID:                keyID,
-		UserID:            params.UserID,
-		ExpiresAt:         database.Now().Add(24 * time.Hour),
+		ID:              keyID,
+		UserID:          params.UserID,
+		LifetimeSeconds: params.LifetimeSeconds,
+		// Make sure in UTC time for common time zone
+		ExpiresAt:         params.ExpiresAt.UTC(),
 		CreatedAt:         database.Now(),
 		UpdatedAt:         database.Now(),
 		HashedSecret:      hashed[:],
diff --git a/coderd/users_test.go b/coderd/users_test.go
@@ -8,11 +8,13 @@ import (
 	"sort"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/database/databasefake"
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/rbac"
@@ -130,6 +132,99 @@ func TestPostLogin(t *testing.T) {
 		})
 		require.NoError(t, err)
 	})
+
+	t.Run("Lifetime&Expire", func(t *testing.T) {
+		t.Parallel()
+		var (
+			ctx = context.Background()
+		)
+		client, api := coderdtest.NewWithAPI(t, nil)
+		admin := coderdtest.CreateFirstUser(t, client)
+
+		split := strings.Split(client.SessionToken, "-")
+		loginKey, err := api.Database.GetAPIKeyByID(ctx, split[0])
+		require.NoError(t, err, "fetch login key")
+		require.Equal(t, int64(86400), loginKey.LifetimeSeconds, "default should be 86400")
+
+		// Generated tokens have a longer life
+		token, err := client.CreateAPIKey(ctx, admin.UserID.String())
+		require.NoError(t, err, "make new api key")
+		split = strings.Split(token.Key, "-")
+		apiKey, err := api.Database.GetAPIKeyByID(ctx, split[0])
+		require.NoError(t, err, "fetch api key")
+
+		require.True(t, apiKey.ExpiresAt.After(time.Now().Add(time.Hour*24*6)), "api key lasts more than 6 days")
+		require.True(t, apiKey.ExpiresAt.After(loginKey.ExpiresAt.Add(time.Hour)), "api key should be longer expires")
+		require.Greater(t, apiKey.LifetimeSeconds, loginKey.LifetimeSeconds, "api key should have longer lifetime")
+	})
+
+	t.Run("APIKeyExtend", func(t *testing.T) {
+		t.Parallel()
+		var (
+			ctx = context.Background()
+		)
+		client, api := coderdtest.NewWithAPI(t, nil)
+		admin := coderdtest.CreateFirstUser(t, client)
+
+		token, err := client.CreateAPIKey(ctx, admin.UserID.String())
+		require.NoError(t, err, "make new api key")
+		client.SessionToken = token.Key
+		split := strings.Split(token.Key, "-")
+
+		apiKey, err := api.Database.GetAPIKeyByID(ctx, split[0])
+		require.NoError(t, err, "fetch api key")
+
+		err = api.Database.UpdateAPIKeyByID(ctx, database.UpdateAPIKeyByIDParams{
+			ID:       apiKey.ID,
+			LastUsed: apiKey.LastUsed,
+			// This should cause a refresh
+			ExpiresAt:         apiKey.ExpiresAt.Add(time.Hour * -2),
+			OAuthAccessToken:  apiKey.OAuthAccessToken,
+			OAuthRefreshToken: apiKey.OAuthRefreshToken,
+			OAuthExpiry:       apiKey.OAuthExpiry,
+		})
+		require.NoError(t, err, "update api key")
+
+		_, err = client.User(ctx, codersdk.Me)
+		require.NoError(t, err, "fetch user")
+
+		apiKey, err = api.Database.GetAPIKeyByID(ctx, split[0])
+		require.NoError(t, err, "fetch refreshed api key")
+		// 1 minute tolerance
+		require.True(t, apiKey.ExpiresAt.After(time.Now().Add(time.Hour*24*7).Add(time.Minute*-1)), "api key lasts 7 days")
+	})
+
+	t.Run("LoginKeyExtend", func(t *testing.T) {
+		t.Parallel()
+		var (
+			ctx = context.Background()
+		)
+		client, api := coderdtest.NewWithAPI(t, nil)
+		_ = coderdtest.CreateFirstUser(t, client)
+		split := strings.Split(client.SessionToken, "-")
+
+		apiKey, err := api.Database.GetAPIKeyByID(ctx, split[0])
+		require.NoError(t, err, "fetch login key")
+
+		err = api.Database.UpdateAPIKeyByID(ctx, database.UpdateAPIKeyByIDParams{
+			ID:       apiKey.ID,
+			LastUsed: apiKey.LastUsed,
+			// This should cause a refresh
+			ExpiresAt:         apiKey.ExpiresAt.Add(time.Hour * -2),
+			OAuthAccessToken:  apiKey.OAuthAccessToken,
+			OAuthRefreshToken: apiKey.OAuthRefreshToken,
+			OAuthExpiry:       apiKey.OAuthExpiry,
+		})
+		require.NoError(t, err, "update login key")
+
+		_, err = client.User(ctx, codersdk.Me)
+		require.NoError(t, err, "fetch user")
+
+		apiKey, err = api.Database.GetAPIKeyByID(ctx, split[0])
+		require.NoError(t, err, "fetch refreshed login key")
+		// 1 minute tolerance
+		require.True(t, apiKey.ExpiresAt.After(time.Now().Add(time.Hour*24).Add(time.Minute*-1)), "login key lasts 24 hrs")
+	})
 }
 
 func TestPostLogout(t *testing.T) {

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+ALTER TABLE api_keys DROP COLUMN lifetime_seconds;`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+-- Default lifetime is 24hours.`
	`2`	`+ALTER TABLE api_keys ADD COLUMN lifetime_seconds bigint default 86400 NOT NULL;`
Original file line number	Diff line number	Diff line change
`@@ -166,7 +166,7 @@ func ExtractAPIKey(db database.Store, oauth *OAuth2Configs) func(http.Handler) h`
`166`	`166`	`}`
`167`	`167`	`// Only update the ExpiresAt once an hour to prevent database spam.`
`168`	`168`	`// We extend the ExpiresAt to reduce re-authentication.`
`169`		`- apiKeyLifetime := 24 * time.Hour`
	`169`	`+ apiKeyLifetime := time.Duration(key.LifetimeSeconds) * time.Second`
`170`	`170`	`if key.ExpiresAt.Sub(now) <= apiKeyLifetime-time.Hour {`
`171`	`171`	`key.ExpiresAt = now.Add(apiKeyLifetime)`
`172`	`172`	`changed = true`