PR review fixes

internal/socks5/socks5.go

@@ -10,6 +10,7 @@ import (
 	"net/netip"
 	"strconv"
 	"strings"
+	"sync"
 )
 
 var (
@@ -227,18 +228,18 @@ func (c *socksConn) handleUDPAssociateRequest(ctx context.Context,
 	associationCtx, associationCancel := context.WithCancel(ctx)
 	defer associationCancel()
 
-	handlerDone := make(chan struct{})
-	go func() {
-		defer close(handlerDone)
-		c.udpRouter.runAssociationHandler(associationCtx, association)
-	}()
+	var wg sync.WaitGroup
 
-	go func() {
+	wg.Go(func() {
+		c.udpRouter.runAssociationHandler(associationCtx, association)
+	})
+
+	wg.Go(func() {
 		_, _ = io.Copy(io.Discard, c.clientConn)
 		associationCancel()
-	}()
+	})
 	<-associationCtx.Done()
-	<-handlerDone
+	wg.Wait()
 	return nil
 }
 

internal/socks5/udp_router.go

@@ -159,8 +159,9 @@ func (r *udpRouter) run(ctx context.Context) error {
 
 func (r *udpRouter) routePacket(sourceAddrPort netip.AddrPort, packet *bytes.Buffer) error {
 	r.mutex.Lock()
-	defer r.mutex.Unlock()
 	association, packetFromClient := r.findClientAssociation(sourceAddrPort)
+	r.mutex.Unlock()
+
 	if !packetFromClient {
 		r.bufferPool.Put(packet)
 		return nil
@@ -330,6 +331,20 @@ func (r *udpRouter) writeClientPacketToDestination(ctx context.Context,
 		return fmt.Errorf("decoding UDP datagram: %w", err)
 	}
 
+	host, portStr, err := net.SplitHostPort(destination)
+	if err != nil {
+		return fmt.Errorf("splitting destination host and port: %w", err)
+	}
+
+	if _, err := netip.ParseAddr(destination); err != nil { // domain name
+		addrs, err := net.DefaultResolver.LookupHost(ctx, host)
+		if err != nil {
+			return fmt.Errorf("resolving destination host: %w", err)
+		}
+
+		destination = net.JoinHostPort(addrs[0], portStr)
+	}
+
 	resolvedDestinationUDPAddress, err := net.ResolveUDPAddr("udp", destination)
 	if err != nil {
 		return fmt.Errorf("resolving destination UDP address: %w", err)