fix: Carry verification Token in Header issue #3

soulteary · Jan 27, 2024 · 3ab0166 · 3ab0166
1 parent 66d5080
commit 3ab0166
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/models/gemini/model.go b/models/gemini/model.go
@@ -11,7 +11,7 @@ import (
 
 type RequestConverter interface {
 	Name() string
-	Convert(req *http.Request, config *define.ModelConfig, payload []byte, openaiPayload define.OpenAI_Payload) (*http.Request, error)
+	Convert(req *http.Request, config *define.ModelConfig, payload []byte, openaiPayload define.OpenAI_Payload, apikey string) (*http.Request, error)
 }
 
 type StripPrefixConverter struct {
@@ -22,7 +22,7 @@ func (c *StripPrefixConverter) Name() string {
 	return "StripPrefix"
 }
 
-func (c *StripPrefixConverter) Convert(req *http.Request, config *define.ModelConfig, payload []byte, openaiPayload define.OpenAI_Payload) (*http.Request, error) {
+func (c *StripPrefixConverter) Convert(req *http.Request, config *define.ModelConfig, payload []byte, openaiPayload define.OpenAI_Payload, apikey string) (*http.Request, error) {
 	req.Host = config.URL.Host
 	req.URL.Scheme = config.URL.Scheme
 	req.URL.Host = config.URL.Host
@@ -36,7 +36,15 @@ func (c *StripPrefixConverter) Convert(req *http.Request, config *define.ModelCo
 	req.URL.RawPath = req.URL.EscapedPath()
 
 	query := req.URL.Query()
-	query.Add("key", config.Key)
+	if config.Key == "" {
+		if apikey == "" {
+			return nil, fmt.Errorf("missing api key")
+		} else {
+			query.Add("key", apikey)
+		}
+	} else {
+		query.Add("key", config.Key)
+	}
 	req.URL.RawQuery = query.Encode()
 	req.Body = io.NopCloser(bytes.NewBuffer(payload))
 	req.ContentLength = int64(len(payload))

diff --git a/models/gemini/proxy.go b/models/gemini/proxy.go
@@ -126,7 +126,7 @@ func getDirector(req *http.Request, body []byte, c *gin.Context, requestConverte
 			network.SendError(c, errors.New("token is empty"))
 			return
 		}
-		req.Header.Set("Authorization", token)
+		req.Header.Del("Authorization")
 
 		repack, err := json.Marshal(payload)
 		if err != nil {
@@ -135,7 +135,7 @@ func getDirector(req *http.Request, body []byte, c *gin.Context, requestConverte
 		}
 
 		originURL := req.URL.String()
-		req, err = requestConverter.Convert(req, deployment, repack, openaiPayload)
+		req, err = requestConverter.Convert(req, deployment, repack, openaiPayload, token)
 		if err != nil {
 			network.SendError(c, errors.Wrap(err, "convert request error"))
 			return