fix bug

ls0f · Feb 5, 2017 · 22bbc1b · 22bbc1b
1 parent 2398da5
commit 22bbc1b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 7 deletions.
diff --git a/src/contorl_test.go b/src/contorl_test.go
@@ -89,4 +89,4 @@ func TestControl_DownloadFile(t *testing.T) {
 	md51, _ := MD5sum("/tmp/test.bin")
 	md52, _ := MD5sum("/tmp/test.bin2")
 	assert.Equal(t, md51, md52)
-}
+}
diff --git a/src/message.go b/src/message.go
@@ -96,10 +96,9 @@ func (m *Message) sendOneMessage(w io.Writer) (n int, err error) {
 func (wrap *MessageWrap) ReadOneMessage() (m *Message, err error) {
 	m = &Message{}
 	if n, ok := wrap.rw.(*net.TCPConn); ok {
-		if wrap.disableReadTimeOut {
-			n.SetReadDeadline(time.Time{})
-		} else {
+		if !wrap.disableReadTimeOut {
 			n.SetReadDeadline(time.Now().Add(time.Second * ReadTimeOut))
+			defer n.SetReadDeadline(time.Time{})
 		}
 	}
 	err = m.readOneMessage(wrap.rw)
@@ -123,10 +122,9 @@ func (wrap *MessageWrap) ReadTheSpecialTypeMessage(msgType uint8) (m *Message, e
 
 func (wrap *MessageWrap) SendOneMessage(m *Message) (err error) {
 	if n, ok := wrap.rw.(*net.TCPConn); ok {
-		if wrap.disableWriteTimeOut {
-			n.SetWriteDeadline(time.Time{})
-		} else {
+		if !wrap.disableWriteTimeOut {
 			n.SetWriteDeadline(time.Now().Add(time.Second * WriteTimeOut))
+			defer n.SetWriteDeadline(time.Time{})
 		}
 	}
 	_, err = m.sendOneMessage(wrap.rw)