ensure the XContentBuilder is always closed in RestBuilderListener

There may be cases where the XContentBuilder is not used and therefore it never gets
closed, which can cause a leak of bytes. This change moves the creation of the builder
into a try with resources block and adds an assertion to verify that we always consume
the bytes in our code; the try-with resources provides protections against memory leaks
caused by plugins, which do not test this.

Author: jaymode

core/src/main/java/org/elasticsearch/common/xcontent/XContentGenerator.java

@@ -94,4 +94,9 @@ public interface XContentGenerator extends Closeable, Flushable {
 
     void copyCurrentStructure(XContentParser parser) throws IOException;
 
+    /**
+     * Returns {@code true} if this XContentGenerator has been closed. A closed generator can not do any more output.
+     */
+    boolean isClosed();
+
 }

core/src/main/java/org/elasticsearch/common/xcontent/json/JsonXContentGenerator.java

@@ -419,4 +419,8 @@ public void close() throws IOException {
         generator.close();
     }
 
+    @Override
+    public boolean isClosed() {
+        return generator.isClosed();
+    }
 }

core/src/main/java/org/elasticsearch/rest/action/RestBuilderListener.java

@@ -34,11 +34,22 @@ public RestBuilderListener(RestChannel channel) {
 
     @Override
     public final RestResponse buildResponse(Response response) throws Exception {
-        return buildResponse(response, channel.newBuilder());
+        try (XContentBuilder builder = channel.newBuilder()) {
+            final RestResponse restResponse = buildResponse(response, builder);
+            assert assertBuilderClosed(builder);
+            return restResponse;
+        }
     }
 
     /**
-     * Builds a response to send back over the channel.
+     * Builds a response to send back over the channel. Implementors should ensure that they close the provided {@link XContentBuilder}
+     * using the {@link XContentBuilder#close()} method.
      */
     public abstract RestResponse buildResponse(Response response, XContentBuilder builder) throws Exception;
+
+    // pkg private method that we can override for testing
+    boolean assertBuilderClosed(XContentBuilder xContentBuilder) {
+        assert xContentBuilder.generator().isClosed() : "callers should ensure the XContentBuilder is closed themselves";
+        return true;
+    }
 }

core/src/test/java/org/elasticsearch/rest/action/RestBuilderListenerTests.java

@@ -0,0 +1,90 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.rest.action;
+
+import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.rest.BytesRestResponse;
+import org.elasticsearch.rest.RestResponse;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.rest.FakeRestChannel;
+import org.elasticsearch.test.rest.FakeRestRequest;
+import org.elasticsearch.transport.TransportResponse;
+import org.elasticsearch.transport.TransportResponse.Empty;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+public class RestBuilderListenerTests extends ESTestCase {
+
+    public void testXContentBuilderClosedInBuildResponse() throws Exception {
+        AtomicReference<XContentBuilder> builderAtomicReference = new AtomicReference<>();
+        RestBuilderListener<TransportResponse.Empty> builderListener =
+            new RestBuilderListener<Empty>(new FakeRestChannel(new FakeRestRequest(), randomBoolean(), 1)) {
+                @Override
+                public RestResponse buildResponse(Empty empty, XContentBuilder builder) throws Exception {
+                    builderAtomicReference.set(builder);
+                    builder.close();
+                    return new BytesRestResponse(RestStatus.OK, BytesRestResponse.TEXT_CONTENT_TYPE, BytesArray.EMPTY);
+                }
+        };
+
+        builderListener.buildResponse(Empty.INSTANCE);
+        assertNotNull(builderAtomicReference.get());
+        assertTrue(builderAtomicReference.get().generator().isClosed());
+    }
+
+    public void testXContentBuilderNotClosedInBuildResponseAssertionsDisabled() throws Exception {
+        AtomicReference<XContentBuilder> builderAtomicReference = new AtomicReference<>();
+        RestBuilderListener<TransportResponse.Empty> builderListener =
+            new RestBuilderListener<Empty>(new FakeRestChannel(new FakeRestRequest(), randomBoolean(), 1)) {
+                @Override
+                public RestResponse buildResponse(Empty empty, XContentBuilder builder) throws Exception {
+                    builderAtomicReference.set(builder);
+                    return new BytesRestResponse(RestStatus.OK, BytesRestResponse.TEXT_CONTENT_TYPE, BytesArray.EMPTY);
+                }
+
+                @Override
+                boolean assertBuilderClosed(XContentBuilder xContentBuilder) {
+                    // don't check the actual builder being closed so we can test auto close
+                    return true;
+                }
+        };
+
+        builderListener.buildResponse(Empty.INSTANCE);
+        assertNotNull(builderAtomicReference.get());
+        assertTrue(builderAtomicReference.get().generator().isClosed());
+    }
+
+    public void testXContentBuilderNotClosedInBuildResponseAssertionsEnabled() throws Exception {
+        assumeTrue("tests are not being run with assertions", RestBuilderListener.class.desiredAssertionStatus());
+
+        RestBuilderListener<TransportResponse.Empty> builderListener =
+            new RestBuilderListener<Empty>(new FakeRestChannel(new FakeRestRequest(), randomBoolean(), 1)) {
+                @Override
+                public RestResponse buildResponse(Empty empty, XContentBuilder builder) throws Exception {
+                    return new BytesRestResponse(RestStatus.OK, BytesRestResponse.TEXT_CONTENT_TYPE, BytesArray.EMPTY);
+                }
+        };
+
+        AssertionError error = expectThrows(AssertionError.class, () -> builderListener.buildResponse(Empty.INSTANCE));
+        assertEquals("callers should ensure the XContentBuilder is closed themselves", error.getMessage());
+    }
+}