diff --git a/common/functions/functions.go b/common/functions/functions.go new file mode 100644 index 000000000..76b0d47b9 --- /dev/null +++ b/common/functions/functions.go @@ -0,0 +1,23 @@ +package functions + +import "v2ray.com/core/common" + +// Task is a function that may return an error. +type Task func() error + +// CloseOnSuccess returns a Task to run a follow task if pre-condition passes, otherwise the error in pre-condition is returned. +func CloseOnSuccess(pre func() error, followup Task) Task { + return func() error { + if err := pre(); err != nil { + return err + } + return followup() + } +} + +// Close returns a Task to close the object. +func Close(obj interface{}) Task { + return func() error { + return common.Close(obj) + } +} diff --git a/proxy/freedom/freedom.go b/proxy/freedom/freedom.go index 4f9b498a0..e8b4a2700 100644 --- a/proxy/freedom/freedom.go +++ b/proxy/freedom/freedom.go @@ -10,6 +10,7 @@ import ( "v2ray.com/core/common" "v2ray.com/core/common/buf" "v2ray.com/core/common/dice" + "v2ray.com/core/common/functions" "v2ray.com/core/common/net" "v2ray.com/core/common/retry" "v2ray.com/core/common/signal" @@ -135,7 +136,7 @@ func (h *Handler) Process(ctx context.Context, link *core.Link, dialer proxy.Dia return nil } - if err := signal.ExecuteParallel(ctx, requestDone, responseDone); err != nil { + if err := signal.ExecuteParallel(ctx, requestDone, functions.CloseOnSuccess(responseDone, functions.Close(output))); err != nil { return newError("connection ends").Base(err) } diff --git a/proxy/shadowsocks/client.go b/proxy/shadowsocks/client.go index 57978cfca..fa89de3b8 100644 --- a/proxy/shadowsocks/client.go +++ b/proxy/shadowsocks/client.go @@ -6,6 +6,7 @@ import ( "v2ray.com/core" "v2ray.com/core/common" "v2ray.com/core/common/buf" + "v2ray.com/core/common/functions" "v2ray.com/core/common/net" "v2ray.com/core/common/protocol" "v2ray.com/core/common/retry" @@ -157,7 +158,7 @@ func (c *Client) Process(ctx context.Context, link *core.Link, dialer proxy.Dial return nil } - if err := signal.ExecuteParallel(ctx, requestDone, responseDone); err != nil { + if err := signal.ExecuteParallel(ctx, requestDone, functions.CloseOnSuccess(responseDone, functions.Close(link.Writer))); err != nil { return newError("connection ends").Base(err) } diff --git a/proxy/socks/client.go b/proxy/socks/client.go index e7fa5afdb..02624962d 100644 --- a/proxy/socks/client.go +++ b/proxy/socks/client.go @@ -7,6 +7,7 @@ import ( "v2ray.com/core" "v2ray.com/core/common" "v2ray.com/core/common/buf" + "v2ray.com/core/common/functions" "v2ray.com/core/common/net" "v2ray.com/core/common/protocol" "v2ray.com/core/common/retry" @@ -129,7 +130,7 @@ func (c *Client) Process(ctx context.Context, link *core.Link, dialer proxy.Dial } } - if err := signal.ExecuteParallel(ctx, requestFunc, responseFunc); err != nil { + if err := signal.ExecuteParallel(ctx, requestFunc, functions.CloseOnSuccess(responseFunc, functions.Close(link.Writer))); err != nil { return newError("connection ends").Base(err) } diff --git a/proxy/vmess/outbound/outbound.go b/proxy/vmess/outbound/outbound.go index e9a52b119..1a6c4f21f 100644 --- a/proxy/vmess/outbound/outbound.go +++ b/proxy/vmess/outbound/outbound.go @@ -11,6 +11,7 @@ import ( "v2ray.com/core" "v2ray.com/core/common" "v2ray.com/core/common/buf" + "v2ray.com/core/common/functions" "v2ray.com/core/common/net" "v2ray.com/core/common/protocol" "v2ray.com/core/common/retry" @@ -159,7 +160,7 @@ func (v *Handler) Process(ctx context.Context, link *core.Link, dialer proxy.Dia return buf.Copy(bodyReader, output, buf.UpdateActivity(timer)) } - if err := signal.ExecuteParallel(ctx, requestDone, responseDone); err != nil { + if err := signal.ExecuteParallel(ctx, requestDone, functions.CloseOnSuccess(responseDone, functions.Close(output))); err != nil { return newError("connection ends").Base(err) }