From 36961ed88220d6b2c76eae238aeb8bdd22d69326 Mon Sep 17 00:00:00 2001
From: maskedeken <52683904+maskedeken@users.noreply.github.com>
Date: Thu, 18 Mar 2021 01:55:51 +0800
Subject: [PATCH] Add remote addr to gRPC transport layer conn (#382)

Co-authored-by: RPRX <63339210+rprx@users.noreply.github.com>
---
 transport/internet/grpc/encoding/hunkconn.go  | 15 +++++++++++++++
 transport/internet/grpc/encoding/multiconn.go | 15 +++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/transport/internet/grpc/encoding/hunkconn.go b/transport/internet/grpc/encoding/hunkconn.go
index 50815144..6cfde6e1 100644
--- a/transport/internet/grpc/encoding/hunkconn.go
+++ b/transport/internet/grpc/encoding/hunkconn.go
@@ -5,12 +5,15 @@ import (
 	"io"
 	"net"
 
+	"google.golang.org/grpc/peer"
+
 	"github.com/xtls/xray-core/common/buf"
 	"github.com/xtls/xray-core/common/net/cnc"
 	"github.com/xtls/xray-core/common/signal/done"
 )
 
 type HunkConn interface {
+	Context() context.Context
 	Send(*Hunk) error
 	Recv() (*Hunk, error)
 	SendMsg(m interface{}) error
@@ -35,11 +38,23 @@ func NewHunkReadWriter(hc HunkConn, cancel context.CancelFunc) *HunkReaderWriter
 }
 
 func NewHunkConn(hc HunkConn, cancel context.CancelFunc) net.Conn {
+	var rAddr net.Addr
+	pr, ok := peer.FromContext(hc.Context())
+	if ok {
+		rAddr = pr.Addr
+	} else {
+		rAddr = &net.TCPAddr{
+			IP:   []byte{0, 0, 0, 0},
+			Port: 0,
+		}
+	}
+
 	wrc := NewHunkReadWriter(hc, cancel)
 	return cnc.NewConnection(
 		cnc.ConnectionInput(wrc),
 		cnc.ConnectionOutput(wrc),
 		cnc.ConnectionOnClose(wrc),
+		cnc.ConnectionRemoteAddr(rAddr),
 	)
 }
 
diff --git a/transport/internet/grpc/encoding/multiconn.go b/transport/internet/grpc/encoding/multiconn.go
index b2b865ab..d247b3eb 100644
--- a/transport/internet/grpc/encoding/multiconn.go
+++ b/transport/internet/grpc/encoding/multiconn.go
@@ -5,12 +5,15 @@ import (
 	"io"
 	"net"
 
+	"google.golang.org/grpc/peer"
+
 	"github.com/xtls/xray-core/common/buf"
 	"github.com/xtls/xray-core/common/net/cnc"
 	"github.com/xtls/xray-core/common/signal/done"
 )
 
 type MultiHunkConn interface {
+	Context() context.Context
 	Send(*MultiHunk) error
 	Recv() (*MultiHunk, error)
 	SendMsg(m interface{}) error
@@ -30,11 +33,23 @@ func NewMultiHunkReadWriter(hc MultiHunkConn, cancel context.CancelFunc) *MultiH
 }
 
 func NewMultiHunkConn(hc MultiHunkConn, cancel context.CancelFunc) net.Conn {
+	var rAddr net.Addr
+	pr, ok := peer.FromContext(hc.Context())
+	if ok {
+		rAddr = pr.Addr
+	} else {
+		rAddr = &net.TCPAddr{
+			IP:   []byte{0, 0, 0, 0},
+			Port: 0,
+		}
+	}
+
 	wrc := NewMultiHunkReadWriter(hc, cancel)
 	return cnc.NewConnection(
 		cnc.ConnectionInputMulti(wrc),
 		cnc.ConnectionOutputMulti(wrc),
 		cnc.ConnectionOnClose(wrc),
+		cnc.ConnectionRemoteAddr(rAddr),
 	)
 }