Repository: dragonflyoss/Dragonfly2
Pull request #1560 (2022-08-12 07:31:03)

chore: add source error metrics
Signed-off-by: Jim Ma <majinjing3@gmail.com>

<!--- Provide a general summary of your changes in the Title above -->

## Description

<!--- Describe your changes in detail -->

## Related Issue

<!--- This project only accepts pull requests related to open issues -->
<!--- If suggesting a new feature or change, please discuss it in an issue first -->
<!--- If fixing a bug, there should be an issue describing it with steps to reproduce -->
<!--- Please link to the issue here: -->

## Motivation and Context

<!--- Why is this change required? What problem does it solve? -->

## Screenshots (if appropriate)

## Types of changes

<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation Update (if none of the other choices apply)

## Checklist

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->

- [ ] My change requires a change to the documentation.
- [ ] I have updated the documentation accordingly.
- [ ] I have read the **CONTRIBUTING** document.
- [ ] I have added tests to cover my changes.



Patch:
diff --git a/client/config/fixme.go b/client/config/fixme.go
index 1ab8f71d..ffef7459 100644
--- a/client/config/fixme.go
+++ b/client/config/fixme.go
@@ -20,12 +20,14 @@ import (
 	"crypto/tls"
 
 	"github.com/johanbrandhorst/certify"
+
+	"d7y.io/dragonfly/v2/pkg/net/ip"
 )
 
 func GetCertificate(certifyClient *certify.Certify) func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
 	return func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
 		// FIXME peers need pure ip cert, certify checks the ServerName, so workaround here
-		hello.ServerName = "peer"
+		hello.ServerName = ip.IPv4
 		return certifyClient.GetCertificate(hello)
 	}
 }
diff --git a/client/config/peerhost.go b/client/config/peerhost.go
index 1c217634..80e6fea4 100644
--- a/client/config/peerhost.go
+++ b/client/config/peerhost.go
@@ -174,10 +174,25 @@ func ConvertPattern(p string, defaultPattern commonv1.Pattern) commonv1.Pattern
 }
 
 type GlobalSecurityOption struct {
-	AutoIssueCert bool                 `mapstructure:"autoIssueCert" yaml:"autoIssueCert"`
-	CACert        serialize.PEMContent `mapstructure:"caCert" yaml:"caCert"`
-	TLSVerify     bool                 `mapstructure:"tlsVerify" yaml:"tlsVerify"`
-}
+	// AutoIssueCert indicates to issue client certificates for all grpc call
+	// if AutoIssueCert is false, any other option in Security will be ignored
+	AutoIssueCert bool `mapstructure:"autoIssueCert" yaml:"autoIssueCert"`
+	// CACert is the root CA certificate for all grpc tls handshake, it can be path or PEM format string
+	CACert serialize.PEMContent `mapstructure:"caCert" yaml:"caCert"`
+	// TLSPrefer indicates to verify client certificates for grpc ServerHandshake
+	TLSVerify bool `mapstructure:"tlsVerify" yaml:"tlsVerify"`
+	// TLSPolicy controls the grpc shandshake behaviors:
+	// force: both ClientHandshake and ServerHandshake are only support tls
+	// prefer: ServerHandshake supports tls and insecure (non-tls), ClientHandshake will only support tls
+	// default or empty: ServerHandshake supports tls and insecure (non-tls), ClientHandshake will only support insecure (non-tls)
+	TLSPolicy string `mapstructure:"tlsPolicy" yaml:"tlsPolicy"`
+}
+
+const (
+	TLSPolicyForce   = "force"
+	TLSPolicyPrefer  = "prefer"
+	TLSPolicyDefault = "default"
+)
 
 type SchedulerOption struct {
 	// Manager is to get the scheduler configuration remotely.
diff --git a/client/config/peerhost_darwin.go b/client/config/peerhost_darwin.go
index 90fd917e..57c8ca33 100644
--- a/client/config/peerhost_darwin.go
+++ b/client/config/peerhost_darwin.go
@@ -81,7 +81,7 @@ var peerHostConfig = func() *DaemonOption {
 			DownloadGRPC: ListenOption{
 				Security: SecurityOption{
 					Insecure:  true,
-					TLSVerify: true,
+					TLSVerify: false,
 				},
 				UnixListen: &UnixListenOption{},
 			},
@@ -176,5 +176,11 @@ var peerHostConfig = func() *DaemonOption {
 				Duration: time.Minute,
 			},
 		},
+		Security: GlobalSecurityOption{
+			AutoIssueCert: false,
+			CACert:        serialize.PEMContent(""),
+			TLSVerify:     false,
+			TLSPolicy:     TLSPolicyDefault,
+		},
 	}
 }
