DragonTCP/core/cmd/dragontcp-server/main.go

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", "", "optional 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", 1048576, "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()
		}
	}
}