6146306117

Committed 11 Sep 2023 12:22PM UTC coverage: 26.069% (-3.9%) from 29.927%

Build # 6146306117

Build Type

Pull #171

github

Committed by

xieydd

Commit Message

chore: Make proxy endpoint port configurable

Signed-off-by: xieydd <xieydd@gmail.com>

Pull Request Pull Request #171: feat: Support BYOC

Run Details

443 of 443 new or added lines in 14 files covered. (100.0%)

939 of 3602 relevant lines covered (26.07%)

1.63 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/agent/pkg/server/proxy_auth.go

package server

import (
        "context"
        "fmt"
        "strings"

        "github.com/gin-gonic/gin"
        "github.com/sirupsen/logrus"
        "github.com/tensorchord/openmodelz/agent/errdefs"
        "github.com/tensorchord/openmodelz/agent/pkg/consts"
)

func (s *Server) proxyAuth(c *gin.Context) (string, string, error) {
        var uid string
        var valid bool

        deployment := c.Param("id")
        if len(deployment) == 0 {
                return "", "", errFromErrDefs(
                        fmt.Errorf("cannot find the deployment name in %s", c.Request.RequestURI), "get-deployment")
        }

        key := c.GetHeader("X-API-Key")
        // Be compatible with the OpenAI API.
        rawKeyStr := c.GetHeader("Authorization")
        logrus.Debug("proxyOther: key: ", key, ", rawKeyStr: ", rawKeyStr)

        if s.validateUnifiedKey(key) {
                // uid 0 means to use unified api key
                uid = "00000000-0000-0000-0000-000000000000"
        } else if len(key) > 0 {
                uid, valid = s.validateAPIKey(key)
                if !valid {
                        return "", "", errdefs.Unauthorized(fmt.Errorf("invalid API key"))
                }
        } else if len(rawKeyStr) > 0 {
                strs := strings.Split(rawKeyStr, " ")
                if len(strs) != 2 {
                        return "", "", errdefs.Unauthorized(fmt.Errorf("invalid Authorization API key"))
                }

                if strs[0] != "Bearer" {
                        return "", "", errdefs.Unauthorized(fmt.Errorf("invalid Authorization API key"))
                }

                uid, valid = s.validateAPIKey(strs[1])
                if !valid {
                        return "", "", errdefs.Unauthorized(fmt.Errorf("invalid Authorization API key"))
                }
        }

        if len(uid) == 0 {
                return "", "", errdefs.Unauthorized(fmt.Errorf("invalid API key"))
        }
        return uid, deployment, nil
}

func (s *Server) proxyNoAuth(c *gin.Context) (string, string, error) {
        deployment := c.Param("id")
        if len(deployment) == 0 {
                return "", "", errdefs.InvalidParameter(
                        fmt.Errorf("cannot find the deployment name in %s", c.Request.RequestURI))
        }

        uid, found := s.getUIDFromDeploymentID(c.Request.Context(), deployment)
        if !found {
                return "", "", errdefs.InvalidParameter(
                        fmt.Errorf("cannot find the user id from the deployment id"))
        }
        return uid, deployment, nil
}

func (s *Server) validateAPIKey(key string) (string, bool) {
        if !strings.HasPrefix(key, consts.APIKEY_PREFIX) {
                return "", false
        }

        apikeys := s.config.ModelZCloud.APIKeys
        uid, exit := apikeys[key]
        if exit {
                return uid, true
        }

        apiServerReady := make(chan struct{})
        go func() {
                if err := s.modelzCloudClient.WaitForAPIServerReady(); err != nil {
                        logrus.Fatalf("failed to wait for apiserver ready: %v", err)
                }
                close(apiServerReady)
        }()
        // Get from apiserver
        apikeys, err := s.modelzCloudClient.GetAPIKeys(context.Background(), apiServerReady, s.config.ModelZCloud.AgentToken, s.config.ModelZCloud.ID)
        if err != nil {
                logrus.Errorf("failed to get apikeys: %v", err)
                return "", false
        }
        uid, exit = apikeys[key]
        if exit {
                return uid, true
        }

        return "", false
}

func (s *Server) validateUnifiedKey(key string) bool {
        if !strings.HasPrefix(key, consts.APIKEY_PREFIX) {
                return false
        }
        if len(s.config.ModelZCloud.UnifiedAPIKey) != 0 && s.config.ModelZCloud.UnifiedAPIKey == key {
                return true
        }
        return false
}

