This commit is contained in:
2026-08-16 19:02:48 -03:00
parent 96fe00eb2b
commit c8e3011f21
31 changed files with 3457 additions and 351 deletions
+96 -23
View File
@@ -17,6 +17,7 @@ import (
"sync"
"time"
"dragontcp/internal/cover"
"dragontcp/internal/protocol"
)
@@ -24,11 +25,12 @@ type chunkSession struct {
id string
target net.Conn
maxChunk int
maxChunks int
maxBuffer int
mu sync.Mutex
notify chan struct{}
chunks map[uint64][]byte
buffered int
nextDown uint64
eof bool
closed bool
@@ -42,14 +44,28 @@ type chunkSession struct {
haveLastUp bool
}
func newChunkSession(id string, target net.Conn, maxChunk, maxChunks int, debug *serverDebug) *chunkSession {
func newChunkSession(id string, target net.Conn, maxChunk, maxBuffer int, debug *serverDebug) *chunkSession {
if maxBuffer < maxChunk {
maxBuffer = maxChunk
}
readSize := maxChunk
if readSize > 64*1024 {
readSize = 64 * 1024
}
mapCapacity := maxBuffer / readSize
if mapCapacity < 1 {
mapCapacity = 1
}
if mapCapacity > 256 {
mapCapacity = 256
}
s := &chunkSession{
id: id,
target: target,
maxChunk: maxChunk,
maxChunks: maxChunks,
maxBuffer: maxBuffer,
notify: make(chan struct{}),
chunks: make(map[uint64][]byte, maxChunks),
chunks: make(map[uint64][]byte, mapCapacity),
lastSeen: time.Now(),
debug: debug,
}
@@ -73,7 +89,12 @@ func (s *chunkSession) touch() {
}
func (s *chunkSession) readTarget() {
buf := make([]byte, s.maxChunk)
ptr := protocol.BufferPool.Get().(*[]byte)
buf := *ptr
defer protocol.BufferPool.Put(ptr)
if s.maxChunk < len(buf) {
buf = buf[:s.maxChunk]
}
for {
n, err := s.target.Read(buf)
@@ -89,10 +110,11 @@ func (s *chunkSession) readTarget() {
s.mu.Unlock()
return
}
if len(s.chunks) < s.maxChunks {
if s.buffered+len(data) <= s.maxBuffer {
seq := s.nextDown
s.nextDown++
s.chunks[seq] = data
s.buffered += len(data)
s.touchLocked()
s.signalLocked()
s.mu.Unlock()
@@ -181,6 +203,7 @@ func (s *chunkSession) pull(want uint64, ack int64, offset, limit int, wait time
removed := false
for seq := range s.chunks {
if seq <= uint64(ack) {
s.buffered -= len(s.chunks[seq])
delete(s.chunks, seq)
removed = true
}
@@ -333,7 +356,8 @@ func decodeWireToken(token string) string {
}
func isChunkCommand(payload []byte) bool {
return bytes.HasPrefix(payload, []byte("COPEN ")) ||
return bytes.HasPrefix(payload, []byte("CPROBE ")) ||
bytes.HasPrefix(payload, []byte("COPEN ")) ||
bytes.HasPrefix(payload, []byte("CPUSH ")) ||
bytes.HasPrefix(payload, []byte("CPULL ")) ||
bytes.HasPrefix(payload, []byte("CCLOSE "))
@@ -349,10 +373,21 @@ func processChunkCommand(
tcpBuffer int,
manager *chunkManager,
maxChunk int,
maxBufferedChunks int,
maxBufferedBytes int,
pollWait time.Duration,
debug *serverDebug,
) error {
if bytes.HasPrefix(payload, []byte("CPROBE ")) {
parts := strings.Fields(string(payload))
if len(parts) != 2 {
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR bad CPROBE"))
}
if !tokenEqual(decodeWireToken(parts[1]), token) {
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
}
return protocol.WriteResponseFrame(conn, requestID, []byte("PROBEOK"))
}
if bytes.HasPrefix(payload, []byte("COPEN ")) {
parts := strings.Fields(string(payload))
if len(parts) != 5 {
@@ -378,7 +413,7 @@ func processChunkCommand(
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR "+err.Error()))
}
session := newChunkSession(sid, target, maxChunk, maxBufferedChunks, debug)
session := newChunkSession(sid, target, maxChunk, maxBufferedBytes, debug)
if err := manager.add(sid, session); err != nil {
session.close()
if debug != nil && debug.enabled {
@@ -521,40 +556,78 @@ func processChunkCommand(
// beyond the magic would be lost.
type prefixedConn struct {
net.Conn
r io.Reader
r io.Reader
headerMask byte
cover cover.Profile
}
func (p *prefixedConn) Read(b []byte) (int, error) { return p.r.Read(b) }
func (p *prefixedConn) Read(b []byte) (int, error) { return p.r.Read(b) }
func (p *prefixedConn) HeaderMask() byte { return p.headerMask }
func (p *prefixedConn) CoverProfile() cover.Profile { return p.cover }
func (p *prefixedConn) ClearPayload() bool { return p.cover.Clear }
// sniffWire reads the two magic bytes and reports whether this connection
// speaks the legacy XOR framing. The returned conn replays them.
func sniffWire(conn net.Conn) (net.Conn, bool, error) {
var magic [2]byte
if _, err := io.ReadFull(conn, magic[:]); err != nil {
return conn, false, err
// sniffWire first checks for the optional self-describing cover preface. If it
// is absent, the bytes are replayed and the legacy/direct B/X classifier is
// used unchanged.
func sniffWire(conn net.Conn) (net.Conn, bool, byte, error) {
var initial [cover.PrefaceSize]byte
if _, err := io.ReadFull(conn, initial[:]); err != nil {
return conn, false, 0, err
}
replayed := &prefixedConn{Conn: conn, r: io.MultiReader(bytes.NewReader(magic[:]), conn)}
return replayed, magic[0] == 'U' && magic[1] == 'P', nil
if profile, ok := cover.DecodePreface(initial); ok {
if profile.Padding > 0 {
padding := make([]byte, int(profile.Padding))
if _, err := io.ReadFull(conn, padding); err != nil {
return conn, false, 0, err
}
}
profiled := &prefixedConn{Conn: conn, r: conn, headerMask: profile.HeaderMask, cover: profile}
return profiled, profile.XOR, profile.HeaderMask, nil
}
magic := initial[:2]
replay := io.MultiReader(bytes.NewReader(initial[:]), conn)
if magic[0]&7 >= 5 {
mask := magic[0] ^ 'U'
if magic[1]^mask != 'P' {
return conn, false, 0, fmt.Errorf("unknown wire header")
}
replayed := &prefixedConn{Conn: conn, r: replay, headerMask: mask}
return replayed, true, mask, nil
}
mask := magic[0] & 0xf8
mode := magic[0] ^ mask
if mode > 4 {
return conn, false, 0, fmt.Errorf("unknown binary mode")
}
replayed := &prefixedConn{Conn: conn, r: replay, headerMask: mask}
return replayed, false, mask, nil
}
// handleXOR serves one connection speaking UP/OK + XOR 0xAD: the v4 chunk
// commands, plus the TUNNEL/TUNNEL2 stream commands.
func handleXOR(
conn net.Conn,
headerMask byte,
token string,
allowPrivate bool,
cache *dnsCache,
tcpBuffer int,
manager *chunkManager,
chunkMax int,
chunkBuffered int,
bufferBytes int,
chunkPollWait time.Duration,
debug *serverDebug,
) {
deadline := newIdleDeadline(conn, 20*time.Second)
for {
_ = conn.SetDeadline(time.Now().Add(20 * time.Second))
if deadline.refresh() != nil {
return
}
requestID, _, payload, err := protocol.ReadRequestFrame(conn)
requestID, _, payload, err := protocol.ReadRequestFrameProfile(conn, headerMask)
if err != nil {
if debug != nil && debug.enabled && err != io.EOF {
debug.errorf("peer=%v read XOR request: %v", conn.RemoteAddr(), err)
@@ -565,7 +638,7 @@ func handleXOR(
if isChunkCommand(payload) {
if err := processChunkCommand(
conn, requestID, payload, token, allowPrivate, cache, tcpBuffer,
manager, chunkMax, chunkBuffered, chunkPollWait, debug,
manager, chunkMax, bufferBytes, chunkPollWait, debug,
); err != nil {
return
}