Skip to content

Commit cc2a61e

Browse files
Genesis929Genesis929
authored
fix(gda): support mTLS and GOOGLE_API_USE_MTLS_ENDPOINT for GDA client (#3460)
## Description The Gemini Data Analytics (GDA) API calls made via custom REST HTTP client lacked dynamic mTLS endpoint routing, violating AIP-4114 and breaking Context Aware Access (CAA). This change fixes the issue by updating the GDA client helper to dynamically select the API endpoint based on standard AIP-4114 environment variables (`GOOGLE_API_USE_CLIENT_CERTIFICATE` and `GOOGLE_API_USE_MTLS_ENDPOINT`). The following tools have been updated to use `NewGDAClient` instead of hardcoded URLs: - `bigquery-conversational-analytics` - `conversational-analytics-ask-data-agent` - `conversational-analytics-get-data-agent-info` - `conversational-analytics-list-accessible-data-agents` - `looker-conversational-analytics` ## Description > Should include a concise description of the changes (bug or feature), it's > impact, along with a summary of the solution ## PR Checklist > Thank you for opening a Pull Request! Before submitting your PR, there are a > few things you can do to make sure it goes smoothly: - [ ] Make sure you reviewed [CONTRIBUTING.md](https://github.com/googleapis/mcp-toolbox/blob/main/CONTRIBUTING.md) - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/mcp-toolbox/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) - [ ] Make sure to add `!` if this involve a breaking change 🛠️ Fixes #
1 parent d6dc5fe commit cc2a61e

8 files changed

Lines changed: 275 additions & 68 deletions

File tree

‎internal/tools/bigquery/bigqueryconversationalanalytics/bigqueryconversationalanalytics.go‎

Lines changed: 19 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,14 @@ import (
3131
"github.com/googleapis/mcp-toolbox/internal/util"
3232
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
3333
"golang.org/x/oauth2"
34+
"google.golang.org/api/option"
3435
)
3536

3637
const resourceType string = "bigquery-conversational-analytics"
3738

38-
const gdaURLFormat = "https://geminidataanalytics.googleapis.com/v1beta/projects/%s/locations/%s:chat"
39+
func getGDAURLFormat() string {
40+
return util.GetGDAEndpoint() + "/v1beta/projects/%s/locations/%s:chat"
41+
}
3942

4043
const instructions = `**INSTRUCTIONS - FOLLOW THESE RULES:**
4144
1. **CONTENT:** Your answer should present the supporting data and then provide a conclusion based on that data.
@@ -155,21 +158,22 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
155158
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
156159
}
157160

158-
var tokenStr string
161+
var tokenSource oauth2.TokenSource
159162

160163
// Get credentials for the API call
161164
if source.UseClientAuthorization() {
162165
// Use client-side access token
163166
if accessToken == "" {
164167
return nil, util.NewClientServerError("tool is configured for client OAuth but no token was provided in the request header", http.StatusUnauthorized, nil)
165168
}
166-
tokenStr, err = accessToken.ParseBearerToken()
169+
tokenStr, err := accessToken.ParseBearerToken()
167170
if err != nil {
168171
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
169172
}
173+
tokenSource = oauth2.StaticTokenSource(&oauth2.Token{AccessToken: tokenStr})
170174
} else {
171175
// Get a token source for the Gemini Data Analytics API.
172-
tokenSource, err := source.BigQueryTokenSourceWithScope(ctx, nil)
176+
tokenSource, err = source.BigQueryTokenSourceWithScope(ctx, nil)
173177
if err != nil {
174178
return nil, util.NewClientServerError("failed to get token source", http.StatusInternalServerError, err)
175179
}
@@ -178,11 +182,6 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
178182
if tokenSource == nil {
179183
return nil, util.NewClientServerError("cloud-platform token source is missing", http.StatusInternalServerError, nil)
180184
}
181-
token, err := tokenSource.Token()
182-
if err != nil {
183-
return nil, util.NewClientServerError("failed to get token from cloud-platform token source", http.StatusInternalServerError, err)
184-
}
185-
tokenStr = token.AccessToken
186185
}
187186

188187
// Extract parameters from the map
@@ -213,12 +212,11 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
213212
if location == "" {
214213
location = "us"
215214
}
216-
caURL := fmt.Sprintf(gdaURLFormat, projectID, location)
215+
caURL := fmt.Sprintf(getGDAURLFormat(), projectID, location)
217216

218217
headers := map[string]string{
219-
source.GetAuthTokenHeaderName(): fmt.Sprintf("Bearer %s", tokenStr),
220-
"Content-Type": "application/json",
221-
"X-Goog-API-Client": util.GDAClientID,
218+
"Content-Type": "application/json",
219+
"X-Goog-API-Client": util.GDAClientID,
222220
}
223221

224222
payload := CAPayload{
@@ -233,8 +231,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
233231
ClientIdEnum: util.GDAClientID,
234232
}
235233

234+
client, err := util.NewGDAClient(ctx, option.WithTokenSource(tokenSource))
235+
if err != nil {
236+
return nil, util.NewClientServerError("failed to create GDA client", http.StatusInternalServerError, err)
237+
}
238+
client.Timeout = 330 * time.Second
239+
236240
// Call the streaming API
237-
response, err := getStream(caURL, payload, headers, source.GetMaxQueryResultRows())
241+
response, err := getStream(client, caURL, payload, headers, source.GetMaxQueryResultRows())
238242
if err != nil {
239243
// getStream wraps network errors or non-200 responses
240244
return nil, util.NewClientServerError("failed to get response from conversational analytics API", http.StatusInternalServerError, err)
@@ -251,7 +255,7 @@ func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (boo
251255
return source.UseClientAuthorization(), nil
252256
}
253257

254-
func getStream(url string, payload CAPayload, headers map[string]string, maxRows int) (string, error) {
258+
func getStream(client *http.Client, url string, payload CAPayload, headers map[string]string, maxRows int) (string, error) {
255259
payloadBytes, err := json.Marshal(payload)
256260
if err != nil {
257261
return "", fmt.Errorf("failed to marshal payload: %w", err)
@@ -265,7 +269,6 @@ func getStream(url string, payload CAPayload, headers map[string]string, maxRows
265269
req.Header.Set(k, v)
266270
}
267271

268-
client := &http.Client{Timeout: 330 * time.Second}
269272
resp, err := client.Do(req)
270273
if err != nil {
271274
return "", fmt.Errorf("failed to send request: %w", err)

‎internal/tools/conversationalanalytics/conversationalanalyticsaskdataagent/conversationalanalyticsaskdataagent.go‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import (
3131
"github.com/googleapis/mcp-toolbox/internal/util"
3232
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
3333
"golang.org/x/oauth2"
34+
"google.golang.org/api/option"
3435
)
3536

3637
const resourceType string = "conversational-analytics-ask-data-agent"
@@ -162,21 +163,22 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
162163
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
163164
}
164165

165-
var tokenStr string
166+
var tokenSource oauth2.TokenSource
166167

167168
// Get credentials for the API call
168169
if source.UseClientAuthorization() {
169170
// Use client-side access token
170171
if accessToken == "" {
171172
return nil, util.NewClientServerError("tool is configured for client OAuth but no token was provided in the request header", http.StatusUnauthorized, nil)
172173
}
173-
tokenStr, err = accessToken.ParseBearerToken()
174+
tokenStr, err := accessToken.ParseBearerToken()
174175
if err != nil {
175176
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
176177
}
178+
tokenSource = oauth2.StaticTokenSource(&oauth2.Token{AccessToken: tokenStr})
177179
} else {
178180
// Get a token source for the Gemini Data Analytics API.
179-
tokenSource, err := source.GoogleCloudTokenSourceWithScope(ctx, "")
181+
tokenSource, err = source.GoogleCloudTokenSourceWithScope(ctx, "")
180182
if err != nil {
181183
return nil, util.NewClientServerError("failed to get token source", http.StatusInternalServerError, err)
182184
}
@@ -185,11 +187,6 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
185187
if tokenSource == nil {
186188
return nil, util.NewClientServerError("cloud-platform token source is missing", http.StatusInternalServerError, nil)
187189
}
188-
token, err := tokenSource.Token()
189-
if err != nil {
190-
return nil, util.NewClientServerError("failed to get token from cloud-platform token source", http.StatusInternalServerError, err)
191-
}
192-
tokenStr = token.AccessToken
193190
}
194191

195192
// Extract parameters from the map
@@ -199,10 +196,9 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
199196

200197
// Construct URL, headers, and payload
201198
projectID := source.GetProjectID()
202-
caURL := fmt.Sprintf("https://geminidataanalytics.googleapis.com/v1beta/projects/%s/locations/%s:chat", projectID, t.Cfg.Location)
199+
caURL := fmt.Sprintf("%s/v1beta/projects/%s/locations/%s:chat", util.GetGDAEndpoint(), projectID, t.Cfg.Location)
203200

204201
headers := map[string]string{
205-
"Authorization": fmt.Sprintf("Bearer %s", tokenStr),
206202
"Content-Type": "application/json",
207203
"X-Goog-API-Client": util.GDAClientID,
208204
}
@@ -218,8 +214,14 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
218214
ClientIdEnum: util.GDAClientID,
219215
}
220216

217+
client, err := util.NewGDAClient(ctx, option.WithTokenSource(tokenSource))
218+
if err != nil {
219+
return nil, util.NewClientServerError("failed to create GDA client", http.StatusInternalServerError, err)
220+
}
221+
client.Timeout = 330 * time.Second
222+
221223
// Call the streaming API
222-
response, err := getStream(caURL, payload, headers, t.Cfg.MaxResults)
224+
response, err := getStream(client, caURL, payload, headers, t.Cfg.MaxResults)
223225
if err != nil {
224226
return nil, util.NewAgentError("failed to get response from conversational analytics API", err)
225227
}
@@ -235,7 +237,7 @@ func (t Tool) RequiresClientAuthorization(resourceMgr tools.SourceProvider) (boo
235237
return source.UseClientAuthorization(), nil
236238
}
237239

238-
func getStream(url string, payload CAPayload, headers map[string]string, maxRows int) (string, error) {
240+
func getStream(client *http.Client, url string, payload CAPayload, headers map[string]string, maxRows int) (string, error) {
239241
payloadBytes, err := json.Marshal(payload)
240242
if err != nil {
241243
return "", fmt.Errorf("failed to marshal payload: %w", err)
@@ -249,7 +251,6 @@ func getStream(url string, payload CAPayload, headers map[string]string, maxRows
249251
req.Header.Set(k, v)
250252
}
251253

252-
client := &http.Client{Timeout: 330 * time.Second}
253254
resp, err := client.Do(req)
254255
if err != nil {
255256
return "", fmt.Errorf("failed to send request: %w", err)

‎internal/tools/conversationalanalytics/conversationalanalyticsgetdataagentinfo/conversationalanalyticsgetdataagentinfo.go‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"io"
2222
"net/http"
2323
"net/url"
24+
"time"
2425

2526
yaml "github.com/goccy/go-yaml"
2627
"github.com/googleapis/mcp-toolbox/internal/sources"
@@ -29,6 +30,7 @@ import (
2930
"github.com/googleapis/mcp-toolbox/internal/util"
3031
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
3132
"golang.org/x/oauth2"
33+
"google.golang.org/api/option"
3234
)
3335

3436
const resourceType string = "conversational-analytics-get-data-agent-info"
@@ -129,21 +131,22 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
129131
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
130132
}
131133

132-
var tokenStr string
134+
var tokenSource oauth2.TokenSource
133135

134136
// Get credentials for the API call
135137
if source.UseClientAuthorization() {
136138
// Use client-side access token
137139
if accessToken == "" {
138140
return nil, util.NewClientServerError("tool is configured for client OAuth but no token was provided in the request header", http.StatusUnauthorized, nil)
139141
}
140-
tokenStr, err = accessToken.ParseBearerToken()
142+
tokenStr, err := accessToken.ParseBearerToken()
141143
if err != nil {
142144
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
143145
}
146+
tokenSource = oauth2.StaticTokenSource(&oauth2.Token{AccessToken: tokenStr})
144147
} else {
145148
// Get a token source for the Gemini Data Analytics API.
146-
tokenSource, err := source.GoogleCloudTokenSourceWithScope(ctx, "")
149+
tokenSource, err = source.GoogleCloudTokenSourceWithScope(ctx, "")
147150
if err != nil {
148151
return nil, util.NewClientServerError("failed to get token source", http.StatusInternalServerError, err)
149152
}
@@ -152,11 +155,6 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
152155
if tokenSource == nil {
153156
return nil, util.NewClientServerError("cloud-platform token source is missing", http.StatusInternalServerError, nil)
154157
}
155-
token, err := tokenSource.Token()
156-
if err != nil {
157-
return nil, util.NewClientServerError("failed to get token from cloud-platform token source", http.StatusInternalServerError, err)
158-
}
159-
tokenStr = token.AccessToken
160158
}
161159

162160
// Extract parameters from the map
@@ -165,16 +163,19 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
165163

166164
// Construct URL
167165
projectID := source.GetProjectID()
168-
caURL := fmt.Sprintf("https://geminidataanalytics.googleapis.com/v1beta/projects/%s/locations/%s/dataAgents/%s", projectID, t.Cfg.Location, url.PathEscape(dataAgentId))
166+
caURL := fmt.Sprintf("%s/v1beta/projects/%s/locations/%s/dataAgents/%s", util.GetGDAEndpoint(), projectID, t.Cfg.Location, url.PathEscape(dataAgentId))
169167

170168
req, err := http.NewRequest("GET", caURL, nil)
171169
if err != nil {
172170
return nil, util.NewClientServerError("failed to create request", http.StatusInternalServerError, err)
173171
}
174-
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", tokenStr))
175172
req.Header.Set("X-Goog-API-Client", util.GDAClientID)
176173

177-
client := &http.Client{}
174+
client, err := util.NewGDAClient(ctx, option.WithTokenSource(tokenSource))
175+
if err != nil {
176+
return nil, util.NewClientServerError("failed to create GDA client", http.StatusInternalServerError, err)
177+
}
178+
client.Timeout = 30 * time.Second
178179
resp, err := client.Do(req)
179180
if err != nil {
180181
return nil, util.NewClientServerError("failed to send request", http.StatusInternalServerError, err)

‎internal/tools/conversationalanalytics/conversationalanalyticslistaccessibledataagents/conversationalanalyticslistaccessibledataagents.go‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"fmt"
2121
"io"
2222
"net/http"
23+
"time"
2324

2425
yaml "github.com/goccy/go-yaml"
2526
"github.com/googleapis/mcp-toolbox/internal/sources"
@@ -28,6 +29,7 @@ import (
2829
"github.com/googleapis/mcp-toolbox/internal/util"
2930
"github.com/googleapis/mcp-toolbox/internal/util/parameters"
3031
"golang.org/x/oauth2"
32+
"google.golang.org/api/option"
3133
)
3234

3335
const resourceType string = "conversational-analytics-list-accessible-data-agents"
@@ -127,21 +129,22 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
127129
return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
128130
}
129131

130-
var tokenStr string
132+
var tokenSource oauth2.TokenSource
131133

132134
// Get credentials for the API call
133135
if source.UseClientAuthorization() {
134136
// Use client-side access token
135137
if accessToken == "" {
136138
return nil, util.NewClientServerError("tool is configured for client OAuth but no token was provided in the request header", http.StatusUnauthorized, nil)
137139
}
138-
tokenStr, err = accessToken.ParseBearerToken()
140+
tokenStr, err := accessToken.ParseBearerToken()
139141
if err != nil {
140142
return nil, util.NewClientServerError("error parsing access token", http.StatusUnauthorized, err)
141143
}
144+
tokenSource = oauth2.StaticTokenSource(&oauth2.Token{AccessToken: tokenStr})
142145
} else {
143146
// Get a token source for the Gemini Data Analytics API.
144-
tokenSource, err := source.GoogleCloudTokenSourceWithScope(ctx, "")
147+
tokenSource, err = source.GoogleCloudTokenSourceWithScope(ctx, "")
145148
if err != nil {
146149
return nil, util.NewClientServerError("failed to get token source", http.StatusInternalServerError, err)
147150
}
@@ -150,25 +153,23 @@ func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, para
150153
if tokenSource == nil {
151154
return nil, util.NewClientServerError("cloud-platform token source is missing", http.StatusInternalServerError, nil)
152155
}
153-
token, err := tokenSource.Token()
154-
if err != nil {
155-
return nil, util.NewClientServerError("failed to get token from cloud-platform token source", http.StatusInternalServerError, err)
156-
}
157-
tokenStr = token.AccessToken
158156
}
159157

160158
// Construct URL
161159
projectID := source.GetProjectID()
162-
caURL := fmt.Sprintf("https://geminidataanalytics.googleapis.com/v1beta/projects/%s/locations/%s/dataAgents:listAccessible", projectID, t.Cfg.Location)
160+
caURL := fmt.Sprintf("%s/v1beta/projects/%s/locations/%s/dataAgents:listAccessible", util.GetGDAEndpoint(), projectID, t.Cfg.Location)
163161

164162
req, err := http.NewRequest("GET", caURL, nil)
165163
if err != nil {
166164
return nil, util.NewClientServerError("failed to create request", http.StatusInternalServerError, err)
167165
}
168-
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", tokenStr))
169166
req.Header.Set("X-Goog-API-Client", util.GDAClientID)
170167

171-
client := &http.Client{}
168+
client, err := util.NewGDAClient(ctx, option.WithTokenSource(tokenSource))
169+
if err != nil {
170+
return nil, util.NewClientServerError("failed to create GDA client", http.StatusInternalServerError, err)
171+
}
172+
client.Timeout = 30 * time.Second
172173
resp, err := client.Do(req)
173174
if err != nil {
174175
return nil, util.NewClientServerError("failed to send request", http.StatusInternalServerError, err)

0 commit comments

Comments
 (0)