chore cleanup

This commit is contained in:
Jason
2019-08-12 15:06:14 +08:00
parent e3b1ff47bd
commit 2254731d3b
3 changed files with 12 additions and 5 deletions

View File

@@ -40,7 +40,7 @@ func (h *udpHandler) fetchUDPInput(conn core.UDPConn, pc *net.UDPConn) {
return return
} }
if _, err = conn.WriteFrom(buf[:n], addr); err != nil { if _, err := conn.WriteFrom(buf[:n], addr); err != nil {
log.Warnf("failed to write UDP data to TUN: %v", err) log.Warnf("failed to write UDP data to TUN: %v", err)
return return
} }
@@ -54,8 +54,11 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
log.Errorf("failed to bind udp address") log.Errorf("failed to bind udp address")
return err return err
} }
h.remoteUDPConnMap.Store(conn, pc) h.remoteUDPConnMap.Store(conn, pc)
go h.fetchUDPInput(conn, pc) go h.fetchUDPInput(conn, pc)
log.Infof("new proxy connection for target: %s:%s", target.Network(), target.String()) log.Infof("new proxy connection for target: %s:%s", target.Network(), target.String())
return nil return nil
} }
@@ -80,9 +83,10 @@ func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr
} }
func (h *udpHandler) Close(conn core.UDPConn) { func (h *udpHandler) Close(conn core.UDPConn) {
conn.Close()
if pc, ok := h.remoteUDPConnMap.Load(conn); ok { if pc, ok := h.remoteUDPConnMap.Load(conn); ok {
pc.(*net.UDPConn).Close() pc.(*net.UDPConn).Close()
h.remoteUDPConnMap.Delete(conn) h.remoteUDPConnMap.Delete(conn)
} }
conn.Close()
} }

View File

@@ -45,8 +45,7 @@ func (h *udpHandler) fetchUDPInput(conn core.UDPConn, pc *net.UDPConn) {
return return
} }
_, err = conn.WriteFrom(buf[:n], addr) if _, err := conn.WriteFrom(buf[:n], addr); err != nil {
if err != nil {
log.Warnf("failed to write UDP data to TUN") log.Warnf("failed to write UDP data to TUN")
return return
} }
@@ -61,9 +60,12 @@ func (h *udpHandler) Connect(conn core.UDPConn, target *net.UDPAddr) error {
return err return err
} }
targetAddr, _ := net.ResolveUDPAddr("udp", h.target) targetAddr, _ := net.ResolveUDPAddr("udp", h.target)
h.remoteAddrMap.Store(conn, targetAddr) h.remoteAddrMap.Store(conn, targetAddr)
h.remoteUDPConnMap.Store(conn, pc) h.remoteUDPConnMap.Store(conn, pc)
go h.fetchUDPInput(conn, pc) go h.fetchUDPInput(conn, pc)
log.Infof("new proxy connection for target: %s:%s", target.Network(), target.String()) log.Infof("new proxy connection for target: %s:%s", target.Network(), target.String())
return nil return nil
} }

View File

@@ -233,6 +233,8 @@ func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr
} }
func (h *udpHandler) Close(conn core.UDPConn) { func (h *udpHandler) Close(conn core.UDPConn) {
conn.Close()
if remoteConn, ok := h.remoteConnMap.Load(conn); ok { if remoteConn, ok := h.remoteConnMap.Load(conn); ok {
remoteConn.(net.Conn).Close() remoteConn.(net.Conn).Close()
h.remoteConnMap.Delete(conn) h.remoteConnMap.Delete(conn)
@@ -243,7 +245,6 @@ func (h *udpHandler) Close(conn core.UDPConn) {
h.remotePacketConnMap.Delete(conn) h.remotePacketConnMap.Delete(conn)
} }
conn.Close()
h.remoteAddrMap.Delete(conn) h.remoteAddrMap.Delete(conn)
if h.sessionStater != nil { if h.sessionStater != nil {