net/rpc: shut down connection if gob has error

The nicest solution would be to buffer the message and only write
it if it encodes correctly, but that adds considerable memory and
CPU overhead for a very rare condition. Instead, we just shut
down the connection if this happens.
Fixes #7689.

LGTM=rsc
R=rsc
CC=golang-codereviews
https://golang.org/cl/146670043

1 files changed, 24 insertions, 1 deletions

diff --git a/src/net/rpc/server.go b/src/net/rpc/server.go
index 6b264b46b8..83728d55a1 100644
--- a/src/net/rpc/server.go
+++ b/src/net/rpc/server.go
@@ -395,6 +395,7 @@ type gobServerCodec struct {
 	dec    *gob.Decoder
 	enc    *gob.Encoder
 	encBuf *bufio.Writer
+	closed bool
 }
 
 func (c *gobServerCodec) ReadRequestHeader(r *Request) error {
@@ -407,15 +408,32 @@ func (c *gobServerCodec) ReadRequestBody(body interface{}) error {
 
 func (c *gobServerCodec) WriteResponse(r *Response, body interface{}) (err error) {
 	if err = c.enc.Encode(r); err != nil {
+		if c.encBuf.Flush() == nil {
+			// Gob couldn't encode the header. Should not happen, so if it does,
+			// shut down the connection to signal that the connection is broken.
+			log.Println("rpc: gob error encoding response:", err)
+			c.Close()
+		}
 		return
 	}
 	if err = c.enc.Encode(body); err != nil {
+		if c.encBuf.Flush() == nil {
+			// Was a gob problem encoding the body but the header has been written.
+			// Shut down the connection to signal that the connection is broken.
+			log.Println("rpc: gob error encoding body:", err)
+			c.Close()
+		}
 		return
 	}
 	return c.encBuf.Flush()
 }
 
 func (c *gobServerCodec) Close() error {
+	if c.closed {
+		// Only call c.rwc.Close once; otherwise the semantics are undefined.
+		return nil
+	}
+	c.closed = true
 	return c.rwc.Close()
 }
 
@@ -426,7 +444,12 @@ func (c *gobServerCodec) Close() error {
 // connection.  To use an alternate codec, use ServeCodec.
 func (server *Server) ServeConn(conn io.ReadWriteCloser) {
 	buf := bufio.NewWriter(conn)
-	srv := &gobServerCodec{conn, gob.NewDecoder(conn), gob.NewEncoder(buf), buf}
+	srv := &gobServerCodec{
+		rwc:    conn,
+		dec:    gob.NewDecoder(conn),
+		enc:    gob.NewEncoder(buf),
+		encBuf: buf,
+	}
 	server.ServeCodec(srv)
 }