net/rpc: shut down connection if gob has error

robpike · robpike · commit cd5b785efe3c · 2014-10-01T13:18:44.000-07:00
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 golang#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
@@ -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)
 }
 
