fixed possible memory leak for UDP stream

This commit is contained in:
yusing 2024-10-14 08:57:46 +08:00
parent 51c6eb4597
commit d3b8cb8cba
2 changed files with 8 additions and 2 deletions

View file

@ -5,6 +5,5 @@ import (
) )
const ( const (
udpBufferSize = 8192
streamStopListenTimeout = 1 * time.Second streamStopListenTimeout = 1 * time.Second
) )

View file

@ -20,6 +20,7 @@ type (
targetAddr *net.UDPAddr targetAddr *net.UDPAddr
} }
UDPConn struct { UDPConn struct {
key string
src *net.UDPConn src *net.UDPConn
dst *net.UDPConn dst *net.UDPConn
U.BidirectionalPipe U.BidirectionalPipe
@ -29,6 +30,8 @@ type (
var NewUDPConnMap = F.NewMap[UDPConnMap] var NewUDPConnMap = F.NewMap[UDPConnMap]
const udpBufferSize = 8192
func NewUDPRoute(base *StreamRoute) StreamImpl { func NewUDPRoute(base *StreamRoute) StreamImpl {
return &UDPRoute{ return &UDPRoute{
StreamRoute: base, StreamRoute: base,
@ -87,6 +90,7 @@ func (route *UDPRoute) Accept() (any, error) {
return nil, err return nil, err
} }
conn = &UDPConn{ conn = &UDPConn{
key,
srcConn, srcConn,
dstConn, dstConn,
U.NewBidirectionalPipe(route.ctx, sourceRWCloser{in, dstConn}, sourceRWCloser{in, srcConn}), U.NewBidirectionalPipe(route.ctx, sourceRWCloser{in, dstConn}, sourceRWCloser{in, srcConn}),
@ -99,7 +103,10 @@ func (route *UDPRoute) Accept() (any, error) {
} }
func (route *UDPRoute) Handle(c any) error { func (route *UDPRoute) Handle(c any) error {
return c.(*UDPConn).Start() conn := c.(*UDPConn)
err := conn.Start()
route.connMap.Delete(conn.key)
return err
} }
func (route *UDPRoute) CloseListeners() { func (route *UDPRoute) CloseListeners() {