V7
This commit is contained in:
@@ -0,0 +1,761 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"dragontcp/internal/protocol"
|
||||
)
|
||||
|
||||
type chunkClientOptions struct {
|
||||
startSize int
|
||||
minSize int
|
||||
maxSize int
|
||||
adaptive bool
|
||||
adaptSuccesses int
|
||||
adaptLog bool
|
||||
pollers int
|
||||
reconnectEvery int
|
||||
pollDelay time.Duration
|
||||
txnTimeout time.Duration
|
||||
tcpBuffer int
|
||||
}
|
||||
|
||||
func wireToken(token string) string {
|
||||
if token == "" {
|
||||
return "-"
|
||||
}
|
||||
return token
|
||||
}
|
||||
|
||||
type adaptiveSizer struct {
|
||||
mu sync.Mutex
|
||||
name string
|
||||
current int
|
||||
min int
|
||||
max int
|
||||
adaptive bool
|
||||
adaptSuccesses int
|
||||
successes int
|
||||
good int
|
||||
bad int
|
||||
logChanges bool
|
||||
}
|
||||
|
||||
func newAdaptiveSizer(name string, opts chunkClientOptions) *adaptiveSizer {
|
||||
start := opts.startSize
|
||||
if start < opts.minSize {
|
||||
start = opts.minSize
|
||||
}
|
||||
if start > opts.maxSize {
|
||||
start = opts.maxSize
|
||||
}
|
||||
return &adaptiveSizer{
|
||||
name: name,
|
||||
current: start,
|
||||
min: opts.minSize,
|
||||
max: opts.maxSize,
|
||||
adaptive: opts.adaptive,
|
||||
adaptSuccesses: opts.adaptSuccesses,
|
||||
logChanges: opts.adaptLog,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *adaptiveSizer) Current() int {
|
||||
s.mu.Lock()
|
||||
n := s.current
|
||||
s.mu.Unlock()
|
||||
return n
|
||||
}
|
||||
|
||||
func (s *adaptiveSizer) Success(attempted int) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if !s.adaptive || s.current >= s.max {
|
||||
return
|
||||
}
|
||||
// Ignore stale successes from records that were already in flight when
|
||||
// another worker changed the shared size.
|
||||
if attempted != s.current {
|
||||
return
|
||||
}
|
||||
|
||||
if attempted > s.good {
|
||||
s.good = attempted
|
||||
}
|
||||
s.successes++
|
||||
|
||||
growAfter := s.adaptSuccesses
|
||||
// When we have converged close to a known failure boundary, stay stable
|
||||
// longer before probing again. This also lets us discover later network
|
||||
// improvements without constantly oscillating around the boundary.
|
||||
if s.bad > 0 && s.bad-s.good <= 32 {
|
||||
growAfter *= 8
|
||||
}
|
||||
if s.successes < growAfter {
|
||||
return
|
||||
}
|
||||
s.successes = 0
|
||||
|
||||
old := s.current
|
||||
var next int
|
||||
if s.bad > old+1 {
|
||||
// Binary-search the gap between known-good and known-bad sizes.
|
||||
next = old + (s.bad-old)/2
|
||||
} else {
|
||||
// Either there is no known ceiling, or we have stayed stable long enough
|
||||
// at it to probe the network again in case conditions improved.
|
||||
if s.bad > 0 {
|
||||
s.bad = 0
|
||||
}
|
||||
step := old / 4
|
||||
if step < 32 {
|
||||
step = 32
|
||||
}
|
||||
next = old + step
|
||||
}
|
||||
|
||||
if next > s.max {
|
||||
next = s.max
|
||||
}
|
||||
if next <= old {
|
||||
return
|
||||
}
|
||||
s.current = next
|
||||
|
||||
if s.logChanges {
|
||||
fmt.Printf("adaptive %s chunk: %d -> %d after stable success\n", s.name, old, next)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *adaptiveSizer) Failure(attempted int) (old, next int) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
old = s.current
|
||||
|
||||
if !s.adaptive {
|
||||
return old, old
|
||||
}
|
||||
// Multiple pollers can fail on the same oversized value at once. Only the
|
||||
// first failure for the current value is allowed to reduce it.
|
||||
if attempted != s.current {
|
||||
return old, old
|
||||
}
|
||||
s.successes = 0
|
||||
|
||||
if s.bad == 0 || attempted < s.bad {
|
||||
s.bad = attempted
|
||||
}
|
||||
|
||||
if s.good > 0 && s.good < attempted {
|
||||
// Return directly to the last size that was proven to work.
|
||||
next = s.good
|
||||
} else {
|
||||
// A previously-good value just failed, so conditions worsened. Forget
|
||||
// the old lower bound and use multiplicative decrease.
|
||||
s.good = 0
|
||||
next = attempted / 2
|
||||
}
|
||||
if next < s.min {
|
||||
next = s.min
|
||||
}
|
||||
if next >= attempted && attempted > s.min {
|
||||
next = attempted - 1
|
||||
}
|
||||
if next < s.min {
|
||||
next = s.min
|
||||
}
|
||||
s.current = next
|
||||
|
||||
if s.logChanges && next != old {
|
||||
fmt.Printf("adaptive %s chunk: %d -> %d after transport failure\n", s.name, old, next)
|
||||
}
|
||||
return old, next
|
||||
}
|
||||
|
||||
type txnLane struct {
|
||||
mu sync.Mutex
|
||||
serverAddr string
|
||||
tcpBuffer int
|
||||
reconnectEvery int
|
||||
timeout time.Duration
|
||||
conn net.Conn
|
||||
count int
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newTxnLane(serverAddr string, tcpBuffer, reconnectEvery int, timeout time.Duration) *txnLane {
|
||||
return &txnLane{
|
||||
serverAddr: serverAddr,
|
||||
tcpBuffer: tcpBuffer,
|
||||
reconnectEvery: reconnectEvery,
|
||||
timeout: timeout,
|
||||
}
|
||||
}
|
||||
|
||||
func (l *txnLane) closeLocked() {
|
||||
if l.conn != nil {
|
||||
_ = l.conn.Close()
|
||||
l.conn = nil
|
||||
}
|
||||
l.count = 0
|
||||
}
|
||||
|
||||
func (l *txnLane) Close() {
|
||||
l.mu.Lock()
|
||||
l.closed = true
|
||||
l.closeLocked()
|
||||
l.mu.Unlock()
|
||||
}
|
||||
|
||||
func (l *txnLane) ensureConn() error {
|
||||
if l.closed {
|
||||
return net.ErrClosed
|
||||
}
|
||||
if l.conn != nil && (l.reconnectEvery <= 0 || l.count < l.reconnectEvery) {
|
||||
return nil
|
||||
}
|
||||
|
||||
l.closeLocked()
|
||||
d := net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
|
||||
conn, err := d.Dial("tcp", l.serverAddr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
protocol.TuneTCP(conn)
|
||||
protocol.TuneTCPBuffer(conn, l.tcpBuffer)
|
||||
l.conn = conn
|
||||
return nil
|
||||
}
|
||||
|
||||
// Do performs exactly one framed transaction. Higher layers decide whether a
|
||||
// failed data record should be retried at a smaller adaptive size.
|
||||
func (l *txnLane) Do(payload []byte) ([]byte, error) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
if err := l.ensureConn(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
timeout := l.timeout
|
||||
if timeout <= 0 {
|
||||
timeout = 5 * time.Second
|
||||
}
|
||||
_ = l.conn.SetDeadline(time.Now().Add(timeout))
|
||||
requestID := requestCounter.Add(1)
|
||||
|
||||
if err := protocol.WriteRequestFrame(l.conn, requestID, payload); err != nil {
|
||||
l.closeLocked()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
responseID, response, err := protocol.ReadResponseFrame(l.conn)
|
||||
if err != nil {
|
||||
l.closeLocked()
|
||||
return nil, err
|
||||
}
|
||||
if responseID != requestID {
|
||||
l.closeLocked()
|
||||
return nil, fmt.Errorf("request ID mismatch")
|
||||
}
|
||||
|
||||
l.count++
|
||||
_ = l.conn.SetDeadline(time.Time{})
|
||||
if l.reconnectEvery > 0 && l.count >= l.reconnectEvery {
|
||||
// For restrictive TCP/53 networks, reconnectEvery=1 must really mean
|
||||
// one request/response per TCP connection. Close immediately after
|
||||
// receiving the response rather than waiting for the next request.
|
||||
l.closeLocked()
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func doControl(lane *txnLane, payload []byte) ([]byte, error) {
|
||||
var lastErr error
|
||||
for attempt := 0; attempt < 5; attempt++ {
|
||||
resp, err := lane.Do(payload)
|
||||
if err == nil {
|
||||
return resp, nil
|
||||
}
|
||||
lastErr = err
|
||||
time.Sleep(time.Duration(attempt+1) * 40 * time.Millisecond)
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
type chunkResult struct {
|
||||
seq uint64
|
||||
data []byte
|
||||
final uint64
|
||||
eof bool
|
||||
err error
|
||||
}
|
||||
|
||||
type chunkConn struct {
|
||||
serverAddr string
|
||||
token string
|
||||
sid string
|
||||
opts chunkClientOptions
|
||||
|
||||
pushLane *txnLane
|
||||
pullLanes []*txnLane
|
||||
|
||||
upSizer *adaptiveSizer
|
||||
downSizer *adaptiveSizer
|
||||
serverMax int
|
||||
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
once sync.Once
|
||||
|
||||
writeMu sync.Mutex
|
||||
upSeq uint64
|
||||
|
||||
claim atomic.Uint64
|
||||
ack atomic.Int64
|
||||
|
||||
results chan chunkResult
|
||||
workers sync.WaitGroup
|
||||
|
||||
readMu sync.Mutex
|
||||
pending map[uint64][]byte
|
||||
nextRead uint64
|
||||
current []byte
|
||||
currentSeq uint64
|
||||
finalKnown bool
|
||||
finalSeq uint64
|
||||
terminalErr error
|
||||
}
|
||||
|
||||
func randomSessionID() (string, error) {
|
||||
var b [16]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(b[:]), nil
|
||||
}
|
||||
|
||||
func openChunkTunnel(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 > protocol.MaxChunkPayload {
|
||||
opts.maxSize = protocol.MaxChunkPayload
|
||||
}
|
||||
if opts.startSize < opts.minSize {
|
||||
opts.startSize = opts.minSize
|
||||
}
|
||||
if opts.startSize > opts.maxSize {
|
||||
opts.startSize = opts.maxSize
|
||||
}
|
||||
if opts.adaptSuccesses < 1 {
|
||||
opts.adaptSuccesses = 64
|
||||
}
|
||||
if opts.pollers < 1 {
|
||||
opts.pollers = 1
|
||||
}
|
||||
if opts.pollers > 128 {
|
||||
opts.pollers = 128
|
||||
}
|
||||
if opts.txnTimeout <= 0 {
|
||||
opts.txnTimeout = 5 * time.Second
|
||||
}
|
||||
|
||||
sid, err := randomSessionID()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
c := &chunkConn{
|
||||
serverAddr: serverAddr,
|
||||
token: token,
|
||||
sid: sid,
|
||||
opts: opts,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
results: make(chan chunkResult, opts.pollers*4),
|
||||
pending: make(map[uint64][]byte, opts.pollers*2),
|
||||
}
|
||||
c.ack.Store(-1)
|
||||
c.upSizer = newAdaptiveSizer("upload", opts)
|
||||
c.downSizer = newAdaptiveSizer("download", opts)
|
||||
|
||||
c.pushLane = newTxnLane(serverAddr, opts.tcpBuffer, opts.reconnectEvery, opts.txnTimeout)
|
||||
|
||||
openPayload := []byte(fmt.Sprintf(
|
||||
"COPEN %s %s %s %d",
|
||||
wireToken(token), sid, targetHost, targetPort,
|
||||
))
|
||||
resp, err := doControl(c.pushLane, openPayload)
|
||||
if err != nil {
|
||||
c.pushLane.Close()
|
||||
cancel()
|
||||
return nil, err
|
||||
}
|
||||
fields := strings.Fields(string(resp))
|
||||
if len(fields) != 2 || fields[0] != "OPENED" {
|
||||
c.pushLane.Close()
|
||||
cancel()
|
||||
return nil, fmt.Errorf("%s", resp)
|
||||
}
|
||||
serverMax, err := strconv.Atoi(fields[1])
|
||||
if err != nil || serverMax < 32 {
|
||||
c.pushLane.Close()
|
||||
cancel()
|
||||
return nil, fmt.Errorf("bad OPENED response: %q", resp)
|
||||
}
|
||||
c.serverMax = serverMax
|
||||
if serverMax < c.opts.maxSize {
|
||||
c.opts.maxSize = serverMax
|
||||
c.upSizer.max = serverMax
|
||||
c.downSizer.max = serverMax
|
||||
if c.upSizer.current > serverMax {
|
||||
c.upSizer.current = serverMax
|
||||
}
|
||||
if c.downSizer.current > serverMax {
|
||||
c.downSizer.current = serverMax
|
||||
}
|
||||
}
|
||||
|
||||
c.pullLanes = make([]*txnLane, opts.pollers)
|
||||
for i := 0; i < opts.pollers; i++ {
|
||||
lane := newTxnLane(serverAddr, opts.tcpBuffer, opts.reconnectEvery, opts.txnTimeout)
|
||||
c.pullLanes[i] = lane
|
||||
c.workers.Add(1)
|
||||
go c.pullWorker(lane)
|
||||
}
|
||||
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func parseDataResponse(resp []byte) (seq uint64, offset int, total int, data []byte, err error) {
|
||||
if len(resp) < 6 || string(resp[:5]) != "DATA " {
|
||||
return 0, 0, 0, nil, fmt.Errorf("not DATA")
|
||||
}
|
||||
|
||||
rest := resp[5:]
|
||||
fields := make([][]byte, 0, 3)
|
||||
start := 0
|
||||
for i := 0; i < len(rest) && len(fields) < 3; i++ {
|
||||
if rest[i] == ' ' {
|
||||
fields = append(fields, rest[start:i])
|
||||
start = i + 1
|
||||
}
|
||||
}
|
||||
if len(fields) != 3 {
|
||||
return 0, 0, 0, nil, fmt.Errorf("bad DATA response")
|
||||
}
|
||||
|
||||
seq, err = strconv.ParseUint(string(fields[0]), 10, 64)
|
||||
if err != nil {
|
||||
return 0, 0, 0, nil, err
|
||||
}
|
||||
offset, err = strconv.Atoi(string(fields[1]))
|
||||
if err != nil || offset < 0 {
|
||||
return 0, 0, 0, nil, fmt.Errorf("bad DATA offset")
|
||||
}
|
||||
total, err = strconv.Atoi(string(fields[2]))
|
||||
if err != nil || total < 0 {
|
||||
return 0, 0, 0, nil, fmt.Errorf("bad DATA total")
|
||||
}
|
||||
|
||||
// start now points immediately after the third separator.
|
||||
return seq, offset, total, rest[start:], nil
|
||||
}
|
||||
|
||||
func (c *chunkConn) pullWorker(lane *txnLane) {
|
||||
defer c.workers.Done()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
seq := c.claim.Add(1) - 1
|
||||
offset := 0
|
||||
var assembled []byte
|
||||
consecutiveMinFailures := 0
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
limit := c.downSizer.Current()
|
||||
ack := c.ack.Load()
|
||||
payload := []byte(fmt.Sprintf(
|
||||
"CPULL %s %s %d %d %d %d",
|
||||
wireToken(c.token), c.sid, ack, seq, offset, limit,
|
||||
))
|
||||
|
||||
resp, err := lane.Do(payload)
|
||||
if err != nil {
|
||||
old, next := c.downSizer.Failure(limit)
|
||||
if next == old && next == c.opts.minSize {
|
||||
consecutiveMinFailures++
|
||||
} else {
|
||||
consecutiveMinFailures = 0
|
||||
}
|
||||
if consecutiveMinFailures >= 8 {
|
||||
select {
|
||||
case c.results <- chunkResult{seq: seq, err: fmt.Errorf("download failed at minimum chunk %d: %w", next, err)}:
|
||||
case <-c.ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
continue
|
||||
}
|
||||
|
||||
if string(resp) == "WAIT" {
|
||||
if c.opts.pollDelay > 0 {
|
||||
select {
|
||||
case <-time.After(c.opts.pollDelay):
|
||||
case <-c.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.HasPrefix(string(resp), "ERR ") {
|
||||
select {
|
||||
case c.results <- chunkResult{seq: seq, err: fmt.Errorf("%s", resp)}:
|
||||
case <-c.ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if strings.HasPrefix(string(resp), "EOF ") {
|
||||
n, err := strconv.ParseUint(strings.TrimSpace(string(resp[4:])), 10, 64)
|
||||
if err != nil {
|
||||
select {
|
||||
case c.results <- chunkResult{seq: seq, err: err}:
|
||||
case <-c.ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
select {
|
||||
case c.results <- chunkResult{seq: seq, eof: true, final: n}:
|
||||
case <-c.ctx.Done():
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
gotSeq, gotOffset, total, fragment, err := parseDataResponse(resp)
|
||||
if err != nil {
|
||||
select {
|
||||
case c.results <- chunkResult{seq: seq, err: err}:
|
||||
case <-c.ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
if gotSeq != seq || gotOffset != offset {
|
||||
select {
|
||||
case c.results <- chunkResult{seq: seq, err: fmt.Errorf("DATA position mismatch")}:
|
||||
case <-c.ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
if total > c.serverMax || total < offset+len(fragment) || len(fragment) == 0 {
|
||||
select {
|
||||
case c.results <- chunkResult{seq: seq, err: fmt.Errorf("invalid DATA fragment size")}:
|
||||
case <-c.ctx.Done():
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if assembled == nil {
|
||||
assembled = make([]byte, 0, total)
|
||||
}
|
||||
assembled = append(assembled, fragment...)
|
||||
offset += len(fragment)
|
||||
consecutiveMinFailures = 0
|
||||
c.downSizer.Success(limit)
|
||||
|
||||
if offset == total {
|
||||
select {
|
||||
case c.results <- chunkResult{seq: seq, data: assembled}:
|
||||
case <-c.ctx.Done():
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *chunkConn) Read(p []byte) (int, error) {
|
||||
c.readMu.Lock()
|
||||
defer c.readMu.Unlock()
|
||||
|
||||
for {
|
||||
if len(c.current) > 0 {
|
||||
n := copy(p, c.current)
|
||||
c.current = c.current[n:]
|
||||
if len(c.current) == 0 {
|
||||
c.nextRead++
|
||||
c.ack.Store(int64(c.currentSeq))
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
if c.terminalErr != nil {
|
||||
return 0, c.terminalErr
|
||||
}
|
||||
|
||||
if c.finalKnown && c.nextRead >= c.finalSeq {
|
||||
return 0, io.EOF
|
||||
}
|
||||
|
||||
if data, ok := c.pending[c.nextRead]; ok {
|
||||
delete(c.pending, c.nextRead)
|
||||
c.current = data
|
||||
c.currentSeq = c.nextRead
|
||||
continue
|
||||
}
|
||||
|
||||
result, ok := <-c.results
|
||||
if !ok {
|
||||
return 0, io.EOF
|
||||
}
|
||||
if result.err != nil {
|
||||
c.terminalErr = result.err
|
||||
return 0, result.err
|
||||
}
|
||||
if result.eof {
|
||||
if !c.finalKnown || result.final < c.finalSeq {
|
||||
c.finalKnown = true
|
||||
c.finalSeq = result.final
|
||||
}
|
||||
continue
|
||||
}
|
||||
if result.seq < c.nextRead {
|
||||
continue
|
||||
}
|
||||
c.pending[result.seq] = result.data
|
||||
}
|
||||
}
|
||||
|
||||
func parseAck(resp []byte, expectedSeq uint64) (int, error) {
|
||||
fields := strings.Fields(string(resp))
|
||||
if len(fields) != 3 || fields[0] != "ACK" {
|
||||
return 0, fmt.Errorf("bad CPUSH response: %q", resp)
|
||||
}
|
||||
seq, err := strconv.ParseUint(fields[1], 10, 64)
|
||||
if err != nil || seq != expectedSeq {
|
||||
return 0, fmt.Errorf("bad CPUSH sequence: %q", resp)
|
||||
}
|
||||
n, err := strconv.Atoi(fields[2])
|
||||
if err != nil || n <= 0 {
|
||||
return 0, fmt.Errorf("bad CPUSH length: %q", resp)
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (c *chunkConn) Write(p []byte) (int, error) {
|
||||
c.writeMu.Lock()
|
||||
defer c.writeMu.Unlock()
|
||||
|
||||
total := 0
|
||||
consecutiveMinFailures := 0
|
||||
|
||||
for len(p) > 0 {
|
||||
size := c.upSizer.Current()
|
||||
n := size
|
||||
if len(p) < n {
|
||||
n = len(p)
|
||||
}
|
||||
|
||||
seq := c.upSeq
|
||||
prefix := []byte(fmt.Sprintf("CPUSH %s %s %d ", wireToken(c.token), c.sid, seq))
|
||||
payload := make([]byte, len(prefix)+n)
|
||||
copy(payload, prefix)
|
||||
copy(payload[len(prefix):], p[:n])
|
||||
|
||||
resp, err := c.pushLane.Do(payload)
|
||||
if err != nil {
|
||||
old, next := c.upSizer.Failure(size)
|
||||
if next == old && next == c.opts.minSize {
|
||||
consecutiveMinFailures++
|
||||
} else {
|
||||
consecutiveMinFailures = 0
|
||||
}
|
||||
if consecutiveMinFailures >= 8 {
|
||||
return total, fmt.Errorf("upload failed at minimum chunk %d: %w", next, err)
|
||||
}
|
||||
time.Sleep(30 * time.Millisecond)
|
||||
continue
|
||||
}
|
||||
|
||||
if strings.HasPrefix(string(resp), "ERR ") {
|
||||
return total, fmt.Errorf("%s", resp)
|
||||
}
|
||||
|
||||
accepted, err := parseAck(resp, seq)
|
||||
if err != nil {
|
||||
return total, err
|
||||
}
|
||||
if accepted > len(p) {
|
||||
return total, fmt.Errorf("server ACK length %d exceeds pending write %d", accepted, len(p))
|
||||
}
|
||||
|
||||
c.upSeq++
|
||||
total += accepted
|
||||
p = p[accepted:]
|
||||
consecutiveMinFailures = 0
|
||||
c.upSizer.Success(size)
|
||||
}
|
||||
|
||||
return total, nil
|
||||
}
|
||||
|
||||
func (c *chunkConn) Close() error {
|
||||
c.once.Do(func() {
|
||||
c.cancel()
|
||||
|
||||
lane := newTxnLane(c.serverAddr, c.opts.tcpBuffer, 1, c.opts.txnTimeout)
|
||||
_, _ = doControl(lane, []byte(fmt.Sprintf("CCLOSE %s %s", wireToken(c.token), c.sid)))
|
||||
lane.Close()
|
||||
|
||||
if c.pushLane != nil {
|
||||
c.pushLane.Close()
|
||||
}
|
||||
for _, lane := range c.pullLanes {
|
||||
lane.Close()
|
||||
}
|
||||
c.workers.Wait()
|
||||
close(c.results)
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *chunkConn) LocalAddr() net.Addr { return dummyAddr("dragontcp-chunk-local") }
|
||||
func (c *chunkConn) RemoteAddr() net.Addr { return dummyAddr("dragontcp-chunk-remote") }
|
||||
func (c *chunkConn) SetDeadline(time.Time) error { return nil }
|
||||
func (c *chunkConn) SetReadDeadline(time.Time) error { return nil }
|
||||
func (c *chunkConn) SetWriteDeadline(time.Time) error { return nil }
|
||||
|
||||
type dummyAddr string
|
||||
|
||||
func (d dummyAddr) Network() string { return "dragontcp-chunk" }
|
||||
func (d dummyAddr) String() string { return string(d) }
|
||||
@@ -0,0 +1,36 @@
|
||||
package main
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestAdaptiveSizerRecoversFromMinimum(t *testing.T) {
|
||||
opts := chunkClientOptions{
|
||||
startSize: 64,
|
||||
minSize: 32,
|
||||
maxSize: 1024,
|
||||
adaptive: true,
|
||||
adaptSuccesses: 2,
|
||||
}
|
||||
s := newAdaptiveSizer("test", opts)
|
||||
_, next := s.Failure(64)
|
||||
if next != 32 {
|
||||
t.Fatalf("failure should reduce 64 -> 32, got %d", next)
|
||||
}
|
||||
|
||||
// When good=32 and bad=64 are adjacent at the controller's probing
|
||||
// granularity, it deliberately waits 8x longer before testing upward.
|
||||
for i := 0; i < 16; i++ {
|
||||
s.Success(32)
|
||||
}
|
||||
if got := s.Current(); got <= 32 {
|
||||
t.Fatalf("adaptive controller remained stuck at minimum: %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWireTokenAllowsEmptyToken(t *testing.T) {
|
||||
if got := wireToken(""); got != "-" {
|
||||
t.Fatalf("empty token wire representation = %q, want '-'", got)
|
||||
}
|
||||
if got := wireToken("secret"); got != "secret" {
|
||||
t.Fatalf("non-empty token changed: %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,478 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"flag"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"dragontcp/internal/protocol"
|
||||
)
|
||||
|
||||
const maxHeader = 128 * 1024
|
||||
|
||||
var requestCounter atomic.Uint32
|
||||
|
||||
func readHTTPHeaders(conn net.Conn) ([]byte, []byte, error) {
|
||||
buf := make([]byte, 0, 8192)
|
||||
tmp := make([]byte, 8192)
|
||||
|
||||
for {
|
||||
n, err := conn.Read(tmp)
|
||||
if n > 0 {
|
||||
buf = append(buf, tmp[:n]...)
|
||||
|
||||
if len(buf) > maxHeader {
|
||||
return nil, nil, fmt.Errorf("HTTP headers too large")
|
||||
}
|
||||
|
||||
if i := bytes.Index(buf, []byte("\r\n\r\n")); i >= 0 {
|
||||
end := i + 4
|
||||
return buf[:end], buf[end:], nil
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func parseHostPort(authority string, defaultPort int) (string, int, error) {
|
||||
authority = strings.TrimSpace(authority)
|
||||
|
||||
if host, portText, err := net.SplitHostPort(authority); err == nil {
|
||||
port, err := strconv.Atoi(portText)
|
||||
return host, port, err
|
||||
}
|
||||
|
||||
// Host without port.
|
||||
if strings.HasPrefix(authority, "[") && strings.HasSuffix(authority, "]") {
|
||||
return strings.Trim(authority, "[]"), defaultPort, nil
|
||||
}
|
||||
|
||||
if strings.Count(authority, ":") == 0 {
|
||||
return authority, defaultPort, nil
|
||||
}
|
||||
|
||||
// Bare IPv6.
|
||||
if ip := net.ParseIP(authority); ip != nil {
|
||||
return authority, defaultPort, nil
|
||||
}
|
||||
|
||||
return "", 0, fmt.Errorf("invalid authority: %s", authority)
|
||||
}
|
||||
|
||||
func rewritePlainHTTPRequest(header []byte) (string, int, []byte, error) {
|
||||
text := string(header)
|
||||
lines := strings.Split(text, "\r\n")
|
||||
if len(lines) == 0 {
|
||||
return "", 0, nil, fmt.Errorf("empty request")
|
||||
}
|
||||
|
||||
parts := strings.SplitN(lines[0], " ", 3)
|
||||
if len(parts) != 3 {
|
||||
return "", 0, nil, fmt.Errorf("invalid request line")
|
||||
}
|
||||
|
||||
method, target, version := parts[0], parts[1], parts[2]
|
||||
|
||||
var (
|
||||
hostHeader string
|
||||
headers []string
|
||||
)
|
||||
|
||||
for _, line := range lines[1:] {
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
k, v, ok := strings.Cut(line, ":")
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
lk := strings.ToLower(strings.TrimSpace(k))
|
||||
|
||||
if lk == "host" {
|
||||
hostHeader = strings.TrimSpace(v)
|
||||
}
|
||||
|
||||
if lk == "connection" ||
|
||||
lk == "proxy-connection" ||
|
||||
lk == "proxy-authorization" {
|
||||
continue
|
||||
}
|
||||
|
||||
headers = append(headers, k+": "+strings.TrimSpace(v))
|
||||
}
|
||||
|
||||
u, err := url.Parse(target)
|
||||
if err != nil {
|
||||
return "", 0, nil, err
|
||||
}
|
||||
|
||||
var host string
|
||||
var port int
|
||||
path := target
|
||||
|
||||
if u.Hostname() != "" {
|
||||
if strings.ToLower(u.Scheme) != "http" {
|
||||
return "", 0, nil, fmt.Errorf("unsupported plain HTTP scheme: %s", u.Scheme)
|
||||
}
|
||||
|
||||
host = u.Hostname()
|
||||
port = 80
|
||||
|
||||
if u.Port() != "" {
|
||||
port, err = strconv.Atoi(u.Port())
|
||||
if err != nil {
|
||||
return "", 0, nil, err
|
||||
}
|
||||
}
|
||||
|
||||
path = u.EscapedPath()
|
||||
if path == "" {
|
||||
path = "/"
|
||||
}
|
||||
if u.RawQuery != "" {
|
||||
path += "?" + u.RawQuery
|
||||
}
|
||||
} else {
|
||||
if hostHeader == "" {
|
||||
return "", 0, nil, fmt.Errorf("missing Host header")
|
||||
}
|
||||
|
||||
host, port, err = parseHostPort(hostHeader, 80)
|
||||
if err != nil {
|
||||
return "", 0, nil, err
|
||||
}
|
||||
if path == "" {
|
||||
path = "/"
|
||||
}
|
||||
}
|
||||
|
||||
var out strings.Builder
|
||||
fmt.Fprintf(&out, "%s %s %s\r\n", method, path, version)
|
||||
|
||||
sawHost := false
|
||||
for _, h := range headers {
|
||||
if strings.HasPrefix(strings.ToLower(h), "host:") {
|
||||
sawHost = true
|
||||
}
|
||||
out.WriteString(h)
|
||||
out.WriteString("\r\n")
|
||||
}
|
||||
|
||||
if !sawHost {
|
||||
if port == 80 {
|
||||
fmt.Fprintf(&out, "Host: %s\r\n", host)
|
||||
} else {
|
||||
fmt.Fprintf(&out, "Host: %s\r\n", net.JoinHostPort(host, strconv.Itoa(port)))
|
||||
}
|
||||
}
|
||||
|
||||
out.WriteString("Connection: close\r\n\r\n")
|
||||
|
||||
return host, port, []byte(out.String()), nil
|
||||
}
|
||||
|
||||
func openDragonTCPTunnel(serverAddr, token, targetHost string, targetPort int, transport string, tcpBuffer int) (net.Conn, error) {
|
||||
d := net.Dialer{
|
||||
Timeout: 10 * time.Second,
|
||||
KeepAlive: 30 * time.Second,
|
||||
}
|
||||
|
||||
conn, err := d.Dial("tcp", serverAddr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
protocol.TuneTCP(conn)
|
||||
protocol.TuneTCPBuffer(conn, tcpBuffer)
|
||||
_ = conn.SetDeadline(time.Now().Add(15 * time.Second))
|
||||
|
||||
// Correlation only; cryptographic randomness is unnecessary here.
|
||||
requestID := requestCounter.Add(1)
|
||||
|
||||
var command []byte
|
||||
if transport == "raw" {
|
||||
command = []byte(fmt.Sprintf("TUNNEL2 %s %s %d RAW", token, targetHost, targetPort))
|
||||
} else {
|
||||
// Legacy XOR command remains compatible with the older server.
|
||||
command = []byte(fmt.Sprintf("TUNNEL %s %s %d", token, targetHost, targetPort))
|
||||
}
|
||||
|
||||
if err := protocol.WriteRequestFrame(conn, requestID, command); err != nil {
|
||||
conn.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
responseID, response, err := protocol.ReadResponseFrame(conn)
|
||||
if err != nil {
|
||||
conn.Close()
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if responseID != requestID {
|
||||
conn.Close()
|
||||
return nil, fmt.Errorf("request ID mismatch")
|
||||
}
|
||||
|
||||
if string(response) != "CONNECTED" {
|
||||
conn.Close()
|
||||
return nil, fmt.Errorf("%s", response)
|
||||
}
|
||||
|
||||
_ = conn.SetDeadline(time.Time{})
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func writeHTTPError(conn net.Conn, code int, reason, detail string) {
|
||||
if detail == "" {
|
||||
detail = reason
|
||||
}
|
||||
|
||||
body := []byte(detail)
|
||||
|
||||
fmt.Fprintf(
|
||||
conn,
|
||||
"HTTP/1.1 %d %s\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: %d\r\nConnection: close\r\n\r\n",
|
||||
code,
|
||||
reason,
|
||||
len(body),
|
||||
)
|
||||
_, _ = conn.Write(body)
|
||||
}
|
||||
|
||||
func handleLocal(conn net.Conn, serverAddr, token, transport string, tcpBuffer int, chunkOpts chunkClientOptions, slots chan struct{}) {
|
||||
defer func() {
|
||||
<-slots
|
||||
_ = conn.Close()
|
||||
}()
|
||||
|
||||
protocol.TuneTCP(conn)
|
||||
protocol.TuneTCPBuffer(conn, tcpBuffer)
|
||||
_ = conn.SetDeadline(time.Now().Add(15 * time.Second))
|
||||
|
||||
header, extra, err := readHTTPHeaders(conn)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
firstLine := strings.SplitN(string(header), "\r\n", 2)[0]
|
||||
parts := strings.SplitN(firstLine, " ", 3)
|
||||
|
||||
if len(parts) != 3 {
|
||||
writeHTTPError(conn, 400, "Bad Request", "invalid HTTP request line")
|
||||
return
|
||||
}
|
||||
|
||||
method, target := parts[0], parts[1]
|
||||
|
||||
if strings.EqualFold(method, "CONNECT") {
|
||||
host, port, err := parseHostPort(target, 443)
|
||||
if err != nil {
|
||||
writeHTTPError(conn, 400, "Bad Request", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
var remote net.Conn
|
||||
if transport == "chunk" {
|
||||
remote, err = openChunkTunnel(serverAddr, token, host, port, chunkOpts)
|
||||
} else {
|
||||
remote, err = openDragonTCPTunnel(serverAddr, token, host, port, transport, tcpBuffer)
|
||||
}
|
||||
if err != nil {
|
||||
writeHTTPError(conn, 502, "Bad Gateway", err.Error())
|
||||
return
|
||||
}
|
||||
defer remote.Close()
|
||||
|
||||
_, _ = conn.Write([]byte(
|
||||
"HTTP/1.1 200 Connection Established\r\n" +
|
||||
"Proxy-Agent: dragontcp-proxy/2.0\r\n\r\n",
|
||||
))
|
||||
|
||||
if len(extra) > 0 {
|
||||
if transport == "xor" {
|
||||
protocol.XorInPlace(extra)
|
||||
}
|
||||
if _, err := remote.Write(extra); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
_ = conn.SetDeadline(time.Time{})
|
||||
if transport == "xor" {
|
||||
protocol.RelayXOR(conn, remote)
|
||||
} else {
|
||||
// raw and chunk connections expose a normal plaintext net.Conn.
|
||||
protocol.RelayRaw(conn, remote)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
host, port, rewritten, err := rewritePlainHTTPRequest(header)
|
||||
if err != nil {
|
||||
writeHTTPError(conn, 400, "Bad Request", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
var remote net.Conn
|
||||
if transport == "chunk" {
|
||||
remote, err = openChunkTunnel(serverAddr, token, host, port, chunkOpts)
|
||||
} else {
|
||||
remote, err = openDragonTCPTunnel(serverAddr, token, host, port, transport, tcpBuffer)
|
||||
}
|
||||
if err != nil {
|
||||
writeHTTPError(conn, 502, "Bad Gateway", err.Error())
|
||||
return
|
||||
}
|
||||
defer remote.Close()
|
||||
|
||||
initial := make([]byte, 0, len(rewritten)+len(extra))
|
||||
initial = append(initial, rewritten...)
|
||||
initial = append(initial, extra...)
|
||||
if transport == "xor" {
|
||||
protocol.XorInPlace(initial)
|
||||
}
|
||||
|
||||
if _, err := remote.Write(initial); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
_ = conn.SetDeadline(time.Time{})
|
||||
if transport == "xor" {
|
||||
protocol.RelayXOR(conn, remote)
|
||||
} else {
|
||||
protocol.RelayRaw(conn, remote)
|
||||
}
|
||||
}
|
||||
|
||||
func main() {
|
||||
var (
|
||||
listenHost = flag.String("listen-host", "127.0.0.1", "local proxy listen host")
|
||||
listenPort = flag.Int("listen-port", 8080, "local proxy listen port")
|
||||
serverHost = flag.String("server-host", "", "remote DragonTCP server host")
|
||||
serverPort = flag.Int("server-port", 53, "remote DragonTCP server port")
|
||||
token = flag.String("token", "", "optional shared token")
|
||||
maxConnections = flag.Int("max-connections", 20000, "max simultaneous proxy connections")
|
||||
transport = flag.String("transport", "chunk", "transport: chunk (mandatory in LiteVPN build)")
|
||||
tcpBuffer = flag.Int("tcp-buffer", 0, "optional TCP read/write buffer bytes; 0 keeps OS autotuning")
|
||||
chunkStart = flag.Int("chunk-start", 1048576, "initial adaptive chunk payload bytes")
|
||||
chunkMin = flag.Int("chunk-min", 32, "minimum adaptive chunk payload bytes")
|
||||
chunkMax = flag.Int("chunk-max", 1048576, "maximum adaptive chunk payload bytes (up to 1 MiB)")
|
||||
chunkAdaptive = flag.Bool("chunk-adaptive", true, "automatically shrink on failures and grow after stable success")
|
||||
chunkSuccesses = flag.Int("chunk-grow-after", 16, "successful data records required before increasing chunk size")
|
||||
chunkAdaptLog = flag.Bool("chunk-adapt-log", true, "print adaptive chunk size changes")
|
||||
chunkSizeLegacy = flag.Int("chunk-size", 0, "legacy fixed chunk size; nonzero disables adaptation")
|
||||
chunkPollers = flag.Int("chunk-pollers", 1, "parallel downstream chunk pollers (LiteVPN default 1)")
|
||||
chunkReconnect = flag.Int("chunk-reconnect-every", 1, "reconnect each transaction lane after N requests; 1 = one request per TCP/53 connection")
|
||||
chunkPollDelay = flag.Duration("chunk-poll-delay", 2*time.Millisecond, "delay after an empty chunk poll")
|
||||
chunkTimeout = flag.Duration("chunk-timeout", 2*time.Second, "per-record transaction timeout before adaptive shrink")
|
||||
)
|
||||
flag.Parse()
|
||||
|
||||
if *serverHost == "" {
|
||||
fmt.Fprintln(os.Stderr, "--server-host is required")
|
||||
os.Exit(2)
|
||||
}
|
||||
|
||||
*transport = strings.ToLower(*transport)
|
||||
if *transport != "chunk" {
|
||||
fmt.Fprintln(os.Stderr, "DragonTCP LiteVPN requires --transport chunk (adaptive XOR-framed TCP/53)")
|
||||
os.Exit(2)
|
||||
}
|
||||
if *chunkSizeLegacy != 0 {
|
||||
if *chunkSizeLegacy < 32 || *chunkSizeLegacy > protocol.MaxChunkPayload {
|
||||
fmt.Fprintf(os.Stderr, "--chunk-size must be between 32 and %d\n", protocol.MaxChunkPayload)
|
||||
os.Exit(2)
|
||||
}
|
||||
*chunkStart = *chunkSizeLegacy
|
||||
*chunkMin = *chunkSizeLegacy
|
||||
*chunkMax = *chunkSizeLegacy
|
||||
*chunkAdaptive = false
|
||||
}
|
||||
if *chunkMin < 32 || *chunkMax > protocol.MaxChunkPayload || *chunkMin > *chunkStart || *chunkStart > *chunkMax {
|
||||
fmt.Fprintf(os.Stderr, "require 32 <= --chunk-min <= --chunk-start <= --chunk-max <= %d\n", protocol.MaxChunkPayload)
|
||||
os.Exit(2)
|
||||
}
|
||||
if *chunkSuccesses < 1 {
|
||||
fmt.Fprintln(os.Stderr, "--chunk-grow-after must be at least 1")
|
||||
os.Exit(2)
|
||||
}
|
||||
if *chunkPollers < 1 || *chunkPollers > 128 {
|
||||
fmt.Fprintln(os.Stderr, "--chunk-pollers must be between 1 and 128")
|
||||
os.Exit(2)
|
||||
}
|
||||
chunkOpts := chunkClientOptions{
|
||||
startSize: *chunkStart,
|
||||
minSize: *chunkMin,
|
||||
maxSize: *chunkMax,
|
||||
adaptive: *chunkAdaptive,
|
||||
adaptSuccesses: *chunkSuccesses,
|
||||
adaptLog: *chunkAdaptLog,
|
||||
pollers: *chunkPollers,
|
||||
reconnectEvery: *chunkReconnect,
|
||||
pollDelay: *chunkPollDelay,
|
||||
txnTimeout: *chunkTimeout,
|
||||
tcpBuffer: *tcpBuffer,
|
||||
}
|
||||
|
||||
listenAddr := net.JoinHostPort(*listenHost, strconv.Itoa(*listenPort))
|
||||
serverAddr := net.JoinHostPort(*serverHost, strconv.Itoa(*serverPort))
|
||||
|
||||
ln, err := net.Listen("tcp", listenAddr)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer ln.Close()
|
||||
|
||||
fmt.Printf("local Go HTTP proxy listening on %s\n", listenAddr)
|
||||
fmt.Printf("remote DragonTCP endpoint=%s\n", serverAddr)
|
||||
fmt.Printf("max_connections=%d transport=%s tcp_buffer=%d\n", *maxConnections, *transport, *tcpBuffer)
|
||||
if *transport == "chunk" {
|
||||
fmt.Printf(
|
||||
"adaptive_chunk=%v start=%d min=%d max=%d grow_after=%d pollers=%d reconnect_every=%d timeout=%s\n",
|
||||
*chunkAdaptive,
|
||||
*chunkStart,
|
||||
*chunkMin,
|
||||
*chunkMax,
|
||||
*chunkSuccesses,
|
||||
*chunkPollers,
|
||||
*chunkReconnect,
|
||||
chunkTimeout.String(),
|
||||
)
|
||||
}
|
||||
|
||||
slots := make(chan struct{}, *maxConnections)
|
||||
|
||||
for {
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, "accept:", err)
|
||||
continue
|
||||
}
|
||||
|
||||
select {
|
||||
case slots <- struct{}{}:
|
||||
go handleLocal(conn, serverAddr, *token, *transport, *tcpBuffer, chunkOpts, slots)
|
||||
default:
|
||||
writeHTTPError(
|
||||
conn,
|
||||
503,
|
||||
"Service Unavailable",
|
||||
"proxy connection limit reached",
|
||||
)
|
||||
_ = conn.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user