client: remove the parameter of 'enableYamux'

For func NewAgentClient(), the parameter of 'enableYamux'
has nothing to do with 'Proxybuiltin' and it can be figured
out by the dialer's address scheme. Thus there is no need to
pass the parameter 'enableYamux'.

Fixes: #659

Signed-off-by: lifupan <lifupan@gmail.com>

protocols/client/client.go

@@ -83,13 +83,13 @@ type dialer func(string, time.Duration) (net.Conn, error)
 //   - hvsock://<path>:<port>. Firecracker implements the virtio-vsock device
 //     model, and mediates communication between AF_UNIX sockets (on the host end)
 //     and AF_VSOCK sockets (on the guest end).
-func NewAgentClient(ctx context.Context, sock string, enableYamux bool) (*AgentClient, error) {
+func NewAgentClient(ctx context.Context, sock string) (*AgentClient, error) {
 	grpcAddr, parsedAddr, err := parse(sock)
 	if err != nil {
 		return nil, err
 	}
 	dialOpts := []grpc.DialOption{grpc.WithInsecure(), grpc.WithBlock()}
-	dialOpts = append(dialOpts, grpc.WithDialer(agentDialer(parsedAddr, enableYamux)))
+	dialOpts = append(dialOpts, grpc.WithDialer(agentDialer(parsedAddr)))
 
 	var tracer opentracing.Tracer
 
@@ -206,8 +206,10 @@ func heartBeat(session *yamux.Session) {
 	}
 }
 
-func agentDialer(addr *url.URL, enableYamux bool) dialer {
+func agentDialer(addr *url.URL) dialer {
 	var d dialer
+	var enableYamux bool
+
 	switch addr.Scheme {
 	case VSockSocketScheme:
 		d = vsockDialer
@@ -217,6 +219,7 @@ func agentDialer(addr *url.URL, enableYamux bool) dialer {
 		fallthrough
 	default:
 		d = unixDialer
+		enableYamux = true
 	}
 
 	if !enableYamux {

protocols/client/client_test.go

@@ -91,13 +91,13 @@ func checkVersion(cli *AgentClient) error {
 	return nil
 }
 
-func agentClientTest(t *testing.T, sock string, success, enableYamux bool, expect string) {
+func agentClientTest(t *testing.T, sock string, success bool, expect string) {
 	dialTimeout := defaultDialTimeout
 	defaultDialTimeout = 1 * time.Second
 	defer func() {
 		defaultDialTimeout = dialTimeout
 	}()
-	cli, err := NewAgentClient(context.Background(), sock, enableYamux)
+	cli, err := NewAgentClient(context.Background(), sock)
 	if success {
 		assert.Nil(t, err, "Failed to create new agent client: %s", err)
 	} else if !success {
@@ -116,17 +116,16 @@ func agentClientTest(t *testing.T, sock string, success, enableYamux bool, expec
 }
 
 func TestNewAgentClient(t *testing.T) {
-	mock, waitCh, err := startMockServer(t, false)
+	mock, waitCh, err := startMockServer(t, true)
 	assert.Nil(t, err, "failed to start mock server: %s", err)
 	defer os.Remove(mockSockAddr)
 
 	cliFunc := func(sock string, success bool, expect string) {
-		agentClientTest(t, sock, success, false, expect)
+		agentClientTest(t, sock, success, expect)
 	}
 
 	// server starts
 	<-waitCh
-	cliFunc(mockSockAddr, true, "")
 	cliFunc(unixMockAddr, true, "")
 	cliFunc(mockBadSchemeAddr, false, "Invalid scheme:")
 	cliFunc(mockBadVsockScheme, false, "Invalid vsock scheme:")
@@ -145,7 +144,7 @@ func TestNewAgentClientWithYamux(t *testing.T) {
 	defer os.Remove(mockSockAddr)
 
 	cliFunc := func(sock string, success bool, expect string) {
-		agentClientTest(t, sock, success, true, expect)
+		agentClientTest(t, sock, success, expect)
 	}
 	// server starts
 	<-waitCh