DragonTCP
This commit is contained in:
@@ -0,0 +1,168 @@
|
|||||||
|
DragonTCP Proxy v6 - Server Debug Build
|
||||||
|
=================================
|
||||||
|
|
||||||
|
This build adds server-side diagnostic logging for the adaptive/chunk transport.
|
||||||
|
The wire protocol and v6 client remain compatible.
|
||||||
|
|
||||||
|
New server flags
|
||||||
|
----------------
|
||||||
|
|
||||||
|
--debug
|
||||||
|
Session/connect/error logging plus periodic aggregate statistics.
|
||||||
|
|
||||||
|
--debug-chunks
|
||||||
|
Logs every COPEN, CPUSH, ACK, CPULL, DATA, WAIT, EOF, and CCLOSE event.
|
||||||
|
This is extremely verbose with 32-byte chunks and can reduce throughput.
|
||||||
|
Enabling --debug-chunks also enables normal debug logging.
|
||||||
|
|
||||||
|
--debug-stats-interval DURATION
|
||||||
|
Aggregate statistics frequency. Default: 5s.
|
||||||
|
Set to 0 to disable periodic statistics.
|
||||||
|
|
||||||
|
Recommended diagnostic command for fixed 32-byte chunks
|
||||||
|
--------------------------------------------------------
|
||||||
|
|
||||||
|
sudo ./dragontcp-server-linux-amd64 --token 'YOUR_SECRET' --chunk-max 32 --chunk-buffered 2048 --chunk-poll-wait 50ms --chunk-session-timeout 5m --max-connections 20000 --tcp-buffer 0 --debug --debug-chunks --debug-stats-interval 5s
|
||||||
|
|
||||||
|
Normal production command with useful low-overhead debug
|
||||||
|
---------------------------------------------------------
|
||||||
|
|
||||||
|
sudo ./dragontcp-server-linux-amd64 --token 'YOUR_SECRET' --chunk-max 32 --chunk-buffered 2048 --chunk-poll-wait 50ms --chunk-session-timeout 5m --max-connections 20000 --tcp-buffer 0 --debug --debug-stats-interval 10s
|
||||||
|
|
||||||
|
Disable all debug logging
|
||||||
|
-------------------------
|
||||||
|
|
||||||
|
Simply omit --debug and --debug-chunks.
|
||||||
|
|
||||||
|
Example debug output
|
||||||
|
--------------------
|
||||||
|
|
||||||
|
[DEBUG] SESSION OPEN id=... target=example.com:443 max_chunk=32 active_sessions=1
|
||||||
|
[CHUNK] CPUSH id=... seq=0 bytes=32 -> ACK accepted=32
|
||||||
|
[CHUNK] CPULL id=... ack=-1 want=0 offset=0 limit=32
|
||||||
|
[CHUNK] DATA id=... seq=0 offset=0 bytes=32 total=32
|
||||||
|
[CHUNK] CPULL id=... want=8 -> WAIT
|
||||||
|
[DEBUG] SESSION CLOSE id=... active_sessions=0
|
||||||
|
[DEBUG] STATS uptime=10s active_connections=8 active_sessions=2 sessions_opened=5 sessions_closed=3 bytes_up=... bytes_down=... push_records=... pull_requests=... data_records=... waits=... errors=0
|
||||||
|
|
||||||
|
Counters
|
||||||
|
--------
|
||||||
|
|
||||||
|
active_connections - currently open DragonTCP TCP connections
|
||||||
|
active_sessions - currently open chunk proxy sessions
|
||||||
|
sessions_opened - total chunk sessions opened
|
||||||
|
sessions_closed - total chunk sessions closed
|
||||||
|
bytes_up - bytes accepted from client and written toward target
|
||||||
|
bytes_down - bytes read from target into chunk buffering
|
||||||
|
push_records - accepted upload CPUSH records
|
||||||
|
pull_requests - CPULL requests received
|
||||||
|
data_records - DATA responses generated
|
||||||
|
waits - WAIT responses because downstream data was not ready yet
|
||||||
|
errors - debug-counted server/protocol errors
|
||||||
|
|
||||||
|
Important performance note
|
||||||
|
--------------------------
|
||||||
|
|
||||||
|
At 32 bytes, --debug-chunks can generate thousands or millions of log lines for
|
||||||
|
large transfers. Use it while diagnosing a failure, then switch to --debug only
|
||||||
|
for normal use.
|
||||||
|
|
||||||
|
Large-chunk update
|
||||||
|
==================
|
||||||
|
|
||||||
|
This is the v6 debug/adaptive-chunk branch with FIXED poller concurrency.
|
||||||
|
It intentionally does NOT include the later adaptive-poller controller.
|
||||||
|
|
||||||
|
Chunk limits
|
||||||
|
------------
|
||||||
|
|
||||||
|
Previous hard limit:
|
||||||
|
|
||||||
|
8192 bytes
|
||||||
|
|
||||||
|
New hard limit:
|
||||||
|
|
||||||
|
1048576 bytes (1 MiB)
|
||||||
|
|
||||||
|
The framed protocol ceiling was increased to 2 MiB so a 1 MiB CPUSH/DATA
|
||||||
|
record plus protocol metadata fits safely.
|
||||||
|
|
||||||
|
New defaults:
|
||||||
|
|
||||||
|
client --chunk-max 65536
|
||||||
|
server --chunk-max 65536
|
||||||
|
|
||||||
|
The adaptive client still begins at:
|
||||||
|
|
||||||
|
--chunk-start 256
|
||||||
|
|
||||||
|
and can grow toward the configured maximum after successful records.
|
||||||
|
|
||||||
|
Use up to 1 MiB adaptive chunks
|
||||||
|
--------------------------------
|
||||||
|
|
||||||
|
Server:
|
||||||
|
|
||||||
|
sudo ./dragontcp-server-linux-amd64 --token 'YOUR_SECRET' --chunk-max 1048576 --debug --debug-stats-interval 10s
|
||||||
|
|
||||||
|
Android ARM64 client with fixed poller count of 8:
|
||||||
|
|
||||||
|
./dragontcp-client-android-arm64 --server-host YOUR_SERVER_IP --token 'YOUR_SECRET' --chunk-start 256 --chunk-min 32 --chunk-max 1048576 --chunk-pollers 8 --chunk-timeout 2s --chunk-adapt-log
|
||||||
|
|
||||||
|
The number of pollers stays exactly at the value passed with --chunk-pollers.
|
||||||
|
Only the chunk size adapts.
|
||||||
|
|
||||||
|
Examples of useful ceilings
|
||||||
|
---------------------------
|
||||||
|
|
||||||
|
--chunk-max 16384 # 16 KiB
|
||||||
|
--chunk-max 32768 # 32 KiB
|
||||||
|
--chunk-max 65536 # 64 KiB (new default maximum)
|
||||||
|
--chunk-max 131072 # 128 KiB
|
||||||
|
--chunk-max 262144 # 256 KiB
|
||||||
|
--chunk-max 524288 # 512 KiB
|
||||||
|
--chunk-max 1048576 # 1 MiB hard maximum
|
||||||
|
|
||||||
|
Fixed-size mode also supports the same range:
|
||||||
|
|
||||||
|
--chunk-size 262144
|
||||||
|
|
||||||
|
Memory note
|
||||||
|
-----------
|
||||||
|
|
||||||
|
Larger server chunk maxima require larger per-session target-read buffers and
|
||||||
|
can increase buffered memory substantially when many sessions are active.
|
||||||
|
For thousands of simultaneous users, do not automatically use 1 MiB unless
|
||||||
|
measurements show that it is useful. Values such as 16-64 KiB are a more
|
||||||
|
reasonable starting point, while the adaptive client can still be configured
|
||||||
|
to probe higher when your network supports it.
|
||||||
|
|
||||||
|
Validation
|
||||||
|
----------
|
||||||
|
|
||||||
|
The updated source and binaries were rebuilt from this v6 debug branch.
|
||||||
|
Validation included:
|
||||||
|
|
||||||
|
* Go builds for Linux amd64, Linux ARM64, Linux ARMv7 client, and Android ARM64.
|
||||||
|
* A protocol round-trip test with a full 1 MiB request and response frame.
|
||||||
|
* An 8 MiB HTTP download through the proxy using fixed 262144-byte (256 KiB)
|
||||||
|
chunk configuration; the downloaded SHA-256 matched the source exactly.
|
||||||
|
|
||||||
|
DragonTCP branding update
|
||||||
|
=========================
|
||||||
|
|
||||||
|
This package was renamed from HOX to DragonTCP.
|
||||||
|
|
||||||
|
Binary names are now:
|
||||||
|
|
||||||
|
dragontcp-server-linux-amd64
|
||||||
|
dragontcp-server-arm64
|
||||||
|
dragontcp-client-linux-amd64
|
||||||
|
dragontcp-client-android-arm64
|
||||||
|
dragontcp-client-arm64
|
||||||
|
dragontcp-client-armv7
|
||||||
|
|
||||||
|
The Go module and command directories were also renamed to DragonTCP.
|
||||||
|
The existing UP/OK wire framing and chunk protocol were intentionally kept
|
||||||
|
unchanged, so this branding change does not break compatibility with the
|
||||||
|
previous protocol implementation.
|
||||||
@@ -0,0 +1,6 @@
|
|||||||
|
1f1058a03099fc04dd00c727656f26a88ebcf43f7ef924d11ecb5675aab96895 bin/dragontcp-client-android-arm64
|
||||||
|
232c95474d90c009f18db9c2d870c910a8248f97185b715e54166eab4baeaebf bin/dragontcp-client-arm64
|
||||||
|
759fc70875d3501c9954e9db99a2b6f51d1ff9f4f39ff3d9b21062b971890e3e bin/dragontcp-client-armv7
|
||||||
|
3b18a4b7dc79819b22ef550b1761c8f36b51793bec35ff077d038c1a7d00ee79 bin/dragontcp-client-linux-amd64
|
||||||
|
4a00acbc2fdeaa1acd34c528b84813c3688a03134d591bed34f94254a46d8d06 bin/dragontcp-server-arm64
|
||||||
|
0bc46c805ceabad0dd149f8b068f98e75eba8cd08b86c9298dc7f34094e9c666 bin/dragontcp-server-linux-amd64
|
||||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,748 @@
|
|||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
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{})
|
||||||
|
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",
|
||||||
|
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",
|
||||||
|
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 ", 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", 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,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", "change-this-token", "shared token")
|
||||||
|
maxConnections = flag.Int("max-connections", 20000, "max simultaneous proxy connections")
|
||||||
|
transport = flag.String("transport", "chunk", "transport: chunk (adaptive framed records), xor, or raw")
|
||||||
|
tcpBuffer = flag.Int("tcp-buffer", 0, "optional TCP read/write buffer bytes; 0 keeps OS autotuning")
|
||||||
|
chunkStart = flag.Int("chunk-start", 256, "initial adaptive chunk payload bytes")
|
||||||
|
chunkMin = flag.Int("chunk-min", 32, "minimum adaptive chunk payload bytes")
|
||||||
|
chunkMax = flag.Int("chunk-max", 65536, "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", 64, "successful data records required before increasing chunk size")
|
||||||
|
chunkAdaptLog = flag.Bool("chunk-adapt-log", false, "print adaptive chunk size changes")
|
||||||
|
chunkSizeLegacy = flag.Int("chunk-size", 0, "legacy fixed chunk size; nonzero disables adaptation")
|
||||||
|
chunkPollers = flag.Int("chunk-pollers", 16, "parallel downstream chunk pollers (1-128)")
|
||||||
|
chunkReconnect = flag.Int("chunk-reconnect-every", 32, "reconnect each transaction lane after N requests; 0 keeps it open")
|
||||||
|
chunkPollDelay = flag.Duration("chunk-poll-delay", 2*time.Millisecond, "delay after an empty chunk poll")
|
||||||
|
chunkTimeout = flag.Duration("chunk-timeout", 5*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 != "raw" && *transport != "xor" && *transport != "chunk" {
|
||||||
|
fmt.Fprintln(os.Stderr, "--transport must be chunk, xor, or raw")
|
||||||
|
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()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,498 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"net"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"dragontcp/internal/protocol"
|
||||||
|
)
|
||||||
|
|
||||||
|
type chunkSession struct {
|
||||||
|
id string
|
||||||
|
target net.Conn
|
||||||
|
maxChunk int
|
||||||
|
maxChunks int
|
||||||
|
|
||||||
|
mu sync.Mutex
|
||||||
|
notify chan struct{}
|
||||||
|
chunks map[uint64][]byte
|
||||||
|
nextDown 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,
|
||||||
|
}
|
||||||
|
go s.readTarget()
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *chunkSession) 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)
|
||||||
|
|
||||||
|
for {
|
||||||
|
n, err := s.target.Read(buf)
|
||||||
|
if n > 0 {
|
||||||
|
data := append([]byte(nil), buf[:n]...)
|
||||||
|
if s.debug != nil && s.debug.enabled {
|
||||||
|
s.debug.bytesDown.Add(uint64(n))
|
||||||
|
}
|
||||||
|
|
||||||
|
for {
|
||||||
|
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
|
||||||
|
}
|
||||||
|
ch := s.notify
|
||||||
|
s.mu.Unlock()
|
||||||
|
<-ch
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
s.touchLocked()
|
||||||
|
s.signalLocked()
|
||||||
|
}
|
||||||
|
s.mu.Unlock()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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))
|
||||||
|
}
|
||||||
|
|
||||||
|
if s.haveLastUp && seq == s.lastUpSeq {
|
||||||
|
s.touch()
|
||||||
|
return s.lastUpLen, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if seq < s.expectedUp {
|
||||||
|
return 0, fmt.Errorf("upload sequence %d is too old", seq)
|
||||||
|
}
|
||||||
|
if seq > s.expectedUp {
|
||||||
|
return 0, fmt.Errorf("unexpected upload sequence %d, expected %d", seq, s.expectedUp)
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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")
|
||||||
|
}
|
||||||
|
|
||||||
|
timer := time.NewTimer(wait)
|
||||||
|
defer timer.Stop()
|
||||||
|
|
||||||
|
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 removed {
|
||||||
|
s.signalLocked()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if chunk, ok := s.chunks[want]; ok {
|
||||||
|
if offset >= len(chunk) {
|
||||||
|
s.mu.Unlock()
|
||||||
|
return nil, len(chunk), false, 0, false, fmt.Errorf("pull offset %d beyond chunk size %d", offset, len(chunk))
|
||||||
|
}
|
||||||
|
end := offset + limit
|
||||||
|
if end > len(chunk) {
|
||||||
|
end = len(chunk)
|
||||||
|
}
|
||||||
|
out := append([]byte(nil), chunk[offset:end]...)
|
||||||
|
total = len(chunk)
|
||||||
|
s.mu.Unlock()
|
||||||
|
return out, total, false, 0, false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if s.eof && want >= s.nextDown {
|
||||||
|
final = s.nextDown
|
||||||
|
s.mu.Unlock()
|
||||||
|
return nil, 0, true, final, false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if s.closed {
|
||||||
|
final = s.nextDown
|
||||||
|
s.mu.Unlock()
|
||||||
|
return nil, 0, true, final, false, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
ch := s.notify
|
||||||
|
s.mu.Unlock()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-ch:
|
||||||
|
continue
|
||||||
|
case <-timer.C:
|
||||||
|
return nil, 0, false, 0, true, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *chunkSession) close() {
|
||||||
|
s.mu.Lock()
|
||||||
|
if s.closed {
|
||||||
|
s.mu.Unlock()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
s.closed = true
|
||||||
|
s.signalLocked()
|
||||||
|
s.mu.Unlock()
|
||||||
|
_ = s.target.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
type chunkManager struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
sessions map[string]*chunkSession
|
||||||
|
timeout time.Duration
|
||||||
|
debug *serverDebug
|
||||||
|
}
|
||||||
|
|
||||||
|
func newChunkManager(timeout time.Duration, debug *serverDebug) *chunkManager {
|
||||||
|
m := &chunkManager{
|
||||||
|
sessions: make(map[string]*chunkSession),
|
||||||
|
timeout: timeout,
|
||||||
|
debug: debug,
|
||||||
|
}
|
||||||
|
go m.cleanupLoop()
|
||||||
|
return m
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *chunkManager) get(id string) *chunkSession {
|
||||||
|
m.mu.RLock()
|
||||||
|
s := m.sessions[id]
|
||||||
|
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 {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
if _, exists := m.sessions[id]; exists {
|
||||||
|
return fmt.Errorf("session already exists")
|
||||||
|
}
|
||||||
|
m.sessions[id] = s
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *chunkManager) remove(id string) {
|
||||||
|
m.mu.Lock()
|
||||||
|
s := m.sessions[id]
|
||||||
|
delete(m.sessions, id)
|
||||||
|
m.mu.Unlock()
|
||||||
|
if s != nil {
|
||||||
|
s.close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *chunkManager) cleanupLoop() {
|
||||||
|
ticker := time.NewTicker(30 * time.Second)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
for range ticker.C {
|
||||||
|
cutoff := time.Now().Add(-m.timeout)
|
||||||
|
var stale []string
|
||||||
|
|
||||||
|
m.mu.RLock()
|
||||||
|
for id, s := range m.sessions {
|
||||||
|
s.mu.Lock()
|
||||||
|
last := s.lastSeen
|
||||||
|
closed := s.closed
|
||||||
|
s.mu.Unlock()
|
||||||
|
if closed || last.Before(cutoff) {
|
||||||
|
stale = append(stale, id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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 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"))
|
||||||
|
}
|
||||||
|
if !tokenEqual(parts[1], token) {
|
||||||
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR 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"))
|
||||||
|
}
|
||||||
|
|
||||||
|
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()))
|
||||||
|
}
|
||||||
|
|
||||||
|
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 {
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
return protocol.WriteResponseFrame(conn, requestID, []byte(fmt.Sprintf("OPENED %d", maxChunk)))
|
||||||
|
}
|
||||||
|
|
||||||
|
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(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)
|
||||||
|
if s == nil {
|
||||||
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR 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 debug != nil {
|
||||||
|
debug.chunkf("CPUSH id=%s seq=%d bytes=%d -> ACK accepted=%d", sid, seq, len(parts[4]), accepted)
|
||||||
|
}
|
||||||
|
return protocol.WriteResponseFrame(conn, requestID, []byte(fmt.Sprintf("ACK %d %d", seq, accepted)))
|
||||||
|
}
|
||||||
|
|
||||||
|
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(parts[1], token) {
|
||||||
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
|
||||||
|
}
|
||||||
|
s := manager.get(parts[2])
|
||||||
|
if s == nil {
|
||||||
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR unknown session"))
|
||||||
|
}
|
||||||
|
ack, err := strconv.ParseInt(parts[3], 10, 64)
|
||||||
|
if err != nil || ack < -1 {
|
||||||
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid ack"))
|
||||||
|
}
|
||||||
|
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"))
|
||||||
|
}
|
||||||
|
if limit > maxChunk {
|
||||||
|
limit = maxChunk
|
||||||
|
}
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
return protocol.WriteResponseFrame(conn, requestID, []byte(fmt.Sprintf("EOF %d", final)))
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
|
||||||
|
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(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"))
|
||||||
|
}
|
||||||
@@ -0,0 +1,83 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"sync/atomic"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
type serverDebug struct {
|
||||||
|
enabled bool
|
||||||
|
chunks bool
|
||||||
|
statsEvery time.Duration
|
||||||
|
started time.Time
|
||||||
|
|
||||||
|
sessionsOpened atomic.Uint64
|
||||||
|
sessionsClosed atomic.Uint64
|
||||||
|
activeSessions atomic.Int64
|
||||||
|
bytesUp atomic.Uint64
|
||||||
|
bytesDown atomic.Uint64
|
||||||
|
pushRecords atomic.Uint64
|
||||||
|
pullRequests atomic.Uint64
|
||||||
|
dataRecords atomic.Uint64
|
||||||
|
waitRecords atomic.Uint64
|
||||||
|
errors atomic.Uint64
|
||||||
|
}
|
||||||
|
|
||||||
|
func newServerDebug(enabled, chunks bool, statsEvery time.Duration) *serverDebug {
|
||||||
|
d := &serverDebug{
|
||||||
|
enabled: enabled || chunks,
|
||||||
|
chunks: chunks,
|
||||||
|
statsEvery: statsEvery,
|
||||||
|
started: time.Now(),
|
||||||
|
}
|
||||||
|
if d.enabled && d.statsEvery > 0 {
|
||||||
|
go d.statsLoop()
|
||||||
|
}
|
||||||
|
return d
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *serverDebug) logf(format string, args ...any) {
|
||||||
|
if d == nil || !d.enabled {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
fmt.Fprintf(os.Stderr, "%s [DEBUG] "+format+"\n", append([]any{time.Now().Format("2006-01-02 15:04:05.000")}, args...)...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *serverDebug) chunkf(format string, args ...any) {
|
||||||
|
if d == nil || !d.chunks {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
fmt.Fprintf(os.Stderr, "%s [CHUNK] "+format+"\n", append([]any{time.Now().Format("2006-01-02 15:04:05.000")}, args...)...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *serverDebug) errorf(format string, args ...any) {
|
||||||
|
if d == nil || !d.enabled {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
d.errors.Add(1)
|
||||||
|
fmt.Fprintf(os.Stderr, "%s [ERROR] "+format+"\n", append([]any{time.Now().Format("2006-01-02 15:04:05.000")}, args...)...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (d *serverDebug) statsLoop() {
|
||||||
|
ticker := time.NewTicker(d.statsEvery)
|
||||||
|
defer ticker.Stop()
|
||||||
|
for range ticker.C {
|
||||||
|
d.logf(
|
||||||
|
"STATS uptime=%s active_connections=%d active_sessions=%d sessions_opened=%d sessions_closed=%d bytes_up=%d bytes_down=%d push_records=%d pull_requests=%d data_records=%d waits=%d errors=%d",
|
||||||
|
time.Since(d.started).Round(time.Second),
|
||||||
|
atomic.LoadInt64(&active),
|
||||||
|
d.activeSessions.Load(),
|
||||||
|
d.sessionsOpened.Load(),
|
||||||
|
d.sessionsClosed.Load(),
|
||||||
|
d.bytesUp.Load(),
|
||||||
|
d.bytesDown.Load(),
|
||||||
|
d.pushRecords.Load(),
|
||||||
|
d.pullRequests.Load(),
|
||||||
|
d.dataRecords.Load(),
|
||||||
|
d.waitRecords.Load(),
|
||||||
|
d.errors.Load(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,367 @@
|
|||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"crypto/subtle"
|
||||||
|
"flag"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
"net/netip"
|
||||||
|
"os"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"dragontcp/internal/protocol"
|
||||||
|
)
|
||||||
|
|
||||||
|
var active int64
|
||||||
|
|
||||||
|
type dnsEntry struct {
|
||||||
|
ips []netip.Addr
|
||||||
|
expires time.Time
|
||||||
|
}
|
||||||
|
|
||||||
|
type dnsCache struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
entries map[string]dnsEntry
|
||||||
|
ttl time.Duration
|
||||||
|
max int
|
||||||
|
}
|
||||||
|
|
||||||
|
func newDNSCache(ttl time.Duration, max int) *dnsCache {
|
||||||
|
return &dnsCache{
|
||||||
|
entries: make(map[string]dnsEntry),
|
||||||
|
ttl: ttl,
|
||||||
|
max: max,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *dnsCache) resolve(ctx context.Context, host string) ([]netip.Addr, error) {
|
||||||
|
if ip, err := netip.ParseAddr(host); err == nil {
|
||||||
|
return []netip.Addr{ip}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
c.mu.RLock()
|
||||||
|
entry, ok := c.entries[host]
|
||||||
|
c.mu.RUnlock()
|
||||||
|
if ok && now.Before(entry.expires) {
|
||||||
|
return entry.ips, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
ips, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
c.mu.Lock()
|
||||||
|
if len(c.entries) >= c.max {
|
||||||
|
// Simple bounded reset keeps the hot cache cheap and prevents growth.
|
||||||
|
c.entries = make(map[string]dnsEntry, c.max)
|
||||||
|
}
|
||||||
|
c.entries[host] = dnsEntry{ips: ips, expires: now.Add(c.ttl)}
|
||||||
|
c.mu.Unlock()
|
||||||
|
|
||||||
|
return ips, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func tokenEqual(a, b string) bool {
|
||||||
|
if len(a) != len(b) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
|
||||||
|
}
|
||||||
|
|
||||||
|
var blockedSpecial = []netip.Prefix{
|
||||||
|
netip.MustParsePrefix("0.0.0.0/8"),
|
||||||
|
netip.MustParsePrefix("100.64.0.0/10"),
|
||||||
|
netip.MustParsePrefix("192.0.0.0/24"),
|
||||||
|
netip.MustParsePrefix("192.0.2.0/24"),
|
||||||
|
netip.MustParsePrefix("198.18.0.0/15"),
|
||||||
|
netip.MustParsePrefix("198.51.100.0/24"),
|
||||||
|
netip.MustParsePrefix("203.0.113.0/24"),
|
||||||
|
netip.MustParsePrefix("240.0.0.0/4"),
|
||||||
|
netip.MustParsePrefix("2001:db8::/32"),
|
||||||
|
}
|
||||||
|
|
||||||
|
func addressAllowed(addr netip.Addr, allowPrivate bool) bool {
|
||||||
|
if addr.IsUnspecified() || addr.IsMulticast() {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
if allowPrivate {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
if !addr.IsGlobalUnicast() ||
|
||||||
|
addr.IsPrivate() ||
|
||||||
|
addr.IsLoopback() ||
|
||||||
|
addr.IsLinkLocalUnicast() {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, prefix := range blockedSpecial {
|
||||||
|
if prefix.Contains(addr) {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
func dialTarget(ctx context.Context, host string, port int, allowPrivate bool, cache *dnsCache, tcpBuffer int) (net.Conn, error) {
|
||||||
|
ips, err := cache.resolve(ctx, host)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
var lastErr error
|
||||||
|
var blocked []string
|
||||||
|
|
||||||
|
d := net.Dialer{
|
||||||
|
Timeout: 10 * time.Second,
|
||||||
|
KeepAlive: 30 * time.Second,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, ip := range ips {
|
||||||
|
if !addressAllowed(ip, allowPrivate) {
|
||||||
|
blocked = append(blocked, ip.String())
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
addr := net.JoinHostPort(ip.String(), strconv.Itoa(port))
|
||||||
|
conn, err := d.DialContext(ctx, "tcp", addr)
|
||||||
|
if err == nil {
|
||||||
|
protocol.TuneTCP(conn)
|
||||||
|
protocol.TuneTCPBuffer(conn, tcpBuffer)
|
||||||
|
return conn, nil
|
||||||
|
}
|
||||||
|
lastErr = err
|
||||||
|
}
|
||||||
|
|
||||||
|
if lastErr != nil {
|
||||||
|
return nil, lastErr
|
||||||
|
}
|
||||||
|
if len(blocked) > 0 {
|
||||||
|
return nil, fmt.Errorf("target resolves only to blocked addresses: %s", strings.Join(blocked, ","))
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("no usable target address")
|
||||||
|
}
|
||||||
|
|
||||||
|
func handle(
|
||||||
|
conn net.Conn,
|
||||||
|
token string,
|
||||||
|
allowPrivate bool,
|
||||||
|
cache *dnsCache,
|
||||||
|
tcpBuffer int,
|
||||||
|
slots chan struct{},
|
||||||
|
manager *chunkManager,
|
||||||
|
chunkMax int,
|
||||||
|
chunkBuffered int,
|
||||||
|
chunkPollWait time.Duration,
|
||||||
|
debug *serverDebug,
|
||||||
|
) {
|
||||||
|
defer func() {
|
||||||
|
<-slots
|
||||||
|
atomic.AddInt64(&active, -1)
|
||||||
|
_ = conn.Close()
|
||||||
|
}()
|
||||||
|
|
||||||
|
protocol.TuneTCP(conn)
|
||||||
|
protocol.TuneTCPBuffer(conn, tcpBuffer)
|
||||||
|
|
||||||
|
for {
|
||||||
|
_ = conn.SetDeadline(time.Now().Add(20 * time.Second))
|
||||||
|
|
||||||
|
requestID, _, payload, err := protocol.ReadRequestFrame(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"),
|
||||||
|
)
|
||||||
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
var (
|
||||||
|
host = flag.String("host", "0.0.0.0", "listen host")
|
||||||
|
port = flag.Int("port", 53, "listen port")
|
||||||
|
token = flag.String("token", "change-this-token", "shared token")
|
||||||
|
maxConnections = flag.Int("max-connections", 20000, "max simultaneous tunnels")
|
||||||
|
allowPrivate = flag.Bool("allow-private", false, "allow private/loopback targets")
|
||||||
|
dnsCacheTTL = flag.Duration("dns-cache-ttl", 30*time.Second, "server DNS cache TTL")
|
||||||
|
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", 65536, "maximum adaptive chunk payload bytes (32 bytes to 1 MiB)")
|
||||||
|
chunkBuffered = flag.Int("chunk-buffered", 256, "maximum buffered destination chunks per 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")
|
||||||
|
debugChunks = flag.Bool("debug-chunks", false, "log every chunk protocol record; very verbose")
|
||||||
|
debugStats = flag.Duration("debug-stats-interval", 5*time.Second, "periodic debug statistics interval; 0 disables")
|
||||||
|
)
|
||||||
|
flag.Parse()
|
||||||
|
|
||||||
|
if *chunkMax < 32 || *chunkMax > protocol.MaxChunkPayload {
|
||||||
|
fmt.Fprintf(os.Stderr, "--chunk-max must be between 32 and %d\n", protocol.MaxChunkPayload)
|
||||||
|
os.Exit(2)
|
||||||
|
}
|
||||||
|
if *chunkBuffered < 8 {
|
||||||
|
fmt.Fprintln(os.Stderr, "--chunk-buffered must be at least 8")
|
||||||
|
os.Exit(2)
|
||||||
|
}
|
||||||
|
|
||||||
|
listenAddr := net.JoinHostPort(*host, strconv.Itoa(*port))
|
||||||
|
ln, err := net.Listen("tcp", listenAddr)
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, err)
|
||||||
|
os.Exit(1)
|
||||||
|
}
|
||||||
|
defer ln.Close()
|
||||||
|
|
||||||
|
fmt.Printf("DragonTCP Go server listening on %s\n", listenAddr)
|
||||||
|
fmt.Printf("max_connections=%d tcp_buffer=%d\n", *maxConnections, *tcpBuffer)
|
||||||
|
|
||||||
|
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())
|
||||||
|
if debug.enabled {
|
||||||
|
fmt.Printf("debug=true debug_chunks=%t stats_interval=%s\n", debug.chunks, debug.statsEvery)
|
||||||
|
}
|
||||||
|
|
||||||
|
for {
|
||||||
|
conn, err := ln.Accept()
|
||||||
|
if err != nil {
|
||||||
|
fmt.Fprintln(os.Stderr, "accept:", err)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case slots <- struct{}{}:
|
||||||
|
atomic.AddInt64(&active, 1)
|
||||||
|
if debug.enabled {
|
||||||
|
debug.logf("ACCEPT peer=%v active_connections=%d", conn.RemoteAddr(), atomic.LoadInt64(&active))
|
||||||
|
}
|
||||||
|
go handle(
|
||||||
|
conn,
|
||||||
|
*token,
|
||||||
|
*allowPrivate,
|
||||||
|
cache,
|
||||||
|
*tcpBuffer,
|
||||||
|
slots,
|
||||||
|
manager,
|
||||||
|
*chunkMax,
|
||||||
|
*chunkBuffered,
|
||||||
|
*chunkPollWait,
|
||||||
|
debug,
|
||||||
|
)
|
||||||
|
default:
|
||||||
|
if debug.enabled {
|
||||||
|
debug.errorf("REJECT peer=%v reason=max-connections", conn.RemoteAddr())
|
||||||
|
}
|
||||||
|
_ = conn.Close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,211 @@
|
|||||||
|
package protocol
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/binary"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
XORKey byte = 0xAD
|
||||||
|
|
||||||
|
// MaxChunkPayload is the hard application-record payload ceiling.
|
||||||
|
// The adaptive chunk protocol may use any size from 32 bytes through 1 MiB.
|
||||||
|
MaxChunkPayload = 1024 * 1024
|
||||||
|
|
||||||
|
// Framed CPUSH/DATA messages include text metadata in addition to chunk
|
||||||
|
// bytes, so keep the frame ceiling comfortably above MaxChunkPayload.
|
||||||
|
MaxHandshake = 2 * 1024 * 1024
|
||||||
|
)
|
||||||
|
|
||||||
|
// 64 KiB balances throughput with memory use at high connection counts.
|
||||||
|
var BufferPool = sync.Pool{
|
||||||
|
New: func() any {
|
||||||
|
b := make([]byte, 64*1024)
|
||||||
|
return &b
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
func ReadRequestFrame(r io.Reader) (uint32, uint32, []byte, error) {
|
||||||
|
var header [14]byte
|
||||||
|
if _, err := io.ReadFull(r, header[:]); err != nil {
|
||||||
|
return 0, 0, nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if header[0] != 'U' || header[1] != 'P' {
|
||||||
|
return 0, 0, nil, errors.New("bad request magic")
|
||||||
|
}
|
||||||
|
|
||||||
|
requestID := binary.BigEndian.Uint32(header[2:6])
|
||||||
|
reserved := binary.BigEndian.Uint32(header[6:10])
|
||||||
|
length := binary.BigEndian.Uint32(header[10:14])
|
||||||
|
|
||||||
|
if length > MaxHandshake {
|
||||||
|
return 0, 0, nil, errors.New("handshake payload too large")
|
||||||
|
}
|
||||||
|
|
||||||
|
payload := make([]byte, int(length))
|
||||||
|
if _, err := io.ReadFull(r, payload); err != nil {
|
||||||
|
return 0, 0, nil, err
|
||||||
|
}
|
||||||
|
XorInPlace(payload)
|
||||||
|
|
||||||
|
return requestID, reserved, payload, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func WriteRequestFrame(w io.Writer, requestID uint32, payload []byte) error {
|
||||||
|
if len(payload) > MaxHandshake {
|
||||||
|
return errors.New("request frame payload too large")
|
||||||
|
}
|
||||||
|
packet := make([]byte, 14+len(payload))
|
||||||
|
packet[0], packet[1] = 'U', 'P'
|
||||||
|
binary.BigEndian.PutUint32(packet[2:6], requestID)
|
||||||
|
binary.BigEndian.PutUint32(packet[6:10], 0)
|
||||||
|
binary.BigEndian.PutUint32(packet[10:14], uint32(len(payload)))
|
||||||
|
copy(packet[14:], payload)
|
||||||
|
XorInPlace(packet[14:])
|
||||||
|
return writeAll(w, packet)
|
||||||
|
}
|
||||||
|
|
||||||
|
func ReadResponseFrame(r io.Reader) (uint32, []byte, error) {
|
||||||
|
var header [10]byte
|
||||||
|
if _, err := io.ReadFull(r, header[:]); err != nil {
|
||||||
|
return 0, nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
if header[0] != 'O' || header[1] != 'K' {
|
||||||
|
return 0, nil, fmt.Errorf("bad response magic: %q", header[:2])
|
||||||
|
}
|
||||||
|
|
||||||
|
requestID := binary.BigEndian.Uint32(header[2:6])
|
||||||
|
length := binary.BigEndian.Uint32(header[6:10])
|
||||||
|
|
||||||
|
if length > MaxHandshake {
|
||||||
|
return 0, nil, errors.New("handshake response too large")
|
||||||
|
}
|
||||||
|
|
||||||
|
payload := make([]byte, int(length))
|
||||||
|
if _, err := io.ReadFull(r, payload); err != nil {
|
||||||
|
return 0, nil, err
|
||||||
|
}
|
||||||
|
XorInPlace(payload)
|
||||||
|
|
||||||
|
return requestID, payload, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func WriteResponseFrame(w io.Writer, requestID uint32, payload []byte) error {
|
||||||
|
if len(payload) > MaxHandshake {
|
||||||
|
return errors.New("response frame payload too large")
|
||||||
|
}
|
||||||
|
packet := make([]byte, 10+len(payload))
|
||||||
|
packet[0], packet[1] = 'O', 'K'
|
||||||
|
binary.BigEndian.PutUint32(packet[2:6], requestID)
|
||||||
|
binary.BigEndian.PutUint32(packet[6:10], uint32(len(payload)))
|
||||||
|
copy(packet[10:], payload)
|
||||||
|
XorInPlace(packet[10:])
|
||||||
|
return writeAll(w, packet)
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeAll(w io.Writer, b []byte) error {
|
||||||
|
for len(b) > 0 {
|
||||||
|
n, err := w.Write(b)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
b = b[n:]
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func CopyXOR(dst net.Conn, src net.Conn) error {
|
||||||
|
ptr := BufferPool.Get().(*[]byte)
|
||||||
|
buf := *ptr
|
||||||
|
defer BufferPool.Put(ptr)
|
||||||
|
|
||||||
|
for {
|
||||||
|
n, err := src.Read(buf)
|
||||||
|
if n > 0 {
|
||||||
|
chunk := buf[:n]
|
||||||
|
XorInPlace(chunk)
|
||||||
|
|
||||||
|
if err2 := writeAll(dst, chunk); err2 != nil {
|
||||||
|
return err2
|
||||||
|
}
|
||||||
|
|
||||||
|
// No restore pass is needed. The next Read overwrites these bytes.
|
||||||
|
}
|
||||||
|
|
||||||
|
if err != nil {
|
||||||
|
if errors.Is(err, io.EOF) {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func relayPair(a, b net.Conn, copier func(net.Conn, net.Conn) error) {
|
||||||
|
done := make(chan struct{}, 2)
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
_ = copier(b, a)
|
||||||
|
if cw, ok := b.(interface{ CloseWrite() error }); ok {
|
||||||
|
_ = cw.CloseWrite()
|
||||||
|
}
|
||||||
|
done <- struct{}{}
|
||||||
|
}()
|
||||||
|
|
||||||
|
go func() {
|
||||||
|
_ = copier(a, b)
|
||||||
|
if cw, ok := a.(interface{ CloseWrite() error }); ok {
|
||||||
|
_ = cw.CloseWrite()
|
||||||
|
}
|
||||||
|
done <- struct{}{}
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Preserve normal TCP half-close semantics. The old implementation set a
|
||||||
|
// 2-second deadline on both connections after the first copy direction
|
||||||
|
// ended, which truncated slow or large responses. Wait for the remaining
|
||||||
|
// direction to drain naturally instead.
|
||||||
|
<-done
|
||||||
|
<-done
|
||||||
|
}
|
||||||
|
|
||||||
|
func RelayXOR(a, b net.Conn) {
|
||||||
|
relayPair(a, b, CopyXOR)
|
||||||
|
}
|
||||||
|
|
||||||
|
// RelayRaw allows Go/Linux to use the optimized TCP io.Copy path. On Linux,
|
||||||
|
// TCP-to-TCP copies can use splice, eliminating the userspace XOR/copy loop.
|
||||||
|
func RelayRaw(a, b net.Conn) {
|
||||||
|
relayPair(a, b, func(dst, src net.Conn) error {
|
||||||
|
_, err := io.Copy(dst, src)
|
||||||
|
return err
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TuneTCP(conn net.Conn) {
|
||||||
|
if tcp, ok := conn.(*net.TCPConn); ok {
|
||||||
|
_ = tcp.SetNoDelay(true)
|
||||||
|
_ = tcp.SetKeepAlive(true)
|
||||||
|
_ = tcp.SetKeepAlivePeriod(30 * time.Second)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TuneTCPBuffer optionally requests larger kernel socket buffers. A value <= 0
|
||||||
|
// leaves Linux/Android autotuning untouched, which is the recommended default
|
||||||
|
// for large connection counts. For a small number of high-BDP mobile links,
|
||||||
|
// values such as 1048576 or 4194304 can improve throughput.
|
||||||
|
func TuneTCPBuffer(conn net.Conn, size int) {
|
||||||
|
if size <= 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if tcp, ok := conn.(*net.TCPConn); ok {
|
||||||
|
_ = tcp.SetReadBuffer(size)
|
||||||
|
_ = tcp.SetWriteBuffer(size)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
//go:build arm || 386
|
||||||
|
|
||||||
|
package protocol
|
||||||
|
|
||||||
|
import "unsafe"
|
||||||
|
|
||||||
|
const xorWordMask32 uint32 = 0xADADADAD
|
||||||
|
|
||||||
|
// XorInPlace is the 32-bit optimized path used by ARMv7/386 builds.
|
||||||
|
// It aligns once, then processes 32 bytes per iteration with native uint32
|
||||||
|
// operations instead of a byte-at-a-time loop.
|
||||||
|
func XorInPlace(b []byte) {
|
||||||
|
n := len(b)
|
||||||
|
if n == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
i := 0
|
||||||
|
for i < n && (uintptr(unsafe.Pointer(&b[i]))&3) != 0 {
|
||||||
|
b[i] ^= XORKey
|
||||||
|
i++
|
||||||
|
}
|
||||||
|
|
||||||
|
for ; i+32 <= n; i += 32 {
|
||||||
|
p := unsafe.Pointer(&b[i])
|
||||||
|
*(*uint32)(unsafe.Add(p, 0)) ^= xorWordMask32
|
||||||
|
*(*uint32)(unsafe.Add(p, 4)) ^= xorWordMask32
|
||||||
|
*(*uint32)(unsafe.Add(p, 8)) ^= xorWordMask32
|
||||||
|
*(*uint32)(unsafe.Add(p, 12)) ^= xorWordMask32
|
||||||
|
*(*uint32)(unsafe.Add(p, 16)) ^= xorWordMask32
|
||||||
|
*(*uint32)(unsafe.Add(p, 20)) ^= xorWordMask32
|
||||||
|
*(*uint32)(unsafe.Add(p, 24)) ^= xorWordMask32
|
||||||
|
*(*uint32)(unsafe.Add(p, 28)) ^= xorWordMask32
|
||||||
|
}
|
||||||
|
|
||||||
|
for ; i+4 <= n; i += 4 {
|
||||||
|
p := (*uint32)(unsafe.Pointer(&b[i]))
|
||||||
|
*p ^= xorWordMask32
|
||||||
|
}
|
||||||
|
|
||||||
|
for ; i < n; i++ {
|
||||||
|
b[i] ^= XORKey
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,50 @@
|
|||||||
|
//go:build amd64 || arm64
|
||||||
|
|
||||||
|
package protocol
|
||||||
|
|
||||||
|
import "unsafe"
|
||||||
|
|
||||||
|
const xorWordMask uint64 = 0xADADADADADADADAD
|
||||||
|
|
||||||
|
// XorInPlace is optimized for 64-bit targets (amd64/arm64).
|
||||||
|
//
|
||||||
|
// It aligns the input once, then XORs 64 bytes per loop iteration using
|
||||||
|
// eight native 64-bit operations. This removes the encoding/binary call
|
||||||
|
// overhead from the hot relay path and lets the compiler generate a tight
|
||||||
|
// load/xor/store loop.
|
||||||
|
func XorInPlace(b []byte) {
|
||||||
|
n := len(b)
|
||||||
|
if n == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
i := 0
|
||||||
|
|
||||||
|
// Align the pointer for native uint64 accesses. This is normally already
|
||||||
|
// aligned for pooled relay buffers, but also makes this safe for subslices.
|
||||||
|
for i < n && (uintptr(unsafe.Pointer(&b[i]))&7) != 0 {
|
||||||
|
b[i] ^= XORKey
|
||||||
|
i++
|
||||||
|
}
|
||||||
|
|
||||||
|
for ; i+64 <= n; i += 64 {
|
||||||
|
p := unsafe.Pointer(&b[i])
|
||||||
|
*(*uint64)(unsafe.Add(p, 0)) ^= xorWordMask
|
||||||
|
*(*uint64)(unsafe.Add(p, 8)) ^= xorWordMask
|
||||||
|
*(*uint64)(unsafe.Add(p, 16)) ^= xorWordMask
|
||||||
|
*(*uint64)(unsafe.Add(p, 24)) ^= xorWordMask
|
||||||
|
*(*uint64)(unsafe.Add(p, 32)) ^= xorWordMask
|
||||||
|
*(*uint64)(unsafe.Add(p, 40)) ^= xorWordMask
|
||||||
|
*(*uint64)(unsafe.Add(p, 48)) ^= xorWordMask
|
||||||
|
*(*uint64)(unsafe.Add(p, 56)) ^= xorWordMask
|
||||||
|
}
|
||||||
|
|
||||||
|
for ; i+8 <= n; i += 8 {
|
||||||
|
p := (*uint64)(unsafe.Pointer(&b[i]))
|
||||||
|
*p ^= xorWordMask
|
||||||
|
}
|
||||||
|
|
||||||
|
for ; i < n; i++ {
|
||||||
|
b[i] ^= XORKey
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,10 @@
|
|||||||
|
//go:build !amd64 && !arm64 && !arm && !386
|
||||||
|
|
||||||
|
package protocol
|
||||||
|
|
||||||
|
// Generic fallback for 32-bit and uncommon architectures.
|
||||||
|
func XorInPlace(b []byte) {
|
||||||
|
for i := range b {
|
||||||
|
b[i] ^= XORKey
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user