TestSecurityContext.java

/**
 * Copyright 2015 StreamSets Inc.
 *
 * Licensed under the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF 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 com.streamsets.datacollector.security;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.streamsets.datacollector.main.RuntimeInfo;
import com.streamsets.datacollector.util.Configuration;
import org.apache.hadoop.minikdc.MiniKdc;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;
import org.mockito.Mockito;

import javax.security.auth.Subject;
import javax.security.auth.kerberos.KerberosPrincipal;
import javax.security.auth.kerberos.KerberosTicket;
import java.io.File;
import java.net.InetAddress;
import java.util.Date;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

public class TestSecurityContext {
  private static File testDir;
  private static File keytabFile;
  private static MiniKdc miniKdc;

  @BeforeClass
  public static void startKdc() throws Exception {
    testDir = new File("target", UUID.randomUUID().toString()).getAbsoluteFile();
    Assert.assertTrue(testDir.mkdirs());
    File kdcDir = new File(testDir, "kdc");
    Assert.assertTrue(kdcDir.mkdirs());
    keytabFile = new File(testDir, "test.keytab");
    miniKdc = new MiniKdc(MiniKdc.createConf(), testDir);
    miniKdc.start();
    miniKdc.createPrincipal(keytabFile, "foo", "bar/localhost");
  }

  @AfterClass
  public static void stopKdc() {
    if (miniKdc != null) {
      miniKdc.stop();
      miniKdc = null;
    }
  }

  private RuntimeInfo getMockRuntimeInfo() {
    RuntimeInfo runtimeInfo = Mockito.mock(RuntimeInfo.class);
    Mockito.when(runtimeInfo.getConfigDir()).thenReturn(testDir.getAbsolutePath());
    Mockito.when(runtimeInfo.getDataDir()).thenReturn(testDir.getAbsolutePath());
    return runtimeInfo;
  }

  @Test
  public void testPrincipalResolution() {
    Configuration conf = new Configuration();
    conf.set(SecurityConfiguration.KERBEROS_ENABLED_KEY, true);
    conf.set(SecurityConfiguration.KERBEROS_KEYTAB_KEY, "test.keytab");

    String hostname = SecurityConfiguration.getLocalHostName();

    Map<String, String> principals = ImmutableMap.<String, String>builder()
                                                 .put("foo", "foo")
                                                 .put("foo/bar","foo/bar")
                                                 .put("foo/bar@REALM", "foo/bar@REALM")
                                                 .put("foo/_HOST", "foo/" + hostname)
                                                 .put("foo/_HOST@REALM", "foo/" + hostname + "@REALM")
                                                 .put("foo/0.0.0.0", "foo/" + hostname)
                                                 .put("foo/0.0.0.0@REALM", "foo/" + hostname + "@REALM").build();

    for (Map.Entry<String, String> entry : principals.entrySet()) {
      conf.set(SecurityConfiguration.KERBEROS_PRINCIPAL_KEY, entry.getKey());
      SecurityContext context = new SecurityContext(getMockRuntimeInfo(), conf);
      Assert.assertEquals(entry.getValue(), context.getSecurityConfiguration().getKerberosPrincipal());
    }
  }

  @Test(expected = RuntimeException.class)
  public void invalidLogin() throws Exception {
    Configuration conf = new Configuration();
    conf.set(SecurityConfiguration.KERBEROS_ENABLED_KEY, true);
    conf.set(SecurityConfiguration.KERBEROS_PRINCIPAL_KEY, "foo");
    conf.set(SecurityConfiguration.KERBEROS_KEYTAB_KEY, "foo.keytab");
    SecurityContext context = new SecurityContext(getMockRuntimeInfo(), conf);
    Assert.assertTrue(context.getSecurityConfiguration().isKerberosEnabled());
    context.login();
  }

  @Test
  public void notLoggedIn() throws Exception {
    Configuration conf = new Configuration();
    conf.set(SecurityConfiguration.KERBEROS_ENABLED_KEY, true);
    conf.set(SecurityConfiguration.KERBEROS_PRINCIPAL_KEY, "foo");
    conf.set(SecurityConfiguration.KERBEROS_KEYTAB_KEY, "test.keytab");
    SecurityContext context = new SecurityContext(getMockRuntimeInfo(), conf);
    Assert.assertNull(context.getSubject());
  }

  @Test
  public void loginKerberosDisabled() throws Exception {
    Configuration conf = new Configuration();
    conf.set(SecurityConfiguration.KERBEROS_ENABLED_KEY, false);
    SecurityContext context = new SecurityContext(getMockRuntimeInfo(), conf);
    Assert.assertFalse(context.getSecurityConfiguration().isKerberosEnabled());
    context.login();
    Subject subject = context.getSubject();
    Assert.assertNotNull(subject);
    System.out.println(subject);
    context.logout();
  }

  @Test
  public void loginFromAbsoluteKeytab() throws Exception {
    loginFromKeytab(keytabFile.getAbsolutePath());
  }

  @Test
  public void loginFromRelativeKeytab() throws Exception {
    loginFromKeytab(keytabFile.getName());
  }

  private void loginFromKeytab(String keytab) throws Exception {
    Configuration conf = new Configuration();
    conf.set(SecurityConfiguration.KERBEROS_ENABLED_KEY, true);
    conf.set(SecurityConfiguration.KERBEROS_PRINCIPAL_KEY, "foo");
    conf.set(SecurityConfiguration.KERBEROS_KEYTAB_KEY, keytab);
    SecurityContext context = new SecurityContext(getMockRuntimeInfo(), conf);
    Assert.assertTrue(context.getSecurityConfiguration().isKerberosEnabled());
    context.login();
    Subject subject = context.getSubject();
    Assert.assertNotNull(subject);
    System.out.println(subject);
    context.logout();
  }

  @Test
  public void testRenewalCalculation() {
    Configuration conf = new Configuration();
    SecurityContext context = new SecurityContext(getMockRuntimeInfo(), conf);
    context = Mockito.spy(context);

    Assert.assertEquals(context.getRenewalWindow(), context.getRenewalWindow(), 0.001);
    Assert.assertTrue(context.getRenewalWindow() >= 0.5);
    Assert.assertTrue(context.getRenewalWindow() < 0.7);

    // brute force test
    for (int i = 0; i < 1000; i++) {
      double window = context.computeRenewalWindow();
      Assert.assertTrue(window >= 0.5);
      Assert.assertTrue(window < 0.7);
    }

    // renewal time
    Mockito.doReturn(0.5D).when(context).getRenewalWindow();
    Assert.assertEquals(10 + (long)(0.5D * (20 - 10)), context.getRenewalTime(10, 20), 0.001);
  }

  private KerberosTicket createMockTGT(String name, Date start, Date end) {
    KerberosPrincipal userPrincipal = new KerberosPrincipal(name);
    KerberosPrincipal serverPrincipal = new KerberosPrincipal("krbtgt/" + userPrincipal.getRealm());
    return new KerberosTicket(new byte[0], userPrincipal,  serverPrincipal, new byte[0], 0, new boolean[0],
        start, start, end, end, new InetAddress[0]);
  }

  @Test
  public void testGetKerberosTicket() {
    long now = System.currentTimeMillis();
    Date expired = new Date(now - TimeUnit.MINUTES.toMillis(30));
    Date valid = new Date(now + TimeUnit.DAYS.toMillis(1));

    KerberosTicket expiredTicket = createMockTGT("short", expired, expired);
    KerberosTicket validTicket = createMockTGT("valid", valid, valid);


    Configuration conf = new Configuration();
    SecurityContext context = new SecurityContext(getMockRuntimeInfo(), conf);
    context = Mockito.spy(context);
    Mockito.doReturn(now).when(context).getTimeNow();

    // 1 ticket, expired, we should return null and remove it from the subject
    Subject subject = new Subject();
    Mockito.doReturn(subject).when(context).getSubject();
    subject.getPrivateCredentials().add(expiredTicket);
    Assert.assertNull(context.getKerberosTicket());
    Assert.assertTrue(subject.getPrivateCredentials().isEmpty());

    // 2 tickets, 1 expired one not, one short other long expired, we should return the non expired one
    // and remove the expired on from subject
    subject = new Subject();
    Mockito.doReturn(subject).when(context).getSubject();
    Mockito.doReturn(now).when(context).getTimeNow();
    subject.getPrivateCredentials().add(expiredTicket);
    subject.getPrivateCredentials().add(validTicket);
    Assert.assertEquals(validTicket, context.getKerberosTicket());
    Assert.assertEquals(ImmutableSet.of(validTicket), subject.getPrivateCredentials());
  }

}