feat: hide sensitive headers

Semior001 · Mar 10, 2024 · f4d23ca · f4d23ca
1 parent 6352185
commit f4d23ca
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 35 deletions.
diff --git a/cmd/groxy/main.go b/cmd/groxy/main.go
@@ -75,7 +75,13 @@ func run(ctx context.Context) error {
 			Delay:         opts.File.Delay,
 		},
 	}}
-	srv := proxy.NewServer(dsvc, proxy.Version(getVersion()))
+
+	proxyOpts := []proxy.Option{proxy.Version(getVersion())}
+	if opts.Debug {
+		proxyOpts = append(proxyOpts, proxy.Debug())
+	}
+
+	srv := proxy.NewServer(dsvc, proxyOpts...)
 
 	ewg, ctx := errgroup.WithContext(ctx)
 	ewg.Go(func() error {

diff --git a/pkg/proxy/middleware/log.go b/pkg/proxy/middleware/log.go
@@ -14,40 +14,76 @@ import (
 )
 
 // Log logs the gRPC requests.
-func Log(next grpc.StreamHandler) grpc.StreamHandler {
-	return func(srv any, stream grpc.ServerStream) (err error) {
-		ctx := stream.Context()
-		ss := &statsStream{ServerStream: stream}
-
-		start := time.Now()
-		defer func() {
-			elapsed := time.Since(start)
-			mtd, ok := grpc.Method(ctx)
-			if !ok {
-				mtd = "unknown"
-			}
-
-			pi, ok := peer.FromContext(ctx)
-			if !ok {
-				pi = &peer.Peer{Addr: &net.IPAddr{IP: net.IPv4zero}}
-			}
-
-			slog.InfoContext(ctx, "request",
-				slog.String("uri", mtd),
-				slog.String("remote", pi.Addr.String()),
-				slog.Duration("elapsed", elapsed),
-				slog.Int("recv_count", ss.recvCount),
-				slog.Int64("recv_size", ss.recvSize),
-				slog.Int("send_count", ss.sendCount),
-				slog.Int64("send_size", ss.sendSize),
-				slogx.Error(err),
-			)
-		}()
-
-		return next(srv, ss)
+func Log(debug bool) func(next grpc.StreamHandler) grpc.StreamHandler {
+	return func(next grpc.StreamHandler) grpc.StreamHandler {
+		return func(srv any, stream grpc.ServerStream) (err error) {
+			ctx := stream.Context()
+			ss := &statsStream{ServerStream: stream}
+
+			start := time.Now()
+			defer func() {
+				elapsed := time.Since(start)
+				mtd, ok := grpc.Method(ctx)
+				if !ok {
+					mtd = "unknown"
+				}
+
+				pi, ok := peer.FromContext(ctx)
+				if !ok {
+					pi = &peer.Peer{Addr: &net.IPAddr{IP: net.IPv4zero}}
+				}
+
+				attrs := []any{
+					slog.String("uri", mtd),
+					slog.String("remote", pi.Addr.String()),
+					slog.Duration("elapsed", elapsed),
+					slog.Int("recv_count", ss.recvCount),
+					slog.Int64("recv_size", ss.recvSize),
+					slog.Int("send_count", ss.sendCount),
+					slog.Int64("send_size", ss.sendSize),
+					slogx.Error(err),
+				}
+
+				if debug {
+					reqHeader, _ := metadata.FromIncomingContext(ctx)
+					attrs = append(attrs,
+						slog.Any("request_header", filterMD(reqHeader)),
+						slog.Any("response_header", filterMD(ss.header)),
+						slog.Any("response_trailer", filterMD(ss.trailer)),
+					)
+				}
+
+				slog.InfoContext(ctx, "request", attrs...)
+			}()
+
+			return next(srv, ss)
+		}
 	}
 }
 
+var hideHeaders = map[string]struct{}{
+	"authorization": {},
+	"cookie":        {},
+	"set-cookie":    {},
+}
+
+func filterMD(md metadata.MD) metadata.MD {
+	if md == nil {
+		return nil
+	}
+
+	out := make(metadata.MD)
+	for k, v := range md {
+		if _, ok := hideHeaders[k]; ok {
+			out[k] = []string{"***"}
+			continue
+		}
+		out[k] = v
+	}
+
+	return out
+}
+
 type statsStream struct {
 	lo      sync.RWMutex
 	header  metadata.MD

diff --git a/pkg/proxy/option.go b/pkg/proxy/option.go
@@ -14,3 +14,6 @@ func Version(v string) Option {
 func WithGRPCServerOptions(opts ...grpc.ServerOption) Option {
 	return func(o *Server) { o.serverOpts = append(o.serverOpts, opts...) }
 }
+
+// Debug sets the debug mode.
+func Debug() Option { return func(s *Server) { s.debug = true } }
diff --git a/pkg/proxy/proxy.go b/pkg/proxy/proxy.go
@@ -25,8 +25,9 @@ type Server struct {
 	defaultResponder func(stream grpc.ServerStream, firstRecv []byte) error
 	matcher          *discovery.Service
 
-	l    net.Listener
-	grpc *grpc.Server
+	debug bool
+	l     net.Listener
+	grpc  *grpc.Server
 }
 
 // NewServer creates a new server.
@@ -55,7 +56,7 @@ func (s *Server) Listen(addr string) (err error) {
 		grpc.UnknownServiceHandler(middleware.Chain(s.handle,
 			middleware.Recoverer,
 			middleware.AppInfo("groxy", "Semior001", s.version),
-			middleware.Log,
+			middleware.Log(s.debug),
 		)),
 		grpc.ForceServerCodec(RawBytesCodec{}),
 	)...)