@@ -92,6 +92,56 @@ func TestServe(t *testing.T) {
9292 }
9393}
9494
95+ func TestServeWithMemoryProtectionMiddleware (t * testing.T ) {
96+ t .Parallel ()
97+
98+ pool , err := dockertest .NewPool ("" )
99+ require .NoError (t , err )
100+
101+ serverToken := "mykey"
102+ serveResource , err := pool .RunWithOptions (& dockertest.RunOptions {
103+ Repository : "authzed/spicedb" ,
104+ Tag : "ci" ,
105+ Cmd : []string {"serve" , "--log-level=debug" , "--grpc-preshared-key" , serverToken , "--telemetry-endpoint=\" \" " },
106+ ExposedPorts : []string {"50051/tcp" },
107+ Env : []string {"GOMEMLIMIT=1B" }, // NOTE: Absurdly low on purpose
108+ }, func (config * docker.HostConfig ) {
109+ config .RestartPolicy = docker.RestartPolicy {
110+ Name : "no" ,
111+ }
112+ })
113+ require .NoError (t , err )
114+ t .Cleanup (func () {
115+ require .NoError (t , pool .Purge (serveResource ))
116+ })
117+
118+ serverPort := serveResource .GetPort ("50051/tcp" )
119+ conn , err := grpc .NewClient (fmt .Sprintf ("localhost:%s" , serverPort ),
120+ grpc .WithTransportCredentials (insecure .NewCredentials ()),
121+ grpcutil .WithInsecureBearerToken (serverToken ),
122+ )
123+
124+ require .NoError (t , err )
125+ t .Cleanup (func () {
126+ _ = conn .Close ()
127+ })
128+
129+ // Health requests bypass the memory middleware
130+ require .Eventually (t , func () bool {
131+ resp , err := healthpb .NewHealthClient (conn ).Check (context .Background (), & healthpb.HealthCheckRequest {Service : "authzed.api.v1.SchemaService" })
132+ return err == nil && resp .GetStatus () == healthpb .HealthCheckResponse_SERVING
133+ }, 5 * time .Second , 1 * time .Second , "server never became healthy" )
134+
135+ // Other requests have the memory middleware
136+ client := v1 .NewSchemaServiceClient (conn )
137+ _ , err = client .WriteSchema (context .Background (), & v1.WriteSchemaRequest {
138+ Schema : `definition user {}` ,
139+ })
140+ s , ok := status .FromError (err )
141+ require .True (t , ok )
142+ require .Equal (t , codes .ResourceExhausted , s .Code ())
143+ }
144+
95145func gracefulShutdown (pool * dockertest.Pool , serveResource * dockertest.Resource ) bool {
96146 closed := make (chan bool , 1 )
97147 go func () {
0 commit comments