84 lines
2.0 KiB
Go
84 lines
2.0 KiB
Go
package client
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"net"
|
|
pb "net-tunnel/pkg/proto"
|
|
"time"
|
|
)
|
|
|
|
func (c *Client) handleUDPData(msg *pb.Message_ProxyData) {
|
|
log.Printf("Received UDP proxy data for connection: %v", msg.ProxyData.ConnId)
|
|
hostPort := net.JoinHostPort(msg.ProxyData.ProxyConfig.LocalIp, fmt.Sprintf("%d", msg.ProxyData.ProxyConfig.LocalPort))
|
|
|
|
// 获取或创建 UDP 连接
|
|
udpConnKey := msg.ProxyData.ConnId
|
|
var conn *net.UDPConn
|
|
|
|
existingConn, ok := c.udpConnections.Load(udpConnKey)
|
|
if ok {
|
|
conn = existingConn.(*net.UDPConn)
|
|
} else {
|
|
udpAddr, err := net.ResolveUDPAddr("udp", hostPort)
|
|
if err != nil {
|
|
log.Printf("Failed to resolve UDP address: %v", err)
|
|
return
|
|
}
|
|
conn, err = net.DialUDP("udp", nil, udpAddr)
|
|
if err != nil {
|
|
log.Printf("Failed to dial UDP address: %v", err)
|
|
return
|
|
}
|
|
c.udpConnections.Store(udpConnKey, conn)
|
|
|
|
go c.readFromUDPConn(msg, conn)
|
|
}
|
|
|
|
_, err := conn.Write(msg.ProxyData.Data)
|
|
if err != nil {
|
|
log.Printf("Failed to write data to UDP connection: %v", err)
|
|
conn.Close()
|
|
c.udpConnections.Delete(udpConnKey)
|
|
return
|
|
}
|
|
}
|
|
|
|
func (c *Client) readFromUDPConn(msg *pb.Message_ProxyData, conn *net.UDPConn) {
|
|
buffer := make([]byte, 4096)
|
|
for {
|
|
err := conn.SetReadDeadline(time.Now().Add(time.Minute * 10))
|
|
if err != nil {
|
|
log.Printf("Failed to set read deadline: %v", err)
|
|
return
|
|
}
|
|
n, _, err := conn.ReadFromUDP(buffer)
|
|
if err != nil {
|
|
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
|
|
continue
|
|
}
|
|
log.Printf("Failed to read from UDP connection: %v", err)
|
|
return
|
|
}
|
|
if n > 0 {
|
|
msg := &pb.Message{
|
|
Content: &pb.Message_ProxyData{
|
|
ProxyData: &pb.ProxyData{
|
|
ConnId: msg.ProxyData.ConnId,
|
|
ProxyConfig: msg.ProxyData.ProxyConfig,
|
|
Data: buffer[:n],
|
|
},
|
|
},
|
|
}
|
|
if err := c.stream.Send(msg); err != nil {
|
|
log.Printf("Failed to send proxy data: %v", err)
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
conn.Close()
|
|
c.udpConnections.Delete(msg.ProxyData.ConnId)
|
|
log.Printf("UDP connection closed: %v", msg.ProxyData.ConnId)
|
|
}
|