From 3627a8e26e6bda1165ed42416186e9d10de0675f Mon Sep 17 00:00:00 2001
From: yuqi01 <yuqi01@baidu.com>
Date: Sat, 6 May 2023 16:30:39 +0800
Subject: [PATCH] =?UTF-8?q?=E6=97=A5=E5=BF=97=E6=89=93=E5=8D=B0=E4=BC=98?=
 =?UTF-8?q?=E5=8C=96,=E5=B0=86=E8=AF=B7=E6=B1=82=E4=B8=AD=E7=9A=84token?=
 =?UTF-8?q?=E4=BD=9C=E4=B8=BA=E7=94=A8=E6=88=B7=E6=A0=87=E8=AF=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.go            | 16 +++++++++++++++-
 pkg/azure/proxy.go |  6 ++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/main.go b/main.go
index 3f1c572..e02c81e 100644
--- a/main.go
+++ b/main.go
@@ -1,9 +1,11 @@
 package main
 
 import (
+	"fmt"
 	"github.com/diemus/azure-openai-proxy/pkg/azure"
 	"github.com/diemus/azure-openai-proxy/pkg/openai"
 	"github.com/gin-gonic/gin"
+	"io"
 	"log"
 	"net/http"
 	"os"
@@ -12,6 +14,7 @@ import (
 var (
 	Address   = "0.0.0.0:8080"
 	ProxyMode = "azure"
+	LogPath   = "log/proxy.log"
 )
 
 func init() {
@@ -22,6 +25,17 @@ func init() {
 	if v := os.Getenv("AZURE_OPENAI_PROXY_MODE"); v != "" {
 		ProxyMode = v
 	}
+	if v := os.Getenv("AZURE_OPENAI_PROXY_LOG_PATH"); v != "" {
+		LogPath = v
+	}
+	logFile, err := os.OpenFile(LogPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
+	if err != nil {
+		fmt.Println("open log file failed, err:", err)
+		return
+	}
+	multi := io.MultiWriter(logFile, os.Stdout)
+	log.SetOutput(multi)
+	log.SetFlags(log.Llongfile | log.Ldate | log.Ltime)
 	log.Printf("loading azure openai proxy address: %s", Address)
 	log.Printf("loading azure openai proxy mode: %s", ProxyMode)
 }
@@ -39,7 +53,7 @@ func main() {
 		router.Any("*path", handleOpenAIProxy)
 	}
 
-	router.Run(Address)
+	_ = router.Run(Address)
 
 }
 
diff --git a/pkg/azure/proxy.go b/pkg/azure/proxy.go
index 15927f2..0fd03ec 100644
--- a/pkg/azure/proxy.go
+++ b/pkg/azure/proxy.go
@@ -76,11 +76,12 @@ func NewOpenAIReverseProxy() *httputil.ReverseProxy {
 		// Replace the Bearer field in the Authorization header with api-key
 		token := ""
 
+		tokenFromReq := strings.ReplaceAll(req.Header.Get("Authorization"), "Bearer ", "")
 		// use the token from the environment variable if it is set
 		if AzureOpenAIToken != "" {
 			token = AzureOpenAIToken
 		} else {
-			token = strings.ReplaceAll(req.Header.Get("Authorization"), "Bearer ", "")
+			token = tokenFromReq
 		}
 
 		req.Header.Set("api-key", token)
@@ -99,7 +100,8 @@ func NewOpenAIReverseProxy() *httputil.ReverseProxy {
 		query.Add("api-version", AzureOpenAIAPIVersion)
 		req.URL.RawQuery = query.Encode()
 
-		log.Printf("proxying request [%s] %s -> %s", model, originURL, req.URL.String())
+		log.Printf("user identity: [%s], proxying request [%s] %s -> %s",
+			tokenFromReq, model, originURL, req.URL.String())
 	}
 	return &httputil.ReverseProxy{Director: director}
 }