github
/
nps
mirror of https://github.com/ehang-io/nps
1
0
Fork 0
Code Issues Releases Wiki Activity

fix:udp buf

pull/413/head
cnlh 2020-02-17 02:30:34 +08:00
parent 1fec6c3649
commit 05c27b9365
1 changed files with 39 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

59
server/proxy/udp.go
View File

@ -1,8 +1,11 @@
package proxy package proxy
import ( import (
"io"
"net" "net"
"strings" "strings"
"sync"
"time"
"ehang.io/nps/bridge" "ehang.io/nps/bridge"
"ehang.io/nps/lib/common" "ehang.io/nps/lib/common"
@ -13,6 +16,7 @@ import (
type UdpModeServer struct { type UdpModeServer struct {
BaseServer BaseServer
addrMap sync.Map
listener *net.UDPConn listener *net.UDPConn
} }
@ -33,8 +37,8 @@ func (s *UdpModeServer) Start() error {
if err != nil { if err != nil {
return err return err
} }
buf := common.BufPoolUdp.Get().([]byte)
for { for {
buf := common.BufPoolUdp.Get().([]byte)
n, addr, err := s.listener.ReadFromUDP(buf) n, addr, err := s.listener.ReadFromUDP(buf)
if err != nil { if err != nil {
if strings.Contains(err.Error(), "use of closed network connection") { if strings.Contains(err.Error(), "use of closed network connection") {
@ -49,28 +53,43 @@ func (s *UdpModeServer) Start() error {
} }
func (s *UdpModeServer) process(addr *net.UDPAddr, data []byte) { func (s *UdpModeServer) process(addr *net.UDPAddr, data []byte) {
if err := s.CheckFlowAndConnNum(s.task.Client); err != nil { if v, ok := s.addrMap.Load(addr.String()); ok {
logs.Warn("client id %d, task id %d,error %s, when udp connection", s.task.Client.Id, s.task.Id, err.Error()) clientConn, ok := v.(io.ReadWriteCloser)
return if ok {
} clientConn.Write(data)
defer s.task.Client.AddConn() s.task.Flow.Add(int64(len(data)), 0)
link := conn.NewLink(common.CONN_UDP, s.task.Target.TargetStr, s.task.Client.Cnf.Crypt, s.task.Client.Cnf.Compress, addr.String(), s.task.Target.LocalProxy) }
if clientConn, err := s.bridge.SendLinkInfo(s.task.Client.Id, link, s.task); err != nil {
return
} else { } else {
target := conn.GetConn(clientConn, s.task.Client.Cnf.Crypt, s.task.Client.Cnf.Compress, nil, true) if err := s.CheckFlowAndConnNum(s.task.Client); err != nil {
defer target.Close() logs.Warn("client id %d, task id %d,error %s, when udp connection", s.task.Client.Id, s.task.Id, err.Error())
s.task.Flow.Add(int64(len(data)), 0) return
buf := common.BufPoolUdp.Get().([]byte) }
defer common.BufPoolUdp.Put(buf) defer s.task.Client.AddConn()
target.Write(data) link := conn.NewLink(common.CONN_UDP, s.task.Target.TargetStr, s.task.Client.Cnf.Crypt, s.task.Client.Cnf.Compress, addr.String(), s.task.Target.LocalProxy)
s.task.Flow.Add(int64(len(data)), 0) if clientConn, err := s.bridge.SendLinkInfo(s.task.Client.Id, link, s.task); err != nil {
if n, err := target.Read(buf); err != nil {
logs.Warn(err)
return return
} else { } else {
s.listener.WriteTo(buf[:n], addr) target := conn.GetConn(clientConn, s.task.Client.Cnf.Crypt, s.task.Client.Cnf.Compress, nil, true)
s.task.Flow.Add(0, int64(n)) s.addrMap.Store(addr.String(), target)
defer target.Close()
target.Write(data)
buf := common.BufPoolUdp.Get().([]byte)
defer common.BufPoolUdp.Put(buf)
s.task.Flow.Add(int64(len(data)), 0)
for {
clientConn.SetReadDeadline(time.Now().Add(time.Minute * 10))
if n, err := target.Read(buf); err != nil {
s.addrMap.Delete(addr.String())
logs.Warn(err)
return
} else {
s.listener.WriteTo(buf[:n], addr)
s.task.Flow.Add(0, int64(n))
}
}
} }
} }
} }