Promote headers to first class citizens on exceptions

This commit merges the pre-existing special exception that
allowed to associate headers with exceptions and the elasticsaerch
base class `ElasticsearchException` This allows for more generic use
of exceptions where plugins can associate meta-data with any elasticsearch
base exception to control behavior etc.

This also addds a generic SecurityException to allow plugins to pass on
information based on the RestStatus.

10 files changed, 168 insertions, 167 deletions

diff --git a/core/src/main/java/org/elasticsearch/ElasticsearchException.java b/core/src/main/java/org/elasticsearch/ElasticsearchException.java
index 9f3dc6f8b7..1fa95d6b1b 100644
--- a/core/src/main/java/org/elasticsearch/ElasticsearchException.java
+++ b/core/src/main/java/org/elasticsearch/ElasticsearchException.java
@@ -19,18 +19,13 @@
 
 package org.elasticsearch;
 
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.stream.NotSerializableExceptionWrapper;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.logging.support.LoggerMessageFormat;
 import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
-import org.elasticsearch.rest.HasRestHeaders;
 import org.elasticsearch.rest.RestStatus;
 
 import java.io.IOException;
@@ -45,6 +40,7 @@ public class ElasticsearchException extends RuntimeException implements ToXConte
 
     public static final String REST_EXCEPTION_SKIP_CAUSE = "rest.exception.skip_cause";
     private static final Map<String, Constructor<? extends ElasticsearchException>> MAPPING;
+    private final Map<String, List<String>> headers = new HashMap<>();
 
     /**
      * Construct a <code>ElasticsearchException</code> with the specified detail message.
@@ -77,6 +73,48 @@ public class ElasticsearchException extends RuntimeException implements ToXConte
     public ElasticsearchException(StreamInput in) throws IOException {
         super(in.readOptionalString(), in.readThrowable());
         readStackTrace(this, in);
+        int numKeys = in.readVInt();
+        for (int i = 0; i < numKeys; i++) {
+            final String key = in.readString();
+            final int numValues = in.readVInt();
+            final ArrayList<String> values = new ArrayList<>(numValues);
+            for (int j = 0; j < numValues; j++) {
+                values.add(in.readString());
+            }
+            headers.put(key, values);
+        }
+    }
+
+    /**
+     * Adds a new header with the given key.
+     * This method will replace existing header if a header with the same key already exists
+     */
+    public void addHeader(String key, String... value) {
+        this.headers.put(key, Arrays.asList(value));
+    }
+
+    /**
+     * Adds a new header with the given key.
+     * This method will replace existing header if a header with the same key already exists
+     */
+    public void addHeader(String key, List<String> value) {
+        this.headers.put(key, value);
+    }
+
+
+    /**
+     * Returns a set of all header keys on this exception
+     */
+    public Set<String> getHeaderKeys() {
+        return headers.keySet();
+    }
+
+    /**
+     * Returns the list of header values for the given key or {@code null} if not header for the
+     * given key exists.
+     */
+    public List<String> getHeader(String key) {
+        return headers.get(key);
     }
 
     /**
@@ -173,6 +211,14 @@ public class ElasticsearchException extends RuntimeException implements ToXConte
         out.writeOptionalString(this.getMessage());
         out.writeThrowable(this.getCause());
         writeStackTraces(this, out);
+        out.writeVInt(headers.size());
+        for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
+            out.writeString(entry.getKey());
+            out.writeVInt(entry.getValue().size());
+            for (String v : entry.getValue()) {
+                out.writeString(v);
+            }
+        }
     }
 
     public static ElasticsearchException readException(StreamInput input, String name) throws IOException {
@@ -198,79 +244,6 @@ public class ElasticsearchException extends RuntimeException implements ToXConte
         return MAPPING.keySet();
     }
 
-    /**
-     * A base class for exceptions that should carry rest headers
-     */
-    @SuppressWarnings("unchecked")
-    public static class WithRestHeadersException extends ElasticsearchException implements HasRestHeaders {
-
-        private final Map<String, List<String>> headers;
-
-        public WithRestHeadersException(String msg, Tuple<String, String[]>... headers) {
-            super(msg);
-            this.headers = headers(headers);
-        }
-
-        protected WithRestHeadersException(String msg, Throwable cause, Map<String, List<String>> headers) {
-            super(msg, cause);
-            this.headers = headers;
-        }
-
-        public WithRestHeadersException(StreamInput in) throws IOException {
-            super(in);
-            int numKeys = in.readVInt();
-            ImmutableMap.Builder<String, List<String>> builder = ImmutableMap.builder();
-            for (int i = 0; i < numKeys; i++) {
-                final String key = in.readString();
-                final int numValues = in.readVInt();
-                final ArrayList<String> headers = new ArrayList<>(numValues);
-                for (int j = 0; j < numValues; j++) {
-                    headers.add(in.readString());
-                }
-                builder.put(key, headers);
-            }
-            headers = builder.build();
-        }
-
-        @Override
-        public void writeTo(StreamOutput out) throws IOException {
-            super.writeTo(out);
-            out.writeVInt(headers.size());
-            for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
-                out.writeString(entry.getKey());
-                out.writeVInt(entry.getValue().size());
-                for (String v : entry.getValue()) {
-                    out.writeString(v);
-                }
-            }
-        }
-
-        @Override
-        public Map<String, List<String>> getHeaders() {
-            return headers;
-        }
-
-        protected static Tuple<String, String[]> header(String name, String... values) {
-            return Tuple.tuple(name, values);
-        }
-
-        private static Map<String, List<String>> headers(Tuple<String, String[]>... headers) {
-            Map<String, List<String>> map = Maps.newHashMap();
-            for (Tuple<String, String[]> header : headers) {
-                List<String> list = map.get(header.v1());
-                if (list == null) {
-                    list = Lists.newArrayList(header.v2());
-                    map.put(header.v1(), list);
-                } else {
-                    for (String value : header.v2()) {
-                        list.add(value);
-                    }
-                }
-            }
-            return ImmutableMap.copyOf(map);
-        }
-    }
-
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         if (this instanceof ElasticsearchWrapperException) {
@@ -559,7 +532,7 @@ public class ElasticsearchException extends RuntimeException implements ToXConte
                 org.elasticsearch.action.PrimaryMissingActionException.class,
                 org.elasticsearch.index.engine.CreateFailedEngineException.class,
                 org.elasticsearch.index.shard.IllegalIndexShardStateException.class,
-                WithRestHeadersException.class,
+                ElasticsearchSecurityException.class,
                 NotSerializableExceptionWrapper.class
         };
         Map<String, Constructor<? extends ElasticsearchException>> mapping = new HashMap<>(exceptions.length);
