V13
This commit is contained in:
+525
-488
File diff suppressed because it is too large
Load Diff
@@ -10,14 +10,11 @@ func TestAdaptiveSizerRecoversFromMinimum(t *testing.T) {
|
||||
adaptive: true,
|
||||
adaptSuccesses: 2,
|
||||
}
|
||||
s := newAdaptiveSizer("test", opts)
|
||||
s := newAdaptiveSizer("test", 64, 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)
|
||||
}
|
||||
@@ -26,11 +23,9 @@ func TestAdaptiveSizerRecoversFromMinimum(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
func TestReconnectZeroMeansPersistent(t *testing.T) {
|
||||
lane := newRequestLane("127.0.0.1:1", 0, 0, 0)
|
||||
if lane.reconnectEvery != 0 {
|
||||
t.Fatalf("reconnectEvery=%d, want 0", lane.reconnectEvery)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -358,25 +358,26 @@ func handleLocal(conn net.Conn, serverAddr, token, transport string, tcpBuffer i
|
||||
|
||||
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")
|
||||
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 (DragonTCP binary adaptive transport)")
|
||||
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, "reserved compatibility setting; binary transport uses one download worker")
|
||||
chunkConcurrency = flag.Int("chunk-concurrency", 1, "maximum adaptive download pipeline depth (1-256); 1 keeps concurrency fixed at one")
|
||||
chunkReconnect = flag.Int("chunk-reconnect-every", 0, "force reconnect after N logical requests; 0 = persistent/automatic")
|
||||
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()
|
||||
|
||||
@@ -387,7 +388,7 @@ func main() {
|
||||
|
||||
*transport = strings.ToLower(*transport)
|
||||
if *transport != "chunk" {
|
||||
fmt.Fprintln(os.Stderr, "DragonTCP LiteVPN requires --transport chunk (adaptive XOR-framed TCP/53)")
|
||||
fmt.Fprintln(os.Stderr, "DragonTCP requires --transport chunk (binary adaptive TCP/53 transport)")
|
||||
os.Exit(2)
|
||||
}
|
||||
if *chunkSizeLegacy != 0 {
|
||||
@@ -412,6 +413,14 @@ func main() {
|
||||
fmt.Fprintln(os.Stderr, "--chunk-pollers must be between 1 and 128")
|
||||
os.Exit(2)
|
||||
}
|
||||
if *chunkConcurrency < 1 || *chunkConcurrency > 256 {
|
||||
fmt.Fprintln(os.Stderr, "--chunk-concurrency must be between 1 and 256")
|
||||
os.Exit(2)
|
||||
}
|
||||
if *chunkReconnect < 0 {
|
||||
fmt.Fprintln(os.Stderr, "--chunk-reconnect-every must be 0 or greater")
|
||||
os.Exit(2)
|
||||
}
|
||||
chunkOpts := chunkClientOptions{
|
||||
startSize: *chunkStart,
|
||||
minSize: *chunkMin,
|
||||
@@ -424,6 +433,7 @@ func main() {
|
||||
pollDelay: *chunkPollDelay,
|
||||
txnTimeout: *chunkTimeout,
|
||||
tcpBuffer: *tcpBuffer,
|
||||
maxPipeline: *chunkConcurrency,
|
||||
}
|
||||
|
||||
listenAddr := net.JoinHostPort(*listenHost, strconv.Itoa(*listenPort))
|
||||
@@ -441,13 +451,14 @@ func main() {
|
||||
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",
|
||||
"adaptive_chunk=%v start=%d min=%d max=%d grow_after=%d pollers=%d concurrency=%d reconnect_every=%d timeout=%s\n",
|
||||
*chunkAdaptive,
|
||||
*chunkStart,
|
||||
*chunkMin,
|
||||
*chunkMax,
|
||||
*chunkSuccesses,
|
||||
*chunkPollers,
|
||||
*chunkConcurrency,
|
||||
*chunkReconnect,
|
||||
chunkTimeout.String(),
|
||||
)
|
||||
|
||||
+330
-328
@@ -3,104 +3,91 @@ package main
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"dragontcp/internal/protocol"
|
||||
"dragontcp/internal/wire"
|
||||
)
|
||||
|
||||
type chunkSession struct {
|
||||
id string
|
||||
target net.Conn
|
||||
maxChunk int
|
||||
maxChunks int
|
||||
type streamSession struct {
|
||||
sid wire.SessionID
|
||||
target net.Conn
|
||||
targetName string
|
||||
maxChunk int
|
||||
maxBuffer int
|
||||
debug *serverDebug
|
||||
|
||||
mu sync.Mutex
|
||||
notify chan struct{}
|
||||
chunks map[uint64][]byte
|
||||
nextDown uint64
|
||||
buf []byte
|
||||
base uint64
|
||||
eof bool
|
||||
closed bool
|
||||
lastSeen time.Time
|
||||
debug *serverDebug
|
||||
|
||||
upMu sync.Mutex
|
||||
expectedUp uint64
|
||||
lastUpSeq uint64
|
||||
lastUpLen int
|
||||
haveLastUp bool
|
||||
}
|
||||
|
||||
func newChunkSession(id string, target net.Conn, maxChunk, maxChunks int, debug *serverDebug) *chunkSession {
|
||||
s := &chunkSession{
|
||||
id: id,
|
||||
target: target,
|
||||
maxChunk: maxChunk,
|
||||
maxChunks: maxChunks,
|
||||
notify: make(chan struct{}),
|
||||
chunks: make(map[uint64][]byte, maxChunks),
|
||||
lastSeen: time.Now(),
|
||||
debug: debug,
|
||||
func newStreamSession(sid wire.SessionID, target net.Conn, targetName string, maxChunk, maxBuffer int, debug *serverDebug) *streamSession {
|
||||
s := &streamSession{
|
||||
sid: sid,
|
||||
target: target,
|
||||
targetName: targetName,
|
||||
maxChunk: maxChunk,
|
||||
maxBuffer: maxBuffer,
|
||||
debug: debug,
|
||||
notify: make(chan struct{}),
|
||||
lastSeen: time.Now(),
|
||||
}
|
||||
go s.readTarget()
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *chunkSession) signalLocked() {
|
||||
func (s *streamSession) signalLocked() {
|
||||
close(s.notify)
|
||||
s.notify = make(chan struct{})
|
||||
}
|
||||
|
||||
func (s *chunkSession) touchLocked() {
|
||||
s.lastSeen = time.Now()
|
||||
}
|
||||
|
||||
func (s *chunkSession) touch() {
|
||||
s.mu.Lock()
|
||||
s.touchLocked()
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *chunkSession) readTarget() {
|
||||
buf := make([]byte, s.maxChunk)
|
||||
func (s *streamSession) touchLocked() { s.lastSeen = time.Now() }
|
||||
|
||||
func (s *streamSession) readTarget() {
|
||||
tmp := make([]byte, 64*1024)
|
||||
for {
|
||||
n, err := s.target.Read(buf)
|
||||
n, err := s.target.Read(tmp)
|
||||
if n > 0 {
|
||||
data := append([]byte(nil), buf[:n]...)
|
||||
if s.debug != nil && s.debug.enabled {
|
||||
s.debug.bytesDown.Add(uint64(n))
|
||||
}
|
||||
|
||||
for {
|
||||
data := append([]byte(nil), tmp[:n]...)
|
||||
for len(data) > 0 {
|
||||
s.mu.Lock()
|
||||
for !s.closed && len(s.buf) >= s.maxBuffer {
|
||||
ch := s.notify
|
||||
s.mu.Unlock()
|
||||
<-ch
|
||||
s.mu.Lock()
|
||||
}
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
if len(s.chunks) < s.maxChunks {
|
||||
seq := s.nextDown
|
||||
s.nextDown++
|
||||
s.chunks[seq] = data
|
||||
s.touchLocked()
|
||||
s.signalLocked()
|
||||
s.mu.Unlock()
|
||||
break
|
||||
room := s.maxBuffer - len(s.buf)
|
||||
take := len(data)
|
||||
if take > room {
|
||||
take = room
|
||||
}
|
||||
ch := s.notify
|
||||
s.buf = append(s.buf, data[:take]...)
|
||||
data = data[take:]
|
||||
s.touchLocked()
|
||||
s.signalLocked()
|
||||
s.mu.Unlock()
|
||||
<-ch
|
||||
if s.debug != nil && s.debug.enabled {
|
||||
s.debug.bytesDown.Add(uint64(take))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if s.debug != nil && s.debug.enabled {
|
||||
s.debug.logf("TARGET EOF session=%s err=%v", s.id, err)
|
||||
}
|
||||
s.mu.Lock()
|
||||
if !s.closed {
|
||||
s.eof = true
|
||||
@@ -113,116 +100,133 @@ func (s *chunkSession) readTarget() {
|
||||
}
|
||||
}
|
||||
|
||||
// push is idempotent for the most recently accepted sequence. This matters
|
||||
// when the server receives a record but the tiny ACK is lost: the client can
|
||||
// retry the same sequence at a smaller adaptive size without duplicating bytes
|
||||
// in the target stream. The ACK reports the length that was actually accepted.
|
||||
func (s *chunkSession) push(seq uint64, data []byte) (int, error) {
|
||||
s.upMu.Lock()
|
||||
defer s.upMu.Unlock()
|
||||
|
||||
if len(data) == 0 || len(data) > s.maxChunk {
|
||||
return 0, fmt.Errorf("upload record size %d is invalid", len(data))
|
||||
func (s *streamSession) ackLocked(offset uint64) {
|
||||
if offset <= s.base {
|
||||
return
|
||||
}
|
||||
|
||||
if s.haveLastUp && seq == s.lastUpSeq {
|
||||
s.touch()
|
||||
return s.lastUpLen, nil
|
||||
end := s.base + uint64(len(s.buf))
|
||||
if offset > end {
|
||||
offset = end
|
||||
}
|
||||
|
||||
if seq < s.expectedUp {
|
||||
return 0, fmt.Errorf("upload sequence %d is too old", seq)
|
||||
drop := int(offset - s.base)
|
||||
if drop <= 0 {
|
||||
return
|
||||
}
|
||||
if seq > s.expectedUp {
|
||||
return 0, fmt.Errorf("unexpected upload sequence %d, expected %d", seq, s.expectedUp)
|
||||
s.buf = s.buf[drop:]
|
||||
s.base = offset
|
||||
if len(s.buf) == 0 {
|
||||
s.buf = nil
|
||||
} else if cap(s.buf) > 4*len(s.buf) && cap(s.buf) > 1024*1024 {
|
||||
compact := append([]byte(nil), s.buf...)
|
||||
s.buf = compact
|
||||
}
|
||||
|
||||
if _, err := s.target.Write(data); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if s.debug != nil && s.debug.enabled {
|
||||
s.debug.bytesUp.Add(uint64(len(data)))
|
||||
s.debug.pushRecords.Add(1)
|
||||
}
|
||||
|
||||
s.lastUpSeq = seq
|
||||
s.lastUpLen = len(data)
|
||||
s.haveLastUp = true
|
||||
s.expectedUp++
|
||||
s.touch()
|
||||
return len(data), nil
|
||||
s.signalLocked()
|
||||
}
|
||||
|
||||
// pull returns at most limit bytes from the requested stored chunk, beginning
|
||||
// at offset. The chunk sequence stays stable while the client retries smaller
|
||||
// fragments, so a large queued chunk can always be recovered after an MTU-like
|
||||
// failure without reopening the proxied destination connection.
|
||||
func (s *chunkSession) pull(want uint64, ack int64, offset, limit int, wait time.Duration) (data []byte, total int, eof bool, final uint64, waitExpired bool, err error) {
|
||||
if offset < 0 || limit <= 0 || limit > s.maxChunk {
|
||||
return nil, 0, false, 0, false, fmt.Errorf("invalid pull offset/limit")
|
||||
}
|
||||
func (s *streamSession) ack(offset uint64) {
|
||||
s.mu.Lock()
|
||||
s.ackLocked(offset)
|
||||
s.touchLocked()
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
timer := time.NewTimer(wait)
|
||||
defer timer.Stop()
|
||||
func (s *streamSession) readAt(offset uint64, limit int, wait time.Duration) ([]byte, byte, error) {
|
||||
if limit < 1 || limit > s.maxChunk {
|
||||
return nil, wire.StatusError, fmt.Errorf("invalid download limit %d", limit)
|
||||
}
|
||||
deadline := time.Now().Add(wait)
|
||||
firstDataAt := time.Time{}
|
||||
|
||||
for {
|
||||
s.mu.Lock()
|
||||
s.touchLocked()
|
||||
|
||||
if ack >= 0 {
|
||||
removed := false
|
||||
for seq := range s.chunks {
|
||||
if seq <= uint64(ack) {
|
||||
delete(s.chunks, seq)
|
||||
removed = true
|
||||
if offset < s.base {
|
||||
s.mu.Unlock()
|
||||
return nil, wire.StatusError, fmt.Errorf("download offset %d was already acknowledged (base=%d)", offset, s.base)
|
||||
}
|
||||
rel64 := offset - s.base
|
||||
if rel64 <= uint64(len(s.buf)) {
|
||||
rel := int(rel64)
|
||||
available := len(s.buf) - rel
|
||||
if available > 0 {
|
||||
if firstDataAt.IsZero() {
|
||||
firstDataAt = time.Now()
|
||||
}
|
||||
}
|
||||
if removed {
|
||||
s.signalLocked()
|
||||
}
|
||||
}
|
||||
|
||||
if chunk, ok := s.chunks[want]; ok {
|
||||
if offset >= len(chunk) {
|
||||
// Coalesce tiny target reads briefly. This prevents a 1-2 byte
|
||||
// producer read from becoming a permanent tiny tunnel record.
|
||||
if available < limit && !s.eof && wait > 0 && time.Since(firstDataAt) < 2*time.Millisecond {
|
||||
ch := s.notify
|
||||
s.mu.Unlock()
|
||||
select {
|
||||
case <-ch:
|
||||
case <-time.After(2 * time.Millisecond):
|
||||
}
|
||||
continue
|
||||
}
|
||||
n := available
|
||||
if n > limit {
|
||||
n = limit
|
||||
}
|
||||
out := append([]byte(nil), s.buf[rel:rel+n]...)
|
||||
s.mu.Unlock()
|
||||
return nil, len(chunk), false, 0, false, fmt.Errorf("pull offset %d beyond chunk size %d", offset, len(chunk))
|
||||
return out, wire.StatusData, nil
|
||||
}
|
||||
end := offset + limit
|
||||
if end > len(chunk) {
|
||||
end = len(chunk)
|
||||
if s.eof || s.closed {
|
||||
s.mu.Unlock()
|
||||
return nil, wire.StatusEOF, nil
|
||||
}
|
||||
out := append([]byte(nil), chunk[offset:end]...)
|
||||
total = len(chunk)
|
||||
} else {
|
||||
s.mu.Unlock()
|
||||
return out, total, false, 0, false, nil
|
||||
return nil, wire.StatusError, fmt.Errorf("download offset %d is beyond buffered stream end %d", offset, s.base+uint64(len(s.buf)))
|
||||
}
|
||||
|
||||
if s.eof && want >= s.nextDown {
|
||||
final = s.nextDown
|
||||
if wait <= 0 || time.Now().After(deadline) {
|
||||
s.mu.Unlock()
|
||||
return nil, 0, true, final, false, nil
|
||||
return nil, wire.StatusWait, nil
|
||||
}
|
||||
|
||||
if s.closed {
|
||||
final = s.nextDown
|
||||
s.mu.Unlock()
|
||||
return nil, 0, true, final, false, nil
|
||||
}
|
||||
|
||||
ch := s.notify
|
||||
remaining := time.Until(deadline)
|
||||
s.mu.Unlock()
|
||||
|
||||
select {
|
||||
case <-ch:
|
||||
continue
|
||||
case <-timer.C:
|
||||
return nil, 0, false, 0, true, nil
|
||||
case <-time.After(remaining):
|
||||
return nil, wire.StatusWait, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *chunkSession) close() {
|
||||
func (s *streamSession) upload(offset uint64, data []byte) error {
|
||||
if len(data) == 0 || len(data) > s.maxChunk {
|
||||
return fmt.Errorf("invalid upload size %d", len(data))
|
||||
}
|
||||
s.upMu.Lock()
|
||||
defer s.upMu.Unlock()
|
||||
|
||||
if offset < s.expectedUp {
|
||||
// Idempotent retry after a lost ACK.
|
||||
if offset+uint64(len(data)) <= s.expectedUp {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("overlapping upload retry at %d", offset)
|
||||
}
|
||||
if offset != s.expectedUp {
|
||||
return fmt.Errorf("upload gap: got %d expected %d", offset, s.expectedUp)
|
||||
}
|
||||
if _, err := s.target.Write(data); err != nil {
|
||||
return err
|
||||
}
|
||||
s.expectedUp += uint64(len(data))
|
||||
s.mu.Lock()
|
||||
s.touchLocked()
|
||||
s.mu.Unlock()
|
||||
if s.debug != nil && s.debug.enabled {
|
||||
s.debug.bytesUp.Add(uint64(len(data)))
|
||||
s.debug.pushRecords.Add(1)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *streamSession) close() {
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.mu.Unlock()
|
||||
@@ -234,272 +238,270 @@ func (s *chunkSession) close() {
|
||||
_ = s.target.Close()
|
||||
}
|
||||
|
||||
type chunkManager struct {
|
||||
type streamManager struct {
|
||||
mu sync.RWMutex
|
||||
sessions map[string]*chunkSession
|
||||
sessions map[string]*streamSession
|
||||
timeout time.Duration
|
||||
debug *serverDebug
|
||||
}
|
||||
|
||||
func newChunkManager(timeout time.Duration, debug *serverDebug) *chunkManager {
|
||||
m := &chunkManager{
|
||||
sessions: make(map[string]*chunkSession),
|
||||
timeout: timeout,
|
||||
debug: debug,
|
||||
}
|
||||
func sidKey(sid wire.SessionID) string { return string(sid[:]) }
|
||||
|
||||
func newStreamManager(timeout time.Duration, debug *serverDebug) *streamManager {
|
||||
m := &streamManager{sessions: make(map[string]*streamSession), timeout: timeout, debug: debug}
|
||||
go m.cleanupLoop()
|
||||
return m
|
||||
}
|
||||
|
||||
func (m *chunkManager) get(id string) *chunkSession {
|
||||
func (m *streamManager) get(sid wire.SessionID) *streamSession {
|
||||
m.mu.RLock()
|
||||
s := m.sessions[id]
|
||||
s := m.sessions[sidKey(sid)]
|
||||
m.mu.RUnlock()
|
||||
return s
|
||||
}
|
||||
|
||||
func (m *chunkManager) count() int {
|
||||
m.mu.RLock()
|
||||
n := len(m.sessions)
|
||||
m.mu.RUnlock()
|
||||
return n
|
||||
}
|
||||
|
||||
func (m *chunkManager) add(id string, s *chunkSession) error {
|
||||
func (m *streamManager) addOrGet(sid wire.SessionID, s *streamSession) (*streamSession, bool) {
|
||||
key := sidKey(sid)
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if _, exists := m.sessions[id]; exists {
|
||||
return fmt.Errorf("session already exists")
|
||||
if old := m.sessions[key]; old != nil {
|
||||
m.mu.Unlock()
|
||||
s.close()
|
||||
return old, false
|
||||
}
|
||||
m.sessions[id] = s
|
||||
return nil
|
||||
m.sessions[key] = s
|
||||
m.mu.Unlock()
|
||||
return s, true
|
||||
}
|
||||
|
||||
func (m *chunkManager) remove(id string) {
|
||||
func (m *streamManager) remove(sid wire.SessionID) {
|
||||
key := sidKey(sid)
|
||||
m.mu.Lock()
|
||||
s := m.sessions[id]
|
||||
delete(m.sessions, id)
|
||||
s := m.sessions[key]
|
||||
delete(m.sessions, key)
|
||||
m.mu.Unlock()
|
||||
if s != nil {
|
||||
s.close()
|
||||
if m.debug != nil && m.debug.enabled {
|
||||
m.debug.sessionsClosed.Add(1)
|
||||
m.debug.activeSessions.Add(-1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (m *chunkManager) cleanupLoop() {
|
||||
func (m *streamManager) count() int { m.mu.RLock(); n := len(m.sessions); m.mu.RUnlock(); return n }
|
||||
|
||||
func (m *streamManager) cleanupLoop() {
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for range ticker.C {
|
||||
cutoff := time.Now().Add(-m.timeout)
|
||||
var stale []string
|
||||
|
||||
var stale []wire.SessionID
|
||||
m.mu.RLock()
|
||||
for id, s := range m.sessions {
|
||||
for _, s := range m.sessions {
|
||||
s.mu.Lock()
|
||||
last := s.lastSeen
|
||||
closed := s.closed
|
||||
sid := s.sid
|
||||
s.mu.Unlock()
|
||||
if closed || last.Before(cutoff) {
|
||||
stale = append(stale, id)
|
||||
stale = append(stale, sid)
|
||||
}
|
||||
}
|
||||
m.mu.RUnlock()
|
||||
|
||||
for _, id := range stale {
|
||||
if m.debug != nil && m.debug.enabled {
|
||||
m.debug.logf("SESSION timeout-close id=%s active_sessions=%d", id, m.count())
|
||||
}
|
||||
m.remove(id)
|
||||
if m.debug != nil && m.debug.enabled {
|
||||
m.debug.sessionsClosed.Add(1)
|
||||
m.debug.activeSessions.Add(-1)
|
||||
}
|
||||
for _, sid := range stale {
|
||||
m.remove(sid)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func decodeWireToken(token string) string {
|
||||
if token == "-" {
|
||||
return ""
|
||||
func parseProbe(payload []byte) (kind byte, value int, token string, err error) {
|
||||
if len(payload) < 11 || !bytes.Equal(payload[:4], wire.ProbeMagic[:]) {
|
||||
return 0, 0, "", fmt.Errorf("bad probe payload")
|
||||
}
|
||||
return token
|
||||
kind = payload[4]
|
||||
tl := int(binary.BigEndian.Uint16(payload[5:7]))
|
||||
value = int(binary.BigEndian.Uint32(payload[7:11]))
|
||||
if 11+tl > len(payload) {
|
||||
return 0, 0, "", fmt.Errorf("bad probe token length")
|
||||
}
|
||||
token = string(payload[11 : 11+tl])
|
||||
return
|
||||
}
|
||||
|
||||
func isChunkCommand(payload []byte) bool {
|
||||
return bytes.HasPrefix(payload, []byte("COPEN ")) ||
|
||||
bytes.HasPrefix(payload, []byte("CPUSH ")) ||
|
||||
bytes.HasPrefix(payload, []byte("CPULL ")) ||
|
||||
bytes.HasPrefix(payload, []byte("CCLOSE "))
|
||||
func probePattern(n int) []byte {
|
||||
out := make([]byte, n)
|
||||
for i := range out {
|
||||
out[i] = byte((i*31 + 17) & 0xff)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func processChunkCommand(
|
||||
conn net.Conn,
|
||||
requestID uint32,
|
||||
payload []byte,
|
||||
token string,
|
||||
allowPrivate bool,
|
||||
cache *dnsCache,
|
||||
tcpBuffer int,
|
||||
manager *chunkManager,
|
||||
maxChunk int,
|
||||
maxBufferedChunks int,
|
||||
pollWait time.Duration,
|
||||
debug *serverDebug,
|
||||
) error {
|
||||
if bytes.HasPrefix(payload, []byte("COPEN ")) {
|
||||
parts := strings.Fields(string(payload))
|
||||
if len(parts) != 5 {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR bad COPEN"))
|
||||
func parseOpen(payload []byte) (token, host string, port int, err error) {
|
||||
if len(payload) < 6 {
|
||||
return "", "", 0, fmt.Errorf("bad OPEN payload")
|
||||
}
|
||||
tl := int(binary.BigEndian.Uint16(payload[0:2]))
|
||||
hl := int(binary.BigEndian.Uint16(payload[2:4]))
|
||||
port = int(binary.BigEndian.Uint16(payload[4:6]))
|
||||
if port < 1 || 6+tl+hl != len(payload) {
|
||||
return "", "", 0, fmt.Errorf("bad OPEN lengths")
|
||||
}
|
||||
token = string(payload[6 : 6+tl])
|
||||
host = string(payload[6+tl:])
|
||||
if host == "" {
|
||||
return "", "", 0, fmt.Errorf("empty target host")
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func processWireRequest(conn net.Conn, req wire.Request, token string, allowPrivate bool, cache *dnsCache, tcpBuffer int, manager *streamManager, maxChunk, maxBuffer int, pollWait time.Duration, debug *serverDebug) error {
|
||||
switch req.Mode {
|
||||
case wire.ModeProbe:
|
||||
kind, value, supplied, err := parseProbe(req.Payload)
|
||||
if err != nil {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
|
||||
}
|
||||
if !tokenEqual(decodeWireToken(parts[1]), token) {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
|
||||
if !tokenEqual(supplied, token) {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("authentication failed"))
|
||||
}
|
||||
sid := parts[2]
|
||||
if len(sid) < 16 || len(sid) > 64 {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid session id"))
|
||||
}
|
||||
host := parts[3]
|
||||
port, err := strconv.Atoi(parts[4])
|
||||
if err != nil || port < 1 || port > 65535 {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid port"))
|
||||
switch kind {
|
||||
case wire.ProbeUpload:
|
||||
if len(req.Payload) > maxChunk {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("probe too large"))
|
||||
}
|
||||
return wire.WriteResponse(conn, wire.StatusOK, nil)
|
||||
case wire.ProbeDownload:
|
||||
if value < 1 || value > maxChunk {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("probe too large"))
|
||||
}
|
||||
return wire.WriteMaskedResponse(conn, wire.StatusData, probePattern(value), req.Session, wire.ModeProbe, req.Seq)
|
||||
case wire.ProbeKeepalive:
|
||||
return wire.WriteResponse(conn, wire.StatusOK, nil)
|
||||
case wire.ProbeBatch:
|
||||
count := value
|
||||
if count < 1 {
|
||||
count = 1
|
||||
}
|
||||
if count > 16 {
|
||||
count = 16
|
||||
}
|
||||
for i := 0; i < count; i++ {
|
||||
data := probePattern(32)
|
||||
if err := wire.WriteMaskedResponse(conn, wire.StatusData, data, req.Session, wire.ModeProbe, req.Seq+uint64(i)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
default:
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("unknown probe kind"))
|
||||
}
|
||||
|
||||
case wire.ModeOpen:
|
||||
supplied, host, port, err := parseOpen(req.Payload)
|
||||
if err != nil {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
|
||||
}
|
||||
if !tokenEqual(supplied, token) {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("authentication failed"))
|
||||
}
|
||||
if old := manager.get(req.Session); old != nil {
|
||||
body := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(body, uint32(maxChunk))
|
||||
return wire.WriteMaskedResponse(conn, wire.StatusOK, body, req.Session, wire.ModeOpen, req.Seq)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
target, err := dialTarget(ctx, host, port, allowPrivate, cache, tcpBuffer)
|
||||
cancel()
|
||||
if err != nil {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR "+err.Error()))
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
|
||||
}
|
||||
|
||||
session := newChunkSession(sid, target, maxChunk, maxBufferedChunks, debug)
|
||||
if err := manager.add(sid, session); err != nil {
|
||||
session.close()
|
||||
if debug != nil && debug.enabled {
|
||||
debug.errorf("COPEN session=%s target=%s:%d failed: %v", sid, host, port, err)
|
||||
}
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR "+err.Error()))
|
||||
}
|
||||
if debug != nil && debug.enabled {
|
||||
session := newStreamSession(req.Session, target, fmt.Sprintf("%s:%d", host, port), maxChunk, maxBuffer, debug)
|
||||
_, created := manager.addOrGet(req.Session, session)
|
||||
if created && debug != nil && debug.enabled {
|
||||
debug.sessionsOpened.Add(1)
|
||||
debug.activeSessions.Add(1)
|
||||
debug.logf("SESSION OPEN id=%s peer=%v target=%s:%d max_chunk=%d active_sessions=%d", sid, conn.RemoteAddr(), host, port, maxChunk, manager.count())
|
||||
debug.chunkf("COPEN id=%s target=%s:%d -> OPENED max=%d", sid, host, port, maxChunk)
|
||||
debug.logf("SESSION OPEN sid=%x target=%s:%d active_sessions=%d", req.Session[:4], host, port, manager.count())
|
||||
}
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte(fmt.Sprintf("OPENED %d", maxChunk)))
|
||||
}
|
||||
body := make([]byte, 4)
|
||||
binary.BigEndian.PutUint32(body, uint32(maxChunk))
|
||||
return wire.WriteMaskedResponse(conn, wire.StatusOK, body, req.Session, wire.ModeOpen, req.Seq)
|
||||
|
||||
if bytes.HasPrefix(payload, []byte("CPUSH ")) {
|
||||
parts := bytes.SplitN(payload, []byte(" "), 5)
|
||||
if len(parts) != 5 {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR bad CPUSH"))
|
||||
}
|
||||
if !tokenEqual(decodeWireToken(string(parts[1])), token) {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
|
||||
}
|
||||
sid := string(parts[2])
|
||||
seq, err := strconv.ParseUint(string(parts[3]), 10, 64)
|
||||
if err != nil {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid sequence"))
|
||||
}
|
||||
s := manager.get(sid)
|
||||
case wire.ModeUpload:
|
||||
s := manager.get(req.Session)
|
||||
if s == nil {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR unknown session"))
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("unknown session"))
|
||||
}
|
||||
accepted, err := s.push(seq, parts[4])
|
||||
if err != nil {
|
||||
if debug != nil && debug.enabled {
|
||||
debug.errorf("CPUSH id=%s seq=%d bytes=%d: %v", sid, seq, len(parts[4]), err)
|
||||
}
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR "+err.Error()))
|
||||
if len(req.Payload) > maxChunk {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("upload too large"))
|
||||
}
|
||||
if debug != nil {
|
||||
debug.chunkf("CPUSH id=%s seq=%d bytes=%d -> ACK accepted=%d", sid, seq, len(parts[4]), accepted)
|
||||
if err := s.upload(req.Seq, req.Payload); err != nil {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
|
||||
}
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte(fmt.Sprintf("ACK %d %d", seq, accepted)))
|
||||
}
|
||||
return wire.WriteResponse(conn, wire.StatusOK, nil)
|
||||
|
||||
if bytes.HasPrefix(payload, []byte("CPULL ")) {
|
||||
parts := strings.Fields(string(payload))
|
||||
if len(parts) != 7 {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR bad CPULL"))
|
||||
}
|
||||
if !tokenEqual(decodeWireToken(parts[1]), token) {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
|
||||
}
|
||||
s := manager.get(parts[2])
|
||||
case wire.ModeDownload:
|
||||
s := manager.get(req.Session)
|
||||
if s == nil {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR unknown session"))
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("unknown session"))
|
||||
}
|
||||
ack, err := strconv.ParseInt(parts[3], 10, 64)
|
||||
if err != nil || ack < -1 {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid ack"))
|
||||
if len(req.Payload) != 14 {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("bad download request"))
|
||||
}
|
||||
want, err := strconv.ParseUint(parts[4], 10, 64)
|
||||
if err != nil {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid want"))
|
||||
}
|
||||
offset, err := strconv.Atoi(parts[5])
|
||||
if err != nil || offset < 0 {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid offset"))
|
||||
}
|
||||
limit, err := strconv.Atoi(parts[6])
|
||||
if err != nil || limit < 1 {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid limit"))
|
||||
ack := binary.BigEndian.Uint64(req.Payload[0:8])
|
||||
limit := int(binary.BigEndian.Uint32(req.Payload[8:12]))
|
||||
count := int(binary.BigEndian.Uint16(req.Payload[12:14]))
|
||||
if limit < 1 {
|
||||
limit = 1
|
||||
}
|
||||
if limit > maxChunk {
|
||||
limit = maxChunk
|
||||
}
|
||||
if count < 1 {
|
||||
count = 1
|
||||
}
|
||||
if count > 256 {
|
||||
count = 256
|
||||
}
|
||||
s.ack(ack)
|
||||
offset := req.Seq
|
||||
if debug != nil && debug.enabled {
|
||||
debug.pullRequests.Add(1)
|
||||
debug.chunkf("CPULL id=%s ack=%d want=%d offset=%d limit=%d", parts[2], ack, want, offset, limit)
|
||||
}
|
||||
|
||||
data, total, eof, final, waitExpired, err := s.pull(want, ack, offset, limit, pollWait)
|
||||
if err != nil {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR "+err.Error()))
|
||||
}
|
||||
if waitExpired {
|
||||
if debug != nil && debug.enabled {
|
||||
debug.waitRecords.Add(1)
|
||||
debug.chunkf("CPULL id=%s want=%d -> WAIT", parts[2], want)
|
||||
for i := 0; i < count; i++ {
|
||||
wait := time.Duration(0)
|
||||
if i == 0 {
|
||||
wait = pollWait
|
||||
}
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("WAIT"))
|
||||
}
|
||||
if eof {
|
||||
if debug != nil {
|
||||
debug.chunkf("CPULL id=%s want=%d -> EOF final=%d", parts[2], want, final)
|
||||
data, status, err := s.readAt(offset, limit, wait)
|
||||
if err != nil {
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte(err.Error()))
|
||||
}
|
||||
switch status {
|
||||
case wire.StatusData:
|
||||
if debug != nil && debug.enabled {
|
||||
debug.dataRecords.Add(1)
|
||||
}
|
||||
if err := wire.WriteMaskedResponse(conn, wire.StatusData, data, req.Session, wire.ModeDownload, offset); err != nil {
|
||||
return err
|
||||
}
|
||||
offset += uint64(len(data))
|
||||
case wire.StatusWait:
|
||||
if debug != nil && debug.enabled {
|
||||
debug.waitRecords.Add(1)
|
||||
}
|
||||
return wire.WriteResponse(conn, wire.StatusWait, nil)
|
||||
case wire.StatusEOF:
|
||||
return wire.WriteResponse(conn, wire.StatusEOF, nil)
|
||||
default:
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("invalid session read status"))
|
||||
}
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte(fmt.Sprintf("EOF %d", final)))
|
||||
}
|
||||
return nil
|
||||
|
||||
if debug != nil && debug.enabled {
|
||||
debug.dataRecords.Add(1)
|
||||
debug.chunkf("DATA id=%s seq=%d offset=%d bytes=%d total=%d", parts[2], want, offset, len(data), total)
|
||||
}
|
||||
prefix := []byte(fmt.Sprintf("DATA %d %d %d ", want, offset, total))
|
||||
out := make([]byte, len(prefix)+len(data))
|
||||
copy(out, prefix)
|
||||
copy(out[len(prefix):], data)
|
||||
return protocol.WriteResponseFrame(conn, requestID, out)
|
||||
case wire.ModeClose:
|
||||
manager.remove(req.Session)
|
||||
return wire.WriteResponse(conn, wire.StatusOK, nil)
|
||||
default:
|
||||
return wire.WriteResponse(conn, wire.StatusError, []byte("unknown mode"))
|
||||
}
|
||||
|
||||
if bytes.HasPrefix(payload, []byte("CCLOSE ")) {
|
||||
parts := strings.Fields(string(payload))
|
||||
if len(parts) != 3 {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR bad CCLOSE"))
|
||||
}
|
||||
if !tokenEqual(decodeWireToken(parts[1]), token) {
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
|
||||
}
|
||||
manager.remove(parts[2])
|
||||
if debug != nil && debug.enabled {
|
||||
debug.sessionsClosed.Add(1)
|
||||
debug.activeSessions.Add(-1)
|
||||
debug.logf("SESSION CLOSE id=%s peer=%v active_sessions=%d", parts[2], conn.RemoteAddr(), manager.count())
|
||||
debug.chunkf("CCLOSE id=%s -> CLOSED", parts[2])
|
||||
}
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("CLOSED"))
|
||||
}
|
||||
|
||||
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR unknown chunk command"))
|
||||
}
|
||||
|
||||
@@ -1,12 +1,22 @@
|
||||
package main
|
||||
|
||||
import "testing"
|
||||
import (
|
||||
"encoding/binary"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDecodeWireTokenAllowsEmptyToken(t *testing.T) {
|
||||
if got := decodeWireToken("-"); got != "" {
|
||||
t.Fatalf("empty wire token decoded as %q", got)
|
||||
func TestParseOpenAllowsEmptyToken(t *testing.T) {
|
||||
host := "example.com"
|
||||
p := make([]byte, 6+len(host))
|
||||
binary.BigEndian.PutUint16(p[0:2], 0)
|
||||
binary.BigEndian.PutUint16(p[2:4], uint16(len(host)))
|
||||
binary.BigEndian.PutUint16(p[4:6], 443)
|
||||
copy(p[6:], host)
|
||||
token, gotHost, port, err := parseOpen(p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := decodeWireToken("secret"); got != "secret" {
|
||||
t.Fatalf("non-empty token changed: %q", got)
|
||||
if token != "" || gotHost != host || port != 443 {
|
||||
t.Fatalf("got token=%q host=%q port=%d", token, gotHost, port)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"crypto/subtle"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/netip"
|
||||
"os"
|
||||
@@ -16,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"dragontcp/internal/protocol"
|
||||
"dragontcp/internal/wire"
|
||||
)
|
||||
|
||||
var active int64
|
||||
@@ -159,9 +159,9 @@ func handle(
|
||||
cache *dnsCache,
|
||||
tcpBuffer int,
|
||||
slots chan struct{},
|
||||
manager *chunkManager,
|
||||
manager *streamManager,
|
||||
chunkMax int,
|
||||
chunkBuffered int,
|
||||
bufferBytes int,
|
||||
chunkPollWait time.Duration,
|
||||
debug *serverDebug,
|
||||
) {
|
||||
@@ -175,110 +175,26 @@ func handle(
|
||||
protocol.TuneTCPBuffer(conn, tcpBuffer)
|
||||
|
||||
for {
|
||||
_ = conn.SetDeadline(time.Now().Add(20 * time.Second))
|
||||
|
||||
requestID, _, payload, err := protocol.ReadRequestFrame(conn)
|
||||
_ = conn.SetDeadline(time.Now().Add(30 * time.Second))
|
||||
req, err := wire.ReadRequest(conn)
|
||||
if err != nil {
|
||||
if debug != nil && debug.enabled && err != io.EOF {
|
||||
debug.errorf("peer=%v read request: %v", conn.RemoteAddr(), err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if isChunkCommand(payload) {
|
||||
if err := processChunkCommand(
|
||||
conn,
|
||||
requestID,
|
||||
payload,
|
||||
token,
|
||||
allowPrivate,
|
||||
cache,
|
||||
tcpBuffer,
|
||||
manager,
|
||||
chunkMax,
|
||||
chunkBuffered,
|
||||
chunkPollWait,
|
||||
debug,
|
||||
); err != nil {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
parts := strings.Fields(string(payload))
|
||||
transport := "xor"
|
||||
|
||||
if len(parts) == 4 && parts[0] == "TUNNEL" {
|
||||
transport = "xor"
|
||||
} else if len(parts) == 5 && parts[0] == "TUNNEL2" {
|
||||
transport = strings.ToLower(parts[4])
|
||||
if transport != "raw" && transport != "xor" {
|
||||
_ = protocol.WriteResponseFrame(conn, requestID, []byte("ERR transport must be RAW or XOR"))
|
||||
return
|
||||
}
|
||||
} else {
|
||||
_ = protocol.WriteResponseFrame(
|
||||
conn,
|
||||
requestID,
|
||||
[]byte("ERR expected TUNNEL, TUNNEL2, or chunk command"),
|
||||
)
|
||||
if err := processWireRequest(
|
||||
conn,
|
||||
req,
|
||||
token,
|
||||
allowPrivate,
|
||||
cache,
|
||||
tcpBuffer,
|
||||
manager,
|
||||
chunkMax,
|
||||
bufferBytes,
|
||||
chunkPollWait,
|
||||
debug,
|
||||
); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
if !tokenEqual(parts[1], token) {
|
||||
_ = protocol.WriteResponseFrame(
|
||||
conn,
|
||||
requestID,
|
||||
[]byte("ERR authentication failed"),
|
||||
)
|
||||
return
|
||||
}
|
||||
|
||||
port, err := strconv.Atoi(parts[3])
|
||||
if err != nil || port < 1 || port > 65535 {
|
||||
_ = protocol.WriteResponseFrame(
|
||||
conn,
|
||||
requestID,
|
||||
[]byte("ERR invalid port"),
|
||||
)
|
||||
return
|
||||
}
|
||||
|
||||
if debug != nil && debug.enabled {
|
||||
debug.logf("TUNNEL peer=%v target=%s:%d transport=%s", conn.RemoteAddr(), parts[2], port, transport)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
target, err := dialTarget(ctx, parts[2], port, allowPrivate, cache, tcpBuffer)
|
||||
cancel()
|
||||
|
||||
if err != nil {
|
||||
if debug != nil && debug.enabled {
|
||||
debug.errorf("TUNNEL target=%s:%d connect failed: %v", parts[2], port, err)
|
||||
}
|
||||
_ = protocol.WriteResponseFrame(
|
||||
conn,
|
||||
requestID,
|
||||
[]byte("ERR "+err.Error()),
|
||||
)
|
||||
return
|
||||
}
|
||||
defer target.Close()
|
||||
|
||||
if err := protocol.WriteResponseFrame(conn, requestID, []byte("CONNECTED")); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
_ = conn.SetDeadline(time.Time{})
|
||||
if transport == "raw" {
|
||||
protocol.RelayRaw(conn, target)
|
||||
} else {
|
||||
protocol.RelayXOR(conn, target)
|
||||
}
|
||||
if debug != nil && debug.enabled {
|
||||
debug.logf("TUNNEL closed peer=%v target=%s:%d transport=%s", conn.RemoteAddr(), parts[2], port, transport)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@@ -293,7 +209,7 @@ func main() {
|
||||
dnsCacheSize = flag.Int("dns-cache-size", 4096, "maximum cached DNS hostnames")
|
||||
tcpBuffer = flag.Int("tcp-buffer", 0, "optional TCP read/write buffer bytes; 0 keeps OS autotuning")
|
||||
chunkMax = flag.Int("chunk-max", 1048576, "maximum adaptive chunk payload bytes (32 bytes to 1 MiB)")
|
||||
chunkBuffered = flag.Int("chunk-buffered", 256, "maximum buffered destination chunks per session")
|
||||
chunkBuffered = flag.Int("chunk-buffered", 256, "compatibility buffer units; 256 = about 16 MiB per active session")
|
||||
chunkPollWait = flag.Duration("chunk-poll-wait", 200*time.Millisecond, "server long-poll wait for chunk data")
|
||||
sessionTimeout = flag.Duration("chunk-session-timeout", 2*time.Minute, "idle chunk session timeout")
|
||||
debugEnabled = flag.Bool("debug", false, "log session/connect/errors and periodic statistics")
|
||||
@@ -325,8 +241,15 @@ func main() {
|
||||
slots := make(chan struct{}, *maxConnections)
|
||||
cache := newDNSCache(*dnsCacheTTL, *dnsCacheSize)
|
||||
debug := newServerDebug(*debugEnabled, *debugChunks, *debugStats)
|
||||
manager := newChunkManager(*sessionTimeout, debug)
|
||||
fmt.Printf("adaptive_chunk_max=%d buffered_chunks=%d poll_wait=%s\n", *chunkMax, *chunkBuffered, chunkPollWait.String())
|
||||
bufferBytes := *chunkBuffered * 65536
|
||||
if bufferBytes < 1024*1024 {
|
||||
bufferBytes = 1024 * 1024
|
||||
}
|
||||
if bufferBytes > 64*1024*1024 {
|
||||
bufferBytes = 64 * 1024 * 1024
|
||||
}
|
||||
manager := newStreamManager(*sessionTimeout, debug)
|
||||
fmt.Printf("binary_transport=true chunk_max=%d buffer_bytes=%d poll_wait=%s\n", *chunkMax, bufferBytes, chunkPollWait.String())
|
||||
if debug.enabled {
|
||||
fmt.Printf("debug=true debug_chunks=%t stats_interval=%s\n", debug.chunks, debug.statsEvery)
|
||||
}
|
||||
@@ -353,7 +276,7 @@ func main() {
|
||||
slots,
|
||||
manager,
|
||||
*chunkMax,
|
||||
*chunkBuffered,
|
||||
bufferBytes,
|
||||
*chunkPollWait,
|
||||
debug,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,177 @@
|
||||
package wire
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
)
|
||||
|
||||
const (
|
||||
RequestHeaderSize = 29
|
||||
ResponseHeaderSize = 5
|
||||
MaxPayload = 2 * 1024 * 1024
|
||||
|
||||
ModeProbe byte = 0
|
||||
ModeOpen byte = 1
|
||||
ModeUpload byte = 2
|
||||
ModeDownload byte = 3
|
||||
ModeClose byte = 4
|
||||
|
||||
StatusOK byte = 0
|
||||
StatusError byte = 1
|
||||
StatusData byte = 2
|
||||
StatusWait byte = 3
|
||||
StatusEOF byte = 4
|
||||
|
||||
ProbeUpload byte = 1
|
||||
ProbeDownload byte = 2
|
||||
ProbeKeepalive byte = 3
|
||||
ProbeBatch byte = 4
|
||||
)
|
||||
|
||||
var ProbeMagic = [4]byte{'D', 'T', 'P', '2'}
|
||||
|
||||
type SessionID [16]byte
|
||||
|
||||
type Request struct {
|
||||
Mode byte
|
||||
Session SessionID
|
||||
Seq uint64
|
||||
Payload []byte
|
||||
}
|
||||
|
||||
func MaskInPlace(data []byte, sid SessionID, mode byte, seq uint64, response bool) {
|
||||
if len(data) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
var seed [30]byte
|
||||
copy(seed[:16], sid[:])
|
||||
seed[16] = mode
|
||||
binary.BigEndian.PutUint64(seed[17:25], seq)
|
||||
if response {
|
||||
seed[25] = 1
|
||||
}
|
||||
|
||||
var counter uint32
|
||||
for off := 0; off < len(data); {
|
||||
binary.BigEndian.PutUint32(seed[26:30], counter)
|
||||
block := sha256.Sum256(seed[:])
|
||||
n := len(data) - off
|
||||
if n > len(block) {
|
||||
n = len(block)
|
||||
}
|
||||
for i := 0; i < n; i++ {
|
||||
data[off+i] ^= block[i]
|
||||
}
|
||||
off += n
|
||||
counter++
|
||||
}
|
||||
}
|
||||
|
||||
func WriteRequest(w io.Writer, mode byte, sid SessionID, seq uint64, plaintext []byte) error {
|
||||
if len(plaintext) > MaxPayload {
|
||||
return fmt.Errorf("request payload too large: %d", len(plaintext))
|
||||
}
|
||||
|
||||
packet := make([]byte, RequestHeaderSize+len(plaintext))
|
||||
packet[0] = mode
|
||||
copy(packet[1:17], sid[:])
|
||||
binary.BigEndian.PutUint64(packet[17:25], seq)
|
||||
binary.BigEndian.PutUint32(packet[25:29], uint32(len(plaintext)))
|
||||
copy(packet[29:], plaintext)
|
||||
MaskInPlace(packet[29:], sid, mode, seq, false)
|
||||
return writeAll(w, packet)
|
||||
}
|
||||
|
||||
func ReadRequest(r io.Reader) (Request, error) {
|
||||
var req Request
|
||||
var header [RequestHeaderSize]byte
|
||||
if _, err := io.ReadFull(r, header[:]); err != nil {
|
||||
return req, err
|
||||
}
|
||||
|
||||
req.Mode = header[0]
|
||||
copy(req.Session[:], header[1:17])
|
||||
req.Seq = binary.BigEndian.Uint64(header[17:25])
|
||||
n := binary.BigEndian.Uint32(header[25:29])
|
||||
if n > MaxPayload {
|
||||
return req, errors.New("request payload too large")
|
||||
}
|
||||
|
||||
if n > 0 {
|
||||
req.Payload = make([]byte, int(n))
|
||||
if _, err := io.ReadFull(r, req.Payload); err != nil {
|
||||
return req, err
|
||||
}
|
||||
MaskInPlace(req.Payload, req.Session, req.Mode, req.Seq, false)
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func WriteResponse(w io.Writer, status byte, body []byte) error {
|
||||
if len(body) > MaxPayload {
|
||||
return fmt.Errorf("response body too large: %d", len(body))
|
||||
}
|
||||
packet := make([]byte, ResponseHeaderSize+len(body))
|
||||
packet[0] = status
|
||||
binary.BigEndian.PutUint32(packet[1:5], uint32(len(body)))
|
||||
copy(packet[5:], body)
|
||||
return writeAll(w, packet)
|
||||
}
|
||||
|
||||
func WriteMaskedResponse(w io.Writer, status byte, body []byte, sid SessionID, mode byte, seq uint64) error {
|
||||
if len(body) > MaxPayload {
|
||||
return fmt.Errorf("response body too large: %d", len(body))
|
||||
}
|
||||
packet := make([]byte, ResponseHeaderSize+len(body))
|
||||
packet[0] = status
|
||||
binary.BigEndian.PutUint32(packet[1:5], uint32(len(body)))
|
||||
copy(packet[5:], body)
|
||||
MaskInPlace(packet[5:], sid, mode, seq, true)
|
||||
return writeAll(w, packet)
|
||||
}
|
||||
|
||||
func ReadResponse(r io.Reader) (byte, []byte, error) {
|
||||
var header [ResponseHeaderSize]byte
|
||||
if _, err := io.ReadFull(r, header[:]); err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
n := binary.BigEndian.Uint32(header[1:5])
|
||||
if n > MaxPayload {
|
||||
return 0, nil, errors.New("response body too large")
|
||||
}
|
||||
var body []byte
|
||||
if n > 0 {
|
||||
body = make([]byte, int(n))
|
||||
if _, err := io.ReadFull(r, body); err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
}
|
||||
return header[0], body, nil
|
||||
}
|
||||
|
||||
func DecodeMaskedResponse(status byte, body []byte, sid SessionID, mode byte, seq uint64) []byte {
|
||||
if len(body) == 0 || status == StatusError {
|
||||
return body
|
||||
}
|
||||
out := append([]byte(nil), body...)
|
||||
MaskInPlace(out, sid, mode, seq, true)
|
||||
return out
|
||||
}
|
||||
|
||||
func writeAll(w io.Writer, b []byte) error {
|
||||
for len(b) > 0 {
|
||||
n, err := w.Write(b)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n <= 0 {
|
||||
return io.ErrShortWrite
|
||||
}
|
||||
b = b[n:]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package wire
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMaskChangesWithSequenceAndRoundTrips(t *testing.T) {
|
||||
var sid SessionID
|
||||
for i := range sid { sid[i] = byte(i+1) }
|
||||
plain := bytes.Repeat([]byte("DragonTCP"), 100)
|
||||
a := append([]byte(nil), plain...)
|
||||
b := append([]byte(nil), plain...)
|
||||
MaskInPlace(a, sid, ModeUpload, 1, false)
|
||||
MaskInPlace(b, sid, ModeUpload, 2, false)
|
||||
if bytes.Equal(a, b) { t.Fatal("different sequences produced identical wire bytes") }
|
||||
MaskInPlace(a, sid, ModeUpload, 1, false)
|
||||
if !bytes.Equal(a, plain) { t.Fatal("mask did not round-trip") }
|
||||
}
|
||||
|
||||
func BenchmarkMask1MiB(b *testing.B) {
|
||||
var sid SessionID
|
||||
data := make([]byte, 1024*1024)
|
||||
b.SetBytes(int64(len(data)))
|
||||
b.ResetTimer()
|
||||
for i:=0;i<b.N;i++ {
|
||||
MaskInPlace(data,sid,ModeUpload,uint64(i),false)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user