2025-02-27 21:39:55

Author: mohanson

daze.go

@@ -22,6 +22,7 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
+	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -419,7 +420,6 @@ func (l *Locale) ServeSocks5UDP(ctx *Context, cli io.ReadWriteCloser) error {
 		bndPort     uint16
 		bnd         *net.UDPConn
 		buf         = make([]byte, 2048)
-		b           bool
 		cpl         = map[string]io.ReadWriteCloser{}
 		dstHost     string
 		dstPort     uint16
@@ -478,10 +478,7 @@ func (l *Locale) ServeSocks5UDP(ctx *Context, cli io.ReadWriteCloser) error {
 		case 0x04:
 			appHeadSize = 22
 		}
-
-		appHead = make([]byte, appHeadSize)
-		copy(appHead, buf[0:appHeadSize])
-
+		appHead = buf[0:appHeadSize]
 		switch appHead[3] {
 		case 0x01:
 			dstHost = net.IP(appHead[4:8]).String()
@@ -495,45 +492,25 @@ func (l *Locale) ServeSocks5UDP(ctx *Context, cli io.ReadWriteCloser) error {
 			dstPort = binary.BigEndian.Uint16(appHead[20:22])
 		}
 		dst = dstHost + ":" + strconv.Itoa(int(dstPort))
-
-		srv, b = cpl[dst]
-		if b {
-			goto send
-		} else {
-			goto init
-		}
-	init:
-		log.Printf("conn: %08x  proto format=socks5", ctx.Cid)
-		srv, err = l.Dialer.Dial(ctx, "udp", dst)
-		if err != nil {
-			log.Printf("conn: %08x  error %s", ctx.Cid, err)
-			continue
-		}
-		cpl[dst] = srv
-		go func(srv io.ReadWriteCloser, appHead []byte, appAddr *net.UDPAddr) error {
-			var (
-				buf = make([]byte, 2048)
-				l   = len(appHead)
-				n   int
-				err error
-			)
-			copy(buf, appHead)
-			for {
-				n, err = srv.Read(buf[l:])
-				if err != nil {
-					break
-				}
-				_, err = bnd.WriteToUDP(buf[:l+n], appAddr)
-				if err != nil {
-					break
-				}
+		srv = cpl[dst]
+		if srv == nil {
+			log.Printf("conn: %08x  proto format=socks5", ctx.Cid)
+			srv, err = l.Dialer.Dial(ctx, "udp", dst)
+			if err != nil {
+				log.Printf("conn: %08x  error %s", ctx.Cid, err)
+				continue
 			}
-			return err
-		}(srv, appHead, appAddr)
-	send:
+			cpl[dst] = srv
+			retHead := slices.Clone(appHead)
+			go ReadCall(srv, func(data []byte) error {
+				return doa.Err(bnd.WriteToUDP(append(retHead, data...), appAddr))
+			})
+		}
 		_, err = srv.Write(buf[appHeadSize:appSize])
 		if err != nil {
 			log.Printf("conn: %08x  error %s", ctx.Cid, err)
+			srv.Close()
+			delete(cpl, dst)
 			continue
 		}
 	}
@@ -1053,6 +1030,25 @@ func (r *RandomReader) Read(p []byte) (int, error) {
 	return len(p), nil
 }
 
+// ReadCall reads data from the given io.Reader and passes it to the provided call function.
+func ReadCall(conn io.Reader, call func([]byte) error) error {
+	var (
+		buf = make([]byte, 2048)
+		err error
+		n   int
+	)
+	for {
+		n, err = conn.Read(buf)
+		if err != nil {
+			break
+		}
+		if err = call(buf[:n]); err != nil {
+			break
+		}
+	}
+	return err
+}
+
 // Salt converts the stupid password passed in by the user to 32-sized byte array.
 func Salt(s string) []byte {
 	h := sha256.Sum256([]byte(s))