feat: SSO Improvement - alter user_sessions table to include access token, implement CRUD ops, GET, POST, PATCH APIs and det token CLIs

master/internal/user/postgres_long_lived_tokens.go

@@ -0,0 +1,125 @@
+package user
+
+import (
+	"context"
+	"time"
+
+	"github.com/o1egl/paseto"
+	"github.com/uptrace/bun"
+	"gopkg.in/guregu/null.v3"
+
+	"github.com/pkg/errors"
+
+	"github.com/determined-ai/determined/master/internal/db"
+	"github.com/determined-ai/determined/master/pkg/model"
+)
+
+const (
+	// TokenExpirationDuration is how long a newly created long lived token is valid.
+	TokenExpirationDuration = 30 * 24 * time.Hour
+)
+
+// LongLivedTokenOption is the return type for WithTokenExpiresAt helper function.
+// It takes a pointer to model.LongLivedToken and modifies it.
+// It’s used to apply optional settings to the LongLivedToken object.
+type LongLivedTokenOption func(f *model.LongLivedToken)
+
+// WithTokenExpiresAt function will add specified expiresAt (if any) to the long lived token table.
+func WithTokenExpiresAt(expiresAt *time.Time) LongLivedTokenOption {
+	return func(s *model.LongLivedToken) {
+		s.ExpiresAt = *expiresAt
+	}
+}
+
+// CreateLongLivedToken creates a row in the long lived token table.
+func CreateLongLivedToken(ctx context.Context, user *model.User, opts ...LongLivedTokenOption) (string, error) {
+	// Populate the default values in the model.
+	longLivedToken := &model.LongLivedToken{
+		UserID:    user.ID,
+		CreatedAt: time.Now(),
+		ExpiresAt: time.Now().Add(TokenExpirationDuration),
+	}
+
+	// Update the optional ExpiresAt field (if passed)
+	for _, opt := range opts {
+		opt(longLivedToken)
+	}
+
+	err := db.Bun().RunInTx(ctx, nil, func(ctx context.Context, tx bun.Tx) error {
+		// A new row is inserted into the long_lived_token table, and the ID of the
+		// inserted row is returned and stored in longLivedToken.ID.
+		_, err := db.Bun().NewInsert().
+			Model(longLivedToken).
+			Column("user_id", "expires_at", "created_at").
+			Returning("id").
+			Exec(ctx, &longLivedToken.ID)
+		if err != nil {
+			return err
+		}
+
+		// A Paseto token is generated using the longLivedToken object and the private key.
+		v2 := paseto.NewV2()
+		privateKey := db.GetTokenKeys().PrivateKey
+		token, err := v2.Sign(privateKey, longLivedToken, nil)
+		if err != nil {
+			return errors.Wrap(err, "failed to generate user authentication token")
+		}
+
+		// The token is hashed using model.HashPassword
+		longLivedToken.TokenValue = token
+		hashedToken, err := model.HashPassword(token)
+		if err != nil {
+			return errors.Wrap(err, "error updating user long lived token")
+		}
+		longLivedToken.TokenValueHash = null.StringFrom(hashedToken)
+
+		// The TokenValueHash is updated in the database.
+		_, err = db.Bun().NewUpdate().
+			Model(longLivedToken).
+			Column("token_value_hash").
+			Where("id = (?)", longLivedToken.ID).
+			Exec(ctx)
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+	if err != nil {
+		return "", err
+	}
+
+	return longLivedToken.TokenValue, nil
+}
+
+// DeleteLongLivenTokenByUserID deletes long lived token if found.
+// If not found, the err will be nil, and the number of affected rows will be zero.
+func DeleteLongLivenTokenByUserID(ctx context.Context, userID model.UserID) error {
+	_, err := db.Bun().NewDelete().
+		Table("long_lived_tokens").
+		Where("user_id = ?", userID).
+		Exec(ctx)
+	return err
+}
+
+// DeleteLongLivenTokenByToken deletes long lived token if found.
+func DeleteLongLivenTokenByToken(ctx context.Context, token string) error {
+	v2 := paseto.NewV2()
+	var longLivedToken model.LongLivedToken
+	// Verification will fail when using external token (Jwt instead of Paseto).
+	// Currently passing Paseto token.
+	if err := v2.Verify(token, db.GetTokenKeys().PublicKey, &longLivedToken, nil); err != nil {
+		return nil //nolint: nilerr
+	}
+	return DeleteLongLivedTokenByTokenID(ctx, longLivedToken.ID)
+}
+
+// DeleteLongLivedTokenByTokenID deletes the long lived token with the given token ID.
+// If not found, the err will be nil, and the number of affected rows will be zero.
+func DeleteLongLivedTokenByTokenID(ctx context.Context, longLivedTokenID model.TokenID) error {
+	_, err := db.Bun().NewDelete().
+		Table("long_lived_tokens").
+		Where("id = ?", longLivedTokenID).
+		Exec(ctx)
+	return err
+}

master/internal/user/postgres_long_lived_tokens_intg_test.go

@@ -0,0 +1,81 @@
+//go:build integration
+// +build integration
+
+package user
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/o1egl/paseto"
+	"github.com/stretchr/testify/require"
+
+	"github.com/determined-ai/determined/master/internal/db"
+	"github.com/determined-ai/determined/master/pkg/model"
+)
+
+func TestCreateLongLivedToken(t *testing.T) {
+	user, err := addTestUser(nil)
+	require.NoError(t, err)
+
+	// Add a LongLivedToken.
+	token, err := CreateLongLivedToken(context.TODO(), user)
+	require.NoError(t, err)
+	require.NotNil(t, token)
+
+	exists, err := db.Bun().NewSelect().Table("long_lived_tokens").
+		Where("user_id = ?", user.ID).Exists(context.TODO())
+	require.True(t, exists)
+	require.NoError(t, err)
+}
+
+func TestCreateLongLivedTokenHasExpiresAt(t *testing.T) {
+	user, err := addTestUser(nil)
+	require.NoError(t, err)
+
+	// Add a LongLivedToken with custom (Now() + 3 Months) Expiry Time.
+	expiresAt := time.Now().Add(TokenExpirationDuration * 3)
+	token, err := CreateLongLivedToken(context.TODO(), user, WithTokenExpiresAt(&expiresAt))
+	require.NoError(t, err)
+	require.NotNil(t, token)
+
+	var restoredSession model.LongLivedToken
+	v2 := paseto.NewV2()
+	err = v2.Verify(token, db.GetTokenKeys().PublicKey, &restoredSession, nil)
+	require.NoError(t, err)
+
+	// Strip monotonic clock readings by using time.Equal
+	require.True(t, restoredSession.ExpiresAt.Equal(expiresAt))
+
+	exists, err := db.Bun().NewSelect().Table("long_lived_tokens").
+		Where("user_id = ?", user.ID).Exists(context.TODO())
+	require.True(t, exists)
+	require.NoError(t, err)
+}
+
+func TestDeleteLongLivenTokenByUserID(t *testing.T) {
+	userID, _, _, err := addTestSession()
+	require.NoError(t, err)
+
+	err = DeleteLongLivenTokenByUserID(context.TODO(), userID)
+	require.NoError(t, err)
+
+	exists, err := db.Bun().NewSelect().Table("long_lived_tokens").
+		Where("user_id = ?", userID).Exists(context.TODO())
+	require.False(t, exists)
+	require.NoError(t, err)
+}
+
+func TestDeleteLongLivenTokenByToken(t *testing.T) {
+	userID, _, token, err := addTestSession()
+	require.NoError(t, err)
+
+	err = DeleteLongLivenTokenByToken(context.TODO(), token)
+	require.NoError(t, err)
+
+	exists, err := db.Bun().NewSelect().Table("long_lived_tokens").
+		Where("user_id = ?", userID).Exists(context.TODO())
+	require.False(t, exists)
+	require.NoError(t, err)
+}

master/pkg/model/user.go

@@ -316,3 +316,17 @@ func HashPassword(password string) (string, error) {
 	}
 	return string(passwordHash), nil
 }
