Переглянути джерело

Add remote addr to gRPC transport layer conn (#382)

Co-authored-by: RPRX <[email protected]>
maskedeken 5 роки тому
батько
коміт
36961ed882

+ 15 - 0
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),
 	)
 }
 

+ 15 - 0
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),
 	)
 }