game/server/cmd/login/main.go

package main

import (
	"crypto/rand"
	"encoding/hex"
	"flag"
	"fmt"
	"log"
	"net"
	"server/internal/db"
	"server/internal/net/packets"
	"sync"
	"time"

	"golang.org/x/crypto/bcrypt"
)

type LoginServer struct {
	conn        *net.UDPConn
	database    *db.Database
	worldHost   string
	worldPort   uint16
	serverID    string
	clients     map[string]*LoginClient // Map of addr string to client
	clientsLock sync.RWMutex
}

type LoginClient struct {
	addr      *net.UDPAddr
	lastSeen  time.Time
	username  string
	authed    bool
}

func main() {
	var (
		port      = flag.String("port", "9998", "Login server port")
		dbDSN     = flag.String("db", "user:password@tcp(localhost:3306)/game", "Database DSN")
		worldHost = flag.String("worldhost", "localhost", "World server host")
		worldPort = flag.Uint("worldport", 9999, "World server port")
	)
	flag.Parse()

	database, err := db.New(*dbDSN)
	if err != nil {
		log.Fatalf("Failed to connect to database: %v", err)
	}
	defer database.Close()

	server := &LoginServer{
		database:  database,
		worldHost: *worldHost,
		worldPort: uint16(*worldPort),
		serverID:  generateServerID(),
		clients:   make(map[string]*LoginClient),
	}

	addr, err := net.ResolveUDPAddr("udp", ":"+*port)
	if err != nil {
		log.Fatalf("Failed to resolve address: %v", err)
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Fatalf("Failed to start login server: %v", err)
	}
	server.conn = conn

	log.Printf("Login server started on UDP port %s (ID: %s)", *port, server.serverID)
	log.Printf("World server configured at %s:%d", server.worldHost, server.worldPort)

	go server.cleanupSessions()
	go server.cleanupClients()

	// Main packet processing loop
	buffer := make([]byte, 1024)
	for {
		n, clientAddr, err := conn.ReadFromUDP(buffer)
		if err != nil {
			log.Printf("Read error: %v", err)
			continue
		}

		go server.handlePacket(buffer[:n], clientAddr)
	}
}

func (s *LoginServer) handlePacket(data []byte, addr *net.UDPAddr) {
	if len(data) == 0 {
		return
	}

	msgType := data[0]

	switch msgType {
	case packets.MSG_LOGIN_REQUEST:
		s.handleLogin(data, addr)
	case packets.MSG_REGISTER_REQUEST:
		s.handleRegister(data, addr)
	}
}

func (s *LoginServer) handleLogin(data []byte, addr *net.UDPAddr) {
	username, password, ok := packets.DecodeLoginRequest(data)
	if !ok {
		s.sendLoginResponse(addr, false, nil, 0, "Invalid request")
		return
	}

	// Update client tracking
	s.clientsLock.Lock()
	s.clients[addr.String()] = &LoginClient{
		addr:     addr,
		lastSeen: time.Now(),
		username: username,
		authed:   false,
	}
	s.clientsLock.Unlock()

	player, err := s.database.GetPlayerByUsername(username)
	if err != nil {
		// Try to register if user doesn't exist
		s.handleRegister(data, addr)
		return
	}

	if err := bcrypt.CompareHashAndPassword([]byte(player.PasswordHash), []byte(password)); err != nil {
		s.sendLoginResponse(addr, false, nil, 0, "Invalid username or password")
		return
	}

	// Generate session token
	token := make([]byte, 32)
	rand.Read(token)
	tokenHex := hex.EncodeToString(token)

	if err := s.database.CreateSession(player.ID, tokenHex, "world", s.serverID, 24*time.Hour); err != nil {
		s.sendLoginResponse(addr, false, nil, 0, "Failed to create session")
		return
	}

	s.database.UpdateLastLogin(player.ID)

	// Mark client as authenticated
	s.clientsLock.Lock()
	if client, exists := s.clients[addr.String()]; exists {
		client.authed = true
	}
	s.clientsLock.Unlock()

	s.sendLoginResponse(addr, true, token, uint32(player.ID), "Login successful")
}

func (s *LoginServer) handleRegister(data []byte, addr *net.UDPAddr) {
	username, password, ok := packets.DecodeRegisterRequest(data)
	if !ok {
		s.sendLoginResponse(addr, false, nil, 0, "Invalid request")
		return
	}

	if len(username) < 3 || len(username) > 20 {
		s.sendLoginResponse(addr, false, nil, 0, "Username must be 3-20 characters")
		return
	}

	if len(password) < 6 {
		s.sendLoginResponse(addr, false, nil, 0, "Password must be at least 6 characters")
		return
	}

	hashedPassword, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		s.sendLoginResponse(addr, false, nil, 0, "Failed to process password")
		return
	}

	playerID, err := s.database.CreatePlayer(username, string(hashedPassword))
	if err != nil {
		// Player might already exist, try login instead
		s.handleLogin(data, addr)
		return
	}

	// Auto-login after registration
	token := make([]byte, 32)
	rand.Read(token)
	tokenHex := hex.EncodeToString(token)

	if err := s.database.CreateSession(playerID, tokenHex, "world", s.serverID, 24*time.Hour); err != nil {
		s.sendLoginResponse(addr, false, nil, 0, "Registration successful but login failed")
		return
	}

	s.sendLoginResponse(addr, true, token, uint32(playerID), "Registration successful")
}

func (s *LoginServer) sendLoginResponse(addr *net.UDPAddr, success bool, token []byte, playerID uint32, message string) {
	response := packets.EncodeLoginResponse(success, token, s.worldHost, s.worldPort, playerID, message)
	s.conn.WriteToUDP(response, addr)
}

func (s *LoginServer) cleanupSessions() {
	ticker := time.NewTicker(1 * time.Hour)
	defer ticker.Stop()

	for range ticker.C {
		if err := s.database.CleanExpiredSessions(); err != nil {
			log.Printf("Failed to clean expired sessions: %v", err)
		}
	}
}

func (s *LoginServer) cleanupClients() {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()

	for range ticker.C {
		s.clientsLock.Lock()
		now := time.Now()
		for addr, client := range s.clients {
			if now.Sub(client.lastSeen) > 5*time.Minute {
				delete(s.clients, addr)
			}
		}
		s.clientsLock.Unlock()
	}
}

func generateServerID() string {
	return fmt.Sprintf("login-%d", time.Now().Unix())
}