feat: managed identity support

Dockerfile

@@ -1,15 +1,15 @@
-FROM golang:1.19 AS builder
+FROM golang:1.22-bookworm AS builder
 
 COPY . /builder
 WORKDIR /builder
 
 RUN make build
 
-FROM alpine:3
+FROM gcr.io/distroless/static-debian12:latest
 
 WORKDIR /app
 
 EXPOSE 8080
 COPY --from=builder /builder/bin .
 
-ENTRYPOINT ["/app/azure-openai-proxy"]
+ENTRYPOINT ["/app/azure-openai-proxy"]

Makefile

@@ -10,4 +10,7 @@ fmt:
 vet:
 	go vet ./...
 
+image:
+	docker build -t azure-openai-proxy:latest .
+
 .PHONY: build fmt vet

README.md

@@ -54,7 +54,7 @@ AZURE_OPENAI_MODEL_MAPPER: gpt-3.5-turbo=gpt-35-turbo
 
 ![Screenshot of the overview UI for an OpenAI Resource in the Azure portal with the endpoint & access keys location circled in red.](assets/images/endpoint.png)
 
-API Key: This value can be found in the **Keys & Endpoint** section when examining your resource from the Azure portal. You can use either `KEY1` or `KEY2`. 
+API Key: This value can be found in the **Keys & Endpoint** section when examining your resource from the Azure portal. You can use either `KEY1` or `KEY2`. If you want to use [Azure Managed Identity](https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/overview), you can leave it blank or use `msi` as the key (if a key is required) and it will use managed identity to authenticate.
 
 ### Proxy
 
@@ -81,7 +81,7 @@ AZURE_OPENAI_SOCKS_PROXY=socks5://127.0.0.1:1080
 ### Use Docker
 
 ````shell
-# config by environment 
+# config by environment
 docker run -d -p 8080:8080 --name=azure-openai-proxy \
   --env AZURE_OPENAI_ENDPOINT=your_azure_endpoint \
   --env AZURE_OPENAI_API_VER=your_azure_api_ver \

azure/init.go

@@ -2,17 +2,19 @@ package azure
 
 import (
 	"fmt"
-	"github.com/spf13/viper"
-	"github.com/stulzq/azure-openai-proxy/constant"
-	"github.com/stulzq/azure-openai-proxy/util"
 	"log"
 	"net/url"
 	"path/filepath"
 	"strings"
+
+	"github.com/spf13/viper"
+	"github.com/stulzq/azure-openai-proxy/constant"
+	"github.com/stulzq/azure-openai-proxy/util"
 )
 
 const (
-	AuthHeaderKey = "api-key"
+	AuthHeaderKey = "Authorization"
+	APIKeyHeaderKey = "api-key"
 )
 
 var (

azure/model.go

@@ -3,7 +3,6 @@ package azure
 import (
 	"bytes"
 	"fmt"
-	"github.com/pkg/errors"
 	"log"
 	"net/http"
 	"net/url"
@@ -74,7 +73,7 @@ func (c *TemplateConverter) Convert(req *http.Request, config *DeploymentConfig)
 	}
 	buff := new(bytes.Buffer)
 	if err := c.Tempalte.Execute(buff, data); err != nil {
-		return req, errors.Wrap(err, "template execute error")
+		return req, fmt.Errorf("template execute error: %w", err)
 	}
 
 	req.Host = config.EndpointUrl.Host

azure/proxy.go

@@ -2,20 +2,27 @@ package azure
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
-	"github.com/stulzq/azure-openai-proxy/util"
 	"io"
 	"log"
 	"net/http"
 	"net/http/httputil"
 	"strings"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/stulzq/azure-openai-proxy/util"
+
 	"github.com/bytedance/sonic"
 	"github.com/gin-gonic/gin"
-	"github.com/pkg/errors"
 )
 
+const cognitiveservicesScope = "https://cognitiveservices.azure.com/.default"
+
 func ProxyWithConverter(requestConverter RequestConverter) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		Proxy(c, requestConverter)
@@ -43,7 +50,7 @@ func ModelProxy(c *gin.Context) {
 			}
 
 			// Set the auth header
