main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"sync/atomic"

	"Wraith/internal/config"
	"Wraith/internal/network"
	"Wraith/internal/protocol"
	"Wraith/internal/protocol/packets"
	"Wraith/internal/proxy"
)

func main() {
	cfgPath := flag.String("config", "config.json", "path to config file")
	flag.Parse()

	cfg, err := config.Load(*cfgPath)
	if err != nil {
		log.Fatalf("loading config: %v", err)
	}

	logger := log.New(os.Stdout, "", log.LstdFlags|log.Lmsgprefix)

	protoReg := protocol.NewRegistry()
	packets.RegisterAll(protoReg)

	handlers := proxy.NewHandlerRegistry()
	proxy.RegisterHandlers(handlers, logger)

	listener, err := net.Listen("tcp", cfg.ListenAddr)
	if err != nil {
		log.Fatalf("listening on %s: %v", cfg.ListenAddr, err)
	}
	logger.Printf("Wraith listening on %s, forwarding to %s (protocol %d / 1.21.4)",
		cfg.ListenAddr, cfg.BackendAddr, packets.ProtocolVersion)

	var nextID int64
	for {
		rawConn, err := listener.Accept()
		if err != nil {
			logger.Printf("accept error: %v", err)
			continue
		}
		id := atomic.AddInt64(&nextID, 1)
		go handleConnection(rawConn, id, cfg, protoReg, handlers, logger)
	}
}

func handleConnection(rawConn net.Conn, id int64, cfg config.Config, protoReg *protocol.Registry, handlers *proxy.HandlerRegistry, logger *log.Logger) {
	sessionID := fmt.Sprintf("#%d(%s)", id, rawConn.RemoteAddr())
	clientConn := network.NewConn(rawConn)
	session := proxy.NewSession(sessionID, clientConn, cfg.BackendAddr, protoReg, handlers, logger)

	logger.Printf("session %s: connected", sessionID)
	if err := session.Run(); err != nil {
		logger.Printf("session %s: ended: %v", sessionID, err)
		return
	}
	logger.Printf("session %s: closed", sessionID)
}