random fixes for preventing memory/conn leak

Signed-off-by: Mark Pashmfouroush <[email protected]>

go.mod

@@ -1,5 +1,7 @@
 module github.com/uoosef/bepass-relay
 
-go 1.19
+go 1.22
 
 require github.com/gaissmai/cidrtree v0.1.4
+
+require github.com/peterbourgon/ff/v4 v4.0.0-alpha.4

go.sum

@@ -1,2 +1,8 @@
 github.com/gaissmai/cidrtree v0.1.4 h1:/aYnv1LIwjtSDHNr1eNN99WJeh6vLrB+Sgr1tRMhHDc=
 github.com/gaissmai/cidrtree v0.1.4/go.mod h1:nrjEeeMZmvoJpLcSvZ3qIVFxw/+9GHKi7wDHHmHKGRI=
+github.com/pelletier/go-toml/v2 v2.0.9 h1:uH2qQXheeefCCkuBBSLi7jCiSmj3VRh2+Goq2N7Xxu0=
+github.com/pelletier/go-toml/v2 v2.0.9/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc=
+github.com/peterbourgon/ff/v4 v4.0.0-alpha.4 h1:aiqS8aBlF9PsAKeMddMSfbwp3smONCn3UO8QfUg0Z7Y=
+github.com/peterbourgon/ff/v4 v4.0.0-alpha.4/go.mod h1:H/13DK46DKXy7EaIxPhk2Y0EC8aubKm35nBjBe8AAGc=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
+gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=

ips.go