diff --git a/core/src/main/java/org/elasticsearch/ElasticsearchSecurityException.java b/core/src/main/java/org/elasticsearch/ElasticsearchSecurityException.java
new file mode 100644
index 0000000000..f4878fe6f9
--- /dev/null
+++ b/core/src/main/java/org/elasticsearch/ElasticsearchSecurityException.java
@@ -0,0 +1,66 @@
+/*
+ * 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;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.rest.RestStatus;
+
+import java.io.IOException;
+
+/**
+ * Generic security exception
+ */
+public class ElasticsearchSecurityException extends ElasticsearchException {
+
+    private final RestStatus status;
+
+    public ElasticsearchSecurityException(String msg, RestStatus status, Throwable cause, Object... args) {
+        super(msg, cause, args);
+        this.status = status ;
+    }
+
+    public ElasticsearchSecurityException(String msg, Throwable cause, Object... args) {
+        this(msg, ExceptionsHelper.status(cause), cause, args);
+    }
+
+    public ElasticsearchSecurityException(String msg, Object... args) {
+        this(msg, RestStatus.INTERNAL_SERVER_ERROR, null, args);
+    }
+
+    public ElasticsearchSecurityException(String msg, RestStatus status, Object... args) {
+        this(msg, status, null, args);
+    }
+
+    public ElasticsearchSecurityException(StreamInput in) throws IOException {
+        super(in);
+        status = RestStatus.readFrom(in);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        RestStatus.writeTo(out, status);
+    }
+
+    @Override
+    public final RestStatus status() {
+        return status;
+    }
+}
diff --git a/core/src/main/java/org/elasticsearch/common/io/stream/NotSerializableExceptionWrapper.java b/core/src/main/java/org/elasticsearch/common/io/stream/NotSerializableExceptionWrapper.java
index 001d8abddf..8252fb6d97 100644
--- a/core/src/main/java/org/elasticsearch/common/io/stream/NotSerializableExceptionWrapper.java
+++ b/core/src/main/java/org/elasticsearch/common/io/stream/NotSerializableExceptionWrapper.java
@@ -21,13 +21,9 @@ package org.elasticsearch.common.io.stream;
 
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
-import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.rest.RestStatus;
 
 import java.io.IOException;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
 
 /**
  * This exception can be used to wrap a given, not serializable exception
@@ -36,27 +32,25 @@ import java.util.Map;
  * the throwable it was created with instead of it's own. The stacktrace has no indication
  * of where this exception was created.
  */
