V7

2026-08-16 03:48:57 -03:00
parent 5621de243a
commit c02b36c83d
48 changed files with 5398 additions and 2221 deletions
@@ -0,0 +1,505 @@
+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 decodeWireToken(token string) string {
+	if token == "-" {
+		return ""
+	}
+	return token
+}
+
+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(decodeWireToken(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(decodeWireToken(string(parts[1])), token) {
+			return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
+		}
+		sid := string(parts[2])
+		seq, err := strconv.ParseUint(string(parts[3]), 10, 64)
+		if err != nil {
+			return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid sequence"))
+		}
+		s := manager.get(sid)
+		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(decodeWireToken(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(decodeWireToken(parts[1]), token) {
+			return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
+		}
+		manager.remove(parts[2])
+		if debug != nil && debug.enabled {
+			debug.sessionsClosed.Add(1)
+			debug.activeSessions.Add(-1)
+			debug.logf("SESSION CLOSE id=%s peer=%v active_sessions=%d", parts[2], conn.RemoteAddr(), manager.count())
+			debug.chunkf("CCLOSE id=%s -> CLOSED", parts[2])
+		}
+		return protocol.WriteResponseFrame(conn, requestID, []byte("CLOSED"))
+	}
+
+	return protocol.WriteResponseFrame(conn, requestID, []byte("ERR unknown chunk command"))
+}