From f933db8ebbaf26a7c97ede5d0eafd2926fa72f2c Mon Sep 17 00:00:00 2001
From: goncalo-rodrigues <goncalo.alfredo.rodrigues@gmail.com>
Date: Thu, 2 Nov 2023 10:46:38 +0100
Subject: [PATCH] [OGE-5650] Return not found status code when db returns not
 found

---
 app.go | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

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
+}