初始化导入 new-api 源码
This commit is contained in:
OpenClaw Task Bot
422
relay/channel/vertex/adaptor.go
Normal file
422
relay/channel/vertex/adaptor.go
Normal file
@@ -0,0 +1,422 @@
|
||||
package vertex
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/claude"
|
||||
"github.com/QuantumNous/new-api/relay/channel/gemini"
|
||||
"github.com/QuantumNous/new-api/relay/channel/openai"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
"github.com/QuantumNous/new-api/setting/reasoning"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/lo"
|
||||
)
|
||||
|
||||
const (
|
||||
RequestModeClaude = 1
|
||||
RequestModeGemini = 2
|
||||
RequestModeOpenSource = 3
|
||||
)
|
||||
|
||||
var claudeModelMap = map[string]string{
|
||||
"claude-3-sonnet-20240229": "claude-3-sonnet@20240229",
|
||||
"claude-3-opus-20240229": "claude-3-opus@20240229",
|
||||
"claude-3-haiku-20240307": "claude-3-haiku@20240307",
|
||||
"claude-3-5-sonnet-20240620": "claude-3-5-sonnet@20240620",
|
||||
"claude-3-5-sonnet-20241022": "claude-3-5-sonnet-v2@20241022",
|
||||
"claude-3-7-sonnet-20250219": "claude-3-7-sonnet@20250219",
|
||||
"claude-sonnet-4-20250514": "claude-sonnet-4@20250514",
|
||||
"claude-opus-4-20250514": "claude-opus-4@20250514",
|
||||
"claude-opus-4-1-20250805": "claude-opus-4-1@20250805",
|
||||
"claude-sonnet-4-5-20250929": "claude-sonnet-4-5@20250929",
|
||||
"claude-haiku-4-5-20251001": "claude-haiku-4-5@20251001",
|
||||
"claude-opus-4-5-20251101": "claude-opus-4-5@20251101",
|
||||
"claude-opus-4-6": "claude-opus-4-6",
|
||||
}
|
||||
|
||||
const anthropicVersion = "vertex-2023-10-16"
|
||||
|
||||
type Adaptor struct {
|
||||
RequestMode int
|
||||
AccountCredentials Credentials
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertGeminiRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeminiChatRequest) (any, error) {
|
||||
// Vertex AI does not support functionResponse.id; keep it stripped here for consistency.
|
||||
if model_setting.GetGeminiSettings().RemoveFunctionResponseIdEnabled {
|
||||
removeFunctionResponseID(request)
|
||||
}
|
||||
geminiAdaptor := gemini.Adaptor{}
|
||||
return geminiAdaptor.ConvertGeminiRequest(c, info, request)
|
||||
}
|
||||
|
||||
func removeFunctionResponseID(request *dto.GeminiChatRequest) {
|
||||
if request == nil {
|
||||
return
|
||||
}
|
||||
|
||||
if len(request.Contents) > 0 {
|
||||
for i := range request.Contents {
|
||||
if len(request.Contents[i].Parts) == 0 {
|
||||
continue
|
||||
}
|
||||
for j := range request.Contents[i].Parts {
|
||||
part := &request.Contents[i].Parts[j]
|
||||
if part.FunctionResponse == nil {
|
||||
continue
|
||||
}
|
||||
if len(part.FunctionResponse.ID) > 0 {
|
||||
part.FunctionResponse.ID = nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(request.Requests) > 0 {
|
||||
for i := range request.Requests {
|
||||
removeFunctionResponseID(&request.Requests[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) {
|
||||
if v, ok := claudeModelMap[info.UpstreamModelName]; ok {
|
||||
c.Set("request_model", v)
|
||||
} else {
|
||||
c.Set("request_model", request.Model)
|
||||
}
|
||||
vertexClaudeReq := copyRequest(request, anthropicVersion)
|
||||
return vertexClaudeReq, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
|
||||
//TODO implement me
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
|
||||
geminiAdaptor := gemini.Adaptor{}
|
||||
return geminiAdaptor.ConvertImageRequest(c, info, request)
|
||||
}
|
||||
|
||||
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
|
||||
if strings.HasPrefix(info.UpstreamModelName, "claude") {
|
||||
a.RequestMode = RequestModeClaude
|
||||
} else if strings.Contains(info.UpstreamModelName, "llama") ||
|
||||
// open source models
|
||||
strings.Contains(info.UpstreamModelName, "-maas") {
|
||||
a.RequestMode = RequestModeOpenSource
|
||||
} else {
|
||||
a.RequestMode = RequestModeGemini
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adaptor) getRequestUrl(info *relaycommon.RelayInfo, modelName, suffix string) (string, error) {
|
||||
region := GetModelRegion(info.ApiVersion, info.OriginModelName)
|
||||
if info.ChannelOtherSettings.VertexKeyType != dto.VertexKeyTypeAPIKey {
|
||||
adc := &Credentials{}
|
||||
if err := common.Unmarshal([]byte(info.ApiKey), adc); err != nil {
|
||||
return "", fmt.Errorf("failed to decode credentials file: %w", err)
|
||||
}
|
||||
a.AccountCredentials = *adc
|
||||
|
||||
if a.RequestMode == RequestModeGemini {
|
||||
if region == "global" {
|
||||
return fmt.Sprintf(
|
||||
"https://aiplatform.googleapis.com/v1/projects/%s/locations/global/publishers/google/models/%s:%s",
|
||||
adc.ProjectID,
|
||||
modelName,
|
||||
suffix,
|
||||
), nil
|
||||
} else {
|
||||
return fmt.Sprintf(
|
||||
"https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/google/models/%s:%s",
|
||||
region,
|
||||
adc.ProjectID,
|
||||
region,
|
||||
modelName,
|
||||
suffix,
|
||||
), nil
|
||||
}
|
||||
} else if a.RequestMode == RequestModeClaude {
|
||||
if region == "global" {
|
||||
return fmt.Sprintf(
|
||||
"https://aiplatform.googleapis.com/v1/projects/%s/locations/global/publishers/anthropic/models/%s:%s",
|
||||
adc.ProjectID,
|
||||
modelName,
|
||||
suffix,
|
||||
), nil
|
||||
} else {
|
||||
return fmt.Sprintf(
|
||||
"https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:%s",
|
||||
region,
|
||||
adc.ProjectID,
|
||||
region,
|
||||
modelName,
|
||||
suffix,
|
||||
), nil
|
||||
}
|
||||
} else if a.RequestMode == RequestModeOpenSource {
|
||||
return fmt.Sprintf(
|
||||
"https://aiplatform.googleapis.com/v1beta1/projects/%s/locations/%s/endpoints/openapi/chat/completions",
|
||||
adc.ProjectID,
|
||||
region,
|
||||
), nil
|
||||
}
|
||||
} else {
|
||||
var keyPrefix string
|
||||
if strings.HasSuffix(suffix, "?alt=sse") {
|
||||
keyPrefix = "&"
|
||||
} else {
|
||||
keyPrefix = "?"
|
||||
}
|
||||
if region == "global" {
|
||||
return fmt.Sprintf(
|
||||
"https://aiplatform.googleapis.com/v1/publishers/google/models/%s:%s%skey=%s",
|
||||
modelName,
|
||||
suffix,
|
||||
keyPrefix,
|
||||
info.ApiKey,
|
||||
), nil
|
||||
} else {
|
||||
return fmt.Sprintf(
|
||||
"https://%s-aiplatform.googleapis.com/v1/publishers/google/models/%s:%s%skey=%s",
|
||||
region,
|
||||
modelName,
|
||||
suffix,
|
||||
keyPrefix,
|
||||
info.ApiKey,
|
||||
), nil
|
||||
}
|
||||
}
|
||||
return "", errors.New("unsupported request mode")
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
suffix := ""
|
||||
if a.RequestMode == RequestModeGemini {
|
||||
if model_setting.GetGeminiSettings().ThinkingAdapterEnabled &&
|
||||
!model_setting.ShouldPreserveThinkingSuffix(info.OriginModelName) {
|
||||
// 新增逻辑：处理 -thinking-<budget> 格式
|
||||
if strings.Contains(info.UpstreamModelName, "-thinking-") {
|
||||
parts := strings.Split(info.UpstreamModelName, "-thinking-")
|
||||
info.UpstreamModelName = parts[0]
|
||||
} else if strings.HasSuffix(info.UpstreamModelName, "-thinking") { // 旧的适配
|
||||
info.UpstreamModelName = strings.TrimSuffix(info.UpstreamModelName, "-thinking")
|
||||
} else if strings.HasSuffix(info.UpstreamModelName, "-nothinking") {
|
||||
info.UpstreamModelName = strings.TrimSuffix(info.UpstreamModelName, "-nothinking")
|
||||
} else if baseModel, level, ok := reasoning.TrimEffortSuffix(info.UpstreamModelName); ok && level != "" {
|
||||
info.UpstreamModelName = baseModel
|
||||
}
|
||||
}
|
||||
|
||||
if info.IsStream {
|
||||
suffix = "streamGenerateContent?alt=sse"
|
||||
} else {
|
||||
suffix = "generateContent"
|
||||
}
|
||||
|
||||
if strings.HasPrefix(info.UpstreamModelName, "imagen") {
|
||||
suffix = "predict"
|
||||
}
|
||||
return a.getRequestUrl(info, info.UpstreamModelName, suffix)
|
||||
} else if a.RequestMode == RequestModeClaude {
|
||||
if info.IsStream {
|
||||
suffix = "streamRawPredict?alt=sse"
|
||||
} else {
|
||||
suffix = "rawPredict"
|
||||
}
|
||||
model := info.UpstreamModelName
|
||||
if v, ok := claudeModelMap[info.UpstreamModelName]; ok {
|
||||
model = v
|
||||
}
|
||||
return a.getRequestUrl(info, model, suffix)
|
||||
} else if a.RequestMode == RequestModeOpenSource {
|
||||
return a.getRequestUrl(info, "", "")
|
||||
}
|
||||
return "", errors.New("unsupported request mode")
|
||||
}
|
||||
|
||||
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
|
||||
channel.SetupApiRequestHeader(info, c, req)
|
||||
if info.ChannelOtherSettings.VertexKeyType != dto.VertexKeyTypeAPIKey {
|
||||
accessToken, err := getAccessToken(a, info)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Set("Authorization", "Bearer "+accessToken)
|
||||
}
|
||||
if a.AccountCredentials.ProjectID != "" {
|
||||
req.Set("x-goog-user-project", a.AccountCredentials.ProjectID)
|
||||
}
|
||||
if strings.Contains(info.UpstreamModelName, "claude") {
|
||||
claude.CommonClaudeHeadersOperation(c, req, info)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
|
||||
if request == nil {
|
||||
return nil, errors.New("request is nil")
|
||||
}
|
||||
if a.RequestMode == RequestModeGemini && strings.HasPrefix(info.UpstreamModelName, "imagen") {
|
||||
prompt := ""
|
||||
for _, m := range request.Messages {
|
||||
if m.Role == "user" {
|
||||
prompt = m.StringContent()
|
||||
if prompt != "" {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if prompt == "" {
|
||||
if p, ok := request.Prompt.(string); ok {
|
||||
prompt = p
|
||||
}
|
||||
}
|
||||
if prompt == "" {
|
||||
return nil, errors.New("prompt is required for image generation")
|
||||
}
|
||||
|
||||
imgReq := dto.ImageRequest{
|
||||
Model: request.Model,
|
||||
Prompt: prompt,
|
||||
N: lo.ToPtr(uint(1)),
|
||||
Size: "1024x1024",
|
||||
}
|
||||
if request.N != nil && *request.N > 0 {
|
||||
imgReq.N = lo.ToPtr(uint(*request.N))
|
||||
}
|
||||
if request.Size != "" {
|
||||
imgReq.Size = request.Size
|
||||
}
|
||||
if len(request.ExtraBody) > 0 {
|
||||
var extra map[string]any
|
||||
if err := json.Unmarshal(request.ExtraBody, &extra); err == nil {
|
||||
if n, ok := extra["n"].(float64); ok && n > 0 {
|
||||
imgReq.N = lo.ToPtr(uint(n))
|
||||
}
|
||||
if size, ok := extra["size"].(string); ok {
|
||||
imgReq.Size = size
|
||||
}
|
||||
// accept aspectRatio in extra body (top-level or under parameters)
|
||||
if ar, ok := extra["aspectRatio"].(string); ok && ar != "" {
|
||||
imgReq.Size = ar
|
||||
}
|
||||
if params, ok := extra["parameters"].(map[string]any); ok {
|
||||
if ar, ok := params["aspectRatio"].(string); ok && ar != "" {
|
||||
imgReq.Size = ar
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
c.Set("request_model", request.Model)
|
||||
return a.ConvertImageRequest(c, info, imgReq)
|
||||
}
|
||||
if a.RequestMode == RequestModeClaude {
|
||||
claudeReq, err := claude.RequestOpenAI2ClaudeMessage(c, *request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
vertexClaudeReq := copyRequest(claudeReq, anthropicVersion)
|
||||
c.Set("request_model", claudeReq.Model)
|
||||
info.UpstreamModelName = claudeReq.Model
|
||||
return vertexClaudeReq, nil
|
||||
} else if a.RequestMode == RequestModeGemini {
|
||||
geminiRequest, err := gemini.CovertOpenAI2Gemini(c, *request, info)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.Set("request_model", request.Model)
|
||||
return geminiRequest, nil
|
||||
} else if a.RequestMode == RequestModeOpenSource {
|
||||
return request, nil
|
||||
}
|
||||
return nil, errors.New("unsupported request mode")
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
|
||||
//TODO implement me
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
|
||||
// TODO implement me
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
|
||||
return channel.DoApiRequest(a, c, info, requestBody)
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
|
||||
claudeAdaptor := claude.Adaptor{}
|
||||
if info.IsStream {
|
||||
switch a.RequestMode {
|
||||
case RequestModeClaude:
|
||||
return claudeAdaptor.DoResponse(c, resp, info)
|
||||
case RequestModeGemini:
|
||||
if info.RelayMode == constant.RelayModeGemini {
|
||||
return gemini.GeminiTextGenerationStreamHandler(c, info, resp)
|
||||
} else {
|
||||
return gemini.GeminiChatStreamHandler(c, info, resp)
|
||||
}
|
||||
case RequestModeOpenSource:
|
||||
return openai.OaiStreamHandler(c, info, resp)
|
||||
}
|
||||
} else {
|
||||
switch a.RequestMode {
|
||||
case RequestModeClaude:
|
||||
return claudeAdaptor.DoResponse(c, resp, info)
|
||||
case RequestModeGemini:
|
||||
if info.RelayMode == constant.RelayModeGemini {
|
||||
return gemini.GeminiTextGenerationHandler(c, info, resp)
|
||||
} else {
|
||||
if strings.HasPrefix(info.UpstreamModelName, "imagen") {
|
||||
return gemini.GeminiImageHandler(c, info, resp)
|
||||
}
|
||||
return gemini.GeminiChatHandler(c, info, resp)
|
||||
}
|
||||
case RequestModeOpenSource:
|
||||
return openai.OpenaiHandler(c, info, resp)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetModelList() []string {
|
||||
var modelList []string
|
||||
for i, s := range ModelList {
|
||||
modelList = append(modelList, s)
|
||||
ModelList[i] = s
|
||||
}
|
||||
for i, s := range claude.ModelList {
|
||||
modelList = append(modelList, s)
|
||||
claude.ModelList[i] = s
|
||||
}
|
||||
for i, s := range gemini.ModelList {
|
||||
modelList = append(modelList, s)
|
||||
gemini.ModelList[i] = s
|
||||
}
|
||||
return modelList
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetChannelName() string {
|
||||
return ChannelName
|
||||
}
|
||||
15
relay/channel/vertex/constants.go
Normal file
15
relay/channel/vertex/constants.go
Normal file
@@ -0,0 +1,15 @@
|
||||
package vertex
|
||||
|
||||
var ModelList = []string{
|
||||
//"claude-3-sonnet-20240229",
|
||||
//"claude-3-opus-20240229",
|
||||
//"claude-3-haiku-20240307",
|
||||
//"claude-3-5-sonnet-20240620",
|
||||
|
||||
//"gemini-1.5-pro-latest", "gemini-1.5-flash-latest",
|
||||
//"gemini-1.5-pro-001", "gemini-1.5-flash-001", "gemini-pro", "gemini-pro-vision",
|
||||
|
||||
"meta/llama3-405b-instruct-maas",
|
||||
}
|
||||
|
||||
var ChannelName = "vertex-ai"
|
||||
42
relay/channel/vertex/dto.go
Normal file
42
relay/channel/vertex/dto.go
Normal file
@@ -0,0 +1,42 @@
|
||||
package vertex
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
)
|
||||
|
||||
type VertexAIClaudeRequest struct {
|
||||
AnthropicVersion string `json:"anthropic_version"`
|
||||
Messages []dto.ClaudeMessage `json:"messages"`
|
||||
System any `json:"system,omitempty"`
|
||||
MaxTokens *uint `json:"max_tokens,omitempty"`
|
||||
StopSequences []string `json:"stop_sequences,omitempty"`
|
||||
Stream *bool `json:"stream,omitempty"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
TopP *float64 `json:"top_p,omitempty"`
|
||||
TopK *int `json:"top_k,omitempty"`
|
||||
Tools any `json:"tools,omitempty"`
|
||||
ToolChoice any `json:"tool_choice,omitempty"`
|
||||
Thinking *dto.Thinking `json:"thinking,omitempty"`
|
||||
OutputConfig json.RawMessage `json:"output_config,omitempty"`
|
||||
//Metadata json.RawMessage `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
func copyRequest(req *dto.ClaudeRequest, version string) *VertexAIClaudeRequest {
|
||||
return &VertexAIClaudeRequest{
|
||||
AnthropicVersion: version,
|
||||
System: req.System,
|
||||
Messages: req.Messages,
|
||||
MaxTokens: req.MaxTokens,
|
||||
Stream: req.Stream,
|
||||
Temperature: req.Temperature,
|
||||
TopP: req.TopP,
|
||||
TopK: req.TopK,
|
||||
StopSequences: req.StopSequences,
|
||||
Tools: req.Tools,
|
||||
ToolChoice: req.ToolChoice,
|
||||
Thinking: req.Thinking,
|
||||
OutputConfig: req.OutputConfig,
|
||||
}
|
||||
}
|
||||
22
relay/channel/vertex/relay-vertex.go
Normal file
22
relay/channel/vertex/relay-vertex.go
Normal file
@@ -0,0 +1,22 @@
|
||||
package vertex
|
||||
|
||||
import "github.com/QuantumNous/new-api/common"
|
||||
|
||||
func GetModelRegion(other string, localModelName string) string {
|
||||
// if other is json string
|
||||
if common.IsJsonObject(other) {
|
||||
m, err := common.StrToMap(other)
|
||||
if err != nil {
|
||||
return other // return original if parsing fails
|
||||
}
|
||||
if m[localModelName] != nil {
|
||||
return m[localModelName].(string)
|
||||
} else {
|
||||
if v, ok := m["default"]; ok {
|
||||
return v.(string)
|
||||
}
|
||||
return "global"
|
||||
}
|
||||
}
|
||||
return other
|
||||
}
|
||||
183
relay/channel/vertex/service_account.go
Normal file
183
relay/channel/vertex/service_account.go
Normal file
@@ -0,0 +1,183 @@
|
||||
package vertex
|
||||
|
||||
import (
|
||||
"crypto/rsa"
|
||||
"crypto/x509"
|
||||
"encoding/json"
|
||||
"encoding/pem"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/bytedance/gopkg/cache/asynccache"
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Credentials struct {
|
||||
ProjectID string `json:"project_id"`
|
||||
PrivateKeyID string `json:"private_key_id"`
|
||||
PrivateKey string `json:"private_key"`
|
||||
ClientEmail string `json:"client_email"`
|
||||
ClientID string `json:"client_id"`
|
||||
}
|
||||
|
||||
var Cache = asynccache.NewAsyncCache(asynccache.Options{
|
||||
RefreshDuration: time.Minute * 35,
|
||||
EnableExpire: true,
|
||||
ExpireDuration: time.Minute * 30,
|
||||
Fetcher: func(key string) (interface{}, error) {
|
||||
return nil, errors.New("not found")
|
||||
},
|
||||
})
|
||||
|
||||
func getAccessToken(a *Adaptor, info *relaycommon.RelayInfo) (string, error) {
|
||||
var cacheKey string
|
||||
if info.ChannelIsMultiKey {
|
||||
cacheKey = fmt.Sprintf("access-token-%d-%d", info.ChannelId, info.ChannelMultiKeyIndex)
|
||||
} else {
|
||||
cacheKey = fmt.Sprintf("access-token-%d", info.ChannelId)
|
||||
}
|
||||
val, err := Cache.Get(cacheKey)
|
||||
if err == nil {
|
||||
return val.(string), nil
|
||||
}
|
||||
|
||||
signedJWT, err := createSignedJWT(a.AccountCredentials.ClientEmail, a.AccountCredentials.PrivateKey)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to create signed JWT: %w", err)
|
||||
}
|
||||
newToken, err := exchangeJwtForAccessToken(signedJWT, info)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to exchange JWT for access token: %w", err)
|
||||
}
|
||||
if err := Cache.SetDefault(cacheKey, newToken); err {
|
||||
return newToken, nil
|
||||
}
|
||||
return newToken, nil
|
||||
}
|
||||
|
||||
func createSignedJWT(email, privateKeyPEM string) (string, error) {
|
||||
|
||||
privateKeyPEM = strings.ReplaceAll(privateKeyPEM, "-----BEGIN PRIVATE KEY-----", "")
|
||||
privateKeyPEM = strings.ReplaceAll(privateKeyPEM, "-----END PRIVATE KEY-----", "")
|
||||
privateKeyPEM = strings.ReplaceAll(privateKeyPEM, "\r", "")
|
||||
privateKeyPEM = strings.ReplaceAll(privateKeyPEM, "\n", "")
|
||||
privateKeyPEM = strings.ReplaceAll(privateKeyPEM, "\\n", "")
|
||||
|
||||
block, _ := pem.Decode([]byte("-----BEGIN PRIVATE KEY-----\n" + privateKeyPEM + "\n-----END PRIVATE KEY-----"))
|
||||
if block == nil {
|
||||
return "", fmt.Errorf("failed to parse PEM block containing the private key")
|
||||
}
|
||||
|
||||
privateKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
rsaPrivateKey, ok := privateKey.(*rsa.PrivateKey)
|
||||
if !ok {
|
||||
return "", fmt.Errorf("not an RSA private key")
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
claims := jwt.MapClaims{
|
||||
"iss": email,
|
||||
"scope": "https://www.googleapis.com/auth/cloud-platform",
|
||||
"aud": "https://www.googleapis.com/oauth2/v4/token",
|
||||
"exp": now.Add(time.Minute * 35).Unix(),
|
||||
"iat": now.Unix(),
|
||||
}
|
||||
|
||||
token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
|
||||
signedToken, err := token.SignedString(rsaPrivateKey)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
return signedToken, nil
|
||||
}
|
||||
|
||||
func exchangeJwtForAccessToken(signedJWT string, info *relaycommon.RelayInfo) (string, error) {
|
||||
|
||||
authURL := "https://www.googleapis.com/oauth2/v4/token"
|
||||
data := url.Values{}
|
||||
data.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
|
||||
data.Set("assertion", signedJWT)
|
||||
|
||||
var client *http.Client
|
||||
var err error
|
||||
if info.ChannelSetting.Proxy != "" {
|
||||
client, err = service.NewProxyHttpClient(info.ChannelSetting.Proxy)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("new proxy http client failed: %w", err)
|
||||
}
|
||||
} else {
|
||||
client = service.GetHttpClient()
|
||||
}
|
||||
|
||||
resp, err := client.PostForm(authURL, data)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var result map[string]interface{}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if accessToken, ok := result["access_token"].(string); ok {
|
||||
return accessToken, nil
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("failed to get access token: %v", result)
|
||||
}
|
||||
|
||||
func AcquireAccessToken(creds Credentials, proxy string) (string, error) {
|
||||
signedJWT, err := createSignedJWT(creds.ClientEmail, creds.PrivateKey)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to create signed JWT: %w", err)
|
||||
}
|
||||
return exchangeJwtForAccessTokenWithProxy(signedJWT, proxy)
|
||||
}
|
||||
|
||||
func exchangeJwtForAccessTokenWithProxy(signedJWT string, proxy string) (string, error) {
|
||||
authURL := "https://www.googleapis.com/oauth2/v4/token"
|
||||
data := url.Values{}
|
||||
data.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
|
||||
data.Set("assertion", signedJWT)
|
||||
|
||||
var client *http.Client
|
||||
var err error
|
||||
if proxy != "" {
|
||||
client, err = service.NewProxyHttpClient(proxy)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("new proxy http client failed: %w", err)
|
||||
}
|
||||
} else {
|
||||
client = service.GetHttpClient()
|
||||
}
|
||||
|
||||
resp, err := client.PostForm(authURL, data)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
var result map[string]interface{}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if accessToken, ok := result["access_token"].(string); ok {
|
||||
return accessToken, nil
|
||||
}
|
||||
return "", fmt.Errorf("failed to get access token: %v", result)
|
||||
}
|
||||
Reference in New Issue
Block a user