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

diff --git a/src/net/rpc/server.go b/src/net/rpc/server.go
index 6b264b46b8e769898c09b08159103a7f449f94af..83728d55a18b4e509a2e9e575ddce41381d2ac16 100644 (file)
--- 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)
 }