@@ -9,6 +9,8 @@ import (
 var (
 	// List of torrent trackers
 	torrentTrackers = []netip.Prefix{
+		netip.MustParsePrefix("127.0.0.0/8"),
+		netip.MustParsePrefix("::1/128"),
 		netip.MustParsePrefix("93.158.213.92/32"),
 		netip.MustParsePrefix("102.223.180.235/32"),
 		netip.MustParsePrefix("23.134.88.6/32"),
@@ -101,7 +103,6 @@ var (
 
 	// List of Cloudflare IP ranges
 	cfRanges = []netip.Prefix{
-		netip.MustParsePrefix("127.0.0.0/8"),
 		netip.MustParsePrefix("103.21.244.0/22"),
 		netip.MustParsePrefix("103.22.200.0/22"),
 		netip.MustParsePrefix("103.31.4.0/22"),
@@ -116,7 +117,6 @@ var (
 		netip.MustParsePrefix("190.93.240.0/20"),
 		netip.MustParsePrefix("197.234.240.0/22"),
 		netip.MustParsePrefix("198.41.128.0/17"),
-		netip.MustParsePrefix("::1/128"),
 		netip.MustParsePrefix("2400:cb00::/32"),
 		netip.MustParsePrefix("2405:8100::/32"),
 		netip.MustParsePrefix("2405:b500::/32"),

relay.go

@@ -3,104 +3,91 @@ package main
 
 import (
 	"bufio"
-	"flag"
+	"context"
+	"errors"
 	"fmt"
 	"io"
-	"log"
+	"log/slog"
 	"net"
 	"net/netip"
+	"os"
+	"os/signal"
 	"strings"
-)
-
-const BUFFER_SIZE = 256 * 1024
-
-type Server struct {
-	host string
-	port string
-}
-
-type Client struct {
-	conn net.Conn
-}
+	"syscall"
 
-type Config struct {
-	Host string
-	Port string
-}
+	"github.com/peterbourgon/ff/v4"
+	"github.com/peterbourgon/ff/v4/ffhelp"
+)
 
-func New(config *Config) *Server {
-	return &Server{
-		host: config.Host,
-		port: config.Port,
-	}
-}
+const BUFFER_SIZE = 2048
 
-func (server *Server) Run() {
-	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%s", server.host, server.port))
+func run(ctx context.Context, l *slog.Logger, bind string) error {
+	listener, err := net.Listen("tcp", bind)
 	if err != nil {
-		log.Fatal(err)
+		return err
 	}
-	defer func() {
-		_ = listener.Close()
-	}()
+	defer listener.Close()
 
 	for {
-		conn, err := listener.Accept()
-		if err != nil {
-			log.Fatal(err)
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+			conn, err := listener.Accept()
+			if err != nil {
+				l.Error("failed to accept connection", "error", err.Error())
+				continue
+			}
+
+			src := netip.MustParseAddrPort(conn.RemoteAddr().String())
+
+			// Check if srcIP is in the whitelist
+			if !connFilter.isSourceAllowed(src.Addr()) {
+				l.Debug("blocked connection", "address", src)
+				conn.Close()
+				continue
+			}
+
+			go handleConnection(l, conn)
 		}
-
-		src, err := netip.ParseAddrPort(conn.RemoteAddr().String())
-		if err != nil {
-			log.Printf("unable to parse host %v", conn.RemoteAddr())
-			_ = conn.Close()
-			continue
-		}
-
-		// Check if srcIP is in the whitelist
-		if !connFilter.isSourceAllowed(src.Addr()) {
-			log.Printf("blocked connection from: %v", src)
-			conn.Close()
-			continue
-		}
-
-		go (&Client{conn: conn}).handleRequest()
 	}
 }
 
-func (client *Client) handleRequest() {
-	defer func() {
-		_ = client.conn.Close()
-	}()
-	reader := bufio.NewReader(client.conn)
+func handleConnection(l *slog.Logger, lConn net.Conn) {
+	reader := bufio.NewReader(lConn)
+
 	header, _ := reader.ReadBytes(byte(13))
 	if len(header) < 1 {
+		lConn.Close()
 		return
 	}
+
 	inputHeader := strings.Split(string(header[:len(header)-1]), "@")
 	if len(inputHeader) < 2 {
+		lConn.Close()
 		return
 	}
+
 	network := "tcp"
 	if inputHeader[0] == "udp" {
 		network = "udp"
 	}
-	address := strings.Replace(inputHeader[1], "$", ":", -1)
-	if strings.Contains(address, "temp-mail.org") {
-		return
-	}
 
+	address := strings.Replace(inputHeader[1], "$", ":", -1)
 	dh, _, err := net.SplitHostPort(address)
 	if err != nil {
+		lConn.Close()
 		return
 	}
+
 	// check if ip is not blocked
 	blockFlag := false
 	addr, err := netip.ParseAddr(dh)
 	if err != nil {
 		// the host may not be an IP, try to resolve it
 		ips, err := net.LookupIP(dh)
 		if err != nil {
+			lConn.Close()
 			return
 		}
 
@@ -112,32 +99,26 @@ func (client *Client) handleRequest() {
 	blockFlag = !addr.IsValid() || !connFilter.isDestinationAllowed(addr)
 
 	if blockFlag {
-		log.Printf("destination host is blocked: %s\n", address)
+		l.Debug("destination host is blocked", "address", address)
+		lConn.Close()
 		return
 	}
 
-	if network == "udp" {
-		handleUDPOverTCP(client.conn, address)
-		return
-	}
-
-	// transmit data
-	log.Printf("%s Dialing to %s...\n", network, address)
+	switch network {
+	case "tcp":
+		rConn, err := net.Dial(network, address)
+		if err != nil {
+			l.Error("failed to dial", "protocol", network, "address", address, "error", err.Error())
+			lConn.Close()
+			return
+		}
 
-	rConn, err := net.Dial(network, address)
+		go handleTCP(lConn, rConn)
 
-	if err != nil {
-		log.Println(fmt.Errorf("failed to connect to socket: %v", err))
-		return
+	case "udp":
+		go handleUDPOverTCP(l, lConn, address)
 	}
-
-	defer func() {
-		_ = rConn.Close()
-	}()
-
-	// transmit data
-	go Copy(client.conn, rConn)
-	Copy(rConn, client.conn)
+	l.Debug("relaying connection", "protocol", network, "address", address)
 }
 
 // Copy reads from src and writes to dst until either EOF is reached on src or
@@ -154,10 +135,33 @@ func Copy(src io.Reader, dst io.Writer) {
 }
 
 func main() {
-	var config Config
-	flag.StringVar(&config.Host, "b", "0.0.0.0", "Server IP address")
-	flag.StringVar(&config.Port, "p", "6666", "Server Port number")
-	flag.Parse()
-	server := New(&config)
-	server.Run()
+	fs := ff.NewFlagSet("bepass-relay")
+	var (
+		verbose = fs.Bool('v', "verbose", "enable verbose logging")
+		bind    = fs.String('b', "bind", "0.0.0.0:6666", "bind address")
+	)
+
+	err := ff.Parse(fs, os.Args[1:])
+	switch {
+	case errors.Is(err, ff.ErrHelp):
+		fmt.Fprintf(os.Stderr, "%s\n", ffhelp.Flags(fs))
+		os.Exit(0)
+	case err != nil:
+		fmt.Fprintf(os.Stderr, "error: %v\n", err)
+		os.Exit(1)
+	}
+
+	l := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
+
+	if *verbose {
+		l = slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug}))
+	}
+
+	ctx, _ := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	if err := run(ctx, l, *bind); err != nil {
+		fmt.Fprintf(os.Stderr, "error: %v\n", err)
+		os.Exit(1)
+	}
+
+	<-ctx.Done()
 }

tcp.go

@@ -0,0 +1,8 @@
+package main
+
+import "net"
+
+func handleTCP(lConn, rConn net.Conn) {
+	go Copy(lConn, rConn)
+	Copy(rConn, lConn)
+}