feat: extract peer address from nginx headers

- Added additional interceptors function to log the address
- Tests to cover the client IP address case
Signed-off-by: Rohit Dandamudi <rohit.dandamudi@siderolabs.com>

cmd/discovery-service/main.go

@@ -124,12 +124,14 @@ func run(ctx context.Context, logger *zap.Logger) error {
 	serverOptions := []grpc.ServerOption{
 		grpc_middleware.WithUnaryServerChain(
 			grpc_ctxtags.UnaryServerInterceptor(grpc_ctxtags.WithFieldExtractor(server.FieldExtractor)),
+			server.AddPeerAddressUnaryServerInterceptor(),
 			grpc_zap.UnaryServerInterceptor(logger),
 			grpc_prometheus.UnaryServerInterceptor,
 			grpc_recovery.UnaryServerInterceptor(recoveryOpt),
 		),
 		grpc_middleware.WithStreamServerChain(
 			grpc_ctxtags.StreamServerInterceptor(grpc_ctxtags.WithFieldExtractor(server.FieldExtractor)),
+			server.AddPeerAddressStreamServerInterceptor(),
 			grpc_zap.StreamServerInterceptor(logger),
 			grpc_prometheus.StreamServerInterceptor,
 			grpc_recovery.StreamServerInterceptor(recoveryOpt),

pkg/server/addr.go

@@ -0,0 +1,39 @@
+// Copyright (c) 2021 Sidero Labs, Inc.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+
+package server
+
+import (
+	"context"
+	"net"
+
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/peer"
+	"inet.af/netaddr"
+)
+
+// PeerAddress is used to extract peer address from the client.
+// it will try to extract the actual client's IP when called via
+// Nginx ingress first if not it will get the nginx or the machine
+// which calls the server, if everything fails returns an empty address.
+func PeerAddress(ctx context.Context) netaddr.IP {
+	if md, ok := metadata.FromIncomingContext(ctx); ok {
+		if vals := md.Get("X-Real-IP"); vals != nil {
+			if ip, err := netaddr.ParseIP(vals[0]); err == nil {
+				return ip
+			}
+		}
+	}
+
+	if peer, ok := peer.FromContext(ctx); ok {
+		if addr, ok := peer.Addr.(*net.TCPAddr); ok {
+			if ip, ok := netaddr.FromStdIP(addr.IP); ok {
+				return ip
+			}
+		}
+	}
+
+	return netaddr.IP{}
+}

pkg/server/logging.go

@@ -6,6 +6,10 @@
 package server
 
 import (
+	"context"
+
+	grpc_ctxtags "github.com/grpc-ecosystem/go-grpc-middleware/tags"
+	"google.golang.org/grpc"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 )
@@ -31,3 +35,31 @@ func FieldExtractor(fullMethod string, req interface{}) map[string]interface{} {
 
 	return nil
 }
+
+// AddPeerAddressUnaryServerInterceptor sets peer.address for logging.
+func AddPeerAddressUnaryServerInterceptor() grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
+		extractPeerAddress(ctx)
+
+		return handler(ctx, req)
+	}
+}
+
+// AddPeerAddressStreamServerInterceptor sets peer.address for logging.
+func AddPeerAddressStreamServerInterceptor() grpc.StreamServerInterceptor {
+	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		ctx := ss.Context()
+
+		extractPeerAddress(ctx)
+
+		return handler(srv, ss)
+	}
+}
+
+func extractPeerAddress(ctx context.Context) {
+	if peerAddress := PeerAddress(ctx); !peerAddress.IsZero() {
+		if tags := grpc_ctxtags.Extract(ctx); tags != nil {
+			tags.Set("peer.address", peerAddress.String())
+		}
+	}
+}

pkg/server/server.go

@@ -9,16 +9,13 @@ package server
 import (
 	"context"
 	"errors"
-	"net"
 	"time"
 
 	prom "github.com/prometheus/client_golang/prometheus"
 	"github.com/talos-systems/discovery-api/api/v1alpha1/server/pb"
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
-	"inet.af/netaddr"
 
 	"github.com/talos-systems/discovery-service/internal/state"
 )
@@ -69,17 +66,8 @@ func (srv *ClusterServer) Hello(ctx context.Context, req *pb.HelloRequest) (*pb.
 
 	resp := &pb.HelloResponse{}
 
-	if peer, ok := peer.FromContext(ctx); ok {
-		if addr, ok := peer.Addr.(*net.TCPAddr); ok {
-			if ip, ok := netaddr.FromStdIP(addr.IP); ok {
-				var err error
-
-				resp.ClientIp, err = ip.MarshalBinary()
-				if err != nil {
-					return nil, err
-				}
-			}
-		}
+	if peerAddress := PeerAddress(ctx); !peerAddress.IsZero() {
+		resp.ClientIp, _ = peerAddress.MarshalBinary() //nolint:errcheck // never fails
 	}
 
 	return resp, nil

pkg/server/server_test.go

@@ -22,6 +22,7 @@ import (
 	"go.uber.org/zap/zaptest"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/durationpb"
@@ -101,6 +102,23 @@ func TestServerAPI(t *testing.T) {
 		assert.Equal(t, []byte{0x7f, 0x0, 0x0, 0x1}, resp.ClientIp) // 127.0.0.1
 	})
 
+	t.Run("HelloWithRealIP", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		ctx = metadata.AppendToOutgoingContext(ctx, "X-Real-IP", "1.2.3.4") // with real IP of client
+
+		resp, err := client.Hello(ctx, &pb.HelloRequest{
+			ClusterId:     "fake",
+			ClientVersion: "v0.12.0",
+		})
+		require.NoError(t, err)
+
+		assert.Equal(t, []byte{0x1, 0x2, 0x3, 0x4}, resp.ClientIp)
+	})
+
 	t.Run("AffiliateUpdate", func(t *testing.T) {
 		t.Parallel()