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

Co-authored-by: RPRX <63339210+rprx@users.noreply.github.com>
This commit is contained in:
maskedeken 2021-03-18 01:55:51 +08:00 committed by GitHub
parent de54d4b08f
commit 36961ed882
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 30 additions and 0 deletions

View file

@ -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),
)
}

View file

@ -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),
)
}