Files
tun2socks/tunnel/tcp.go
LanceLi 2d0bd1d219 Improve: tunnel/tcp pipe (#219)
Co-authored-by: xjasonlyu <xjasonlyu@gmail.com>
2023-04-03 19:15:24 +08:00

76 lines
1.9 KiB
Go

package tunnel
import (
"io"
"net"
"sync"
"time"
"github.com/xjasonlyu/tun2socks/v2/common/pool"
"github.com/xjasonlyu/tun2socks/v2/core/adapter"
"github.com/xjasonlyu/tun2socks/v2/log"
M "github.com/xjasonlyu/tun2socks/v2/metadata"
"github.com/xjasonlyu/tun2socks/v2/proxy"
"github.com/xjasonlyu/tun2socks/v2/tunnel/statistic"
)
const (
// tcpWaitTimeout implements a TCP half-close timeout.
tcpWaitTimeout = 60 * time.Second
)
func handleTCPConn(originConn adapter.TCPConn) {
defer originConn.Close()
id := originConn.ID()
metadata := &M.Metadata{
Network: M.TCP,
SrcIP: net.IP(id.RemoteAddress),
SrcPort: id.RemotePort,
DstIP: net.IP(id.LocalAddress),
DstPort: id.LocalPort,
}
remoteConn, err := proxy.Dial(metadata)
if err != nil {
log.Warnf("[TCP] dial %s: %v", metadata.DestinationAddress(), err)
return
}
metadata.MidIP, metadata.MidPort = parseAddr(remoteConn.LocalAddr())
remoteConn = statistic.DefaultTCPTracker(remoteConn, metadata)
defer remoteConn.Close()
log.Infof("[TCP] %s <-> %s", metadata.SourceAddress(), metadata.DestinationAddress())
pipe(originConn, remoteConn)
}
// pipe copies copy data to & from provided net.Conn(s) bidirectionally.
func pipe(origin, remote net.Conn) {
wg := sync.WaitGroup{}
wg.Add(2)
go unidirectionalStream(remote, origin, "origin->remote", &wg)
go unidirectionalStream(origin, remote, "remote->origin", &wg)
wg.Wait()
}
func unidirectionalStream(dst, src net.Conn, dir string, wg *sync.WaitGroup) {
defer wg.Done()
buf := pool.Get(pool.RelayBufferSize)
if _, err := io.CopyBuffer(dst, src, buf); err != nil {
log.Debugf("[TCP] copy data for %s: %v", dir, err)
}
pool.Put(buf)
// Do the upload/download side TCP half-close.
if cr, ok := src.(interface{ CloseRead() error }); ok {
cr.CloseRead()
}
if cw, ok := dst.(interface{ CloseWrite() error }); ok {
cw.CloseWrite()
}
// Set TCP half-close timeout.
dst.SetReadDeadline(time.Now().Add(tcpWaitTimeout))
}