From d9f79b0ecc76a0ea465da1c554ae7399c53c1113 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 10 Dec 2013 16:49:14 -0800
Subject: [PATCH] packer/rpc: hard close all streams when underlying conn
 closes

---
 packer/rpc/muxconn.go      | 13 ++++++++-----
 packer/rpc/muxconn_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/packer/rpc/muxconn.go b/packer/rpc/muxconn.go
index 08c28bb25..34b5988c3 100644
--- a/packer/rpc/muxconn.go
+++ b/packer/rpc/muxconn.go
@@ -234,7 +234,13 @@ func (m *MuxConn) openStream(id uint32) (*Stream, error) {
 }
 
 func (m *MuxConn) loop() {
-	defer m.Close()
+	defer func() {
+		m.mu.Lock()
+		defer m.mu.Unlock()
+		for _, w := range m.streams {
+			w.remoteClose()
+		}
+	}()
 
 	var id uint32
 	var packetType muxPacketType
@@ -384,16 +390,13 @@ func (s *Stream) Close() error {
 		return fmt.Errorf("Stream in bad state: %d", s.state)
 	}
 
-	if _, err := s.mux.write(s.id, muxPacketFin, nil); err != nil {
-		return err
-	}
-
 	if s.state == streamStateEstablished {
 		s.setState(streamStateFinWait1)
 	} else {
 		s.remoteClose()
 	}
 
+	s.mux.write(s.id, muxPacketFin, nil)
 	return nil
 }
 
diff --git a/packer/rpc/muxconn_test.go b/packer/rpc/muxconn_test.go
index 392ee6b5a..f4ba59db1 100644
--- a/packer/rpc/muxconn_test.go
+++ b/packer/rpc/muxconn_test.go
@@ -114,6 +114,32 @@ func TestMuxConn(t *testing.T) {
 	<-doneCh
 }
 
+func TestMuxConn_socketClose(t *testing.T) {
+	client, server := testMux(t)
+	defer client.Close()
+	defer server.Close()
+
+	go func() {
+		_, err := server.Accept(0)
+		if err != nil {
+			t.Fatalf("err: %s", err)
+		}
+
+		server.rwc.Close()
+	}()
+
+	s0, err := client.Dial(0)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	var data [1024]byte
+	_, err = s0.Read(data[:])
+	if err != io.EOF {
+		t.Fatalf("err: %s", err)
+	}
+}
+
 func TestMuxConn_clientClosesStreams(t *testing.T) {
 	client, server := testMux(t)
 	defer client.Close()