jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-06-02 00:57:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: potential use of mismatched tokens (#1092)

Browse Source 
Co-authored-by: Kent Dong <ch3cho@qq.com>
This commit is contained in:
pepesi
2024-08-26 15:40:55 +08:00
committed by GitHub
parent 496346fe95
commit f5b8341f7f
2 changed files with 16 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
plugins/wasm-go/extensions/ai-proxy/provider/moonshot.go
View File

@@ -89,7 +89,8 @@ func (m *moonshotProvider) OnRequestBody(ctx wrapper.HttpContext, apiName ApiNam
return types.ActionContinue, replaceJsonRequestBody(request, log) return types.ActionContinue, replaceJsonRequestBody(request, log)
} }
err := m.getContextContent(func(content string, err error) { apiKey := m.config.GetOrSetTokenWithContext(ctx)
err := m.getContextContent(apiKey, func(content string, err error) {
defer func() { defer func() {
_ = proxywasm.ResumeHttpRequest() _ = proxywasm.ResumeHttpRequest()
}() }()
@@ -114,13 +115,13 @@ func (m *moonshotProvider) performChatCompletion(ctx wrapper.HttpContext, fileCo
return replaceJsonRequestBody(request, log) return replaceJsonRequestBody(request, log)
} }
func (m *moonshotProvider) getContextContent(callback func(string, error), log wrapper.Log) error { func (m *moonshotProvider) getContextContent(apiKey string, callback func(string, error), log wrapper.Log) error {
if m.config.moonshotFileId != "" { if m.config.moonshotFileId != "" {
if m.fileContent != "" { if m.fileContent != "" {
callback(m.fileContent, nil) callback(m.fileContent, nil)
return nil return nil
} }
return m.sendRequest(http.MethodGet, "/v1/files/"+m.config.moonshotFileId+"/content", "", return m.sendRequest(http.MethodGet, "/v1/files/"+m.config.moonshotFileId+"/content", "", apiKey,
func(statusCode int, responseHeaders http.Header, responseBody []byte) { func(statusCode int, responseHeaders http.Header, responseBody []byte) {
responseString := string(responseBody) responseString := string(responseBody)
if statusCode != http.StatusOK { if statusCode != http.StatusOK {
@@ -141,13 +142,13 @@ func (m *moonshotProvider) getContextContent(callback func(string, error), log w
return errors.New("both moonshotFileId and context are not configured") return errors.New("both moonshotFileId and context are not configured")
} }
func (m *moonshotProvider) sendRequest(method, path string, body string, callback wrapper.ResponseCallback) error { func (m *moonshotProvider) sendRequest(method, path, body, apiKey string, callback wrapper.ResponseCallback) error {
switch method { switch method {
case http.MethodGet: case http.MethodGet:
headers := util.CreateHeaders("Authorization", "Bearer "+m.config.GetRandomToken()) headers := util.CreateHeaders("Authorization", "Bearer "+apiKey)
return m.client.Get(path, headers, callback, m.config.timeout) return m.client.Get(path, headers, callback, m.config.timeout)
case http.MethodPost: case http.MethodPost:
headers := util.CreateHeaders("Authorization", "Bearer "+m.config.GetRandomToken(), "Content-Type", "application/json") headers := util.CreateHeaders("Authorization", "Bearer "+apiKey, "Content-Type", "application/json")
return m.client.Post(path, headers, []byte(body), callback, m.config.timeout) return m.client.Post(path, headers, []byte(body), callback, m.config.timeout)
default: default:
return errors.New("unsupported method: " + method) return errors.New("unsupported method: " + method)

9
plugins/wasm-go/extensions/ai-proxy/provider/provider.go
View File

@@ -290,6 +290,15 @@ func (c *ProviderConfig) Validate() error {
return nil return nil
} }
func (c *ProviderConfig) GetOrSetTokenWithContext(ctx wrapper.HttpContext) string {
ctxApiKey := ctx.GetContext(ctxKeyApiName)
if ctxApiKey == nil {
ctxApiKey = c.GetRandomToken()
ctx.SetContext(ctxKeyApiName, ctxApiKey)
}
return ctxApiKey.(string)
}
func (c *ProviderConfig) GetRandomToken() string { func (c *ProviderConfig) GetRandomToken() string {
apiTokens := c.apiTokens apiTokens := c.apiTokens
count := len(apiTokens) count := len(apiTokens)