Fix: properly store renewed access token into grant table

Signed-off-by: Daishan Peng <[email protected]>

Author: StrongMonkey

database/database.go

@@ -499,6 +499,52 @@ func (d *Database) StoreGrant(grant *Grant) error {
 	return err
 }
 
+// UpdateGrant updates an existing grant's properties
+func (d *Database) UpdateGrant(grant *Grant) error {
+	var query string
+	if d.dbType == "postgres" {
+		query = `
+			UPDATE grants 
+			SET scope = $1, metadata = $2, props = $3, expires_at = $4
+			WHERE id = $5 AND user_id = $6
+		`
+	} else {
+		query = `
+			UPDATE grants 
+			SET scope = ?, metadata = ?, props = ?, expires_at = ?
+			WHERE id = ? AND user_id = ?
+		`
+	}
+
+	scope, _ := json.Marshal(grant.Scope)
+	metadata, _ := json.Marshal(grant.Metadata)
+	props, _ := json.Marshal(grant.Props)
+
+	result, err := d.db.Exec(query,
+		scope,
+		metadata,
+		props,
+		grant.ExpiresAt,
+		grant.ID,
+		grant.UserID,
+	)
+	if err != nil {
+		return err
+	}
+
+	// Check if any rows were affected
+	rowsAffected, err := result.RowsAffected()
+	if err != nil {
+		return err
+	}
+
+	if rowsAffected == 0 {
+		return fmt.Errorf("grant not found: id=%s, user_id=%s", grant.ID, grant.UserID)
+	}
+
+	return nil
+}
+
 // GetGrant retrieves a grant by ID and user ID
 func (d *Database) GetGrant(grantID, userID string) (*Grant, error) {
 	var query string

main.go

@@ -327,6 +327,80 @@ func (p *OAuthProxy) decryptPropsIfNeeded(props map[string]interface{}) (map[str
 	return result, nil
 }
 
+// updateGrant updates a grant with new token information
+func (p *OAuthProxy) updateGrant(grantID, userID string, newTokenInfo *providers.TokenInfo) error {
+	// Get the existing grant
+	grant, err := p.db.GetGrant(grantID, userID)
+	if err != nil {
+		return fmt.Errorf("failed to get grant: %w", err)
+	}
+
+	// Prepare sensitive props data
+	sensitiveProps := map[string]interface{}{
+		"access_token":  newTokenInfo.AccessToken,
+		"refresh_token": newTokenInfo.RefreshToken,
+		"expires_at":    newTokenInfo.ExpireAt,
+	}
+
+	// Add existing user info if available
+	if grant.Props != nil {
+		if email, ok := grant.Props["email"].(string); ok {
+			sensitiveProps["email"] = email
+		}
+		if name, ok := grant.Props["name"].(string); ok {
+			sensitiveProps["name"] = name
+		}
+		if userID, ok := grant.Props["user_id"].(string); ok {
+			sensitiveProps["user_id"] = userID
+		}
+	}
+
+	// Initialize props map
+	props := make(map[string]interface{})
+
+	// Check if encryption is enabled
+	if p.encryptionKey != "" {
+		// Decode the encryption key from base64
+		encryptionKey, err := base64.StdEncoding.DecodeString(p.encryptionKey)
+		if err != nil {
+			return fmt.Errorf("failed to decode encryption key: %w", err)
+		}
+
+		// Validate key length (must be 32 bytes for AES-256)
+		if len(encryptionKey) != 32 {
+			return fmt.Errorf("invalid encryption key length: %d bytes (expected 32)", len(encryptionKey))
+		}
+
+		// Encrypt the sensitive props data
+		encryptedProps, err := encryptData(sensitiveProps, encryptionKey)
+		if err != nil {
+			return fmt.Errorf("failed to encrypt props data: %w", err)
+		}
+
+		// Store encrypted data
+		props["encrypted_data"] = encryptedProps.Data
+		props["iv"] = encryptedProps.IV
+		props["algorithm"] = encryptedProps.Algorithm
+		props["encrypted"] = true
+	} else {
+		// Store data in plain text if no encryption key is provided
+		for key, value := range sensitiveProps {
+			props[key] = value
+		}
+		props["encrypted"] = false
+	}
+
+	// Update the grant with new props
+	grant.Props = props
+
+	// Update the grant in the database
+	if err := p.db.UpdateGrant(grant); err != nil {
+		return fmt.Errorf("failed to update grant: %w", err)
+	}
+
+	return nil
+}
+
 // databaseAdapter adapts the database to the tokens.Database interface
 type databaseAdapter struct {
 	db *database.Database
@@ -984,7 +1058,17 @@ func (p *OAuthProxy) mcpProxyHandler(c *gin.Context) {
 						return
 					}
 
-					// Update the token info with the new access token
+					// Update the grant with new token information
+					if err := p.updateGrant(tokenInfo.GrantID, tokenInfo.UserID, newTokenInfo); err != nil {
+						log.Printf("Failed to update grant: %v", err)
+						c.JSON(http.StatusInternalServerError, gin.H{
+							"error":             "server_error",
+							"error_description": "Failed to update grant with new token",
+						})
+						return
+					}
+
+					// Update the token info with the new access token for the current request
 					tokenInfo.Props["access_token"] = newTokenInfo.AccessToken
 					if newTokenInfo.RefreshToken != "" {
 						tokenInfo.Props["refresh_token"] = newTokenInfo.RefreshToken
@@ -1042,6 +1126,7 @@ func (p *OAuthProxy) mcpProxyHandler(c *gin.Context) {
 		},
 		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
 			log.Printf("Proxy error: %v", err)
+			c.Abort()
 		},
 	}
 

tokens/jwt.go

@@ -23,6 +23,7 @@ type Database interface {
 
 type TokenClaims struct {
 	UserID    string                 `json:"user_id"`
+	GrantID   string                 `json:"grant_id"`
 	Props     map[string]interface{} `json:"props,omitempty"`
 	ExpiresAt time.Time              `json:"expires_at"`
 }
@@ -95,6 +96,7 @@ func (tm *TokenManager) ValidateAccessToken(tokenString string) (*TokenClaims, e
 	// Create TokenClaims with the grant's props
 	claims := &TokenClaims{
 		UserID:    userID,
+		GrantID:   grantID,
 		Props:     grant.Props,
 		ExpiresAt: tokenData.ExpiresAt,
 	}
@@ -111,6 +113,7 @@ func (tm *TokenManager) GetTokenInfo(tokenString string) (*TokenInfo, error) {
 
 	return &TokenInfo{
 		UserID:    claims.UserID,
+		GrantID:   claims.GrantID,
 		Props:     claims.Props,
 		ExpiresAt: claims.ExpiresAt,
 	}, nil
@@ -119,6 +122,7 @@ func (tm *TokenManager) GetTokenInfo(tokenString string) (*TokenInfo, error) {
 // TokenInfo represents token information
 type TokenInfo struct {
 	UserID    string
+	GrantID   string
 	Props     map[string]interface{}
 	ExpiresAt time.Time
 }