Fix auth service monitor wrapper to maintain EmailInviter implementat…

…ion when required (#7916)
treeverse · Jun 25, 2024 · 7fac5bb · 7fac5bb
1 parent 847faf1
commit 7fac5bb
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 45 deletions.
diff --git a/.github/workflows/esti.yaml b/.github/workflows/esti.yaml
@@ -70,7 +70,7 @@ jobs:
         if: steps.restore-cache.outputs.cache-hit != 'true'
         run: |
           make -j3 gen-api gen-code gen-ui VERSION=${{ steps.version.outputs.tag }}
-          tar -czf /tmp/generated.tar.gz ./webui/dist ./pkg/auth/{client,wrapper}.gen.go ./pkg/authentication/apiclient/client.gen.go ./pkg/permissions/actions.gen.go ./pkg/api/apigen/lakefs.gen.go
+          tar -czf /tmp/generated.tar.gz ./webui/dist ./pkg/auth/{client,service_wrapper,service_inviter_wrapper}.gen.go ./pkg/authentication/apiclient/client.gen.go ./pkg/permissions/actions.gen.go ./pkg/api/apigen/lakefs.gen.go
 
         # must upload artifact in order to download generated later
       - name: Store generated code

diff --git a/Makefile b/Makefile
@@ -294,7 +294,8 @@ validate-permissions-gen: gen-code
 
 .PHONY: validate-wrapper
 validate-wrapper: gen-code
-	git diff --quiet -- pkg/auth/wrapper.gen.go || (echo "Modification verification failed! pkg/auth/wrapper.gen.go"; false)
+	git diff --quiet -- pkg/auth/service_wrapper.gen.go || (echo "Modification verification failed! pkg/auth/service_wrapper.gen.go"; false)
+	git diff --quiet -- pkg/auth/service_inviter_wrapper.gen.go || (echo "Modification verification failed! pkg/auth/service_inviter_wrapper.gen.go"; false)
 
 .PHONY: validate-wrapgen-testcode
 validate-wrapgen-testcode: gen-code

diff --git a/cmd/lakefs/cmd/run.go b/cmd/lakefs/cmd/run.go
@@ -66,6 +66,38 @@ func checkAuthModeSupport(cfg *config.Config) error {
 	return nil
 }
 
+func NewAuthService(ctx context.Context, cfg *config.Config, logger logging.Logger, kvStore kv.Store) auth.Service {
+	if err := checkAuthModeSupport(cfg); err != nil {
+		logger.WithError(err).Fatal("Unsupported auth mode")
+	}
+	if cfg.IsAuthTypeAPI() {
+		apiService, err := auth.NewAPIAuthService(
+			cfg.Auth.API.Endpoint,
+			cfg.Auth.API.Token.SecureValue(),
+			cfg.Auth.AuthenticationAPI.ExternalPrincipalsEnabled,
+			crypt.NewSecretStore([]byte(cfg.Auth.Encrypt.SecretKey)),
+			authparams.ServiceCache(cfg.Auth.Cache),
+			logger.WithField("service", "auth_api"),
+		)
+		if err != nil {
+			logger.WithError(err).Fatal("failed to create authentication service")
+		}
+		if !cfg.Auth.API.SkipHealthCheck {
+			if err := apiService.CheckHealth(ctx, logger, cfg.Auth.API.HealthCheckTimeout); err != nil {
+				logger.WithError(err).Fatal("Auth API health check failed")
+			}
+		}
+		return auth.NewMonitoredAuthServiceAndInviter(apiService)
+	}
+	authService := auth.NewAuthService(
+		kvStore,
+		crypt.NewSecretStore([]byte(cfg.Auth.Encrypt.SecretKey)),
+		authparams.ServiceCache(cfg.Auth.Cache),
+		logger.WithField("service", "auth_service"),
+	)
+	return auth.NewMonitoredAuthService(authService)
+}
+
 var runCmd = &cobra.Command{
 	Use:   "run",
 	Short: "Run lakeFS",
@@ -111,40 +143,7 @@ var runCmd = &cobra.Command{
 		authMetadataManager := auth.NewKVMetadataManager(version.Version, cfg.Installation.FixedID, cfg.Database.Type, kvStore)
 		idGen := &actions.DecreasingIDGenerator{}
 
-		// initialize authorization service
-		var authService auth.Service
-
-		if err := checkAuthModeSupport(cfg); err != nil {
-			logger.WithError(err).Fatal("Unsupported auth mode")
-		}
-		if cfg.IsAuthTypeAPI() {
-			apiService, err := auth.NewAPIAuthService(
-				cfg.Auth.API.Endpoint,
-				cfg.Auth.API.Token.SecureValue(),
-				cfg.Auth.AuthenticationAPI.ExternalPrincipalsEnabled,
-				crypt.NewSecretStore([]byte(cfg.Auth.Encrypt.SecretKey)),
-				authparams.ServiceCache(cfg.Auth.Cache),
-				logger.WithField("service", "auth_api"),
-			)
-			if err != nil {
-				logger.WithError(err).Fatal("failed to create authentication service")
-			}
-			authService = apiService
-			if !cfg.Auth.API.SkipHealthCheck {
-				if err := apiService.CheckHealth(ctx, logger, cfg.Auth.API.HealthCheckTimeout); err != nil {
-					logger.WithError(err).Fatal("Auth API health check failed")
-				}
-			}
-		} else {
-			authService = auth.NewAuthService(
-				kvStore,
-				crypt.NewSecretStore([]byte(cfg.Auth.Encrypt.SecretKey)),
-				authparams.ServiceCache(cfg.Auth.Cache),
-				logger.WithField("service", "auth_service"),
-			)
-		}
-		authService = auth.NewMonitoredAuthService(authService)
-
+		authService := NewAuthService(ctx, cfg, logger, kvStore)
 		// initialize authentication service
 		var authenticationService authentication.Service
 		if cfg.IsAuthenticationTypeAPI() {

diff --git a/cmd/lakefs/cmd/run_test.go b/cmd/lakefs/cmd/run_test.go
@@ -0,0 +1,32 @@
+package cmd_test
+
+import (
+	"context"
+	"github.com/treeverse/lakefs/cmd/lakefs/cmd"
+	"github.com/treeverse/lakefs/pkg/auth"
+	"github.com/treeverse/lakefs/pkg/config"
+	"github.com/treeverse/lakefs/pkg/logging"
+	"testing"
+)
+
+func TestGetAuthService(t *testing.T) {
+	t.Run("maintain_inviter", func(t *testing.T) {
+		cfg := &config.Config{}
+		cfg.Auth.API.Endpoint = "http://localhost:8000"
+		cfg.Auth.API.SkipHealthCheck = true
+		service := cmd.NewAuthService(context.Background(), cfg, logging.ContextUnavailable(), nil)
+		_, ok := service.(auth.EmailInviter)
+		if !ok {
+			t.Fatalf("expected Service to be of type EmailInviter")
+		}
+	})
+	t.Run("maintain_service", func(t *testing.T) {
+		cfg := &config.Config{}
+		cfg.Auth.UIConfig.RBAC = config.AuthRBACSimplified
+		service := cmd.NewAuthService(context.Background(), cfg, logging.ContextUnavailable(), nil)
+		_, ok := service.(auth.EmailInviter)
+		if ok {
+			t.Fatalf("expected Service to not be of type EmailInviter")
+		}
+	})
+}
diff --git a/pkg/auth/monitored.go b/pkg/auth/monitored.go
@@ -16,15 +16,24 @@ var (
 		}, []string{"operation", "success"})
 )
 
+func ObserveDuration(operation string, duration time.Duration, success bool) {
+	status := "failure"
+	if success {
+		status = "success"
+	}
+	authDurationSecs.WithLabelValues(operation, status).Observe(duration.Seconds())
+}
+
+func NewMonitoredAuthServiceAndInviter(service ServiceAndInviter) *MonitoredServiceAndInviter {
+	return &MonitoredServiceAndInviter{
+		Wrapped: service,
+		Observe: ObserveDuration,
+	}
+}
+
 func NewMonitoredAuthService(service Service) *MonitoredService {
 	return &MonitoredService{
 		Wrapped: service,
-		Observe: func(op string, duration time.Duration, success bool) {
-			status := "failure"
-			if success {
-				status = "success"
-			}
-			authDurationSecs.WithLabelValues(op, status).Observe(duration.Seconds())
-		},
+		Observe: ObserveDuration,
 	}
 }
diff --git a/pkg/auth/service.go b/pkg/auth/service.go
@@ -1,9 +1,11 @@
 package auth
 
-//go:generate go run github.com/treeverse/lakefs/tools/wrapgen --package auth --output ./wrapper.gen.go --interface Service ./service.go
+//go:generate go run github.com/treeverse/lakefs/tools/wrapgen --package auth --output ./service_inviter_wrapper.gen.go --interface ServiceAndInviter ./service.go
+//go:generate go run github.com/treeverse/lakefs/tools/wrapgen --package auth --output ./service_wrapper.gen.go --interface Service ./service.go
 
 // Must run goimports after wrapgen: it adds unused imports.
-//go:generate go run golang.org/x/tools/cmd/goimports@latest -w ./wrapper.gen.go
+//go:generate go run golang.org/x/tools/cmd/goimports@latest -w ./service_inviter_wrapper.gen.go
+//go:generate go run golang.org/x/tools/cmd/goimports@latest -w ./service_wrapper.gen.go
 
 //go:generate go run github.com/deepmap/oapi-codegen/cmd/[email protected] -package auth -generate "types,client" -o client.gen.go ../../api/authorization.yml
 //go:generate go run github.com/golang/mock/[email protected] -package=mock -destination=mock/mock_auth_client.go github.com/treeverse/lakefs/pkg/auth ClientWithResponsesInterface
@@ -93,6 +95,11 @@ type ExternalPrincipalsService interface {
 	ListUserExternalPrincipals(ctx context.Context, userID string, params *model.PaginationParams) ([]*model.ExternalPrincipal, *model.Paginator, error)
 }
 
+type ServiceAndInviter interface {
+	Service
+	EmailInviter
+}
+
 type Service interface {
 	SecretStore() crypt.SecretStore
 	Cache() Cache