Extract peer real IP from Load Balancer when possible

go.mod

@@ -27,7 +27,7 @@ require (
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 	golang.zx2c4.com/wireguard/windows v0.5.3
 	google.golang.org/grpc v1.56.3
-	google.golang.org/protobuf v1.30.0
+	google.golang.org/protobuf v1.31.0
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 )
 

go.sum

@@ -940,8 +940,9 @@ google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlba
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
-google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng=
 google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
+google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
+google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

management/server/grpcserver.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"fmt"
+	"net/netip"
 	"strings"
 	"time"
 
@@ -11,6 +12,7 @@ import (
 	log "github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	gRPCPeer "google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
 
@@ -109,17 +111,45 @@ func (s *GRPCServer) GetServerKey(ctx context.Context, req *proto.Empty) (*proto
 	}, nil
 }
 
+// getRealIP returns real peer IP from load balance or fallback to IP of the GRPC peer
+func getRealIP(ctx context.Context) string {
+	if headers, ok := metadata.FromIncomingContext(ctx); ok {
+		value := headers.Get("x-forwarded-for")
+		if len(value) == 0 || value[0] == "" {
+			value = headers.Get("x-real-ip")
+		}
+
+		if len(value) > 0 && value[0] != "" {
+			ips := strings.Split(value[0], ",")
+			ipStr := strings.TrimSpace(ips[len(ips)-1])
+			ip, err := netip.ParseAddr(ipStr)
+			if err == nil {
+				return ip.String()
+			}
+		}
+
+		if grpcPeer, ok := gRPCPeer.FromContext(ctx); ok {
+			addrPort, err := netip.ParseAddrPort(grpcPeer.Addr.String())
+			if err == nil {
+				return addrPort.Addr().String()
+			}
+			log.Errorf("when resolve ip address of GRPC peer: %s", err)
+		}
+	}
+
+	// not happen normally as we should be able to resolve ip address above
+	return ""
+}
+
 // Sync validates the existence of a connecting peer, sends an initial state (all available for the connecting peers) and
 // notifies the connected peer of any updates (e.g. new peers under the same account)
 func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_SyncServer) error {
 	reqStart := time.Now()
 	if s.appMetrics != nil {
 		s.appMetrics.GRPCMetrics().CountSyncRequest()
 	}
-	p, ok := gRPCPeer.FromContext(srv.Context())
-	if ok {
-		log.Debugf("Sync request from peer [%s] [%s]", req.WgPubKey, p.Addr.String())
-	}
+	realIP := getRealIP(srv.Context())
+	log.Debugf("Sync request from peer [%s] [%s]", req.WgPubKey, realIP)
 
 	syncReq := &proto.SyncRequest{}
 	peerKey, err := s.parseRequest(req, syncReq)
@@ -290,10 +320,8 @@ func (s *GRPCServer) Login(ctx context.Context, req *proto.EncryptedMessage) (*p
 	if s.appMetrics != nil {
 		s.appMetrics.GRPCMetrics().CountLoginRequest()
 	}
-	p, ok := gRPCPeer.FromContext(ctx)
-	if ok {
-		log.Debugf("Login request from peer [%s] [%s]", req.WgPubKey, p.Addr.String())
-	}
+	realIP := getRealIP(ctx)
+	log.Debugf("Login request from peer [%s] [%s]", req.WgPubKey, realIP)
 
 	loginReq := &proto.LoginRequest{}
 	peerKey, err := s.parseRequest(req, loginReq)
@@ -303,8 +331,7 @@ func (s *GRPCServer) Login(ctx context.Context, req *proto.EncryptedMessage) (*p
 
 	if loginReq.GetMeta() == nil {
 		msg := status.Errorf(codes.FailedPrecondition,
-			"peer system meta has to be provided to log in. Peer %s, remote addr %s", peerKey.String(),
-			p.Addr.String())
+			"peer system meta has to be provided to log in. Peer %s, remote addr %s", peerKey.String(), realIP)
 		log.Warn(msg)
 		return nil, msg
 	}