diff --git a/client/config/peerhost_linux.go b/client/config/peerhost_linux.go
index 8ee1bcf4..89e86296 100644
--- a/client/config/peerhost_linux.go
+++ b/client/config/peerhost_linux.go
@@ -29,6 +29,7 @@ import (
 	"d7y.io/dragonfly/v2/pkg/dfnet"
 	"d7y.io/dragonfly/v2/pkg/net/fqdn"
 	"d7y.io/dragonfly/v2/pkg/net/ip"
+	"d7y.io/dragonfly/v2/pkg/serialize"
 )
 
 var peerHostConfig = func() *DaemonOption {
@@ -81,7 +82,7 @@ var peerHostConfig = func() *DaemonOption {
 			DownloadGRPC: ListenOption{
 				Security: SecurityOption{
 					Insecure:  true,
-					TLSVerify: true,
+					TLSVerify: false,
 				},
 				UnixListen: &UnixListenOption{},
 			},
@@ -175,5 +176,11 @@ var peerHostConfig = func() *DaemonOption {
 				Duration: time.Minute,
 			},
 		},
+		Security: GlobalSecurityOption{
+			AutoIssueCert: false,
+			CACert:        serialize.PEMContent(""),
+			TLSVerify:     false,
+			TLSPolicy:     TLSPolicyDefault,
+		},
 	}
 }
diff --git a/client/daemon/daemon.go b/client/daemon/daemon.go
index 7976ed60..f339e09d 100644
--- a/client/daemon/daemon.go
+++ b/client/daemon/daemon.go
@@ -145,13 +145,15 @@ func New(opt *config.DaemonOption, d dfpath.Dfpath) (Daemon, error) {
 			}
 
 			// issue a certificate to reduce first time delay
