refactor: use remote service when it is available

1 files changed, 14 insertions, 15 deletions

diff --git a/pkg/authz/server.go b/pkg/authz/server.go
index 6eedcca..4ab7c17 100644
--- a/pkg/authz/server.go
+++ b/pkg/authz/server.go
@@ -27,25 +27,24 @@ func New(ctx context.Context, options ...grpc.ServerOption) *Server {
 		grpc.StreamInterceptor(pls.LogGRPCStream(logger)),
 	)...)
 
-	connection := Connection.From(ctx)
+	auth.RegisterAuthorizationServer(server, authorizationServiceFor(ctx))
+	reflection.Register(server)
+
+	return &Server{
+		Server: server,
+	}
+}
 
-	if x.IsZero(connection) {
-		auth.RegisterAuthorizationServer(server, NewLocalCheckService())
-	} else {
+func authorizationServiceFor(ctx context.Context) auth.AuthorizationServer {
+	connection := Connection.From(ctx)
+	svcs := []auth.AuthorizationServer{}
+	if x.IsPresent(connection) {
 		pls.LogNow(ctx, log.Fields{"authzd": map[string]string{
 			"target": connection.CanonicalTarget(),
 			"state":  connection.GetState().String(),
 		}})
-		auth.RegisterAuthorizationServer(
-			server,
-			NewRemoteCheckService(
-				auth.NewAuthorizationClient(connection),
-			),
-		)
-	}
-	reflection.Register(server)
-
-	return &Server{
-		Server: server,
+		svcs = append(svcs, NewRemoteCheckService(auth.NewAuthorizationClient(connection)))
 	}
+	svcs = append(svcs, NewLocalCheckService())
+	return NewCheckService(svcs)
 }