QR2: Wait for connections to stop before shutdown

This commit is contained in:
mkwcat 2024-05-07 21:19:56 -04:00
parent 74e5c6a910
commit 832c54acee
No known key found for this signature in database
GPG Key ID: 7A505679CE9E7AA9

View File

@ -3,6 +3,7 @@ package qr2
import ( import (
"encoding/binary" "encoding/binary"
"net" "net"
"sync"
"time" "time"
"wwfc/common" "wwfc/common"
"wwfc/logging" "wwfc/logging"
@ -26,7 +27,11 @@ const (
ClientExploitReply = 0x10 ClientExploitReply = 0x10
) )
var masterConn net.PacketConn var (
masterConn net.PacketConn
inShutdown = false
waitGroup = sync.WaitGroup{}
)
func StartServer(reload bool) { func StartServer(reload bool) {
// Get config // Get config
@ -39,53 +44,91 @@ func StartServer(reload bool) {
} }
masterConn = conn masterConn = conn
inShutdown = false
if reload { if reload {
loadSessions() err := loadSessions()
loadLogins() if err != nil {
loadGroups() panic(err)
logging.Info("QR2", "Loaded", aurora.Cyan(len(sessions)), "sessions") }
logging.Notice("QR2", "Loaded", aurora.Cyan(len(sessions)), "sessions")
err = loadLogins()
if err != nil {
panic(err)
}
logging.Notice("QR2", "Loaded", aurora.Cyan(len(logins)), "logins")
err = loadGroups()
if err != nil {
panic(err)
}
logging.Notice("QR2", "Loaded", aurora.Cyan(len(groups)), "groups")
} }
waitGroup.Add(1)
go func() { go func() {
defer waitGroup.Done()
// Close the listener when the application closes. // Close the listener when the application closes.
defer conn.Close() defer conn.Close()
logging.Notice("QR2", "Listening on", aurora.BrightCyan(address)) logging.Notice("QR2", "Listening on", aurora.BrightCyan(address))
for { for {
if inShutdown {
return
}
buf := make([]byte, 1024) buf := make([]byte, 1024)
_, addr, err := conn.ReadFrom(buf) n, addr, err := conn.ReadFrom(buf)
if err != nil { if err != nil || n == 0 {
continue continue
} }
waitGroup.Add(1)
go handleConnection(conn, *addr.(*net.UDPAddr), buf) go handleConnection(conn, *addr.(*net.UDPAddr), buf)
} }
}() }()
} }
func Shutdown() { func Shutdown() {
inShutdown = true
masterConn.Close()
waitGroup.Wait()
mutex.Lock()
defer mutex.Unlock()
err := saveSessions() err := saveSessions()
if err != nil { if err != nil {
logging.Error("QR2", "Failed to save sessions:", err) logging.Error("QR2", "Failed to save sessions:", err)
} }
logging.Info("QR2", "Saved", aurora.Cyan(len(sessions)), "sessions") logging.Notice("QR2", "Saved", aurora.Cyan(len(sessions)), "sessions")
err = saveLogins() err = saveLogins()
if err != nil { if err != nil {
logging.Error("QR2", "Failed to save logins:", err) logging.Error("QR2", "Failed to save logins:", err)
} }
logging.Info("QR2", "Saved", aurora.Cyan(len(logins)), "logins") logging.Notice("QR2", "Saved", aurora.Cyan(len(logins)), "logins")
err = saveGroups() err = saveGroups()
if err != nil { if err != nil {
logging.Error("QR2", "Failed to save groups:", err) logging.Error("QR2", "Failed to save groups:", err)
} }
logging.Notice("QR2", "Saved", aurora.Cyan(len(groups)), "groups")
} }
func handleConnection(conn net.PacketConn, addr net.UDPAddr, buffer []byte) { func handleConnection(conn net.PacketConn, addr net.UDPAddr, buffer []byte) {
defer waitGroup.Done()
packetType := buffer[0] packetType := buffer[0]
moduleName := "QR2:" + addr.String() moduleName := "QR2:" + addr.String()