feat: block instances (#7129)

* docs: fix init description typos

* feat: block instances using limits

* translate

* unit tests

* fix translations

* redirect /ui/login

* fix http interceptor

* cleanup

* fix http interceptor

* fix: delete cookies on gateway 200

* add integration tests

* add command test

* docs

* fix integration tests

* add bulk api and integration test

* optimize bulk set limits

* unit test bulk limits

* fix broken link

* fix assets middleware

* fix broken link

* validate instance id format

* Update internal/eventstore/search_query.go

Co-authored-by: Livio Spring <livio.a@gmail.com>

* remove support for owner bulk limit commands

* project limits to instances

* migrate instances projection

* Revert "migrate instances projection"

This reverts commit 214218732a.

* join limits, remove owner

* remove todo

* use optional bool

* normally validate instance ids

* use 302

* cleanup

* cleanup

* Update internal/api/grpc/system/limits_converter.go

Co-authored-by: Livio Spring <livio.a@gmail.com>

* remove owner

* remove owner from reset

---------

Co-authored-by: Livio Spring <livio.a@gmail.com>

internal/api/grpc/server/gateway.go

@@ -185,7 +185,7 @@ func addInterceptors(
 	handler = http_mw.ActivityHandler(handler)
 	// For some non-obvious reason, the exhaustedCookieInterceptor sends the SetCookie header
 	// only if it follows the http_mw.DefaultTelemetryHandler
-	handler = exhaustedCookieInterceptor(handler, accessInterceptor, queries)
+	handler = exhaustedCookieInterceptor(handler, accessInterceptor)
 	handler = http_mw.DefaultMetricsHandler(handler)
 	return handler
 }
@@ -205,14 +205,12 @@ func http1Host(next http.Handler, http1HostName string) http.Handler {
 func exhaustedCookieInterceptor(
 	next http.Handler,
 	accessInterceptor *http_mw.AccessInterceptor,
-	queries *query.Queries,
 ) http.Handler {
 	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
 		next.ServeHTTP(&cookieResponseWriter{
 			ResponseWriter:    writer,
 			accessInterceptor: accessInterceptor,
 			request:           request,
-			queries:           queries,
 		}, request)
 	})
 }
