Fix Logging in Collectsub Server (#995)

* fixing csub server logger Signed-off-by: Soham Arora <arorasoham9@gmail.com> * mod and sum Signed-off-by: Soham Arora <arorasoham9@gmail.com> * lowercase leading letter for aded functions Signed-off-by: Soham Arora <arorasoham9@gmail.com> --------- Signed-off-by: Soham Arora <arorasoham9@gmail.com>
guacsec · Jul 25, 2023 · 093d702 · 093d702
1 parent 9368f3a
commit 093d702
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 3 deletions.
diff --git a/go.mod b/go.mod
@@ -174,6 +174,7 @@ require (
 	github.com/golang/mock v1.6.0
 	github.com/google/go-github/v50 v50.2.0
 	github.com/google/osv-scanner v1.3.6
+	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
 	github.com/jedib0t/go-pretty/v6 v6.4.6
 	github.com/jeremywohl/flatten v1.0.1
 	github.com/manifoldco/promptui v0.9.0

diff --git a/go.sum b/go.sum
@@ -1571,6 +1571,7 @@ github.com/grafana/regexp v0.0.0-20221122212121-6b5c0a4cb7fd/go.mod h1:M5qHK+eWf
 github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.0/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
 github.com/grpc-ecosystem/go-grpc-middleware v1.0.1-0.20190118093823-f849b5445de4/go.mod h1:FiyG127CGDf3tlThmgyCl78X/SZQqEOJBCDaAfeWzPs=
+github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 h1:+9834+KizmvFV7pXQGSXQTsaWhq2GjuNUt0aUU0YBYw=
 github.com/grpc-ecosystem/go-grpc-middleware v1.3.0/go.mod h1:z0ButlSOZa5vEBq9m2m2hlwIgKw+rp3sdCBRoJY+30Y=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
 github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=

diff --git a/pkg/collectsub/server/server.go b/pkg/collectsub/server/server.go
@@ -22,11 +22,16 @@ import (
 	"sync"
 	"time"
 
+	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
+	grpc_zap "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap"
+	"github.com/grpc-ecosystem/go-grpc-middleware/logging/zap/ctxzap"
 	pb "github.com/guacsec/guac/pkg/collectsub/collectsub"
 	"github.com/guacsec/guac/pkg/collectsub/server/db/simpledb"
 	db "github.com/guacsec/guac/pkg/collectsub/server/db/types"
 	"github.com/guacsec/guac/pkg/logging"
+	"go.uber.org/zap"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 )
 
 type server struct {
@@ -50,7 +55,7 @@ func NewServer(port int) (*server, error) {
 }
 
 func (s *server) AddCollectEntries(ctx context.Context, in *pb.AddCollectEntriesRequest) (*pb.AddCollectEntriesResponse, error) {
-	logger := logging.FromContext(ctx)
+	logger := ctxzap.Extract(ctx).Sugar()
 	logger.Infof("AddCollectEntries called with entries: %v", in.Entries)
 
 	err := s.Db.AddCollectEntries(ctx, in.Entries)
@@ -64,7 +69,7 @@ func (s *server) AddCollectEntries(ctx context.Context, in *pb.AddCollectEntries
 }
 
 func (s *server) GetCollectEntries(ctx context.Context, in *pb.GetCollectEntriesRequest) (*pb.GetCollectEntriesResponse, error) {
-	logger := logging.FromContext(ctx)
+	logger := ctxzap.Extract(ctx).Sugar()
 	logger.Infof("GetCollectEntries called with filters: %v", in.Filters)
 
 	ret, err := s.Db.GetCollectEntries(ctx, in.Filters, in.SinceTime)
@@ -78,13 +83,52 @@ func (s *server) GetCollectEntries(ctx context.Context, in *pb.GetCollectEntries
 	}, nil
 }
 
+func contextPropagationUnaryServerInterceptor() grpc.UnaryServerInterceptor {
+	return func(
+		ctx context.Context,
+		req interface{},
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (interface{}, error) {
+		if md, ok := metadata.FromIncomingContext(ctx); ok {
+			ctx = metadata.NewOutgoingContext(ctx, md)
+		}
+		return handler(ctx, req)
+	}
+}
+
+func contextToZapFieldsUnaryServerInterceptor() grpc.UnaryServerInterceptor {
+	return func(
+		ctx context.Context,
+		req interface{},
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (interface{}, error) {
+		if md, ok := metadata.FromIncomingContext(ctx); ok {
+			for k, v := range md {
+				ctxzap.AddFields(ctx, zap.Strings(k, v))
+			}
+		}
+		return handler(ctx, req)
+	}
+}
+
 func (s *server) Serve(ctx context.Context) error {
 	logger := logging.FromContext(ctx)
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", s.port))
 	if err != nil {
 		return fmt.Errorf("error opening port %d when starting csub server: %w", s.port, err)
 	}
-	gs := grpc.NewServer()
+	opts := []grpc.ServerOption{
+		grpc.UnaryInterceptor(
+			grpc_middleware.ChainUnaryServer(
+				contextPropagationUnaryServerInterceptor(),
+				grpc_zap.UnaryServerInterceptor(logger.Desugar()),
+				contextToZapFieldsUnaryServerInterceptor(),
+			)),
+	}
+	gs := grpc.NewServer(opts...)
+
 	pb.RegisterColectSubscriberServiceServer(gs, s)
 
 	var wg sync.WaitGroup