From dfcdb6a96a9ab600885aada1b952318c8863c2e7 Mon Sep 17 00:00:00 2001
From: onyx-and-iris <code@onyxandiris.online>
Date: Wed, 20 Mar 2024 18:02:55 +0000
Subject: [PATCH] move packet header check methods into Session

log invalid packets at error level
---
 pkg/udpproxy/session.go  | 27 +++++++++++++++++++++++----
 pkg/udpproxy/udpproxy.go | 20 ++------------------
 2 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/pkg/udpproxy/session.go b/pkg/udpproxy/session.go
index 4c804bb..e5afc77 100644
--- a/pkg/udpproxy/session.go
+++ b/pkg/udpproxy/session.go
@@ -1,6 +1,7 @@
 package udpproxy
 
 import (
+	"errors"
 	"net"
 	"strings"
 	"time"
@@ -37,6 +38,14 @@ func (s *Session) isRconPacket(buf []byte) bool {
 	return string(buf[:8]) == "\xff\xff\xff\xffrcon"
 }
 
+func (s *Session) isQueryPacket(buf []byte) bool {
+	return string(buf[:13]) == "\xff\xff\xff\xffgetstatus" || string(buf[:11]) == "\xff\xff\xff\xffgetinfo"
+}
+
+func (s *Session) isValidPacket(buf []byte) bool {
+	return s.isRconPacket(buf) || s.isQueryPacket(buf)
+}
+
 func (s *Session) isResponsePacket(buf []byte) bool {
 	return string(buf[:9]) == "\xff\xff\xff\xffprint"
 }
@@ -46,7 +55,7 @@ func (s *Session) listen() error {
 		buf := make([]byte, 2048)
 		n, err := s.serverConn.Read(buf)
 		if err != nil {
-			log.Println(err)
+			log.Error(err)
 			continue
 		}
 
@@ -55,14 +64,18 @@ func (s *Session) listen() error {
 }
 
 func (s *Session) proxyFrom(buf []byte) error {
+	if !s.isResponsePacket(buf) {
+		return nil
+	}
+
 	s.updateTime = time.Now()
 	_, err := s.proxyConn.WriteToUDP(buf, s.caddr)
 	if err != nil {
-		log.Println(err)
+		log.Error(err)
 		return err
 	}
 
-	if s.isResponsePacket(buf) {
+	if log.GetLevel() == log.DebugLevel {
 		parts := strings.Split(string(buf[10:]), " ")
 		log.Debugf("Response: %s", strings.Join(parts, " "))
 	}
@@ -71,10 +84,16 @@ func (s *Session) proxyFrom(buf []byte) error {
 }
 
 func (s *Session) proxyTo(buf []byte) error {
+	if !s.isValidPacket(buf) {
+		err := errors.New("not a rcon or query packet")
+		log.Error(err.Error())
+		return err
+	}
+
 	s.updateTime = time.Now()
 	_, err := s.serverConn.Write(buf)
 	if err != nil {
-		log.Println(err)
+		log.Error(err)
 		return err
 	}
 
diff --git a/pkg/udpproxy/udpproxy.go b/pkg/udpproxy/udpproxy.go
index 24e5220..dadadb2 100644
--- a/pkg/udpproxy/udpproxy.go
+++ b/pkg/udpproxy/udpproxy.go
@@ -37,18 +37,6 @@ func New(port, target string) (*Client, error) {
 	}, nil
 }
 
-func (c *Client) isRconPacket(buf []byte) bool {
-	return string(buf[:8]) == "\xff\xff\xff\xffrcon"
-}
-
-func (c *Client) isQueryPacket(buf []byte) bool {
-	return string(buf[:13]) == "\xff\xff\xff\xffgetstatus" || string(buf[:11]) == "\xff\xff\xff\xffgetinfo"
-}
-
-func (c *Client) isValidPacket(buf []byte) bool {
-	return c.isRconPacket(buf) || c.isQueryPacket(buf)
-}
-
 func (c *Client) ListenAndServe() error {
 	var err error
 	c.proxyConn, err = net.ListenUDP("udp", c.laddr)
@@ -62,18 +50,14 @@ func (c *Client) ListenAndServe() error {
 		buf := make([]byte, 2048)
 		n, caddr, err := c.proxyConn.ReadFromUDP(buf)
 		if err != nil {
-			log.Println(err)
-		}
-
-		if !c.isValidPacket(buf[:n]) {
-			continue
+			log.Error(err)
 		}
 
 		session, found := c.sessions[caddr.String()]
 		if !found {
 			session, err = newSession(caddr, c.raddr, c.proxyConn)
 			if err != nil {
-				log.Println(err)
+				log.Error(err)
 				continue
 			}