-public final class NotSerializableExceptionWrapper extends ElasticsearchException.WithRestHeadersException {
+public final class NotSerializableExceptionWrapper extends ElasticsearchException {
 
     private final String name;
     private final RestStatus status;
 
-    public NotSerializableExceptionWrapper(Throwable other, Map<String, List<String>> headers) {
-        super(other.getMessage(), other.getCause(), headers);
+    public NotSerializableExceptionWrapper(Throwable other) {
+        super(other.getMessage(), other.getCause());
         this.name = ElasticsearchException.getExceptionName(other);
         this.status = ExceptionsHelper.status(other);
         setStackTrace(other.getStackTrace());
         for (Throwable otherSuppressed : other.getSuppressed()) {
             addSuppressed(otherSuppressed);
         }
-    }
-
-    public NotSerializableExceptionWrapper(WithRestHeadersException other) {
-        this(other, other.getHeaders());
-    }
-
-    public NotSerializableExceptionWrapper(Throwable other) {
-        this(other, Collections.EMPTY_MAP);
+        if (other instanceof ElasticsearchException) {
+            ElasticsearchException ex = (ElasticsearchException) other;
+            for (String key : ex.getHeaderKeys()) {
+                this.addHeader(key, ex.getHeader(key));
+            }
+        }
     }
 
     public NotSerializableExceptionWrapper(StreamInput in) throws IOException {
diff --git a/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java b/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
index 75b0d1de5b..6670e36fc6 100644
--- a/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
+++ b/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
@@ -511,9 +511,6 @@ public abstract class StreamOutput extends OutputStream {
                 final String name = throwable.getClass().getName();
                 if (throwable instanceof ElasticsearchException && ElasticsearchException.isRegistered(name)) {
                     ex = (ElasticsearchException) throwable;
-                } else if (throwable instanceof ElasticsearchException.WithRestHeadersException) {
-                    // ensure we transport also the headers
-                    ex = new NotSerializableExceptionWrapper((ElasticsearchException.WithRestHeadersException)throwable);
                 } else {
                     ex = new NotSerializableExceptionWrapper(throwable);
                 }
diff --git a/core/src/main/java/org/elasticsearch/rest/BytesRestResponse.java b/core/src/main/java/org/elasticsearch/rest/BytesRestResponse.java
index 02820158a6..b64c9c5873 100644
--- a/core/src/main/java/org/elasticsearch/rest/BytesRestResponse.java
+++ b/core/src/main/java/org/elasticsearch/rest/BytesRestResponse.java
@@ -21,6 +21,7 @@ package org.elasticsearch.rest;
 
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.bootstrap.Elasticsearch;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.collect.Tuple;
@@ -94,8 +95,8 @@ public class BytesRestResponse extends RestResponse {
             this.content = builder.bytes();
             this.contentType = builder.contentType().restContentType();
         }
-        if (t instanceof HasRestHeaders) {
-            addHeaders(((HasRestHeaders) t).getHeaders());
+        if (t instanceof ElasticsearchException) {
+            copyHeaders(((ElasticsearchException) t));
         }
     }
 
diff --git a/core/src/main/java/org/elasticsearch/rest/HasRestHeaders.java b/core/src/main/java/org/elasticsearch/rest/HasRestHeaders.java
deleted file mode 100644
index 2eaeb89dee..0000000000
--- a/core/src/main/java/org/elasticsearch/rest/HasRestHeaders.java
+++ /dev/null
@@ -1,40 +0,0 @@
-/*
- * 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;
-
-import org.elasticsearch.ElasticsearchException;
-
-import java.util.List;
-import java.util.Map;
-
-/**
- * Classes that carry rest headers should implement this interface. Specifically, exceptions that
- * get translated to a rest response, can implement this interface and the headers will be added
- * the the response.
- *
- * @see ElasticsearchException.WithRestHeadersException
- */
-public interface HasRestHeaders {
-
-    /**
-     * @return The rest headers
-     */
-    Map<String, List<String>> getHeaders();
-}
diff --git a/core/src/main/java/org/elasticsearch/rest/RestResponse.java b/core/src/main/java/org/elasticsearch/rest/RestResponse.java
index 72da75a0fc..80ad0e16a3 100644
--- a/core/src/main/java/org/elasticsearch/rest/RestResponse.java
+++ b/core/src/main/java/org/elasticsearch/rest/RestResponse.java
@@ -20,19 +20,16 @@
 package org.elasticsearch.rest;
 
 import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.Nullable;
 import org.elasticsearch.common.bytes.BytesReference;
 
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
 /**
  *
  */
-public abstract class RestResponse implements HasRestHeaders {
+public abstract class RestResponse {
 
     protected Map<String, List<String>> customHeaders;
 
@@ -53,17 +50,18 @@ public abstract class RestResponse implements HasRestHeaders {
      */
     public abstract RestStatus status();
 
-    public void addHeaders(Map<String, List<String>> headers) {
+    public void copyHeaders(ElasticsearchException ex) {
+        Set<String> headerKeySet = ex.getHeaderKeys();
         if (customHeaders == null) {
-            customHeaders = new HashMap<>(headers.size());
+            customHeaders = new HashMap<>(headerKeySet.size());
         }
-        for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
-            List<String> values = customHeaders.get(entry.getKey());
+        for (String key : headerKeySet) {
+            List<String> values = customHeaders.get(key);
             if (values == null) {
                 values = Lists.newArrayList();
-                customHeaders.put(entry.getKey(), values);
+                customHeaders.put(key, values);
             }
-            values.addAll(entry.getValue());
+            values.addAll(ex.getHeader(key));
         }
     }
 
@@ -85,7 +83,6 @@ public abstract class RestResponse implements HasRestHeaders {
     /**
      * Returns custom headers that have been added, or null if none have been set.
      */
-    @Override
     @Nullable
     public Map<String, List<String>> getHeaders() {
         return customHeaders;
diff --git a/core/src/test/java/org/elasticsearch/ElasticsearchExceptionTests.java b/core/src/test/java/org/elasticsearch/ElasticsearchExceptionTests.java
index 699c3ca616..34bc1178b0 100644
--- a/core/src/test/java/org/elasticsearch/ElasticsearchExceptionTests.java
+++ b/core/src/test/java/org/elasticsearch/ElasticsearchExceptionTests.java
@@ -265,7 +265,7 @@ public class ElasticsearchExceptionTests extends ElasticsearchTestCase {
                 new IllegalArgumentException("alalaal"),
                 new NullPointerException("boom"),
                 new EOFException("dadada"),
-                new SecurityException("nono!"),
+                new ElasticsearchSecurityException("nono!"),
                 new NumberFormatException("not a number"),
                 new CorruptIndexException("baaaam", "this is my resource"),
                 new IndexFormatTooNewException("tooo new", 1, 1, 1),
diff --git a/core/src/test/java/org/elasticsearch/ExceptionSerializationTests.java b/core/src/test/java/org/elasticsearch/ExceptionSerializationTests.java
index 1ff70f0cf4..0c2c838e61 100644
--- a/core/src/test/java/org/elasticsearch/ExceptionSerializationTests.java
+++ b/core/src/test/java/org/elasticsearch/ExceptionSerializationTests.java
@@ -33,7 +33,6 @@ import org.elasticsearch.cluster.metadata.SnapshotId;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.*;
 import org.elasticsearch.common.breaker.CircuitBreakingException;
-import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.PathUtils;
 import org.elasticsearch.common.io.stream.*;
 import org.elasticsearch.common.transport.LocalTransportAddress;
@@ -583,30 +582,35 @@ public class ExceptionSerializationTests extends ElasticsearchTestCase {
     }
 
     public void testWithRestHeadersException() throws IOException {
-        ElasticsearchException.WithRestHeadersException ex = serialize(new ElasticsearchException.WithRestHeadersException("msg", new Tuple("foo", new String[]{"foo", "bar"})));
+        ElasticsearchException ex = new ElasticsearchException("msg");
+        ex.addHeader("foo", "foo", "bar");
+        ex = serialize(ex);
         assertEquals("msg", ex.getMessage());
-        assertEquals(2, ex.getHeaders().get("foo").size());
-        assertEquals("foo", ex.getHeaders().get("foo").get(0));
-        assertEquals("bar", ex.getHeaders().get("foo").get(1));
+        assertEquals(2, ex.getHeader("foo").size());
+        assertEquals("foo", ex.getHeader("foo").get(0));
+        assertEquals("bar", ex.getHeader("foo").get(1));
 
         RestStatus status = randomFrom(RestStatus.values());
         // ensure we are carrying over the headers even if not serialized
-        ElasticsearchException serialize = serialize((ElasticsearchException) new UnknownHeaderException("msg", status, new Tuple("foo", new String[]{"foo", "bar"})));
+        UnknownHeaderException uhe = new UnknownHeaderException("msg", status);
+        uhe.addHeader("foo", "foo", "bar");
+
+        ElasticsearchException serialize = serialize((ElasticsearchException)uhe);
         assertTrue(serialize instanceof NotSerializableExceptionWrapper);
         NotSerializableExceptionWrapper e = (NotSerializableExceptionWrapper) serialize;
         assertEquals("msg", e.getMessage());
-        assertEquals(2, e.getHeaders().get("foo").size());
-        assertEquals("foo", e.getHeaders().get("foo").get(0));
-        assertEquals("bar", e.getHeaders().get("foo").get(1));
+        assertEquals(2, e.getHeader("foo").size());
+        assertEquals("foo", e.getHeader("foo").get(0));
+        assertEquals("bar", e.getHeader("foo").get(1));
         assertSame(status, e.status());
 
     }
 
-    public static class UnknownHeaderException extends ElasticsearchException.WithRestHeadersException {
+    public static class UnknownHeaderException extends ElasticsearchException {
         private final RestStatus status;
 
-        public UnknownHeaderException(String msg, RestStatus status, Tuple<String, String[]>... headers) {
-            super(msg, headers);
+        public UnknownHeaderException(String msg, RestStatus status) {
+            super(msg);
             this.status = status;
         }
 
@@ -615,4 +619,11 @@ public class ExceptionSerializationTests extends ElasticsearchTestCase {
             return status;
         }
     }
+
+    public void testElasticsearchSecurityException() throws IOException {
+        ElasticsearchSecurityException ex = new ElasticsearchSecurityException("user [{}] is not allowed", RestStatus.UNAUTHORIZED, "foo");
+        ElasticsearchSecurityException e = serialize(ex);
+        assertEquals(ex.status(), e.status());
+        assertEquals(RestStatus.UNAUTHORIZED, e.status());
+    }
 }
diff --git a/core/src/test/java/org/elasticsearch/rest/BytesRestResponseTests.java b/core/src/test/java/org/elasticsearch/rest/BytesRestResponseTests.java
index f60b178c40..478d77b663 100644
--- a/core/src/test/java/org/elasticsearch/rest/BytesRestResponseTests.java
+++ b/core/src/test/java/org/elasticsearch/rest/BytesRestResponseTests.java
@@ -152,10 +152,12 @@ public class BytesRestResponseTests extends ElasticsearchTestCase {
         assertEquals(expected.trim(), text.trim());
     }
 
-    public static class WithHeadersException extends ElasticsearchException.WithRestHeadersException {
+    public static class WithHeadersException extends ElasticsearchException {
 
         WithHeadersException() {
-            super("", header("n1", "v11", "v12"), header("n2", "v21", "v22"));
+            super("");
+            this.addHeader("n1", "v11", "v12");
+            this.addHeader("n2", "v21", "v22");
         }
     }