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: kata-containers#300

Signed-off-by: Peng Tao <bergwolf@gmail.com>

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
 	}
 }