1	package server
2
3	import (
4	"context"
5	"fmt"
6	"strings"
7
8	"github.com/gin-gonic/gin"
9	"github.com/sirupsen/logrus"
10	"github.com/tensorchord/openmodelz/agent/errdefs"
11	"github.com/tensorchord/openmodelz/agent/pkg/consts"
12	)
13
14	func (s Server) proxyAuth(c gin.Context) (string, string, error) {	×
15	var uid string	×
16	var valid bool	×
17		×
18	deployment := c.Param("id")	×
19	if len(deployment) == 0 {	×
20	return "", "", errFromErrDefs(	×
21	fmt.Errorf("cannot find the deployment name in %s", c.Request.RequestURI), "get-deployment")	×
22	}	×
23
24	key := c.GetHeader("X-API-Key")	×
25	// Be compatible with the OpenAI API.	×
26	rawKeyStr := c.GetHeader("Authorization")	×
27	logrus.Debug("proxyOther: key: ", key, ", rawKeyStr: ", rawKeyStr)	×
28		×
29	if s.validateUnifiedKey(key) {	×
30	// uid 0 means to use unified api key	×
31	uid = "00000000-0000-0000-0000-000000000000"	×
32	} else if len(key) > 0 {	×
33	uid, valid = s.validateAPIKey(key)	×
34	if !valid {	×
35	return "", "", errdefs.Unauthorized(fmt.Errorf("invalid API key"))	×
36	}	×
37	} else if len(rawKeyStr) > 0 {	×
38	strs := strings.Split(rawKeyStr, " ")	×
39	if len(strs) != 2 {	×
40	return "", "", errdefs.Unauthorized(fmt.Errorf("invalid Authorization API key"))	×
41	}	×
42
43	if strs[0] != "Bearer" {	×
44	return "", "", errdefs.Unauthorized(fmt.Errorf("invalid Authorization API key"))	×
45	}	×
46
47	uid, valid = s.validateAPIKey(strs[1])	×
48	if !valid {	×
49	return "", "", errdefs.Unauthorized(fmt.Errorf("invalid Authorization API key"))	×
50	}	×
51	}
52
53	if len(uid) == 0 {	×
54	return "", "", errdefs.Unauthorized(fmt.Errorf("invalid API key"))	×
55	}	×
56	return uid, deployment, nil	×
57	}
58
59	func (s Server) proxyNoAuth(c gin.Context) (string, string, error) {	×
60	deployment := c.Param("id")	×
61	if len(deployment) == 0 {	×
62	return "", "", errdefs.InvalidParameter(	×
63	fmt.Errorf("cannot find the deployment name in %s", c.Request.RequestURI))	×
64	}	×
65
66	uid, found := s.getUIDFromDeploymentID(c.Request.Context(), deployment)	×
67	if !found {	×
68	return "", "", errdefs.InvalidParameter(	×
69	fmt.Errorf("cannot find the user id from the deployment id"))	×
70	}	×
71	return uid, deployment, nil	×
72	}
73
74	func (s *Server) validateAPIKey(key string) (string, bool) {	×
75	if !strings.HasPrefix(key, consts.APIKEY_PREFIX) {	×
76	return "", false	×
77	}	×
78
79	apikeys := s.config.ModelZCloud.APIKeys	×
80	uid, exit := apikeys[key]	×
81	if exit {	×
82	return uid, true	×
83	}	×
84
85	apiServerReady := make(chan struct{})	×
86	go func() {	×
87	if err := s.modelzCloudClient.WaitForAPIServerReady(); err != nil {	×
88	logrus.Fatalf("failed to wait for apiserver ready: %v", err)	×
89	}	×
90	close(apiServerReady)	×
91	}()
92	// Get from apiserver
93	apikeys, err := s.modelzCloudClient.GetAPIKeys(context.Background(), apiServerReady, s.config.ModelZCloud.AgentToken, s.config.ModelZCloud.ID)	×
94	if err != nil {	×
95	logrus.Errorf("failed to get apikeys: %v", err)	×
96	return "", false	×
97	}	×
98	uid, exit = apikeys[key]	×
99	if exit {	×
100	return uid, true	×
101	}	×
102
103	return "", false	×
104	}
105
106	func (s *Server) validateUnifiedKey(key string) bool {	×
107	if !strings.HasPrefix(key, consts.APIKEY_PREFIX) {	×
108	return false	×
109	}	×
110	if len(s.config.ModelZCloud.UnifiedAPIKey) != 0 && s.config.ModelZCloud.UnifiedAPIKey == key {	×
111	return true	×
112	}	×
113	return false	×
114	}

tensorchord / openmodelz / 6146306117

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous