redirect to desired page after authenticating

internal/domains/authentication/auth.go

@@ -100,27 +100,28 @@ func (a *Authenticator) RunBackgroundCleanup(ctx context.Context) error {
 }
 
 // Exchange exchanges an auth code for an access token.
-func (a *Authenticator) Exchange(ctx context.Context, state, code string) (accessToken string, expiration time.Time, err error) {
+func (a *Authenticator) Exchange(ctx context.Context, state, code string) (accessToken, targetURI string, expiration time.Time, err error) {
 	// validate state
 
-	if exp, err := a.GetStateExpiration(ctx, state); errors.Is(err, consts.ErrNotFound) {
-		return "", time.Time{}, fmt.Errorf("invalid state: %w", consts.ErrNotFound)
+	exp, targetURI, err := a.GetStateExpirationAndURL(ctx, state)
+	if errors.Is(err, consts.ErrNotFound) {
+		return "", "", time.Time{}, fmt.Errorf("invalid state: %w", consts.ErrNotFound)
 	} else if err != nil {
-		return "", time.Time{}, fmt.Errorf("failed to load state expiration: %w", err)
+		return "", "", time.Time{}, fmt.Errorf("failed to load state expiration: %w", err)
 	} else if exp.Before(time.Now()) {
-		return "", time.Time{}, fmt.Errorf("invalid state: state expired")
+		return "", "", time.Time{}, fmt.Errorf("invalid state: state expired")
 	}
 
 	// obtain token and profile
 
 	tkn, err := a.Config.Exchange(ctx, code)
 	if err != nil {
-		return "", time.Time{}, fmt.Errorf("failed to exchange an authorization code for a token: %w", err)
+		return "", "", time.Time{}, fmt.Errorf("failed to exchange an authorization code for a token: %w", err)
 	}
 
 	idToken, claims, err := a.verifyIDTokenAndClaimsFromToken(ctx, tkn)
 	if err != nil {
-		return "", time.Time{}, err
+		return "", "", time.Time{}, err
 	}
 
 	claimsJSON, _ := json.Marshal(claims)
@@ -200,10 +201,10 @@ func (a *Authenticator) Exchange(ctx context.Context, state, code string) (acces
 			"claims": json.RawMessage(claimsJSON),
 		},
 	); err != nil {
-		return "", time.Time{}, fmt.Errorf("failed to perform query to save tokens: %w", err)
+		return "", "", time.Time{}, fmt.Errorf("failed to perform query to save tokens: %w", err)
 	}
 
-	return tkn.AccessToken, tkn.Expiry.UTC(), nil
+	return tkn.AccessToken, targetURI, tkn.Expiry.UTC(), nil
 }
 
 // VerifyIDToken verifies that an *oauth2.Token is a valid *oidc.IDToken.

internal/domains/authentication/store.go

@@ -10,10 +10,11 @@ import (
 	"time"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgtype"
 )
 
 // NewState creates a new state for logging in, saving it in the database.
-func (a *Authenticator) NewState(ctx context.Context) ([32]byte, error) {
+func (a *Authenticator) NewState(ctx context.Context, targetURI string) ([32]byte, error) {
 	state, err := generateRandomState()
 	if err != nil {
 		return state, fmt.Errorf("failed to generate random state: %w", err)
@@ -21,9 +22,20 @@ func (a *Authenticator) NewState(ctx context.Context) ([32]byte, error) {
 
 	if _, err = a.db.Exec(
 		ctx,
-		"INSERT INTO oauth_login_states (state) VALUES (@state)",
+		`
+			INSERT INTO
+				oauth_login_states (
+					state,
+					target_uri
+				)
+			VALUES (
+				@state,
+				@target_uri
+			)
+		`,
 		pgx.NamedArgs{
-			"state": state[:],
+			"state":      state[:],
+			"target_uri": targetURI,
 		},
 	); err != nil {
 		return state, fmt.Errorf("failed to execute query: %w", err)
@@ -38,31 +50,41 @@ func generateRandomState() ([32]byte, error) {
 	return b, err
 }
 
-// GetStateExpiration get's the oauth state's expiration
-// func (a *Authenticator) GetStateExpiration(ctx context.Context, state [32]byte) (time.Time, error) {
-func (a *Authenticator) GetStateExpiration(ctx context.Context, state string) (time.Time, error) {
+// GetStateExpirationAndURL get's the oauth state's expiration
+func (a *Authenticator) GetStateExpirationAndURL(ctx context.Context, state string) (time.Time, string, error) {
 	rows, err := a.db.Query(
 		ctx,
-		`SELECT expiration from oauth_login_states WHERE state = ('\x' || @state)::BYTEA`,
+		`
+			SELECT
+				expiration,
+				target_uri
+			FROM
+				oauth_login_states
+			WHERE
+				state = ('\x' || @state)::BYTEA`,
 		pgx.NamedArgs{
-			//"state": state[:],
 			"state": state,
 		},
 	)
 	if err != nil {
-		return time.Time{}, fmt.Errorf("failed to perform query: %w", err)
+		return time.Time{}, "", fmt.Errorf("failed to perform query: %w", err)
 	}
 
-	exp, err := pgx.CollectExactlyOneRow(rows, pgx.RowTo[time.Time])
+	type Row struct {
+		Expiration time.Time
+		Target_uri pgtype.Text
+	}
+
+	r, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByNameLax[Row])
 	if err != nil {
 		if errors.Is(err, pgx.ErrNoRows) {
-			return time.Time{}, consts.ErrNotFound
+			return time.Time{}, "", consts.ErrNotFound
 		}
 
-		return time.Time{}, fmt.Errorf("failed to scan row: %w", err)
+		return time.Time{}, "", fmt.Errorf("failed to scan row: %w", err)
 	}
 
-	return exp, nil
+	return r.Expiration, r.Target_uri.String, nil
 }
 
 // TODO: need to automatically clean up expired tokens