diff --git a/app.go b/app.go index 9ce4b122256..1802d4f902f 100644 --- a/app.go +++ b/app.go @@ -610,7 +610,15 @@ func (h *Headscale) Serve() error { } // Start the local gRPC server without TLS and without authentication - grpcSocket := grpc.NewServer(zerolog.UnaryInterceptor()) + grpcOptions := []grpc.ServerOption{ + grpc.UnaryInterceptor( + grpcMiddleware.ChainUnaryServer( + serverErrorInterceptor, + zerolog.NewUnaryServerInterceptor(), + ), + ), + } + grpcSocket := grpc.NewServer(grpcOptions...) v1.RegisterHeadscaleServiceServer(grpcSocket, newHeadscaleV1APIServer(h)) reflection.Register(grpcSocket) @@ -650,6 +658,7 @@ func (h *Headscale) Serve() error { grpc.UnaryInterceptor( grpcMiddleware.ChainUnaryServer( h.grpcAuthenticationInterceptor, + serverErrorInterceptor, zerolog.NewUnaryServerInterceptor(), ), ), @@ -991,3 +1000,13 @@ func readOrCreatePrivateKey(path string) (*key.MachinePrivate, error) { return &machineKey, nil } + +func serverErrorInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + resp, err := handler(ctx, req) + + if errors.Is(err, gorm.ErrRecordNotFound) { + return resp, status.Error(codes.NotFound, err.Error()) + } + + return resp, err +}