AMBARI-10071. Kerberos service check doesn't work is new HTTP session is created (rlevas)

ambari-server/src/main/java/org/apache/ambari/server/controller/KerberosHelper.java

@@ -1730,14 +1730,10 @@ public class KerberosHelper {
    * @param value   the Kerberos service check identifier to store or null to clear any previously set value
    */
   private void setKerberosServiceCheckIdentifier(Cluster cluster, String value) {
-    Map<String, Object> sessionAttributes = cluster.getSessionAttributes();
-
-    if (sessionAttributes != null) {
-      if (value == null) {
-        sessionAttributes.remove(SERVICE_CHECK_IDENTIFIER);
-      } else {
-        sessionAttributes.put(SERVICE_CHECK_IDENTIFIER, value);
-      }
+    if (value == null) {
+      cluster.removeSessionAttribute(SERVICE_CHECK_IDENTIFIER);
+    } else {
+      cluster.setSessionAttribute(SERVICE_CHECK_IDENTIFIER, value);
     }
   }
 

ambari-server/src/main/java/org/apache/ambari/server/state/Cluster.java

@@ -476,6 +476,21 @@ public interface Cluster {
    */
   public void addSessionAttributes(Map<String, Object> attributes);
 
+  /**
+   * Sets or adds an attribute in the session for this cluster
+   *
+   * @param key   the name of the key which identifies the attribute in the map
+   * @param value the value to set
+   */
+  public void setSessionAttribute(String key, Object value);
+
+  /**
+   * Removes an attribute from the session for this cluster
+   *
+   * @param key the name of the key which identifies the attribute in the map
+   */
+  public void removeSessionAttribute(String key);
+
   /**
    * Get the map of session attributes for this cluster.
    *

ambari-server/src/main/java/org/apache/ambari/server/state/cluster/ClusterImpl.java

@@ -2530,23 +2530,34 @@ public class ClusterImpl implements Cluster {
   @Override
   public void addSessionAttributes(Map<String, Object> attributes) {
     if (attributes != null && !attributes.isEmpty()) {
-
       Map<String, Object>  sessionAttributes = new HashMap<String, Object>(getSessionAttributes());
-
       sessionAttributes.putAll(attributes);
+      setSessionAttributes(attributes);
+    }
+  }
 
-      String attributeName = CLUSTER_SESSION_ATTRIBUTES_PREFIX + getClusterName();
+  @Override
+  public void setSessionAttribute(String key, Object value){
+    if (key != null && !key.isEmpty()) {
+      Map<String, Object> sessionAttributes = new HashMap<String, Object>(getSessionAttributes());
+      sessionAttributes.put(key, value);
+      setSessionAttributes(sessionAttributes);
+    }
+  }
 
-      getSessionManager().setAttribute(attributeName, sessionAttributes);
+  @Override
+  public void removeSessionAttribute(String key) {
+    if (key != null && !key.isEmpty()) {
+      Map<String, Object> sessionAttributes = new HashMap<String, Object>(getSessionAttributes());
+      sessionAttributes.remove(key);
+      setSessionAttributes(sessionAttributes);
     }
   }
 
   @Override
   public Map<String, Object> getSessionAttributes() {
-    String attributeName = CLUSTER_SESSION_ATTRIBUTES_PREFIX + getClusterName();
-
     Map<String, Object>  attributes =
-        (Map<String, Object>) getSessionManager().getAttribute(attributeName);
+        (Map<String, Object>) getSessionManager().getAttribute(getClusterSessionAttributeName());
 
     return attributes == null ? Collections.<String, Object>emptyMap() : attributes;
   }
@@ -2559,4 +2570,25 @@ public class ClusterImpl implements Cluster {
   protected AmbariSessionManager getSessionManager() {
     return sessionManager;
   }
+
+  /**
+   * Set the map of session attributes for this cluster.
+   * <p/>
+   * This is a private method so that it may be used as a utility for add and update operations.
+   *
+   * @param sessionAttributes the map of session attributes for this cluster; never null
+   */
+  private void setSessionAttributes(Map<String, Object> sessionAttributes) {
+    getSessionManager().setAttribute(getClusterSessionAttributeName(), sessionAttributes);
+  }
+
+  /**
+   * Generates and returns the cluster-specific attribute name to use to set and get cluster-specific
+   * session attributes.
+   *
+   * @return the name of the cluster-specific session attribute
+   */
+  private String getClusterSessionAttributeName() {
+    return CLUSTER_SESSION_ATTRIBUTES_PREFIX + getClusterName();
+  }
 }

ambari-server/src/test/java/org/apache/ambari/server/controller/KerberosHelperTest.java

@@ -363,6 +363,11 @@ public class KerberosHelperTest extends EasyMockSupport {
     testCreateTestIdentity(new KerberosCredential("principal", "password", "keytab"));
   }
 
+  @Test (expected = KerberosMissingAdminCredentialsException.class)
+  public void testCreateTestIdentityNoCredentials() throws Exception {
+    testCreateTestIdentity(null);
+  }
+
   @Test
   public void testDeleteTestIdentity() throws Exception {
     testDeleteTestIdentity(new KerberosCredential("principal", "password", "keytab"));
@@ -1965,6 +1970,13 @@ public class KerberosHelperTest extends EasyMockSupport {
     final Config krb5ConfConfig = createNiceMock(Config.class);
     expect(krb5ConfConfig.getProperties()).andReturn(krb5ConfProperties).anyTimes();
 
+    final Map<String,Object> attributeMap = new HashMap<String, Object>();
+    if (kerberosCredential != null) {
+      attributeMap.put("kerberos_admin/" + KerberosCredential.KEY_NAME_PRINCIPAL, kerberosCredential.getPrincipal());
+      attributeMap.put("kerberos_admin/" + KerberosCredential.KEY_NAME_PASSWORD, kerberosCredential.getPassword());
+      attributeMap.put("kerberos_admin/" + KerberosCredential.KEY_NAME_KEYTAB, kerberosCredential.getKeytab());
+    }
+
     final Cluster cluster = createNiceMock(Cluster.class);
     expect(cluster.getDesiredConfigByType("krb5-conf")).andReturn(krb5ConfConfig).anyTimes();
     expect(cluster.getDesiredConfigByType("kerberos-env")).andReturn(kerberosEnvConfig).anyTimes();
@@ -1991,13 +2003,17 @@ public class KerberosHelperTest extends EasyMockSupport {
     expect(cluster.getCurrentStackVersion())
         .andReturn(new StackId("HDP", "2.2"))
         .anyTimes();
-    expect(cluster.getSessionAttributes()).andReturn(new HashMap<String, Object>() {{
-      if (kerberosCredential != null) {
-        put("kerberos_admin/" + KerberosCredential.KEY_NAME_PRINCIPAL, kerberosCredential.getPrincipal());
-        put("kerberos_admin/" + KerberosCredential.KEY_NAME_PASSWORD, kerberosCredential.getPassword());
-        put("kerberos_admin/" + KerberosCredential.KEY_NAME_KEYTAB, kerberosCredential.getKeytab());
+    expect(cluster.getSessionAttributes()).andReturn(attributeMap).anyTimes();
+
+    cluster.setSessionAttribute(anyObject(String.class), anyObject());
+    expectLastCall().andAnswer(new IAnswer<Object>() {
+      @Override
+      public Object answer() throws Throwable {
+        Object[] args = getCurrentArguments();
+        attributeMap.put((String) args[0], args[1]);
+        return null;
       }
-    }}).anyTimes();
+    }).anyTimes();
 
     final Clusters clusters = injector.getInstance(Clusters.class);
     expect(clusters.getHostsForCluster("c1"))

ambari-server/src/test/java/org/apache/ambari/server/state/cluster/ClusterImplTest.java

@@ -26,11 +26,7 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
-import static org.easymock.EasyMock.createMock;
-import static org.easymock.EasyMock.createMockBuilder;
-import static org.easymock.EasyMock.expect;
-import static org.easymock.EasyMock.replay;
-import static org.easymock.EasyMock.verify;
+import static org.easymock.EasyMock.*;
 import static org.junit.Assert.*;
 
 public class ClusterImplTest {
@@ -61,6 +57,80 @@ public class ClusterImplTest {
     verify(sessionManager, cluster);
   }
 
+  @Test
+  public void testSetSessionAttribute() throws Exception {
+    Map<String, Object> attributes = new HashMap<String, Object>();
+    attributes.put("foo", "bar");
+    attributes.put("foo2", "bar2");
+
+    Map<String, Object> updatedAttributes = new HashMap<String, Object>(attributes);
+    updatedAttributes.put("foo2", "updated value");
+
+    Map<String, Object> addedAttributes = new HashMap<String, Object>(updatedAttributes);
+    updatedAttributes.put("foo3", "added value");
+
+    AmbariSessionManager sessionManager = createMock(AmbariSessionManager.class);
+
+    ClusterImpl cluster =
+      createMockBuilder(ClusterImpl.class).
+        addMockedMethod("getSessionManager").
+        addMockedMethod("getClusterName").
+        addMockedMethod("getSessionAttributes").
+        createMock();
+
+    expect(cluster.getSessionManager()).andReturn(sessionManager);
+    expect(cluster.getClusterName()).andReturn("c1");
+    expect(cluster.getSessionAttributes()).andReturn(attributes);
+
+    sessionManager.setAttribute("cluster_session_attributes:c1", updatedAttributes);
+    expectLastCall().once();
+
+    expect(cluster.getSessionManager()).andReturn(sessionManager);
+    expect(cluster.getClusterName()).andReturn("c1");
+    expect(cluster.getSessionAttributes()).andReturn(updatedAttributes);
+
+    sessionManager.setAttribute("cluster_session_attributes:c1", addedAttributes);
+    expectLastCall().once();
+
+    replay(sessionManager, cluster);
+
+    cluster.setSessionAttribute("foo2", "updated value");
+    cluster.setSessionAttribute("foo3", "added value");
+
+    verify(sessionManager, cluster);
+  }
+
+  @Test
+  public void testRemoveSessionAttribute() throws Exception {
+    Map<String, Object> attributes = new HashMap<String, Object>();
+    attributes.put("foo", "bar");
+    attributes.put("foo2", "bar2");
+
+    Map<String, Object> trimmedAttributes = new HashMap<String, Object>(attributes);
+    trimmedAttributes.remove("foo2");
+
+    AmbariSessionManager sessionManager = createMock(AmbariSessionManager.class);
+
+    ClusterImpl cluster =
+      createMockBuilder(ClusterImpl.class).
+        addMockedMethod("getSessionManager").
+        addMockedMethod("getClusterName").
+        addMockedMethod("getSessionAttributes").
+        createMock();
+
+    expect(cluster.getSessionManager()).andReturn(sessionManager);
+    expect(cluster.getClusterName()).andReturn("c1");
+    expect(cluster.getSessionAttributes()).andReturn(attributes);
+    sessionManager.setAttribute("cluster_session_attributes:c1", trimmedAttributes);
+    expectLastCall().once();
+
+    replay(sessionManager, cluster);
+
+    cluster.removeSessionAttribute("foo2");
+
+    verify(sessionManager, cluster);
+  }
+
   @Test
   public void testGetSessionAttributes() throws Exception {
     Map<String, Object> attributes = new HashMap<String, Object>();