-			req.Header.Set(AuthHeaderKey, deployment.ApiKey)
+			req.Header.Set(APIKeyHeaderKey, deployment.ApiKey)
 
 			// Send the request
 			client := &http.Client{}
@@ -69,14 +76,14 @@ func ModelProxy(c *gin.Context) {
 			}
 
 			// Parse the response body as JSON
-			var deplotmentInfo DeploymentInfo
-			err = json.Unmarshal(body, &deplotmentInfo)
+			var deploymentInfo DeploymentInfo
+			err = json.Unmarshal(body, &deploymentInfo)
 			if err != nil {
 				log.Printf("error parsing response body for deployment %s: %v", deployment.DeploymentName, err)
 				results <- nil
 				return
 			}
-			results <- deplotmentInfo.Data
+			results <- deploymentInfo.Data
 		}(deployment)
 	}
 
@@ -134,12 +141,12 @@ func Proxy(c *gin.Context, requestConverter RequestConverter) {
 		if model == "" {
 			_model, err := sonic.Get(body, "model")
 			if err != nil {
-				util.SendError(c, errors.Wrap(err, "get model error"))
+				util.SendError(c, fmt.Errorf("get model error: %w", err))
 				return
 			}
 			_modelStr, err := _model.String()
 			if err != nil {
-				util.SendError(c, errors.Wrap(err, "get model name error"))
+				util.SendError(c, fmt.Errorf("get model name error: %w", err))
 				return
 			}
 			model = _modelStr
@@ -152,23 +159,45 @@ func Proxy(c *gin.Context, requestConverter RequestConverter) {
 			return
 		}
 
-		// get auth token from header or deployemnt config
+		// get auth token from header or deployment config
 		token := deployment.ApiKey
-		if token == "" {
+		tokenFound := false
+		if token == "" && token != "msi" {
 			rawToken := req.Header.Get("Authorization")
 			token = strings.TrimPrefix(rawToken, "Bearer ")
+			req.Header.Set(APIKeyHeaderKey, token)
+			req.Header.Del("Authorization")
+			tokenFound = true
+		}
+		// get azure token using managed identity
+		var azureToken azcore.AccessToken
+		if token == "" || token == "msi" {
+			cred, err := azidentity.NewManagedIdentityCredential(nil)
+			if err != nil {
+				util.SendError(c, fmt.Errorf("failed to create managed identity credential: %w", err))
+			}
+
+			azureToken, err = cred.GetToken(context.TODO(), policy.TokenRequestOptions{
+				Scopes: []string{cognitiveservicesScope},
+			})
+			if err != nil {
+				util.SendError(c, fmt.Errorf("failed to get token: %w", err))
+			}
+
+			req.Header.Del(APIKeyHeaderKey)
+			req.Header.Set(AuthHeaderKey, "Bearer "+azureToken.Token)
+			tokenFound = true
 		}
-		if token == "" {
+
+		if !tokenFound {
 			util.SendError(c, errors.New("token is empty"))
 			return
 		}
-		req.Header.Set(AuthHeaderKey, token)
-		req.Header.Del("Authorization")
 
 		originURL := req.URL.String()
 		req, err = requestConverter.Convert(req, deployment)
 		if err != nil {
-			util.SendError(c, errors.Wrap(err, "convert request error"))
+			util.SendError(c, fmt.Errorf("convert request error: %w", err))
 			return
 		}
 		log.Printf("proxying request [%s] %s -> %s", model, originURL, req.URL.String())
@@ -177,7 +206,7 @@ func Proxy(c *gin.Context, requestConverter RequestConverter) {
 	proxy := &httputil.ReverseProxy{Director: director}
 	transport, err := util.NewProxyFromEnv()
 	if err != nil {
-		util.SendError(c, errors.Wrap(err, "get proxy error"))
+		util.SendError(c, fmt.Errorf("get proxy error: %w", err))
 		return
 	}
 	if transport != nil {
@@ -201,7 +230,7 @@ func Proxy(c *gin.Context, requestConverter RequestConverter) {
 func GetDeploymentByModel(model string) (*DeploymentConfig, error) {
 	deploymentConfig, exist := ModelDeploymentConfig[model]
 	if !exist {
-		return nil, errors.New(fmt.Sprintf("deployment config for %s not found", model))
+		return nil, fmt.Errorf("deployment config for %s not found", model)
 	}
 	return &deploymentConfig, nil
 }

cmd/main.go

@@ -3,17 +3,17 @@ package main
 import (
 	"context"
 	"fmt"
-	"github.com/spf13/pflag"
-	"github.com/spf13/viper"
-	"github.com/stulzq/azure-openai-proxy/azure"
 	"log"
 	"net/http"
 	"os"
 	"os/signal"
 	"syscall"
 
+	"github.com/spf13/pflag"
+	"github.com/spf13/viper"
+	"github.com/stulzq/azure-openai-proxy/azure"
+
 	"github.com/gin-gonic/gin"
-	"github.com/pkg/errors"
 )
 
 var (
@@ -74,7 +74,7 @@ func runServer(srv *http.Server) {
 	go func() {
 		log.Printf("Server listening at %s\n", srv.Addr)
 		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
-			panic(errors.Errorf("listen: %s\n", err))
+			panic(fmt.Errorf("listen: %w", err))
 		}
 	}()
 

go.mod

@@ -3,36 +3,44 @@ module github.com/stulzq/azure-openai-proxy
 go 1.21
 
 require (
-	github.com/bytedance/sonic v1.10.2
-	github.com/gin-gonic/gin v1.9.1
-	github.com/pkg/errors v0.9.1
+	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.11.1
+	github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.5.2
+	github.com/bytedance/sonic v1.11.8
+	github.com/gin-gonic/gin v1.10.0
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.18.2
-	github.com/stretchr/testify v1.8.4
-	golang.org/x/net v0.19.0
+	github.com/stretchr/testify v1.9.0
+	golang.org/x/net v0.25.0
 )
 
 require (
-	github.com/chenzhuoyu/base64x v0.0.0-20230717121745-296ad89f973d // indirect
-	github.com/chenzhuoyu/iasm v0.9.1 // indirect
+	github.com/Azure/azure-sdk-for-go/sdk/internal v1.5.2 // indirect
+	github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2 // indirect
+	github.com/bytedance/sonic/loader v0.1.1 // indirect
+	github.com/cloudwego/base64x v0.1.4 // indirect
+	github.com/cloudwego/iasm v0.2.0 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/fsnotify/fsnotify v1.7.0 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.3 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect
 	github.com/go-playground/universal-translator v0.18.1 // indirect
-	github.com/go-playground/validator/v10 v10.16.0 // indirect
+	github.com/go-playground/validator/v10 v10.20.0 // indirect
 	github.com/goccy/go-json v0.10.2 // indirect
+	github.com/golang-jwt/jwt/v5 v5.2.1 // indirect
+	github.com/google/uuid v1.6.0 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
-	github.com/klauspost/cpuid/v2 v2.2.6 // indirect
-	github.com/leodido/go-urn v1.2.4 // indirect
+	github.com/klauspost/cpuid/v2 v2.2.7 // indirect
+	github.com/kylelemons/godebug v1.1.0 // indirect
+	github.com/leodido/go-urn v1.4.0 // indirect
 	github.com/magiconair/properties v1.8.7 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
-	github.com/pelletier/go-toml/v2 v2.1.1 // indirect
+	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
+	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/sagikazarmark/locafero v0.4.0 // indirect
 	github.com/sagikazarmark/slog-shim v0.1.0 // indirect
@@ -43,12 +51,12 @@ require (
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.2.12 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/arch v0.6.0 // indirect
-	golang.org/x/crypto v0.17.0 // indirect
+	golang.org/x/arch v0.8.0 // indirect
+	golang.org/x/crypto v0.23.0 // indirect
 	golang.org/x/exp v0.0.0-20231214170342-aacd6d4b4611 // indirect
-	golang.org/x/sys v0.15.0 // indirect
-	golang.org/x/text v0.14.0 // indirect
-	google.golang.org/protobuf v1.31.0 // indirect
+	golang.org/x/sys v0.20.0 // indirect
+	golang.org/x/text v0.15.0 // indirect
+	google.golang.org/protobuf v1.34.1 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )