535 lines
13 KiB
Go
535 lines
13 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"dragontcp/internal/cover"
|
|
"dragontcp/internal/protocol"
|
|
"dragontcp/internal/wire"
|
|
)
|
|
|
|
const (
|
|
bpModeProbe byte = 0
|
|
bpModeUpload byte = 1
|
|
bpModeDownload byte = 2
|
|
bpModeBatchDownload byte = 3
|
|
bpModeACK byte = 4
|
|
bpHeaderSize = 29
|
|
)
|
|
|
|
var bpOpenMagic = [4]byte{'D', 'O', 'P', '1'}
|
|
var bpCloseMagic = [4]byte{'D', 'C', 'L', '1'}
|
|
|
|
type bpPhysicalConn struct {
|
|
conn net.Conn
|
|
requests int
|
|
}
|
|
|
|
type bpLane struct {
|
|
mu sync.Mutex
|
|
serverAddr string
|
|
tcpBuffer int
|
|
reconnectEvery int
|
|
timeout time.Duration
|
|
coverProfile cover.Profile
|
|
autoReconnect bool
|
|
pc *bpPhysicalConn
|
|
closed bool
|
|
}
|
|
|
|
func newBPLane(serverAddr string, tcpBuffer, reconnectEvery int, timeout time.Duration, coverProfile cover.Profile) *bpLane {
|
|
autoReconnect := reconnectEvery == 1
|
|
if autoReconnect {
|
|
reconnectEvery = 0
|
|
}
|
|
return &bpLane{
|
|
serverAddr: serverAddr,
|
|
tcpBuffer: tcpBuffer,
|
|
reconnectEvery: reconnectEvery,
|
|
timeout: timeout,
|
|
coverProfile: coverProfile,
|
|
autoReconnect: autoReconnect,
|
|
}
|
|
}
|
|
|
|
func (l *bpLane) transportFailureLocked(reused bool) {
|
|
if l.autoReconnect && reused {
|
|
l.reconnectEvery = 1
|
|
}
|
|
l.discardLocked()
|
|
}
|
|
|
|
func (l *bpLane) discardLocked() {
|
|
if l.pc != nil {
|
|
_ = l.pc.conn.Close()
|
|
l.pc = nil
|
|
}
|
|
}
|
|
|
|
func (l *bpLane) closeAfterLocked() {
|
|
if l.pc != nil && l.reconnectEvery > 0 && l.pc.requests >= l.reconnectEvery {
|
|
l.discardLocked()
|
|
}
|
|
}
|
|
|
|
func (l *bpLane) ensureLocked() error {
|
|
if l.closed {
|
|
return net.ErrClosed
|
|
}
|
|
if l.pc != nil {
|
|
if l.reconnectEvery <= 0 || l.pc.requests < l.reconnectEvery {
|
|
return nil
|
|
}
|
|
l.discardLocked()
|
|
}
|
|
d := net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
|
|
conn, err := d.Dial("tcp", l.serverAddr)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := cover.WritePreface(conn, l.coverProfile); err != nil {
|
|
_ = conn.Close()
|
|
return err
|
|
}
|
|
protocol.TuneTCP(conn)
|
|
protocol.TuneTCPBuffer(conn, l.tcpBuffer)
|
|
l.pc = &bpPhysicalConn{conn: conn}
|
|
return nil
|
|
}
|
|
|
|
func (l *bpLane) Close() {
|
|
l.mu.Lock()
|
|
l.closed = true
|
|
l.discardLocked()
|
|
l.mu.Unlock()
|
|
}
|
|
|
|
func writeBPRequest(w io.Writer, mode byte, sid wire.SessionID, seq uint64, payload []byte, downloadHint uint32, headerMask byte, clear bool) error {
|
|
n := uint32(len(payload))
|
|
if mode == bpModeDownload {
|
|
n = downloadHint
|
|
payload = nil
|
|
}
|
|
if len(payload) > wire.MaxPayload {
|
|
return fmt.Errorf("BP payload too large: %d", len(payload))
|
|
}
|
|
var header [bpHeaderSize]byte
|
|
header[0] = mode ^ headerMask
|
|
copy(header[1:17], sid[:])
|
|
binary.BigEndian.PutUint64(header[17:25], seq)
|
|
binary.BigEndian.PutUint32(header[25:29], n)
|
|
if clear {
|
|
buffers := net.Buffers{header[:], payload}
|
|
_, err := buffers.WriteTo(w)
|
|
return err
|
|
}
|
|
packet := make([]byte, bpHeaderSize+len(payload))
|
|
copy(packet[:bpHeaderSize], header[:])
|
|
copy(packet[bpHeaderSize:], payload)
|
|
wire.MaskInPlace(packet[bpHeaderSize:], sid, mode, seq, false)
|
|
for len(packet) > 0 {
|
|
written, err := w.Write(packet)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if written <= 0 {
|
|
return io.ErrShortWrite
|
|
}
|
|
packet = packet[written:]
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func readBPResponse(r io.Reader, sid wire.SessionID, mode byte, seq uint64, headerMask byte, clear bool) (byte, []byte, error) {
|
|
status, body, err := wire.ReadResponseProfile(r, headerMask)
|
|
if err == nil && status != wire.StatusError && len(body) > 0 && !clear {
|
|
wire.MaskInPlace(body, sid, mode, seq, true)
|
|
}
|
|
return status, body, err
|
|
}
|
|
|
|
func (l *bpLane) single(mode byte, sid wire.SessionID, seq uint64, payload []byte, downloadHint uint32) (byte, []byte, error) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
timeout := l.timeout
|
|
if timeout <= 0 {
|
|
timeout = 5 * time.Second
|
|
}
|
|
var lastErr error
|
|
for attempt := 0; attempt < 2; attempt++ {
|
|
if err := l.ensureLocked(); err != nil {
|
|
if isTransportTimeout(err) {
|
|
return 0, nil, err
|
|
}
|
|
lastErr = err
|
|
continue
|
|
}
|
|
reused := l.pc.requests > 0
|
|
_ = l.pc.conn.SetDeadline(time.Now().Add(timeout))
|
|
if err := writeBPRequest(l.pc.conn, mode, sid, seq, payload, downloadHint, l.coverProfile.HeaderMask, l.coverProfile.Clear); err != nil {
|
|
l.transportFailureLocked(reused)
|
|
if isTransportTimeout(err) {
|
|
return 0, nil, err
|
|
}
|
|
lastErr = err
|
|
continue
|
|
}
|
|
status, body, err := readBPResponse(l.pc.conn, sid, mode, seq, l.coverProfile.HeaderMask, l.coverProfile.Clear)
|
|
if err != nil {
|
|
l.transportFailureLocked(reused)
|
|
if isTransportTimeout(err) {
|
|
return 0, nil, err
|
|
}
|
|
lastErr = err
|
|
continue
|
|
}
|
|
l.pc.requests++
|
|
_ = l.pc.conn.SetDeadline(time.Time{})
|
|
l.closeAfterLocked()
|
|
return status, body, nil
|
|
}
|
|
return 0, nil, fmt.Errorf("BP request failed after reconnect: %w", lastErr)
|
|
}
|
|
|
|
func decodeBPData(body []byte) ([]byte, error) {
|
|
if len(body) < 4 {
|
|
return nil, fmt.Errorf("short BP DATA body")
|
|
}
|
|
n := int(binary.BigEndian.Uint32(body[:4]))
|
|
if n < 0 || n > len(body)-4 {
|
|
return nil, fmt.Errorf("bad BP DATA length")
|
|
}
|
|
return append([]byte(nil), body[4:4+n]...), nil
|
|
}
|
|
|
|
func (l *bpLane) download(sid wire.SessionID, offset uint64, maxChunk, count int) ([][]byte, byte, error) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
timeout := l.timeout
|
|
if timeout <= 0 {
|
|
timeout = 5 * time.Second
|
|
}
|
|
mode := bpModeDownload
|
|
payload := []byte(nil)
|
|
hint := uint32(maxChunk)
|
|
if count > 1 {
|
|
mode = bpModeBatchDownload
|
|
payload = make([]byte, 6)
|
|
binary.BigEndian.PutUint32(payload[:4], uint32(maxChunk))
|
|
binary.BigEndian.PutUint16(payload[4:6], uint16(count))
|
|
hint = 0
|
|
}
|
|
responses := 1
|
|
if mode == bpModeBatchDownload {
|
|
responses = count
|
|
}
|
|
var lastErr error
|
|
for attempt := 0; attempt < 2; attempt++ {
|
|
if err := l.ensureLocked(); err != nil {
|
|
if isTransportTimeout(err) {
|
|
return nil, 0, err
|
|
}
|
|
lastErr = err
|
|
continue
|
|
}
|
|
reused := l.pc.requests > 0
|
|
_ = l.pc.conn.SetDeadline(time.Now().Add(timeout))
|
|
if err := writeBPRequest(l.pc.conn, mode, sid, offset, payload, hint, l.coverProfile.HeaderMask, l.coverProfile.Clear); err != nil {
|
|
l.transportFailureLocked(reused)
|
|
if isTransportTimeout(err) {
|
|
return nil, 0, err
|
|
}
|
|
lastErr = err
|
|
continue
|
|
}
|
|
|
|
out := make([][]byte, 0, responses)
|
|
lastStatus := wire.StatusOK
|
|
for i := 0; i < responses; i++ {
|
|
status, body, err := readBPResponse(l.pc.conn, sid, mode, offset, l.coverProfile.HeaderMask, l.coverProfile.Clear)
|
|
if err != nil {
|
|
l.transportFailureLocked(reused)
|
|
if isTransportTimeout(err) {
|
|
return nil, 0, err
|
|
}
|
|
lastErr = err
|
|
goto retry
|
|
}
|
|
lastStatus = status
|
|
switch status {
|
|
case wire.StatusData:
|
|
data, err := decodeBPData(body)
|
|
if err != nil {
|
|
l.discardLocked()
|
|
return out, status, err
|
|
}
|
|
if len(data) > 0 {
|
|
out = append(out, data)
|
|
}
|
|
case wire.StatusOK, wire.StatusWait:
|
|
case wire.StatusEOF:
|
|
case wire.StatusError:
|
|
l.discardLocked()
|
|
return out, status, fmt.Errorf("%s", string(body))
|
|
default:
|
|
l.discardLocked()
|
|
return out, status, fmt.Errorf("unexpected BP download status %d", status)
|
|
}
|
|
}
|
|
l.pc.requests++
|
|
_ = l.pc.conn.SetDeadline(time.Time{})
|
|
l.closeAfterLocked()
|
|
return out, lastStatus, nil
|
|
retry:
|
|
}
|
|
return nil, 0, fmt.Errorf("BP download request failed after reconnect: %w", lastErr)
|
|
}
|
|
|
|
type bpConn struct {
|
|
sid wire.SessionID
|
|
opts chunkClientOptions
|
|
uploadLane *bpLane
|
|
downloadLane *bpLane
|
|
upSizer *adaptiveSizer
|
|
downSizer *adaptiveSizer
|
|
|
|
writeMu sync.Mutex
|
|
upOffset uint64
|
|
|
|
readMu sync.Mutex
|
|
readBuf []byte
|
|
downloadOffset uint64
|
|
consumedOffset uint64
|
|
lastAck uint64
|
|
eof bool
|
|
pipeline int
|
|
|
|
closeOnce sync.Once
|
|
dead atomic.Bool
|
|
}
|
|
|
|
func openBPTunnel(serverAddr, token, targetHost string, targetPort int, opts chunkClientOptions) (net.Conn, error) {
|
|
if opts.minSize < 32 {
|
|
opts.minSize = 32
|
|
}
|
|
if opts.maxSize < opts.minSize {
|
|
opts.maxSize = opts.minSize
|
|
}
|
|
if opts.maxSize > 1024*1024 {
|
|
opts.maxSize = 1024 * 1024
|
|
}
|
|
if opts.startSize < opts.minSize || opts.startSize > opts.maxSize {
|
|
opts.startSize = opts.maxSize
|
|
}
|
|
if opts.txnTimeout <= 0 {
|
|
opts.txnTimeout = 5 * time.Second
|
|
}
|
|
if opts.maxPipeline < 1 {
|
|
opts.maxPipeline = 1
|
|
}
|
|
if opts.maxPipeline > 256 {
|
|
opts.maxPipeline = 256
|
|
}
|
|
if opts.minPipeline < 1 {
|
|
opts.minPipeline = 1
|
|
}
|
|
if opts.minPipeline > opts.maxPipeline {
|
|
opts.minPipeline = opts.maxPipeline
|
|
}
|
|
reconnect := opts.reconnectEvery
|
|
if reconnect < 0 {
|
|
reconnect = 0
|
|
}
|
|
|
|
sid, err := randomSessionID()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
uploadLane := newBPLane(serverAddr, opts.tcpBuffer, reconnect, opts.txnTimeout, opts.coverProfile)
|
|
status, body, err := uploadLane.single(bpModeUpload, sid, 0, nil, 0)
|
|
if err != nil {
|
|
uploadLane.Close()
|
|
return nil, err
|
|
}
|
|
if status == wire.StatusError {
|
|
uploadLane.Close()
|
|
return nil, fmt.Errorf("%s", string(body))
|
|
}
|
|
if status != wire.StatusOK {
|
|
uploadLane.Close()
|
|
return nil, fmt.Errorf("bad BP registration response %d", status)
|
|
}
|
|
openPayload, err := encodeOpen(token, targetHost, targetPort)
|
|
if err != nil {
|
|
uploadLane.Close()
|
|
return nil, err
|
|
}
|
|
openPayload = append(append([]byte(nil), bpOpenMagic[:]...), openPayload...)
|
|
status, body, err = uploadLane.single(bpModeUpload, sid, 1, openPayload, 0)
|
|
if err != nil {
|
|
uploadLane.Close()
|
|
return nil, err
|
|
}
|
|
if status == wire.StatusError {
|
|
uploadLane.Close()
|
|
return nil, fmt.Errorf("%s", string(body))
|
|
}
|
|
if status != wire.StatusOK {
|
|
uploadLane.Close()
|
|
return nil, fmt.Errorf("bad BP OPEN response %d", status)
|
|
}
|
|
|
|
c := &bpConn{
|
|
sid: sid,
|
|
opts: opts,
|
|
uploadLane: uploadLane,
|
|
downloadLane: newBPLane(serverAddr, opts.tcpBuffer, reconnect, opts.txnTimeout, opts.coverProfile),
|
|
pipeline: opts.maxPipeline,
|
|
}
|
|
c.upSizer = newAdaptiveSizer("BP upload", opts.startSize, opts)
|
|
c.downSizer = newAdaptiveSizer("BP download", opts.startSize, opts)
|
|
return c, nil
|
|
}
|
|
|
|
func (c *bpConn) fillReadBuffer() error {
|
|
if c.eof {
|
|
return io.EOF
|
|
}
|
|
for len(c.readBuf) == 0 && !c.eof {
|
|
if c.consumedOffset > c.lastAck {
|
|
status, body, err := c.downloadLane.single(bpModeACK, c.sid, c.consumedOffset, nil, 0)
|
|
if err != nil {
|
|
if isTransportTimeout(err) {
|
|
c.dead.Store(true)
|
|
}
|
|
return err
|
|
}
|
|
if status == wire.StatusError {
|
|
return fmt.Errorf("%s", string(body))
|
|
}
|
|
c.lastAck = c.consumedOffset
|
|
}
|
|
chunk := c.downSizer.Current()
|
|
count := c.pipeline
|
|
if count < c.opts.minPipeline {
|
|
count = c.opts.minPipeline
|
|
}
|
|
if count > c.opts.maxPipeline {
|
|
count = c.opts.maxPipeline
|
|
}
|
|
data, status, err := c.downloadLane.download(c.sid, c.downloadOffset, chunk, count)
|
|
if err != nil {
|
|
if isTransportTimeout(err) {
|
|
c.dead.Store(true)
|
|
return err
|
|
}
|
|
old, next := c.downSizer.FailureReason(chunk, err)
|
|
if old == next && next == c.opts.minSize {
|
|
return err
|
|
}
|
|
time.Sleep(30 * time.Millisecond)
|
|
continue
|
|
}
|
|
c.readBuf = appendChunkParts(c.readBuf, data)
|
|
for _, part := range data {
|
|
c.downloadOffset += uint64(len(part))
|
|
}
|
|
if len(data) > 0 {
|
|
c.downSizer.Success(chunk)
|
|
if c.pipeline < c.opts.maxPipeline {
|
|
c.pipeline++
|
|
}
|
|
}
|
|
if status == wire.StatusEOF {
|
|
c.eof = true
|
|
}
|
|
if len(c.readBuf) == 0 && !c.eof {
|
|
delay := c.opts.pollDelay
|
|
if delay <= 0 {
|
|
delay = 5 * time.Millisecond
|
|
}
|
|
time.Sleep(delay)
|
|
}
|
|
}
|
|
if c.eof && len(c.readBuf) == 0 {
|
|
return io.EOF
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *bpConn) Read(p []byte) (int, error) {
|
|
c.readMu.Lock()
|
|
defer c.readMu.Unlock()
|
|
if len(p) == 0 {
|
|
return 0, nil
|
|
}
|
|
if len(c.readBuf) == 0 {
|
|
if err := c.fillReadBuffer(); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
n := copy(p, c.readBuf)
|
|
c.readBuf = c.readBuf[n:]
|
|
c.consumedOffset += uint64(n)
|
|
return n, nil
|
|
}
|
|
|
|
func (c *bpConn) Write(p []byte) (int, error) {
|
|
c.writeMu.Lock()
|
|
defer c.writeMu.Unlock()
|
|
total := 0
|
|
for len(p) > 0 {
|
|
size := c.upSizer.Current()
|
|
n := minInt(size, len(p))
|
|
status, body, err := c.uploadLane.single(bpModeUpload, c.sid, c.upOffset+2, p[:n], 0)
|
|
if err != nil {
|
|
if isTransportTimeout(err) {
|
|
c.dead.Store(true)
|
|
return total, err
|
|
}
|
|
old, next := c.upSizer.FailureReason(size, err)
|
|
if old == next && next == c.opts.minSize {
|
|
return total, err
|
|
}
|
|
time.Sleep(30 * time.Millisecond)
|
|
continue
|
|
}
|
|
if status == wire.StatusError {
|
|
return total, fmt.Errorf("%s", string(body))
|
|
}
|
|
if status != wire.StatusOK {
|
|
return total, fmt.Errorf("unexpected BP upload status %d", status)
|
|
}
|
|
c.upOffset += uint64(n)
|
|
total += n
|
|
p = p[n:]
|
|
c.upSizer.Success(size)
|
|
}
|
|
return total, nil
|
|
}
|
|
|
|
func (c *bpConn) Close() error {
|
|
c.closeOnce.Do(func() {
|
|
// A timeout is a hard tunnel boundary. Do not create a fresh physical
|
|
// connection only to close a session that has already been abandoned.
|
|
if !c.dead.Load() {
|
|
_, _, _ = c.downloadLane.single(bpModeACK, c.sid, c.consumedOffset, bpCloseMagic[:], 0)
|
|
}
|
|
c.uploadLane.Close()
|
|
c.downloadLane.Close()
|
|
})
|
|
return nil
|
|
}
|
|
|
|
func (c *bpConn) LocalAddr() net.Addr { return dummyAddr("dragontcp-bp-local") }
|
|
func (c *bpConn) RemoteAddr() net.Addr { return dummyAddr("dragontcp-bp-remote") }
|
|
func (c *bpConn) SetDeadline(time.Time) error { return nil }
|
|
func (c *bpConn) SetReadDeadline(time.Time) error { return nil }
|
|
func (c *bpConn) SetWriteDeadline(time.Time) error { return nil }
|