Pārlūkot izejas kodu

support user whitelist

Alexander Belanger 4 gadi atpakaļ
vecāks
revīzija
caf641334e

+ 4 - 3
api/server/handlers/project/get_usage.go

@@ -31,9 +31,10 @@ func (p *ProjectGetUsageHandler) ServeHTTP(w http.ResponseWriter, r *http.Reques
 	res := &types.GetProjectUsageResponse{}
 
 	currUsage, limit, usageCache, err := usage.GetUsage(&usage.GetUsageOpts{
-		Project: proj,
-		DOConf:  p.Config().DOConf,
-		Repo:    p.Repo(),
+		Project:          proj,
+		DOConf:           p.Config().DOConf,
+		Repo:             p.Repo(),
+		WhitelistedUsers: p.Config().WhitelistedUsers,
 	})
 
 	if err != nil {

+ 4 - 3
api/server/router/middleware/usage.go

@@ -28,9 +28,10 @@ func (b *UsageMiddleware) Middleware(next http.Handler) http.Handler {
 
 		// get the project usage limits
 		currentUsage, limit, _, err := usage.GetUsage(&usage.GetUsageOpts{
-			Project: proj,
-			DOConf:  b.config.DOConf,
-			Repo:    b.config.Repo,
+			Project:          proj,
+			DOConf:           b.config.DOConf,
+			Repo:             b.config.Repo,
+			WhitelistedUsers: b.config.WhitelistedUsers,
 		})
 
 		if err != nil {

+ 3 - 0
api/server/shared/config/config.go

@@ -89,6 +89,9 @@ type Config struct {
 
 	// BillingManager manages billing for Porter instances with billing enabled
 	BillingManager billing.BillingManager
+
+	// WhitelistedUsers do not count toward usage limits
+	WhitelistedUsers map[uint]uint
 }
 
 type ConfigLoader interface {

+ 1 - 0
api/server/shared/config/env/envconfs.go

@@ -51,6 +51,7 @@ type ServerConf struct {
 
 	IronPlansAPIKey    string `env:"IRON_PLANS_API_KEY"`
 	IronPlansServerURL string `env:"IRON_PLANS_SERVER_URL"`
+	WhitelistedUsers   []uint `env:"WHITELISTED_USERS"`
 
 	DOClientID                 string `env:"DO_CLIENT_ID"`
 	DOClientSecret             string `env:"DO_CLIENT_SECRET"`

+ 9 - 0
api/server/shared/config/loader/loader.go

@@ -180,6 +180,15 @@ func (e *EnvConfigLoader) LoadConfig() (res *config.Config, err error) {
 		},
 	}
 
+	// construct the whitelisted users map
+	wlUsers := make(map[uint]uint)
+
+	for _, userID := range sc.WhitelistedUsers {
+		wlUsers[userID] = userID
+	}
+
+	res.WhitelistedUsers = wlUsers
+
 	res.URLCache = urlcache.Init(sc.DefaultApplicationHelmRepoURL, sc.DefaultAddonHelmRepoURL)
 
 	provAgent, err := getProvisionerAgent(sc)

+ 18 - 5
internal/usage/usage.go

@@ -14,9 +14,10 @@ import (
 )
 
 type GetUsageOpts struct {
-	Repo    repository.Repository
-	DOConf  *oauth2.Config
-	Project *models.Project
+	Repo             repository.Repository
+	DOConf           *oauth2.Config
+	Project          *models.Project
+	WhitelistedUsers map[uint]uint
 }
 
 // GetUsage gets a project's current usage and usage limit
@@ -45,6 +46,14 @@ func GetUsage(opts *GetUsageOpts) (
 		return nil, nil, nil, err
 	}
 
+	countedRoles := make([]models.Role, 0)
+
+	for _, role := range roles {
+		if _, exists := opts.WhitelistedUsers[role.UserID]; !exists {
+			countedRoles = append(countedRoles, role)
+		}
+	}
+
 	usageCache, err := opts.Repo.ProjectUsage().ReadProjectUsageCache(opts.Project.ID)
 	isCacheFound := true
 
@@ -72,7 +81,7 @@ func GetUsage(opts *GetUsageOpts) (
 			usageCache.ResourceMemory = memory
 		}
 
-		isExceeded := isUsageExceeded(usageCache, limit, uint(len(roles)), uint(len(clusters)))
+		isExceeded := isUsageExceeded(usageCache, limit, uint(len(countedRoles)), uint(len(clusters)))
 
 		if !usageCache.Exceeded && isExceeded {
 			// update the usage cache with a time exceeded
@@ -89,11 +98,15 @@ func GetUsage(opts *GetUsageOpts) (
 		}
 	}
 
+	// we check whether it's currently exceeded based on the cache every time, since
+	// it's an inexpensive operation and involves no further DB lookups
+	usageCache.Exceeded = isUsageExceeded(usageCache, limit, uint(len(countedRoles)), uint(len(clusters)))
+
 	return &types.ProjectUsage{
 		ResourceCPU:    usageCache.ResourceCPU,
 		ResourceMemory: usageCache.ResourceMemory,
 		Clusters:       uint(len(clusters)),
-		Users:          uint(len(roles)),
+		Users:          uint(len(countedRoles)),
 	}, limit, usageCache, nil
 }
 

+ 15 - 12
services/usage/usage.go

@@ -17,17 +17,19 @@ import (
 )
 
 type UsageTracker struct {
-	db     *gorm.DB
-	repo   repository.Repository
-	doConf *oauth2.Config
+	db               *gorm.DB
+	repo             repository.Repository
+	doConf           *oauth2.Config
+	whitelistedUsers map[uint]uint
 }
 
 type UsageTrackerOpts struct {
-	DBConf         *env.DBConf
-	DOClientID     string
-	DOClientSecret string
-	DOScopes       []string
-	ServerURL      string
+	DBConf           *env.DBConf
+	DOClientID       string
+	DOClientSecret   string
+	DOScopes         []string
+	ServerURL        string
+	WhitelistedUsers map[uint]uint
 }
 
 const stepSize = 100
@@ -54,7 +56,7 @@ func NewUsageTracker(opts *UsageTrackerOpts) (*UsageTracker, error) {
 		BaseURL:      opts.ServerURL,
 	})
 
-	return &UsageTracker{db, repo, doConf}, nil
+	return &UsageTracker{db, repo, doConf, opts.WhitelistedUsers}, nil
 }
 
 type UsageTrackerResponse struct {
@@ -89,9 +91,10 @@ func (u *UsageTracker) GetProjectUsage() (map[uint]*UsageTrackerResponse, error)
 		// go through each project
 		for _, project := range projects {
 			_, limit, cache, err := usage.GetUsage(&usage.GetUsageOpts{
-				Repo:    u.repo,
-				DOConf:  u.doConf,
-				Project: project,
+				Repo:             u.repo,
+				DOConf:           u.doConf,
+				Project:          project,
+				WhitelistedUsers: u.whitelistedUsers,
 			})
 
 			if err != nil {