fix race problem in udp when passing the pointer to channel

bepass-org · Aug 29, 2023 · a5b51b8 · a5b51b8
1 parent ab592b9
commit a5b51b8
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 16 deletions.
diff --git a/transport/transport.go b/transport/transport.go
@@ -22,7 +22,7 @@ type UdpBind struct {
 	SocksWriter   io.Writer
 	SocksReq      *socks5.Request
 	AssociateBind *net.UDPConn
-	RecvChan      chan *UDPPacket
+	RecvChan      chan UDPPacket
 }
 
 type UDPConf struct {
@@ -132,7 +132,7 @@ func (t *Transport) TunnelUDP(w io.Writer, req *socks5.Request) error {
 		return err
 	}
 
-	bindWriteChannel := make(chan *UDPPacket)
+	bindWriteChannel := make(chan UDPPacket)
 	tunnelWriteChannel, channelIndex, err := t.Tunnel.PersistentDial(tunnelEndpoint, bindWriteChannel)
 	if err != nil {
 		t.Logger.Errorf("Unable to get or create tunnel for udpBindWriteChannel %v\r\n", err)
@@ -165,15 +165,14 @@ func (t *Transport) TunnelUDP(w io.Writer, req *socks5.Request) error {
 			if err != nil {
 				continue
 			}
-			tunnelWriteChannel <- &UDPPacket{
+			tunnelWriteChannel <- UDPPacket{
 				Channel: channelIndex,
 				Data:    pk.Data,
 			}
 		}
 	}()
 	for {
 		datagram := <-udpBind.RecvChan
-		fmt.Print("injam")
 		pkb, err := statute.NewDatagram(req.RawDestAddr.String(), datagram.Data)
 		if err != nil {
 			continue

diff --git a/transport/ws.go b/transport/ws.go
@@ -6,7 +6,6 @@ import (
 	"bepass/wsconnadapter"
 	"context"
 	"encoding/binary"
-	"fmt"
 	"github.com/gorilla/websocket"
 	"golang.org/x/net/proxy"
 	"net"
@@ -15,8 +14,8 @@ import (
 )
 
 type EstablishedTunnel struct {
-	tunnelWriteChannel chan *UDPPacket
-	bindWriteChannels  map[uint16]chan *UDPPacket
+	tunnelWriteChannel chan UDPPacket
+	bindWriteChannels  map[uint16]chan UDPPacket
 	channelIndex       uint16
 }
 
@@ -56,18 +55,18 @@ func (w *WSTunnel) Dial(endpoint string) (*websocket.Conn, error) {
 	return conn, err
 }
 
-func (w *WSTunnel) PersistentDial(tunnelEndpoint string, bindWriteChannel chan *UDPPacket) (chan *UDPPacket, uint16, error) {
+func (w *WSTunnel) PersistentDial(tunnelEndpoint string, bindWriteChannel chan UDPPacket) (chan UDPPacket, uint16, error) {
 	if tunnel, ok := w.EstablishedTunnels[tunnelEndpoint]; ok {
 		tunnel.channelIndex = tunnel.channelIndex + 1
 		tunnel.bindWriteChannels[tunnel.channelIndex] = bindWriteChannel
 		return tunnel.tunnelWriteChannel, tunnel.channelIndex, nil
 	}
 
-	tunnelWriteChannel := make(chan *UDPPacket)
+	tunnelWriteChannel := make(chan UDPPacket)
 
 	w.EstablishedTunnels[tunnelEndpoint] = &EstablishedTunnel{
 		tunnelWriteChannel: tunnelWriteChannel,
-		bindWriteChannels:  make(map[uint16]chan *UDPPacket),
+		bindWriteChannels:  make(map[uint16]chan UDPPacket),
 		channelIndex:       1,
 	}
 
@@ -162,15 +161,10 @@ func (w *WSTunnel) PersistentDial(tunnelEndpoint string, bindWriteChannel chan *
 							continue
 						}
 
-						fmt.Println(rawPacket[:n])
-
 						// first 2 packets of response is channel id
-
 						channelID := binary.BigEndian.Uint16(rawPacket[:2])
 
-						fmt.Println(channelID)
-
-						pkt := &UDPPacket{
+						pkt := UDPPacket{
 							channelID,
 							rawPacket[2:n],
 						}