-			_, err := certifyClient.GetCertificate(&tls.ClientHelloInfo{
+			cert, err := certifyClient.GetCertificate(&tls.ClientHelloInfo{
 				ServerName: ip.IPv4,
 			})
 			if err != nil {
 				logger.Errorf("issue certificate error: %s", err.Error())
 				return nil, err
 			}
+			logger.Debugf("request cert from manager, common name: %s, issuer: %s",
+				cert.Leaf.Subject.CommonName, cert.Leaf.Issuer.CommonName)
 		}
 
 		// New dynconfig manager client.
@@ -337,6 +339,7 @@ func loadGPRCTLSCredentials(opt config.SecurityOption, certifyClient *certify.Ce
 	}
 
 	opt.TLSConfig.ClientCAs = certPool
+	opt.TLSConfig.RootCAs = certPool
 
 	// Load server's certificate and private key
 	if certifyClient == nil {
@@ -347,7 +350,6 @@ func loadGPRCTLSCredentials(opt config.SecurityOption, certifyClient *certify.Ce
 		opt.TLSConfig.Certificates = []tls.Certificate{serverCert}
 	} else {
 		// enable auto issue certificate
-		opt.TLSConfig.Certificates = nil
 		opt.TLSConfig.GetCertificate = config.GetCertificate(certifyClient)
 		opt.TLSConfig.GetClientCertificate = certifyClient.GetClientCertificate
 	}
@@ -356,7 +358,15 @@ func loadGPRCTLSCredentials(opt config.SecurityOption, certifyClient *certify.Ce
 		opt.TLSConfig.ClientAuth = tls.RequireAndVerifyClientCert
 	}
 
-	return credentials.NewTLS(opt.TLSConfig), nil
+	switch security.TLSPolicy {
+	case config.TLSPolicyDefault, config.TLSPolicyPrefer:
+		return rpc.NewMuxTransportCredentials(opt.TLSConfig,
+			rpc.WithTLSPreferClientHandshake(security.TLSPolicy == config.TLSPolicyPrefer)), nil
+	case config.TLSPolicyForce:
+		return credentials.NewTLS(opt.TLSConfig), nil
+	default:
+		return nil, fmt.Errorf("invalid tlsPolicy: %s", security.TLSPolicy)
+	}
 }
 
 func loadGlobalGPRCTLSCredentials(certifyClient *certify.Certify, security config.GlobalSecurityOption) (credentials.TransportCredentials, error) {
@@ -370,17 +380,26 @@ func loadGlobalGPRCTLSCredentials(certifyClient *certify.Certify, security confi
 		return nil, fmt.Errorf("failed to add global CA's certificate")
 	}
 
-	config := &tls.Config{
+	tlsConfig := &tls.Config{
 		ClientCAs:            certPool,
+		RootCAs:              certPool,
 		GetCertificate:       config.GetCertificate(certifyClient),
 		GetClientCertificate: certifyClient.GetClientCertificate,
 	}
 
 	if security.TLSVerify {
-		config.ClientAuth = tls.RequireAndVerifyClientCert
+		tlsConfig.ClientAuth = tls.RequireAndVerifyClientCert
 	}
 
-	return credentials.NewTLS(config), nil
+	switch security.TLSPolicy {
+	case config.TLSPolicyDefault, config.TLSPolicyPrefer:
+		return rpc.NewMuxTransportCredentials(tlsConfig,
+			rpc.WithTLSPreferClientHandshake(security.TLSPolicy == config.TLSPolicyPrefer)), nil
+	case config.TLSPolicyForce:
+		return credentials.NewTLS(tlsConfig), nil
+	default:
+		return nil, fmt.Errorf("invalid tlsPolicy: %s", security.TLSPolicy)
+	}
 }
 
 func (*clientDaemon) prepareTCPListener(opt config.ListenOption, withTLS bool) (net.Listener, int, error) {
diff --git a/client/daemon/peer/piece_manager.go b/client/daemon/peer/piece_manager.go
index 237797c1..2bbf0920 100644
--- a/client/daemon/peer/piece_manager.go
+++ b/client/daemon/peer/piece_manager.go
@@ -485,7 +485,7 @@ func (pm *pieceManager) downloadKnownLengthSource(ctx context.Context, pt Task,
 
 		pt.ReportPieceResult(request, result, nil)
 		pt.PublishPieceInfo(pieceNum, uint32(result.Size))
-		if supportConcurrent {
+		if supportConcurrent && pieceNum+2 < maxPieceNum {
 			// the time unit of FinishTime and BeginTime is ns
 			speed := float64(pieceSize) / float64((result.FinishTime-result.BeginTime)/1000000)
 			if speed < float64(pm.concurrentOption.ThresholdSpeed) {
diff --git a/client/daemon/rpcserver/rpcserver.go b/client/daemon/rpcserver/rpcserver.go
index f8ade718..c0fada39 100644
--- a/client/daemon/rpcserver/rpcserver.go
+++ b/client/daemon/rpcserver/rpcserver.go
@@ -38,6 +38,8 @@ import (
 	"go.uber.org/atomic"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials"
+	grpcpeer "google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/emptypb"
 
@@ -211,7 +213,28 @@ func (s *server) sendFirstPieceTasks(
 	return sendExistPieces(sync.Context(), log, get, request, sync, sentMap, false)
 }
 
+func printAuthInfo(ctx context.Context) {
+	if peerInfo, ok := grpcpeer.FromContext(ctx); ok {
+		if tlsInfo, ok := peerInfo.AuthInfo.(credentials.TLSInfo); ok {
+			for i, pc := range tlsInfo.State.PeerCertificates {
+				logger.Debugf("peer cert depth %d, issuer: %#v", i, pc.Issuer.CommonName)
+				logger.Debugf("peer cert depth %d, common name: %#v", i, pc.Subject.CommonName)
+				if len(pc.IPAddresses) > 0 {
+					logger.Debugf("peer cert depth %d, ip: %#v", i, pc.IPAddresses)
+				}
+				if len(pc.DNSNames) > 0 {
+					logger.Debugf("peer cert depth %d, dns: %#v", i, pc.DNSNames)
+				}
+			}
+		}
+	}
+}
+
 func (s *server) SyncPieceTasks(sync dfdaemonv1.Daemon_SyncPieceTasksServer) error {
+	if logger.IsDebug() {
+		printAuthInfo(sync.Context())
+	}
+
 	request, err := sync.Recv()
 	if err != nil {
 		logger.Errorf("receive first sync piece tasks request error: %s", err.Error())
diff --git a/client/daemon/rpcserver/seeder.go b/client/daemon/rpcserver/seeder.go
index 7a18ded6..9865375b 100644
--- a/client/daemon/rpcserver/seeder.go
+++ b/client/daemon/rpcserver/seeder.go
@@ -52,6 +52,10 @@ func (s *seeder) SyncPieceTasks(tasksServer cdnsystemv1.Seeder_SyncPieceTasksSer
 }
 
 func (s *seeder) ObtainSeeds(seedRequest *cdnsystemv1.SeedRequest, seedsServer cdnsystemv1.Seeder_ObtainSeedsServer) error {
+	if logger.IsDebug() {
+		printAuthInfo(seedsServer.Context())
+	}
+
 	metrics.SeedPeerConcurrentDownloadGauge.Inc()
 	defer metrics.SeedPeerConcurrentDownloadGauge.Dec()
 	metrics.SeedPeerDownloadCount.Add(1)
diff --git a/internal/dflog/logger.go b/internal/dflog/logger.go
index 27a72abb..d6292bf2 100644
--- a/internal/dflog/logger.go
+++ b/internal/dflog/logger.go
@@ -215,6 +215,10 @@ func (log *SugaredLoggerOnWith) Debug(args ...any) {
 	CoreLogger.Debugw(fmt.Sprint(args...), log.withArgs...)
 }
 
+func (log *SugaredLoggerOnWith) IsDebug() bool {
+	return coreLogLevelEnabler.Enabled(zap.DebugLevel)
+}
+
 func Infof(template string, args ...any) {
 	CoreLogger.Infof(template, args...)
 }
@@ -243,6 +247,14 @@ func Debugf(template string, args ...any) {
 	CoreLogger.Debugf(template, args...)
 }
 
+func Debug(args ...any) {
+	CoreLogger.Debug(args...)
+}
+
+func IsDebug() bool {
+	return coreLogLevelEnabler.Enabled(zap.DebugLevel)
+}
+
 func Fatalf(template string, args ...any) {
 	CoreLogger.Fatalf(template, args...)
 }
diff --git a/pkg/rpc/mux.go b/pkg/rpc/mux.go
new file mode 100644
index 00000000..56ae8d0f
--- /dev/null
+++ b/pkg/rpc/mux.go
@@ -0,0 +1,145 @@
+/*
+ *     Copyright 2022 The Dragonfly Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package rpc
+
+import (
+	"bytes"
+	"context"
+	"crypto/tls"
+	"fmt"
+	"net"
+
+	"github.com/soheilhy/cmux"
+	"google.golang.org/grpc/credentials"
+)
+
+const (
+	// cmux's TLS matcher read at least 3 + 1 bytes
+	tlsRecordPrefix = 4
+)
+
+type muxTransportCredentials struct {
+	credentials credentials.TransportCredentials
+	tlsMatcher  cmux.Matcher
+	tlsPrefer   bool
+}
+
+func WithTLSPreferClientHandshake(prefer bool) func(m *muxTransportCredentials) {
+	return func(m *muxTransportCredentials) {
+		m.tlsPrefer = prefer
+	}
+}
+
+func NewMuxTransportCredentials(c *tls.Config, opts ...func(m *muxTransportCredentials)) credentials.TransportCredentials {
+	m := &muxTransportCredentials{
+		tlsMatcher:  cmux.TLS(),
+		credentials: credentials.NewTLS(c),
+	}
+
+	for _, opt := range opts {
+		opt(m)
+	}
+
+	return m
+}
+
+func (m *muxTransportCredentials) ClientHandshake(ctx context.Context, s string, conn net.Conn) (net.Conn, credentials.AuthInfo, error) {
+	if m.tlsPrefer {
+		return m.credentials.ClientHandshake(ctx, s, conn)
+	}
+	return conn, info{credentials.CommonAuthInfo{SecurityLevel: credentials.NoSecurity}}, nil
+}
+
+func (m *muxTransportCredentials) ServerHandshake(conn net.Conn) (net.Conn, credentials.AuthInfo, error) {
+	var prefix = make([]byte, tlsRecordPrefix)
+
+	n, err := conn.Read(prefix)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	if n != tlsRecordPrefix {
+		_ = conn.Close()
+		return nil, nil, fmt.Errorf("short read handshake")
+	}
+
+	conn = &muxConn{
+		Conn: conn,
+		buf:  prefix,
+	}
+
+	// tls
+	if m.tlsMatcher(bytes.NewReader(prefix)) {
+		return m.credentials.ServerHandshake(conn)
+	}
+
+	// non-tls
+	return conn, info{credentials.CommonAuthInfo{SecurityLevel: credentials.NoSecurity}}, nil
+}
+
+func (m *muxTransportCredentials) Info() credentials.ProtocolInfo {
+	return credentials.ProtocolInfo{
+		ProtocolVersion:  "",
+		SecurityProtocol: "mux",
+		ServerName:       "",
+	}
+}
+
+func (m *muxTransportCredentials) Clone() credentials.TransportCredentials {
+	return &muxTransportCredentials{
+		tlsMatcher:  cmux.TLS(),
+		credentials: m.credentials.Clone(),
+	}
+}
+
+func (m *muxTransportCredentials) OverrideServerName(s string) error {
+	return m.credentials.OverrideServerName(s)
+}
+
+// info contains the auth information for an insecure connection.
+// It implements the AuthInfo interface.
+type info struct {
+	credentials.CommonAuthInfo
+}
+
+// AuthType returns the type of info as a string.
+func (info) AuthType() string {
+	return "insecure"
+}
+
+type muxConn struct {
+	net.Conn
+	buf []byte
+}
+
+func (m *muxConn) Read(b []byte) (int, error) {
+	if len(m.buf) == 0 {
+		return m.Conn.Read(b)
+	}
+
+	wn := copy(b, m.buf)
+	if wn < len(m.buf) {
+		m.buf = m.buf[wn:]
+		return wn, nil
+	}
+
+	m.buf = nil
+	b = b[wn:]
+
+	n, err := m.Conn.Read(b)
+	return n + wn, err
+}
diff --git a/scheduler/metrics/metrics.go b/scheduler/metrics/metrics.go
index fa9bf999..cceb8a87 100644
--- a/scheduler/metrics/metrics.go
+++ b/scheduler/metrics/metrics.go
@@ -142,6 +142,13 @@ var (
 		Help:      "Counter of the number of peer task.",
 	}, []string{"tag", "app", "type"})
 
+	PeerTaskSourceErrorCounter = promauto.NewCounterVec(prometheus.CounterOpts{
+		Namespace: constants.MetricsNamespace,
+		Subsystem: constants.SchedulerMetricsName,
+		Name:      "peer_task_source_error_total",
+		Help:      "Counter of the source error code number of peer task.",
+	}, []string{"tag", "app", "protocol", "code"})
+
 	PeerTaskDownloadDuration = promauto.NewHistogramVec(prometheus.HistogramOpts{
 		Namespace: constants.MetricsNamespace,
 		Subsystem: constants.SchedulerMetricsName,
diff --git a/scheduler/service/service.go b/scheduler/service/service.go
index 37b5958d..60b6d3ac 100644
--- a/scheduler/service/service.go
+++ b/scheduler/service/service.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"net/url"
 	"time"
 
 	"go.opentelemetry.io/otel/trace"
@@ -729,7 +730,7 @@ func (s *Service) handlePeerSuccess(ctx context.Context, peer *resource.Peer) {
 	}
 
 	// If the peer type is tiny and back-to-source,
-	// it need to directly download the tiny file and store the data in task DirectPiece.
+	// it needs to directly download the tiny file and store the data in task DirectPiece.
 	if sizeScope == commonv1.SizeScope_TINY && len(peer.Task.DirectPiece) == 0 {
 		data, err := peer.DownloadTinyFile()
 		if err != nil {
@@ -762,7 +763,7 @@ func (s *Service) handlePeerFail(ctx context.Context, peer *resource.Peer) {
 }
 
 // handleLegacySeedPeer handles seed server's task has left,
-// but did not notify the schduler to leave the task.
+// but did not notify the scheduler to leave the task.
 func (s *Service) handleLegacySeedPeer(ctx context.Context, peer *resource.Peer) {
 	if err := peer.FSM.Event(resource.PeerEventLeave); err != nil {
 		peer.Log.Errorf("peer fsm event failed: %s", err.Error())
@@ -798,7 +799,7 @@ func (s *Service) handleTaskSuccess(ctx context.Context, task *resource.Task, re
 }
 
 // Conditions for the task to switch to the TaskStateSucceeded are:
-// 1. Seed peer downloads the resource falied.
+// 1. Seed peer downloads the resource failed.
 // 2. Dfdaemon back-to-source to download failed.
 func (s *Service) handleTaskFail(ctx context.Context, task *resource.Task, backToSourceErr *errordetailsv1.SourceError, seedPeerErr error) {
 	// If peer back-to-source fails due to an unrecoverable error,
@@ -822,8 +823,19 @@ func (s *Service) handleTaskFail(ctx context.Context, task *resource.Task, backT
 			for _, detail := range st.Details() {
 				switch d := detail.(type) {
 				case *errordetailsv1.SourceError:
+					var proto = "unknown"
+					if u, err := url.Parse(task.URL); err == nil {
+						proto = u.Scheme
+					}
+					// TODO currently, metrics.PeerTaskSourceErrorCounter is only updated for seed peer source error, need update for normal peer
 					if d.Metadata != nil {
 						task.Log.Infof("source error: %d/%s", d.Metadata.StatusCode, d.Metadata.Status)
+						metrics.PeerTaskSourceErrorCounter.WithLabelValues(
+							task.URLMeta.Tag, task.URLMeta.Application, proto, fmt.Sprintf("%d", d.Metadata.StatusCode)).Inc()
+					} else {
+						task.Log.Warn("source error, but no metadata found")
+						metrics.PeerTaskSourceErrorCounter.WithLabelValues(
+							task.URLMeta.Tag, task.URLMeta.Application, proto, "0").Inc()
 					}
 					if !d.Temporary {
 						task.Log.Infof("source error is not temporary, notify other peers task aborted")


Test patch:
diff --git a/pkg/rpc/mux_test.go b/pkg/rpc/mux_test.go
new file mode 100644
index 00000000..9cb9a476
--- /dev/null
+++ b/pkg/rpc/mux_test.go
@@ -0,0 +1,69 @@
+/*
+ *     Copyright 2022 The Dragonfly Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package rpc
+
+import (
+	"bytes"
+	"io/ioutil"
+	"net"
+	"testing"
+
+	testifyassert "github.com/stretchr/testify/assert"
+)
+
+type testConn struct {
+	net.Conn
+	buf *bytes.Buffer
+}
+
+func (conn *testConn) Read(b []byte) (int, error) {
+	return conn.buf.Read(b)
+}
+
+func Test_muxConn(t *testing.T) {
+	assert := testifyassert.New(t)
+
+	testCases := []struct {
+		name    string
+		bufSize int
+		data    []byte
+	}{
+		{
+			name:    "buf size equal data size",
+			bufSize: 4,
+			data:    []byte("hell"),
+		},
+		{
+			name:    "buf size less than data size",
+			bufSize: 4,
+			data:    []byte("hello world"),
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			mc := &muxConn{
+				buf:  tc.data[:tc.bufSize],
+				Conn: &testConn{Conn: nil, buf: bytes.NewBuffer(tc.data[tc.bufSize:])},
+			}
+
+			data, err := ioutil.ReadAll(mc)
+			assert.Nil(err, "read all should ok")
+			assert.Equal(tc.data, data, "data shloud be same")
+		})
+	}
+}