@Override
  public DataEncryptionKey newDataEncryptionKey() {
    if (encryptDataTransfer) {
      synchronized (this) {
        if (encryptionKey == null ||
            encryptionKey.expiryDate < timer.now()) {
          // Encryption Key (EK) is generated from Block Key (BK).
          // Check if EK is expired, and generate a new one using the current BK
          // if so, otherwise continue to use the previously generated EK.
          //
          // It's important to make sure that when EK is not expired, the BK
          // used to generate the EK is not expired and removed, because
          // the same BK will be used to re-generate the EK
          // by BlockTokenSecretManager.
          //
          // The current implementation ensures that when an EK is not expired
          // (within tokenLifetime), the BK that's used to generate it
          // still has at least "keyUpdateInterval" of life time before
          // the BK gets expired and removed.
          // See BlockTokenSecretManager for details.
          LOG.debug("Generating new data encryption key because current key "
              + (encryptionKey == null ?
              "is null." : "expired on " + encryptionKey.expiryDate));
          encryptionKey = blockTokenSecretManager.generateDataEncryptionKey();
        }
        return encryptionKey;
      }
    } else {
      return null;
    }
  }
@Test
  public void testNewDataEncryptionKey() throws Exception {
    final Configuration conf = new HdfsConfiguration();
    // Enable data transport encryption and access token
    conf.setBoolean(DFSConfigKeys.DFS_ENCRYPT_DATA_TRANSFER_KEY, true);
    conf.setBoolean(DFSConfigKeys.DFS_BLOCK_ACCESS_TOKEN_ENABLE_KEY, true);

    final long keyUpdateInterval = 2 * 1000;
    final long tokenLifeTime = keyUpdateInterval;
    final String blockPoolId = "bp-foo";
    FakeTimer fakeTimer = new FakeTimer();
    BlockTokenSecretManager btsm = new BlockTokenSecretManager(
        keyUpdateInterval, tokenLifeTime, 0, 1, blockPoolId, null, false);
    Whitebox.setInternalState(btsm, "timer", fakeTimer);

    // When KeyManager asks for block keys, return them from btsm directly
    NamenodeProtocol namenode = mock(NamenodeProtocol.class);
    when(namenode.getBlockKeys()).thenReturn(btsm.exportKeys());

    // Instantiate a KeyManager instance and get data encryption key.
    KeyManager keyManager = new KeyManager(blockPoolId, namenode,
        true, conf);
    Whitebox.setInternalState(keyManager, "timer", fakeTimer);
    Whitebox.setInternalState(
        Whitebox.getInternalState(keyManager, "blockTokenSecretManager"),
        "timer", fakeTimer);
    final DataEncryptionKey dek = keyManager.newDataEncryptionKey();
    final long remainingTime = dek.expiryDate - fakeTimer.now();
    assertEquals("KeyManager dataEncryptionKey should expire in 2 seconds",
        keyUpdateInterval, remainingTime);
    // advance the timer to expire the block key and data encryption key
    fakeTimer.advance(keyUpdateInterval + 1);

    // After the initial data encryption key expires, KeyManager should
    // regenerate a valid data encryption key using the current block key.
    final DataEncryptionKey dekAfterExpiration =
        keyManager.newDataEncryptionKey();
    assertNotEquals("KeyManager should generate a new data encryption key",
        dek, dekAfterExpiration);
    assertTrue("KeyManager has an expired DataEncryptionKey!",
        dekAfterExpiration.expiryDate > fakeTimer.now());
  }