New
This commit is contained in:
@@ -0,0 +1,499 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
"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 {
|
||||
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 {
|
||||
lastErr = err
|
||||
l.transportFailureLocked(reused)
|
||||
continue
|
||||
}
|
||||
status, body, err := readBPResponse(l.pc.conn, sid, mode, seq, l.coverProfile.HeaderMask, l.coverProfile.Clear)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
l.transportFailureLocked(reused)
|
||||
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 {
|
||||
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 {
|
||||
lastErr = err
|
||||
l.transportFailureLocked(reused)
|
||||
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 {
|
||||
lastErr = err
|
||||
l.transportFailureLocked(reused)
|
||||
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
|
||||
}
|
||||
|
||||
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 {
|
||||
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 {
|
||||
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 {
|
||||
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() {
|
||||
_, _, _ = 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 }
|
||||
Reference in New Issue
Block a user