From 1967166b6e13d3934dda2dd50cc9fb0d61968f8d Mon Sep 17 00:00:00 2001 From: Daishan Peng Date: Thu, 18 Sep 2025 12:33:11 -0700 Subject: [PATCH] Feat: add salesforce oauth and header support Signed-off-by: Daishan Peng --- pkg/oauth/callback/callback.go | 11 +++++++++++ pkg/proxy/proxy.go | 7 +++++++ 2 files changed, 18 insertions(+) diff --git a/pkg/oauth/callback/callback.go b/pkg/oauth/callback/callback.go index 504c030..08dfea1 100644 --- a/pkg/oauth/callback/callback.go +++ b/pkg/oauth/callback/callback.go @@ -14,6 +14,7 @@ import ( "github.com/obot-platform/mcp-oauth-proxy/pkg/mcpui" "github.com/obot-platform/mcp-oauth-proxy/pkg/providers" "github.com/obot-platform/mcp-oauth-proxy/pkg/types" + "golang.org/x/oauth2" ) type Store interface { @@ -195,6 +196,8 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { "expires_at": tokenInfo.Expiry.Unix(), } + extractSpecialProps(tokenInfo, sensitiveProps) + // Only add user info if we have it if needsUserInfo { sensitiveProps["email"] = userInfo.Email @@ -328,3 +331,11 @@ func (p *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { w.Header().Set("Location", parsedURL.String()) w.WriteHeader(http.StatusFound) } + +func extractSpecialProps(tokenInfo *oauth2.Token, props map[string]any) { + // Salesforce instance URL + url := tokenInfo.Extra("instance_url") + if v, ok := url.(string); ok && v != "" { + props["instance_url"] = v + } +} diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go index e027c0d..119994b 100644 --- a/pkg/proxy/proxy.go +++ b/pkg/proxy/proxy.go @@ -492,6 +492,13 @@ func setHeaders(header http.Header, props map[string]any) { } else { header.Del("X-Forwarded-Access-Token") } + + // Salesforce instance URL + if instanceURL, ok := props["instance_url"].(string); ok { + header.Set("X-Forwarded-Instance-URL", instanceURL) + } else { + header.Del("X-Forwarded-Instance-URL") + } } // updateGrant updates a grant with new token information