Xray-core/transport/internet/grpc/encoding/hunkconn.go
maskedeken 36961ed882
Add remote addr to gRPC transport layer conn (#382)
Co-authored-by: RPRX <63339210+rprx@users.noreply.github.com>
2021-03-17 17:55:51 +00:00

139 lines
2.5 KiB
Go

package encoding
import (
"context"
"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
RecvMsg(m interface{}) error
}
type StreamCloser interface {
CloseSend() error
}
type HunkReaderWriter struct {
hc HunkConn
cancel context.CancelFunc
done *done.Instance
buf []byte
index int
}
func NewHunkReadWriter(hc HunkConn, cancel context.CancelFunc) *HunkReaderWriter {
return &HunkReaderWriter{hc, cancel, done.New(), nil, 0}
}
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),
)
}
func (h *HunkReaderWriter) forceFetch() error {
hunk, err := h.hc.Recv()
if err != nil {
if err == io.EOF {
return err
}
return newError("failed to fetch hunk from gRPC tunnel").Base(err)
}
h.buf = hunk.Data
h.index = 0
return nil
}
func (h *HunkReaderWriter) Read(buf []byte) (int, error) {
if h.done.Done() {
return 0, io.EOF
}
if h.index >= len(h.buf) {
if err := h.forceFetch(); err != nil {
return 0, err
}
}
n := copy(buf, h.buf[h.index:])
h.index += n
return n, nil
}
func (h *HunkReaderWriter) ReadMultiBuffer() (buf.MultiBuffer, error) {
if h.done.Done() {
return nil, io.EOF
}
if h.index >= len(h.buf) {
if err := h.forceFetch(); err != nil {
return nil, err
}
}
if cap(h.buf) == buf.Size {
b := h.buf
h.index = len(h.buf)
return buf.MultiBuffer{buf.NewExisted(b)}, nil
}
b := buf.New()
_, err := b.ReadFrom(h)
if err != nil {
return nil, err
}
return buf.MultiBuffer{b}, nil
}
func (h *HunkReaderWriter) Write(buf []byte) (int, error) {
if h.done.Done() {
return 0, io.ErrClosedPipe
}
err := h.hc.Send(&Hunk{Data: buf[:]})
if err != nil {
return 0, newError("failed to send data over gRPC tunnel").Base(err)
}
return len(buf), nil
}
func (h *HunkReaderWriter) Close() error {
if h.cancel != nil {
h.cancel()
}
if sc, match := h.hc.(StreamCloser); match {
return sc.CloseSend()
}
return h.done.Close()
}