+
+// TokenID is the type for user token IDs.
+type TokenID int
+
+// LongLivedToken corresponds to a row in the "long_lived_tokens" DB table.
+type LongLivedToken struct {
+	bun.BaseModel  `bun:"table:long_lived_tokens"`
+	ID             TokenID     `db:"id" json:"id"`
+	UserID         UserID      `db:"user_id" json:"user_id"`
+	TokenValue     string      `bun:"-"`
+	TokenValueHash null.String `db:"token_value_hash" json:"token_value_hash"`
+	ExpiresAt      time.Time   `db:"expires_at" json:"expires_at"`
+	CreatedAt      time.Time   `db:"created_at" json:"created_at"`
+}

master/static/migrations/20240828013336_add-long-lived-tokens-table.tx.up.sql

@@ -0,0 +1,11 @@
+CREATE TABLE public.long_lived_tokens (
+    id integer GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
+    user_id integer NOT NULL,
+    token_value_hash text, -- Hash of the token value for secure storage
+    expires_at TIMESTAMP NOT NULL,
+    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    FOREIGN KEY(user_id) REFERENCES public.users(id) ON DELETE CASCADE
+);
+
+-- Adding an index on token_value_hash for faster lookups
+CREATE INDEX idx_token_value_hash ON long_lived_tokens(token_value_hash);