From df5a8f6d3b2e218e4b4b619488aa3f242bca188a 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    | 28 +++++++++++++++++++++-------
 flake.nix |  2 +-
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/app.go b/app.go
index 9ce4b12225..d58f0db886 100644
--- a/app.go
+++ b/app.go
@@ -19,7 +19,6 @@ import (
 
 	"github.com/coreos/go-oidc/v3/oidc"
 	"github.com/gorilla/mux"
-	grpcMiddleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
 	v1 "github.com/ori-edge/headscale/gen/go/headscale/v1"
 	"github.com/patrickmn/go-cache"
@@ -610,7 +609,13 @@ func (h *Headscale) Serve() error {
 	}
 
 	// Start the local gRPC server without TLS and without authentication
-	grpcSocket := grpc.NewServer(zerolog.UnaryInterceptor())
+	grpcOptions := []grpc.ServerOption{
+		grpc.ChainUnaryInterceptor(
+			serverErrorInterceptor,
+			zerolog.NewUnaryServerInterceptor(),
+		),
+	}
+	grpcSocket := grpc.NewServer(grpcOptions...)
 
 	v1.RegisterHeadscaleServiceServer(grpcSocket, newHeadscaleV1APIServer(h))
 	reflection.Register(grpcSocket)
@@ -647,11 +652,10 @@ func (h *Headscale) Serve() error {
 		log.Info().Msgf("Enabling remote gRPC at %s", h.cfg.GRPCAddr)
 
 		grpcOptions := []grpc.ServerOption{
-			grpc.UnaryInterceptor(
-				grpcMiddleware.ChainUnaryServer(
-					h.grpcAuthenticationInterceptor,
-					zerolog.NewUnaryServerInterceptor(),
-				),
+			grpc.ChainUnaryInterceptor(
+				h.grpcAuthenticationInterceptor,
+				serverErrorInterceptor,
+				zerolog.NewUnaryServerInterceptor(),
 			),
 		}
 
@@ -991,3 +995,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
+}
diff --git a/flake.nix b/flake.nix
index 7e26bb8309..a58d9de592 100644
--- a/flake.nix
+++ b/flake.nix
@@ -36,7 +36,7 @@
 
             # When updating go.mod or go.sum, a new sha will need to be calculated,
             # update this if you have a mismatch after doing a change to thos files.
-            vendorSha256 = "sha256-46lWrzQpqSSh3riujKVw+1V1HYaWOdgbT9adqvWp6U4=";
+            vendorSha256 = "sha256-hSlmgjRpulu+RRg26udtOq4lQj6JVJK+3lnvu5gKyso=";
 
             ldflags = [ "-s" "-w" "-X github.com/ori-edge/headscale/cmd/headscale/cli.Version=v${version}" ];
           };