From c89828e1fb0060fd644c166fb68de30498bfb3e0 Mon Sep 17 00:00:00 2001 From: SAPikachu Date: Sun, 27 Jan 2019 20:14:13 +0800 Subject: [PATCH] Fix potential memory leak --- proxy/dokodemo/dokodemo.go | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go index f4e26203c..fca02b712 100644 --- a/proxy/dokodemo/dokodemo.go +++ b/proxy/dokodemo/dokodemo.go @@ -117,7 +117,7 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn in return nil } - var tReader buf.Reader + var tConn net.Conn responseDone := func() error { defer timer.SetTimeout(plcy.Timeouts.UplinkOnly) @@ -136,14 +136,16 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn in sockopt.BindAddress = dest.Address.IP() sockopt.BindPort = uint32(dest.Port) } - tConn, err := internet.DialSystem(ctx, net.DestinationFromAddr(conn.RemoteAddr()), sockopt) + var err error + tConn, err = internet.DialSystem(ctx, net.DestinationFromAddr(conn.RemoteAddr()), sockopt) if err != nil { return err } writer = &buf.SequentialWriter{Writer: tConn} - tReader = buf.NewReader(tConn) + tReader := buf.NewReader(tConn) go func() { defer tConn.Close() + defer common.Close(link.Writer) if err := buf.Copy(tReader, link.Writer, buf.UpdateActivity(timer)); err != nil { newError("failed to transport request (TPROXY conn)").Base(err).WriteToLog() } @@ -151,6 +153,11 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn in } } + defer func() { + if tConn != nil { + tConn.Close() + } + }() if err := buf.Copy(link.Reader, writer, buf.UpdateActivity(timer)); err != nil { return newError("failed to transport response").Base(err) } @@ -161,8 +168,8 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn in if err := task.Run(ctx, task.OnSuccess(requestDone, task.Close(link.Writer)), responseDone); err != nil { common.Interrupt(link.Reader) common.Interrupt(link.Writer) - if tReader != nil { - common.Interrupt(tReader) + if tConn != nil { + tConn.Close() } return newError("connection ends").Base(err) }