From 40d79063f5364eb78f3507e555d478dfd58edab1 Mon Sep 17 00:00:00 2001
From: onyx-and-iris <code@onyxandiris.online>
Date: Wed, 7 May 2025 15:41:34 +0100
Subject: [PATCH] add validation logic to ports-mapping flag

---
 cmd/q3rcon-proxy/main.go | 28 +++++++++++++++++++++++++---
 1 file changed, 25 insertions(+), 3 deletions(-)

diff --git a/cmd/q3rcon-proxy/main.go b/cmd/q3rcon-proxy/main.go
index 5abe1e2..cd958c8 100644
--- a/cmd/q3rcon-proxy/main.go
+++ b/cmd/q3rcon-proxy/main.go
@@ -4,10 +4,11 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
-	"github.com/onyx-and-iris/q3rcon-proxy/pkg/udpproxy"
+	udpproxy "github.com/onyx-and-iris/q3rcon-proxy"
 	log "github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v3"
 )
@@ -41,6 +42,27 @@ func main() {
 				Usage:    "Proxy and target ports (proxy:target)",
 				Sources:  cli.EnvVars("Q3RCON_PORTS_MAPPING"),
 				Required: true,
+				Action: func(ctx context.Context, cmd *cli.Command, v string) error {
+					// Validate the ports mapping
+					for mapping := range strings.SplitSeq(v, ";") {
+						ports := strings.Split(mapping, ":")
+						if len(ports) != 2 {
+							return fmt.Errorf("invalid ports mapping: %s", mapping)
+						}
+						proxyPort, err := strconv.Atoi(ports[0])
+						if err != nil || proxyPort < 1 || proxyPort > 65535 {
+							return fmt.Errorf("invalid proxy port: %s", ports[0])
+						}
+						targetPort, err := strconv.Atoi(ports[1])
+						if err != nil || targetPort < 1 || targetPort > 65535 {
+							return fmt.Errorf("invalid target port: %s", ports[1])
+						}
+						if proxyPort == targetPort {
+							return fmt.Errorf("proxy and target ports cannot be the same: %s", mapping)
+						}
+					}
+					return nil
+				},
 			},
 			&cli.IntFlag{
 				Name:    "session-timeout",
@@ -99,7 +121,7 @@ func initProxy(cfg proxyConfig, errChan chan error) {
 	hostAddr := fmt.Sprintf("%s:%s", cfg.proxyHost, proxyPort)
 	proxyAddr := fmt.Sprintf("%s:%s", cfg.targetHost, targetPort)
 
-	c, err := udpproxy.New(
+	server, err := udpproxy.New(
 		hostAddr, proxyAddr,
 		udpproxy.WithSessionTimeout(time.Duration(cfg.sessionTimeout)*time.Minute))
 	if err != nil {
@@ -109,5 +131,5 @@ func initProxy(cfg proxyConfig, errChan chan error) {
 
 	log.Printf("q3rcon-proxy initialized: [proxy] (%s) [target] (%s)", hostAddr, proxyAddr)
 
-	errChan <- c.ListenAndServe()
+	errChan <- server.ListenAndServe()
 }