From 8abefd18bc79c7ad185d10b3222824e64d998b17 Mon Sep 17 00:00:00 2001 From: Peng Tao Date: Wed, 18 Jul 2018 02:34:30 +0800 Subject: [PATCH] protocols/client: close yamux session when closing the stream For yamux builtin client, we should close the yamux session when closing the stream otherwise the session keeps open and can confuse the server. Fixes: #300 Signed-off-by: Peng Tao --- protocols/client/client.go | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/protocols/client/client.go b/protocols/client/client.go index df92d8e510..8a11618820 100644 --- a/protocols/client/client.go +++ b/protocols/client/client.go @@ -37,6 +37,15 @@ type AgentClient struct { conn *grpc.ClientConn } +type yamuxSessionStream struct { + net.Conn + session *yamux.Session +} + +func (y *yamuxSessionStream) Close() error { + return y.session.Close() +} + type dialer func(string, time.Duration) (net.Conn, error) // NewAgentClient creates a new agent gRPC client and handles both unix and vsock addresses. @@ -161,7 +170,12 @@ func agentDialer(addr *url.URL, enableYamux bool) dialer { return nil, err } - return stream, nil + y := &yamuxSessionStream{ + Conn: stream.(net.Conn), + session: session, + } + + return y, nil } }