🐛 fix: fix vertex AI credential acquisition failure when using a proxy (

#300)
MartialBE · Jul 16, 2024 · 5ded841 · 5ded841
1 parent 7a6e928
commit 5ded841
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/common/utils/proxy.go b/common/utils/proxy.go
@@ -16,6 +16,7 @@ type ContextKey string
 
 const ProxyHTTPAddrKey ContextKey = "proxyHttpAddr"
 const ProxySock5AddrKey ContextKey = "proxySock5Addr"
+const ProxyAddrKey ContextKey = "proxyAddr"
 
 func ProxyFunc(req *http.Request) (*url.URL, error) {
 	proxyAddr := req.Context().Value(ProxyHTTPAddrKey)

diff --git a/providers/vertexai/base.go b/providers/vertexai/base.go
@@ -6,10 +6,13 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
+	"net/url"
 	"one-api/common/cache"
 	"one-api/common/logger"
 	"one-api/common/requester"
+	"one-api/common/utils"
 	"one-api/model"
 	"one-api/providers/base"
 	"one-api/providers/vertexai/category"
@@ -19,7 +22,9 @@ import (
 
 	credentials "cloud.google.com/go/iam/credentials/apiv1"
 	"cloud.google.com/go/iam/credentials/apiv1/credentialspb"
+	"golang.org/x/net/proxy"
 	"google.golang.org/api/option"
+	"google.golang.org/grpc"
 )
 
 const TokenCacheKey = "api_token:vertexai"
@@ -104,7 +109,11 @@ func (p *VertexAIProvider) GetToken() (string, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
-	client, err := credentials.NewIamCredentialsClient(ctx, option.WithCredentialsJSON([]byte(p.Channel.Key)))
+	if p.Channel.Proxy != nil && *p.Channel.Proxy != "" {
+		ctx = context.WithValue(ctx, utils.ProxyAddrKey, *p.Channel.Proxy)
+	}
+
+	client, err := credentials.NewIamCredentialsClient(ctx, option.WithCredentialsJSON([]byte(p.Channel.Key)), option.WithGRPCDialOption(grpc.WithContextDialer(customDialer)))
 	if err != nil {
 		return "", fmt.Errorf("failed to create IAM credentials client: %w", err)
 	}
@@ -163,3 +172,24 @@ func errorHandle(vertexaiError *VertexaiError) *types.OpenAIError {
 		Code:    vertexaiError.Error.Code,
 	}
 }
+
+func customDialer(ctx context.Context, addr string) (net.Conn, error) {
+	proxyAddress, ok := ctx.Value(utils.ProxyAddrKey).(string)
+	if !ok || proxyAddress == "" {
+		return net.Dial("tcp", addr)
+	}
+
+	proxyURL, err := url.Parse(proxyAddress)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing proxy address: %w", err)
+	}
+
+	dialer := &net.Dialer{}
+
+	dialerProxy, err := proxy.FromURL(proxyURL, dialer)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create HTTP dialer: %v", err)
+	}
+
+	return dialerProxy.Dial("tcp", addr)
+}