From bf73ba366528ba2bab14bc96a3d32120f94b78a5 Mon Sep 17 00:00:00 2001
From: uboness <uboness@gmail.com>
Date: Fri, 8 Aug 2014 18:02:52 +0200
Subject: [PATCH] Updated the code base to use the latest changes in
 TransportMessage on es-core

Now the UsernamePasswordToken, once resolved, is placed in the (new) request context as is (avoids the need to re-parse headers)

Original commit: elastic/x-pack-elasticsearch@f898a941577c75af2294a8db52e250e13ec313ce
---
 .../shield/authc/support/UsernamePasswordToken.java   | 11 ++++++-----
 .../authc/support/UsernamePasswordTokenTests.java     |  4 ++++
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/elasticsearch/shield/authc/support/UsernamePasswordToken.java b/src/main/java/org/elasticsearch/shield/authc/support/UsernamePasswordToken.java
index 5e34adcfb5a..b60935cb9ee 100644
--- a/src/main/java/org/elasticsearch/shield/authc/support/UsernamePasswordToken.java
+++ b/src/main/java/org/elasticsearch/shield/authc/support/UsernamePasswordToken.java
@@ -19,7 +19,8 @@ import java.util.regex.Pattern;
  */
 public class UsernamePasswordToken implements AuthenticationToken {
 
-    private static final String TOKEN_HEADER = "X-ES-UsernamePasswordToken";
+    private static final String TOKEN_KEY = "X-ES-UsernamePasswordToken";
+
     static final String BASIC_AUTH_HEADER = "Authorization";
     private static final Pattern BASIC_AUTH_PATTERN = Pattern.compile("Basic\\s(.+)");
 
@@ -42,17 +43,17 @@ public class UsernamePasswordToken implements AuthenticationToken {
     }
 
     public static UsernamePasswordToken extractToken(TransportRequest request, UsernamePasswordToken defaultToken) {
-        UsernamePasswordToken token = (UsernamePasswordToken) request.getHeader(TOKEN_HEADER);
+        UsernamePasswordToken token = (UsernamePasswordToken) request.context().get(TOKEN_KEY);
         if (token != null) {
             return token;
         }
 
-        String authStr = (String) request.getHeader(BASIC_AUTH_HEADER);
+        String authStr = request.getHeader(BASIC_AUTH_HEADER);
         if (authStr == null) {
             if (defaultToken == null) {
                 return null;
             }
-            request.putHeader(TOKEN_HEADER, defaultToken);
+            request.context().put(TOKEN_KEY, defaultToken);
             return defaultToken;
         }
 
@@ -64,7 +65,7 @@ public class UsernamePasswordToken implements AuthenticationToken {
         String userpasswd = new String(Base64.decodeBase64(matcher.group(1)), Charsets.UTF_8);
         int i = userpasswd.indexOf(':');
         token = new UsernamePasswordToken(userpasswd.substring(0, i), userpasswd.substring(i+1).toCharArray());
-        request.putHeader(TOKEN_HEADER, token);
+        request.context().put(TOKEN_KEY, token);
         return token;
     }
 
diff --git a/src/test/java/org/elasticsearch/shield/authc/support/UsernamePasswordTokenTests.java b/src/test/java/org/elasticsearch/shield/authc/support/UsernamePasswordTokenTests.java
index 6060a81595b..11125c23134 100644
--- a/src/test/java/org/elasticsearch/shield/authc/support/UsernamePasswordTokenTests.java
+++ b/src/test/java/org/elasticsearch/shield/authc/support/UsernamePasswordTokenTests.java
@@ -44,5 +44,9 @@ public class UsernamePasswordTokenTests extends ElasticsearchTestCase {
         assertThat(token, notNullValue());
         assertThat(token.principal(), equalTo("user1"));
         assertThat(new String(token.credentials()), equalTo("test123"));
+
+        // making sure that indeed, once resolved the instance is reused across multiple resolve calls
+        UsernamePasswordToken token2 = UsernamePasswordToken.extractToken(request, null);
+        assertThat(token, is(token2));
     }
 }