From a4605be73f5a6d0beea91d604fae897e5edcb713 Mon Sep 17 00:00:00 2001
From: Nicholas Molnar <65710+neekolas@users.noreply.github.com>
Date: Thu, 21 Mar 2024 09:33:52 -0700
Subject: [PATCH] Add proxy protocol

---
 go.mod                 | 1 +
 go.sum                 | 2 ++
 pkg/api/interceptor.go | 2 ++
 pkg/api/server.go      | 4 +++-
 4 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index 8a974cf6..83044d23 100644
--- a/go.mod
+++ b/go.mod
@@ -135,6 +135,7 @@ require (
 	github.com/opencontainers/runtime-spec v1.0.3-0.20210326190908-1c3f411f0417 // indirect
 	github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 // indirect
 	github.com/philhofer/fwd v1.1.1 // indirect
+	github.com/pires/go-proxyproto v0.7.0
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/prometheus/client_model v0.4.0 // indirect
 	github.com/prometheus/common v0.42.0 // indirect
diff --git a/go.sum b/go.sum
index 774f616a..c40fd15d 100644
--- a/go.sum
+++ b/go.sum
@@ -910,6 +910,8 @@ github.com/pierrec/lz4 v0.0.0-20190327172049-315a67e90e41/go.mod h1:3/3N9NVKO0je
 github.com/pierrec/lz4 v2.0.5+incompatible/go.mod h1:pdkljMzZIN41W+lC3N2tnIh5sFi+IEE17M5jbnwPHcY=
 github.com/pierrec/lz4 v2.5.2+incompatible/go.mod h1:pdkljMzZIN41W+lC3N2tnIh5sFi+IEE17M5jbnwPHcY=
 github.com/pierrec/lz4/v4 v4.1.14/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
+github.com/pires/go-proxyproto v0.7.0 h1:IukmRewDQFWC7kfnb66CSomk2q/seBuilHBYFwyq0Hs=
+github.com/pires/go-proxyproto v0.7.0/go.mod h1:Vz/1JPY/OACxWGQNIRY2BeyDmpoaWmEP40O9LbuiFR4=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
diff --git a/pkg/api/interceptor.go b/pkg/api/interceptor.go
index 88af2697..7915853b 100644
--- a/pkg/api/interceptor.go
+++ b/pkg/api/interceptor.go
@@ -149,6 +149,8 @@ func (wa *WalletAuthorizer) applyLimits(ctx context.Context, fullMethod string,
 		// requests without an IP address are bucketed together as "ip_unknown"
 		ip = "ip_unknown"
 	}
+	// TODO: Remove this noisy log line
+	wa.Log.Info("got peer address", logging.String("client_ip", ip))
 
 	// with no wallet apply regular limits
 	var isPriority bool
diff --git a/pkg/api/server.go b/pkg/api/server.go
index 96e12636..ff1139c7 100644
--- a/pkg/api/server.go
+++ b/pkg/api/server.go
@@ -29,6 +29,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 
+	"github.com/pires/go-proxyproto"
 	messagev1 "github.com/xmtp/xmtp-node-go/pkg/api/message/v1"
 	apicontext "github.com/xmtp/xmtp-node-go/pkg/api/message/v1/context"
 	mlsv1 "github.com/xmtp/xmtp-node-go/pkg/mls/api/v1"
@@ -86,7 +87,8 @@ func New(config *Config) (*Server, error) {
 func (s *Server) startGRPC() error {
 	var err error
 
-	s.grpcListener, err = net.Listen("tcp", addrString(s.GRPCAddress, s.GRPCPort))
+	grpcListener, err := net.Listen("tcp", addrString(s.GRPCAddress, s.GRPCPort))
+	s.grpcListener = &proxyproto.Listener{Listener: grpcListener}
 	if err != nil {
 		return errors.Wrap(err, "creating grpc listener")
 	}