@@ -221,7 +219,7 @@ type cookieResponseWriter struct {
 	http.ResponseWriter
 	accessInterceptor *http_mw.AccessInterceptor
 	request           *http.Request
-	queries           *query.Queries
+	headerWritten     bool
 }
 
 func (r *cookieResponseWriter) WriteHeader(status int) {
@@ -231,9 +229,18 @@ func (r *cookieResponseWriter) WriteHeader(status int) {
 	if status == http.StatusTooManyRequests {
 		r.accessInterceptor.SetExhaustedCookie(r.ResponseWriter, r.request)
 	}
+	r.headerWritten = true
 	r.ResponseWriter.WriteHeader(status)
 }
 
+func (r *cookieResponseWriter) Write(bytes []byte) (int, error) {
+	if !r.headerWritten {
+		// If no header was written before the data, the status code is 200 and we can delete the cookie
+		r.accessInterceptor.DeleteExhaustedCookie(r.ResponseWriter)
+	}
+	return r.ResponseWriter.Write(bytes)
+}
+
 func grpcCredentials(tlsConfig *tls.Config) credentials.TransportCredentials {
 	creds := insecure.NewCredentials()
 	if tlsConfig != nil {

internal/api/grpc/server/middleware/instance_interceptor_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"reflect"
 	"testing"
+	"time"
 
 	"golang.org/x/text/language"
 	"google.golang.org/grpc"
@@ -164,6 +165,14 @@ func (m *mockInstanceVerifier) InstanceByID(context.Context) (authz.Instance, er
 
 type mockInstance struct{}
 
+func (m *mockInstance) Block() *bool {
+	panic("shouldn't be called here")
+}
+
+func (m *mockInstance) AuditLogRetention() *time.Duration {
+	panic("shouldn't be called here")
+}
+
 func (m *mockInstance) InstanceID() string {
 	return "instanceID"
 }

internal/api/grpc/server/middleware/limits_interceptor.go

@@ -0,0 +1,31 @@
+package middleware
+
+import (
+	"context"
+	"strings"
+
+	"google.golang.org/grpc"
+
+	"github.com/zitadel/zitadel/internal/api/authz"
+	"github.com/zitadel/zitadel/internal/zerrors"
+)
+
+func LimitsInterceptor(ignoreService ...string) grpc.UnaryServerInterceptor {
+	for idx, service := range ignoreService {
+		if !strings.HasPrefix(service, "/") {
+			ignoreService[idx] = "/" + service
+		}
+	}
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (_ interface{}, err error) {
+		for _, service := range ignoreService {
+			if strings.HasPrefix(info.FullMethod, service) {
+				return handler(ctx, req)
+			}
+		}
+		instance := authz.GetInstance(ctx)
+		if block := instance.Block(); block != nil && *block {
+			return nil, zerrors.ThrowResourceExhausted(nil, "LIMITS-molsj", "Errors.Limits.Instance.Blocked")
+		}
+		return handler(ctx, req)
+	}
+}

internal/api/grpc/server/server.go

@@ -53,9 +53,10 @@ func CreateServer(
 				middleware.InstanceInterceptor(queries, hostHeaderName, system_pb.SystemService_ServiceDesc.ServiceName, healthpb.Health_ServiceDesc.ServiceName),
 				middleware.AccessStorageInterceptor(accessSvc),
 				middleware.ErrorHandler(),
+				middleware.LimitsInterceptor(system_pb.SystemService_ServiceDesc.ServiceName),
 				middleware.AuthorizationInterceptor(verifier, authConfig),
-				middleware.QuotaExhaustedInterceptor(accessSvc, system_pb.SystemService_ServiceDesc.ServiceName),
 				middleware.TranslationHandler(),
+				middleware.QuotaExhaustedInterceptor(accessSvc, system_pb.SystemService_ServiceDesc.ServiceName),
 				middleware.ValidationHandler(),
 				middleware.ServiceHandler(),
 				middleware.ActivityInterceptor(),

internal/api/grpc/system/limits.go

@@ -4,15 +4,12 @@ import (
 	"context"
 
 	"github.com/zitadel/zitadel/internal/api/grpc/object"
+	objectpb "github.com/zitadel/zitadel/pkg/grpc/object"
 	"github.com/zitadel/zitadel/pkg/grpc/system"
 )
 
 func (s *Server) SetLimits(ctx context.Context, req *system.SetLimitsRequest) (*system.SetLimitsResponse, error) {
-	details, err := s.command.SetLimits(
-		ctx,
-		req.GetInstanceId(),
-		instanceLimitsPbToCommand(req),
-	)
+	details, err := s.command.SetLimits(ctx, setInstanceLimitsPbToCommand(req))
 	if err != nil {
 		return nil, err
 	}
@@ -21,8 +18,23 @@ func (s *Server) SetLimits(ctx context.Context, req *system.SetLimitsRequest) (*
 	}, nil
 }
 
-func (s *Server) ResetLimits(ctx context.Context, req *system.ResetLimitsRequest) (*system.ResetLimitsResponse, error) {
-	details, err := s.command.ResetLimits(ctx, req.GetInstanceId())
+func (s *Server) BulkSetLimits(ctx context.Context, req *system.BulkSetLimitsRequest) (*system.BulkSetLimitsResponse, error) {
+	details, targetDetails, err := s.command.SetInstanceLimitsBulk(ctx, bulkSetInstanceLimitsPbToCommand(req))
+	if err != nil {
+		return nil, err
+	}
+	resp := &system.BulkSetLimitsResponse{
+		Details:       object.AddToDetailsPb(details.Sequence, details.EventDate, details.ResourceOwner),
+		TargetDetails: make([]*objectpb.ObjectDetails, len(targetDetails)),
+	}
+	for i := range targetDetails {
+		resp.TargetDetails[i] = object.AddToDetailsPb(targetDetails[i].Sequence, targetDetails[i].EventDate, targetDetails[i].ResourceOwner)
+	}
+	return resp, nil
+}
+
+func (s *Server) ResetLimits(ctx context.Context, _ *system.ResetLimitsRequest) (*system.ResetLimitsResponse, error) {
+	details, err := s.command.ResetLimits(ctx)
 	if err != nil {
 		return nil, err
 	}

internal/api/grpc/system/limits_converter.go

@@ -7,10 +7,23 @@ import (
 	"github.com/zitadel/zitadel/pkg/grpc/system"
 )
 
-func instanceLimitsPbToCommand(req *system.SetLimitsRequest) *command.SetLimits {
+func setInstanceLimitsPbToCommand(req *system.SetLimitsRequest) *command.SetLimits {
 	var setLimits = new(command.SetLimits)
 	if req.AuditLogRetention != nil {
 		setLimits.AuditLogRetention = gu.Ptr(req.AuditLogRetention.AsDuration())
 	}
+	setLimits.Block = req.Block
 	return setLimits
 }
+
+func bulkSetInstanceLimitsPbToCommand(req *system.BulkSetLimitsRequest) []*command.SetInstanceLimitsBulk {
+	cmds := make([]*command.SetInstanceLimitsBulk, len(req.Limits))
+	for i := range req.Limits {
+		setLimitsReq := req.Limits[i]
+		cmds[i] = &command.SetInstanceLimitsBulk{
+			InstanceID: setLimitsReq.GetInstanceId(),
+			SetLimits:  *setInstanceLimitsPbToCommand(req.Limits[i]),
+		}
+	}
+	return cmds
+}

internal/api/grpc/system/limits_integration_block_test.go

@@ -0,0 +1,287 @@
+//go:build integration
+
+package system_test
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/muhlemmer/gu"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/durationpb"
+
+	"github.com/zitadel/zitadel/internal/api/authz"
+	"github.com/zitadel/zitadel/internal/query"
+	"github.com/zitadel/zitadel/pkg/grpc/admin"
+	"github.com/zitadel/zitadel/pkg/grpc/system"
+)
+
+func TestServer_Limits_Block(t *testing.T) {
+	domain, instanceID, iamOwnerCtx := Tester.UseIsolatedInstance(t, CTX, SystemCTX)
+	tests := []*test{
+		publicAPIBlockingTest(domain),
+		{
+			name: "mutating API",
+			testGrpc: func(tt assert.TestingT, expectBlocked bool) {
+				randomGrpcIdpName := randomString("idp-grpc", 5)
+				_, err := Tester.Client.Admin.AddGitHubProvider(iamOwnerCtx, &admin.AddGitHubProviderRequest{
+					Name:         randomGrpcIdpName,
+					ClientId:     "client-id",
+					ClientSecret: "client-secret",
+				})
+				assertGrpcError(tt, err, expectBlocked)
+				//nolint:contextcheck
+				idpExists := idpExistsCondition(tt, instanceID, randomGrpcIdpName)
+				if expectBlocked {
+					// We ensure that the idp really is not created
+					assert.Neverf(tt, idpExists, 5*time.Second, 1*time.Second, "idp should never be created")
+				} else {
+					assert.Eventuallyf(tt, idpExists, 5*time.Second, 1*time.Second, "idp should be created")
+				}
+			},
+			testHttp: func(tt assert.TestingT) (*http.Request, error, func(assert.TestingT, *http.Response, bool)) {
+				randomHttpIdpName := randomString("idp-http", 5)
+				req, err := http.NewRequestWithContext(
+					CTX,
+					"POST",
+					fmt.Sprintf("http://%s/admin/v1/idps/github", net.JoinHostPort(domain, "8080")),
+					strings.NewReader(`{
+	"name": "`+randomHttpIdpName+`",
+	"clientId": "client-id",
+	"clientSecret": "client-secret"
+}`),
+				)
+				if err != nil {
+					return nil, err, nil
+				}
+				req.Header.Set("Authorization", Tester.BearerToken(iamOwnerCtx))
+				return req, nil, func(ttt assert.TestingT, response *http.Response, expectBlocked bool) {
+					assertLimitResponse(ttt, response, expectBlocked)
+					assertSetLimitingCookie(ttt, response, expectBlocked)
+				}
+			},
+		}, {
+			name: "discovery",
+			testHttp: func(tt assert.TestingT) (*http.Request, error, func(assert.TestingT, *http.Response, bool)) {
+				req, err := http.NewRequestWithContext(
+					CTX,
+					"GET",
+					fmt.Sprintf("http://%s/.well-known/openid-configuration", net.JoinHostPort(domain, "8080")),
+					nil,
+				)
+				return req, err, func(ttt assert.TestingT, response *http.Response, expectBlocked bool) {
+					assertLimitResponse(ttt, response, expectBlocked)
+					assertSetLimitingCookie(ttt, response, expectBlocked)
+				}
+			},
+		}, {
+			name: "login",
+			testHttp: func(tt assert.TestingT) (*http.Request, error, func(assert.TestingT, *http.Response, bool)) {
+				req, err := http.NewRequestWithContext(
+					CTX,
+					"GET",
+					fmt.Sprintf("http://%s/ui/login/login/externalidp/callback", net.JoinHostPort(domain, "8080")),
+					nil,
+				)
+				return req, err, func(ttt assert.TestingT, response *http.Response, expectBlocked bool) {
+					// the login paths should return a redirect if the instance is blocked
+					if expectBlocked {
+						assert.Equal(ttt, http.StatusFound, response.StatusCode)
+					} else {
+						assertLimitResponse(ttt, response, false)
+					}
+					assertSetLimitingCookie(ttt, response, expectBlocked)
+				}
+			},
+		}, {
+			name: "console",
+			testHttp: func(tt assert.TestingT) (*http.Request, error, func(assert.TestingT, *http.Response, bool)) {
+				req, err := http.NewRequestWithContext(
+					CTX,
+					"GET",
+					fmt.Sprintf("http://%s/ui/console/", net.JoinHostPort(domain, "8080")),
+					nil,
+				)
+				return req, err, func(ttt assert.TestingT, response *http.Response, expectBlocked bool) {
+					// the console is not blocked so we can render a link to an instance management portal.
+					// A CDN can cache these assets easily
+					// We also don't care about a cookie because the environment.json already takes care of that.
+					assertLimitResponse(ttt, response, false)
+				}
+			},
+		}, {
+			name: "environment.json",
+			testHttp: func(tt assert.TestingT) (*http.Request, error, func(assert.TestingT, *http.Response, bool)) {
+				req, err := http.NewRequestWithContext(
+					CTX,
+					"GET",
+					fmt.Sprintf("http://%s/ui/console/assets/environment.json", net.JoinHostPort(domain, "8080")),
+					nil,
+				)
+				return req, err, func(ttt assert.TestingT, response *http.Response, expectBlocked bool) {
+					// the environment.json should always return successfully
+					assertLimitResponse(ttt, response, false)
+					assertSetLimitingCookie(ttt, response, expectBlocked)
+					body, err := io.ReadAll(response.Body)
+					assert.NoError(ttt, err)
+					var compFunc assert.ComparisonAssertionFunc = assert.NotContains
+					if expectBlocked {
+						compFunc = assert.Contains
+					}
+					compFunc(ttt, string(body), `"exhausted":true`)
+				}
+			},
+		}}
+	_, err := Tester.Client.System.SetLimits(SystemCTX, &system.SetLimitsRequest{
+		InstanceId: instanceID,
+		Block:      gu.Ptr(true),
+	})
+	require.NoError(t, err)
+	// The following call ensures that an undefined bool is not deserialized to false
+	_, err = Tester.Client.System.SetLimits(SystemCTX, &system.SetLimitsRequest{
+		InstanceId:        instanceID,
+		AuditLogRetention: durationpb.New(time.Hour),
+	})
+	require.NoError(t, err)
+	for _, tt := range tests {
+		var isFirst bool
+		t.Run(tt.name+" with blocking", func(t *testing.T) {
+			isFirst = isFirst || !t.Skipped()
+			testBlockingAPI(t, tt, true, isFirst)
+		})
+	}
+	_, err = Tester.Client.System.SetLimits(SystemCTX, &system.SetLimitsRequest{
+		InstanceId: instanceID,
+		Block:      gu.Ptr(false),
+	})
+	require.NoError(t, err)
+	for _, tt := range tests {
+		var isFirst bool
+		t.Run(tt.name+" without blocking", func(t *testing.T) {
+			isFirst = isFirst || !t.Skipped()
+			testBlockingAPI(t, tt, false, isFirst)
+		})
+	}
+}
+
+type test struct {
+	name     string
+	testHttp func(t assert.TestingT) (req *http.Request, err error, assertResponse func(t assert.TestingT, response *http.Response, expectBlocked bool))
+	testGrpc func(t assert.TestingT, expectBlocked bool)
+}
+
+func testBlockingAPI(t *testing.T, tt *test, expectBlocked bool, isFirst bool) {
+	req, err, assertResponse := tt.testHttp(t)
+	require.NoError(t, err)
+	testHTTP := func(tt assert.TestingT) {
+		resp, err := (&http.Client{
+			// Don't follow redirects
+			CheckRedirect: func(req *http.Request, via []*http.Request) error {
+				return http.ErrUseLastResponse
+			},
+		}).Do(req)
+		defer func() {
+			require.NoError(t, resp.Body.Close())
+		}()
+		require.NoError(t, err)
+		assertResponse(t, resp, expectBlocked)
+	}
+	if isFirst {
+		// limits are eventually consistent, so we need to wait for the blocking to be set on the first test
+		assert.EventuallyWithT(t, func(c *assert.CollectT) {
+			testHTTP(c)
+		}, 15*time.Second, time.Second, "wait for blocking to be set")
+	} else {
+		testHTTP(t)
+	}
+	if tt.testGrpc != nil {
+		tt.testGrpc(t, expectBlocked)
+	}
+}
+
+func publicAPIBlockingTest(domain string) *test {
+	return &test{
+		name: "public API",
+		testGrpc: func(tt assert.TestingT, expectBlocked bool) {
+			conn, err := grpc.DialContext(CTX, net.JoinHostPort(domain, "8080"),
+				grpc.WithBlock(),
+				grpc.WithTransportCredentials(insecure.NewCredentials()),
+			)
+			assert.NoError(tt, err)
+			_, err = admin.NewAdminServiceClient(conn).Healthz(CTX, &admin.HealthzRequest{})
+			assertGrpcError(tt, err, expectBlocked)
+		},
+		testHttp: func(tt assert.TestingT) (*http.Request, error, func(assert.TestingT, *http.Response, bool)) {
+			req, err := http.NewRequestWithContext(
+				CTX,
+				"GET",
+				fmt.Sprintf("http://%s/admin/v1/healthz", net.JoinHostPort(domain, "8080")),
+				nil,
+			)
+			return req, err, func(ttt assert.TestingT, response *http.Response, expectBlocked bool) {
+				assertLimitResponse(ttt, response, expectBlocked)
+				assertSetLimitingCookie(ttt, response, expectBlocked)
+			}
+		},
+	}
+}
+
+// If expectSet is true, we expect the cookie to be set
+// If expectSet is false, we expect the cookie to be deleted
+func assertSetLimitingCookie(t assert.TestingT, response *http.Response, expectSet bool) {
+	for _, cookie := range response.Cookies() {
+		if cookie.Name == "zitadel.quota.exhausted" {
+			if expectSet {
+				assert.Greater(t, cookie.MaxAge, 0)
+			} else {
+				assert.LessOrEqual(t, cookie.MaxAge, 0)
+			}
+			return
+		}
+	}
+	assert.FailNow(t, "cookie not found")
+}
+
+func assertGrpcError(t assert.TestingT, err error, expectBlocked bool) {
+	if expectBlocked {
+		assert.Equal(t, codes.ResourceExhausted, status.Convert(err).Code())
+		return
+	}
+	assert.NoError(t, err)
+}
+
+func assertLimitResponse(t assert.TestingT, response *http.Response, expectBlocked bool) {
+	if expectBlocked {
+		assert.Equal(t, http.StatusTooManyRequests, response.StatusCode)
+		return
+	}
+	assert.GreaterOrEqual(t, response.StatusCode, 200)
+	assert.Less(t, response.StatusCode, 300)
+}
+
+func idpExistsCondition(t assert.TestingT, instanceID, idpName string) func() bool {
+	return func() bool {
+		nameQuery, err := query.NewIDPTemplateNameSearchQuery(query.TextEquals, idpName)
+		assert.NoError(t, err)
+		instanceQuery, err := query.NewIDPTemplateResourceOwnerSearchQuery(instanceID)
+		assert.NoError(t, err)
+		idps, err := Tester.Queries.IDPTemplates(authz.WithInstanceID(CTX, instanceID), &query.IDPTemplateSearchQueries{
+			Queries: []query.SearchQuery{
+				instanceQuery,
+				nameQuery,
+			},
+		}, false)
+		assert.NoError(t, err)
+		return len(idps.Templates) > 0
+	}
+}

internal/api/grpc/system/limits_integration_bulk_test.go

@@ -0,0 +1,75 @@
+//go:build integration
+
+package system_test
+
+import (
+	"testing"
+
+	"github.com/muhlemmer/gu"
+	"github.com/stretchr/testify/require"
+
+	"github.com/zitadel/zitadel/internal/integration"
+	"github.com/zitadel/zitadel/pkg/grpc/system"
+)
+
+func TestServer_Limits_Bulk(t *testing.T) {
+	const len = 5
+	type instance struct{ domain, id string }
+	instances := make([]*instance, len)
+	for i := 0; i < len; i++ {
+		domain := integration.RandString(5) + ".integration.localhost"
+		resp, err := Tester.Client.System.CreateInstance(SystemCTX, &system.CreateInstanceRequest{
+			InstanceName: "testinstance",
+			CustomDomain: domain,
+			Owner: &system.CreateInstanceRequest_Machine_{
+				Machine: &system.CreateInstanceRequest_Machine{
+					UserName: "owner",
+					Name:     "owner",
+				},
+			},
+		})
+		require.NoError(t, err)
+		instances[i] = &instance{domain, resp.GetInstanceId()}
+	}
+	resp, err := Tester.Client.System.BulkSetLimits(SystemCTX, &system.BulkSetLimitsRequest{
+		Limits: []*system.SetLimitsRequest{{
+			InstanceId: instances[0].id,
+			Block:      gu.Ptr(true),
+		}, {
+			InstanceId: instances[1].id,
+			Block:      gu.Ptr(false),
+		}, {
+			InstanceId: instances[2].id,
+			Block:      gu.Ptr(true),
+		}, {
+			InstanceId: instances[3].id,
+			Block:      gu.Ptr(false),
+		}, {
+			InstanceId: instances[4].id,
+			Block:      gu.Ptr(true),
+		}},
+	})
+	require.NoError(t, err)
+	details := resp.GetTargetDetails()
+	require.Len(t, details, len)
+	t.Run("the first instance is blocked", func(t *testing.T) {
+		require.Equal(t, instances[0].id, details[0].GetResourceOwner(), "resource owner must be instance id")
+		testBlockingAPI(t, publicAPIBlockingTest(instances[0].domain), true, true)
+	})
+	t.Run("the second instance isn't blocked", func(t *testing.T) {
+		require.Equal(t, instances[1].id, details[1].GetResourceOwner(), "resource owner must be instance id")
+		testBlockingAPI(t, publicAPIBlockingTest(instances[1].domain), false, true)
+	})
+	t.Run("the third instance is blocked", func(t *testing.T) {
+		require.Equal(t, instances[2].id, details[2].GetResourceOwner(), "resource owner must be instance id")
+		testBlockingAPI(t, publicAPIBlockingTest(instances[2].domain), true, true)
+	})
+	t.Run("the fourth instance isn't blocked", func(t *testing.T) {
+		require.Equal(t, instances[3].id, details[3].GetResourceOwner(), "resource owner must be instance id")
+		testBlockingAPI(t, publicAPIBlockingTest(instances[3].domain), false, true)
+	})
+	t.Run("the fifth instance is blocked", func(t *testing.T) {
+		require.Equal(t, instances[4].id, details[4].GetResourceOwner(), "resource owner must be instance id")
+		testBlockingAPI(t, publicAPIBlockingTest(instances[4].domain), true, true)
+	})
+}

internal/api/grpc/system/quotas_enabled/quota_integration_test.go

@@ -1,6 +1,6 @@
 //go:build integration
 
-package system_test
+package quotas_enabled_test
 
 import (
 	"bytes"
@@ -23,13 +23,12 @@ import (
 var callURL = "http://localhost:" + integration.PortQuotaServer
 
 func TestServer_QuotaNotification_Limit(t *testing.T) {
-	_, instanceID, iamOwnerCtx := Tester.UseIsolatedInstance(t, CTX, SystemCTX)
 	amount := 10
 	percent := 50
 	percentAmount := amount * percent / 100
 
 	_, err := Tester.Client.System.SetQuota(SystemCTX, &system.SetQuotaRequest{
-		InstanceId:    instanceID,
+		InstanceId:    Tester.Instance.InstanceID(),
 		Unit:          quota_pb.Unit_UNIT_REQUESTS_ALL_AUTHENTICATED,
 		From:          timestamppb.Now(),
 		ResetInterval: durationpb.New(time.Minute * 5),
@@ -51,23 +50,23 @@ func TestServer_QuotaNotification_Limit(t *testing.T) {
 	require.NoError(t, err)
 
 	for i := 0; i < percentAmount; i++ {
-		_, err := Tester.Client.Admin.GetDefaultOrg(iamOwnerCtx, &admin.GetDefaultOrgRequest{})
+		_, err := Tester.Client.Admin.GetDefaultOrg(IAMOwnerCTX, &admin.GetDefaultOrgRequest{})
 		require.NoErrorf(t, err, "error in %d call of %d", i, percentAmount)
 	}
 	awaitNotification(t, Tester.QuotaNotificationChan, quota.RequestsAllAuthenticated, percent)
 
 	for i := 0; i < (amount - percentAmount); i++ {
-		_, err := Tester.Client.Admin.GetDefaultOrg(iamOwnerCtx, &admin.GetDefaultOrgRequest{})
+		_, err := Tester.Client.Admin.GetDefaultOrg(IAMOwnerCTX, &admin.GetDefaultOrgRequest{})
 		require.NoErrorf(t, err, "error in %d call of %d", i, percentAmount)
 	}
 	awaitNotification(t, Tester.QuotaNotificationChan, quota.RequestsAllAuthenticated, 100)
 
-	_, limitErr := Tester.Client.Admin.GetDefaultOrg(iamOwnerCtx, &admin.GetDefaultOrgRequest{})
+	_, limitErr := Tester.Client.Admin.GetDefaultOrg(IAMOwnerCTX, &admin.GetDefaultOrgRequest{})
 	require.Error(t, limitErr)
 }
 
 func TestServer_QuotaNotification_NoLimit(t *testing.T) {
-	_, instanceID, iamOwnerCtx := Tester.UseIsolatedInstance(t, CTX, SystemCTX)
+	_, instanceID, IAMOwnerCTX := Tester.UseIsolatedInstance(t, CTX, SystemCTX)
 	amount := 10
 	percent := 50
 	percentAmount := amount * percent / 100
@@ -95,24 +94,24 @@ func TestServer_QuotaNotification_NoLimit(t *testing.T) {
 	require.NoError(t, err)
 
 	for i := 0; i < percentAmount; i++ {
-		_, err := Tester.Client.Admin.GetDefaultOrg(iamOwnerCtx, &admin.GetDefaultOrgRequest{})
+		_, err := Tester.Client.Admin.GetDefaultOrg(IAMOwnerCTX, &admin.GetDefaultOrgRequest{})
 		require.NoErrorf(t, err, "error in %d call of %d", i, percentAmount)
 	}
 	awaitNotification(t, Tester.QuotaNotificationChan, quota.RequestsAllAuthenticated, percent)
 
 	for i := 0; i < (amount - percentAmount); i++ {
-		_, err := Tester.Client.Admin.GetDefaultOrg(iamOwnerCtx, &admin.GetDefaultOrgRequest{})
+		_, err := Tester.Client.Admin.GetDefaultOrg(IAMOwnerCTX, &admin.GetDefaultOrgRequest{})
 		require.NoErrorf(t, err, "error in %d call of %d", i, percentAmount)
 	}
 	awaitNotification(t, Tester.QuotaNotificationChan, quota.RequestsAllAuthenticated, 100)
 
 	for i := 0; i < amount; i++ {
-		_, err := Tester.Client.Admin.GetDefaultOrg(iamOwnerCtx, &admin.GetDefaultOrgRequest{})
+		_, err := Tester.Client.Admin.GetDefaultOrg(IAMOwnerCTX, &admin.GetDefaultOrgRequest{})
 		require.NoErrorf(t, err, "error in %d call of %d", i, percentAmount)
 	}
 	awaitNotification(t, Tester.QuotaNotificationChan, quota.RequestsAllAuthenticated, 200)
 
-	_, limitErr := Tester.Client.Admin.GetDefaultOrg(iamOwnerCtx, &admin.GetDefaultOrgRequest{})
+	_, limitErr := Tester.Client.Admin.GetDefaultOrg(IAMOwnerCTX, &admin.GetDefaultOrgRequest{})
 	require.NoError(t, limitErr)
 }
 

internal/api/grpc/system/quotas_enabled/server_integration_test.go

@@ -0,0 +1,37 @@
+//go:build integration
+
+package quotas_enabled_test
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/zitadel/zitadel/internal/integration"
+)
+
+var (
+	CTX         context.Context
+	SystemCTX   context.Context
+	IAMOwnerCTX context.Context
+	Tester      *integration.Tester
+)
+
+func TestMain(m *testing.M) {
+	os.Exit(func() int {
+		ctx, _, cancel := integration.Contexts(5 * time.Minute)
+		defer cancel()
+		CTX = ctx
+
+		Tester = integration.NewTester(ctx, `
+Quotas:
+  Access:
+    Enabled: true
+`)
+		defer Tester.Done()
+		SystemCTX = Tester.WithAuthorization(ctx, integration.SystemUser)
+		IAMOwnerCTX = Tester.WithAuthorization(ctx, integration.IAMOwner)
+		return m.Run()
+	}())
+}