diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/DataXceiver.java b/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/DataXceiver.java index a29fab9734eb42..ff7359fa265dc0 100644 --- a/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/DataXceiver.java +++ b/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/DataXceiver.java @@ -36,6 +36,7 @@ import org.apache.hadoop.hdfs.protocol.datatransfer.BlockPinningException; import org.apache.hadoop.hdfs.protocol.datatransfer.DataTransferProtoUtil; import org.apache.hadoop.hdfs.protocol.datatransfer.IOStreamPair; +import org.apache.hadoop.hdfs.protocol.datatransfer.InvalidEncryptionKeyException; import org.apache.hadoop.hdfs.protocol.datatransfer.Op; import org.apache.hadoop.hdfs.protocol.datatransfer.Receiver; import org.apache.hadoop.hdfs.protocol.datatransfer.Sender; @@ -797,7 +798,6 @@ public void writeBlock(final ExtendedBlock block, mirrorNode = targets[0].getXferAddr(connectToDnViaHostname); LOG.debug("Connecting to datanode {}", mirrorNode); mirrorTarget = NetUtils.createSocketAddr(mirrorNode); - mirrorSock = datanode.newSocket(); try { DataNodeFaultInjector.get().failMirrorConnection(); @@ -806,32 +806,52 @@ public void writeBlock(final ExtendedBlock block, (HdfsConstants.READ_TIMEOUT_EXTENSION * targets.length); int writeTimeout = dnConf.socketWriteTimeout + (HdfsConstants.WRITE_TIMEOUT_EXTENSION * targets.length); - NetUtils.connect(mirrorSock, mirrorTarget, timeoutValue); - mirrorSock.setTcpNoDelay(dnConf.getDataTransferServerTcpNoDelay()); - mirrorSock.setSoTimeout(timeoutValue); - mirrorSock.setKeepAlive(true); - if (dnConf.getTransferSocketSendBufferSize() > 0) { - mirrorSock.setSendBufferSize( - dnConf.getTransferSocketSendBufferSize()); - } - - OutputStream unbufMirrorOut = NetUtils.getOutputStream(mirrorSock, - writeTimeout); - InputStream unbufMirrorIn = NetUtils.getInputStream(mirrorSock); DataEncryptionKeyFactory keyFactory = datanode.getDataEncryptionKeyFactoryForBlock(block); - SecretKey secretKey = null; - if (dnConf.overwriteDownstreamDerivedQOP) { - String bpid = block.getBlockPoolId(); - BlockKey blockKey = datanode.blockPoolTokenSecretManager - .get(bpid).getCurrentKey(); - secretKey = blockKey.getKey(); + OutputStream unbufMirrorOut; + InputStream unbufMirrorIn; + int encryptionKeyRetryCount = 0; + while (true) { + try { + mirrorSock = datanode.newSocket(); + NetUtils.connect(mirrorSock, mirrorTarget, timeoutValue); + mirrorSock.setTcpNoDelay( + dnConf.getDataTransferServerTcpNoDelay()); + mirrorSock.setSoTimeout(timeoutValue); + mirrorSock.setKeepAlive(true); + if (dnConf.getTransferSocketSendBufferSize() > 0) { + mirrorSock.setSendBufferSize( + dnConf.getTransferSocketSendBufferSize()); + } + + unbufMirrorOut = NetUtils.getOutputStream(mirrorSock, + writeTimeout); + unbufMirrorIn = NetUtils.getInputStream(mirrorSock); + SecretKey secretKey = null; + if (dnConf.overwriteDownstreamDerivedQOP) { + String bpid = block.getBlockPoolId(); + BlockKey blockKey = datanode.blockPoolTokenSecretManager + .get(bpid).getCurrentKey(); + secretKey = blockKey.getKey(); + } + IOStreamPair saslStreams = datanode.saslClient.socketSend( + mirrorSock, unbufMirrorOut, unbufMirrorIn, keyFactory, + blockToken, targets[0], secretKey); + unbufMirrorOut = saslStreams.out; + unbufMirrorIn = saslStreams.in; + break; + } catch (InvalidEncryptionKeyException e) { + IOUtils.closeSocket(mirrorSock); + mirrorSock = null; + if (!prepareRetryAfterInvalidEncryptionKey(keyFactory, + ++encryptionKeyRetryCount)) { + throw e; + } + LOG.info("Retrying connection to mirror {} for block {} after " + + "InvalidEncryptionKeyException", + targets[0], block, e); + } } - IOStreamPair saslStreams = datanode.saslClient.socketSend( - mirrorSock, unbufMirrorOut, unbufMirrorIn, keyFactory, - blockToken, targets[0], secretKey); - unbufMirrorOut = saslStreams.out; - unbufMirrorIn = saslStreams.in; mirrorOut = new DataOutputStream(new BufferedOutputStream(unbufMirrorOut, smallBufferSize)); mirrorIn = new DataInputStream(unbufMirrorIn); @@ -1211,21 +1231,40 @@ public void replaceBlock(final ExtendedBlock block, final String dnAddr = proxySource.getXferAddr(connectToDnViaHostname); LOG.debug("Connecting to datanode {}", dnAddr); InetSocketAddress proxyAddr = NetUtils.createSocketAddr(dnAddr); - proxySock = datanode.newSocket(); - NetUtils.connect(proxySock, proxyAddr, dnConf.socketTimeout); - proxySock.setTcpNoDelay(dnConf.getDataTransferServerTcpNoDelay()); - proxySock.setSoTimeout(dnConf.socketTimeout); - proxySock.setKeepAlive(true); - - OutputStream unbufProxyOut = NetUtils.getOutputStream(proxySock, - dnConf.socketWriteTimeout); - InputStream unbufProxyIn = NetUtils.getInputStream(proxySock); DataEncryptionKeyFactory keyFactory = datanode.getDataEncryptionKeyFactoryForBlock(block); - IOStreamPair saslStreams = datanode.saslClient.socketSend(proxySock, - unbufProxyOut, unbufProxyIn, keyFactory, blockToken, proxySource); - unbufProxyOut = saslStreams.out; - unbufProxyIn = saslStreams.in; + OutputStream unbufProxyOut; + InputStream unbufProxyIn; + int encryptionKeyRetryCount = 0; + while (true) { + try { + proxySock = datanode.newSocket(); + NetUtils.connect(proxySock, proxyAddr, dnConf.socketTimeout); + proxySock.setTcpNoDelay(dnConf.getDataTransferServerTcpNoDelay()); + proxySock.setSoTimeout(dnConf.socketTimeout); + proxySock.setKeepAlive(true); + + unbufProxyOut = NetUtils.getOutputStream(proxySock, + dnConf.socketWriteTimeout); + unbufProxyIn = NetUtils.getInputStream(proxySock); + IOStreamPair saslStreams = datanode.saslClient.socketSend( + proxySock, unbufProxyOut, unbufProxyIn, keyFactory, blockToken, + proxySource); + unbufProxyOut = saslStreams.out; + unbufProxyIn = saslStreams.in; + break; + } catch (InvalidEncryptionKeyException e) { + IOUtils.closeSocket(proxySock); + proxySock = null; + if (!prepareRetryAfterInvalidEncryptionKey(keyFactory, + ++encryptionKeyRetryCount)) { + throw e; + } + LOG.info("Retrying connection to proxy {} for block {} after " + + "InvalidEncryptionKeyException", + proxySource, block, e); + } + } proxyOut = new DataOutputStream(new BufferedOutputStream(unbufProxyOut, smallBufferSize)); @@ -1313,6 +1352,15 @@ public void replaceBlock(final ExtendedBlock block, datanode.metrics.addReplaceBlockOp(elapsed()); } + private static boolean prepareRetryAfterInvalidEncryptionKey( + DataEncryptionKeyFactory keyFactory, int retryCount) { + if (retryCount > 1) { + return false; + } + keyFactory.clearDataEncryptionKey(); + return true; + } + /** * Separated for testing. diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/erasurecode/StripedBlockWriter.java b/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/erasurecode/StripedBlockWriter.java index 5b8b22c3c11362..d614e896908c1e 100644 --- a/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/erasurecode/StripedBlockWriter.java +++ b/hadoop-hdfs-project/hadoop-hdfs/src/main/java/org/apache/hadoop/hdfs/server/datanode/erasurecode/StripedBlockWriter.java @@ -27,6 +27,7 @@ import org.apache.hadoop.hdfs.protocol.ExtendedBlock; import org.apache.hadoop.hdfs.protocol.datatransfer.BlockConstructionStage; import org.apache.hadoop.hdfs.protocol.datatransfer.IOStreamPair; +import org.apache.hadoop.hdfs.protocol.datatransfer.InvalidEncryptionKeyException; import org.apache.hadoop.hdfs.protocol.datatransfer.Sender; import org.apache.hadoop.hdfs.protocol.datatransfer.sasl.DataEncryptionKeyFactory; import org.apache.hadoop.hdfs.security.token.block.BlockTokenIdentifier; @@ -90,6 +91,15 @@ class StripedBlockWriter { init(); } + static boolean prepareRetryAfterInvalidEncryptionKey( + DataEncryptionKeyFactory keyFactory, int retryCount) { + if (retryCount > 1) { + return false; + } + keyFactory.clearDataEncryptionKey(); + return true; + } + ByteBuffer getTargetBuffer() { return targetBuffer; } @@ -113,12 +123,6 @@ private void init() throws IOException { try { InetSocketAddress targetAddr = stripedWriter.getSocketAddress4Transfer(target); - socket = datanode.newSocket(); - NetUtils.connect(socket, targetAddr, - datanode.getDnConf().getSocketTimeout()); - socket.setTcpNoDelay( - datanode.getDnConf().getDataTransferServerTcpNoDelay()); - socket.setSoTimeout(datanode.getDnConf().getSocketTimeout()); DataNodeFaultInjector.get().stripedBlockWriterInit(targetBuffer); Token blockToken = @@ -127,15 +131,35 @@ private void init() throws IOException { new StorageType[]{storageType}, new String[]{storageId}); long writeTimeout = datanode.getDnConf().getSocketWriteTimeout(); - OutputStream unbufOut = NetUtils.getOutputStream(socket, writeTimeout); - InputStream unbufIn = NetUtils.getInputStream(socket); DataEncryptionKeyFactory keyFactory = datanode.getDataEncryptionKeyFactoryForBlock(block); - IOStreamPair saslStreams = datanode.getSaslClient().socketSend( - socket, unbufOut, unbufIn, keyFactory, blockToken, target); - - unbufOut = saslStreams.out; - unbufIn = saslStreams.in; + OutputStream unbufOut; + InputStream unbufIn; + int encryptionKeyRetryCount = 0; + while (true) { + try { + socket = datanode.newSocket(); + NetUtils.connect(socket, targetAddr, + datanode.getDnConf().getSocketTimeout()); + socket.setTcpNoDelay( + datanode.getDnConf().getDataTransferServerTcpNoDelay()); + socket.setSoTimeout(datanode.getDnConf().getSocketTimeout()); + unbufOut = NetUtils.getOutputStream(socket, writeTimeout); + unbufIn = NetUtils.getInputStream(socket); + IOStreamPair saslStreams = datanode.getSaslClient().socketSend( + socket, unbufOut, unbufIn, keyFactory, blockToken, target); + unbufOut = saslStreams.out; + unbufIn = saslStreams.in; + break; + } catch (InvalidEncryptionKeyException e) { + IOUtils.closeSocket(socket); + socket = null; + if (!prepareRetryAfterInvalidEncryptionKey(keyFactory, + ++encryptionKeyRetryCount)) { + throw e; + } + } + } out = new DataOutputStream(new BufferedOutputStream(unbufOut, DFSUtilClient.getSmallBufferSize(conf))); diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/TestDataXceiverEncryptionKey.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/TestDataXceiverEncryptionKey.java new file mode 100644 index 00000000000000..b36508a2712819 --- /dev/null +++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/TestDataXceiverEncryptionKey.java @@ -0,0 +1,274 @@ +/** + * Licensed to 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 org.apache.hadoop.hdfs.server.datanode; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.InputStream; +import java.io.OutputStream; +import java.net.Socket; +import java.net.SocketAddress; +import java.nio.channels.SocketChannel; +import java.util.concurrent.atomic.AtomicInteger; + +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.StorageType; +import org.apache.hadoop.hdfs.DFSTestUtil; +import org.apache.hadoop.hdfs.net.Peer; +import org.apache.hadoop.hdfs.net.PeerServer; +import org.apache.hadoop.hdfs.protocol.DatanodeID; +import org.apache.hadoop.hdfs.protocol.DatanodeInfo; +import org.apache.hadoop.hdfs.protocol.ExtendedBlock; +import org.apache.hadoop.hdfs.protocol.datatransfer.BlockConstructionStage; +import org.apache.hadoop.hdfs.protocol.datatransfer.IOStreamPair; +import org.apache.hadoop.hdfs.protocol.datatransfer.InvalidEncryptionKeyException; +import org.apache.hadoop.hdfs.protocol.datatransfer.sasl.DataEncryptionKeyFactory; +import org.apache.hadoop.hdfs.protocol.datatransfer.sasl.SaslDataTransferClient; +import org.apache.hadoop.hdfs.security.token.block.BlockTokenIdentifier; +import org.apache.hadoop.hdfs.security.token.block.DataEncryptionKey; +import org.apache.hadoop.hdfs.server.datanode.fsdataset.FsDatasetSpi; +import org.apache.hadoop.hdfs.server.datanode.fsdataset.FsVolumeSpi; +import org.apache.hadoop.hdfs.server.datanode.metrics.DataNodeMetrics; +import org.apache.hadoop.hdfs.server.protocol.DatanodeRegistration; +import org.apache.hadoop.io.Text; +import org.apache.hadoop.security.token.Token; +import org.apache.hadoop.util.DataChecksum; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +/** + * Test DataXceiver handling of InvalidEncryptionKeyException. + */ +@Timeout(60) +public class TestDataXceiverEncryptionKey { + + @Test + public void testWriteBlockRetriesInvalidEncryptionKeyToMirror() + throws Exception { + Peer peer = createPeer(); + Configuration conf = new Configuration(); + CountingKeyFactory keyFactory = new CountingKeyFactory(); + RetryDataNode dataNode = new RetryDataNode(conf, keyFactory); + AtomicInteger socketSendCount = new AtomicInteger(); + org.mockito.Mockito.doAnswer(invocation -> { + if (socketSendCount.getAndIncrement() == 0) { + throw new InvalidEncryptionKeyException("test invalid key"); + } + return new IOStreamPair( + new ByteArrayInputStream(new byte[0]), new ByteArrayOutputStream()); + }).when(dataNode.saslClient).socketSend( + any(Socket.class), any(OutputStream.class), any(InputStream.class), + any(DataEncryptionKeyFactory.class), any(), + any(DatanodeID.class), any()); + + DataXceiverServer server = new DataXceiverServer( + mock(PeerServer.class), conf, dataNode); + DataXceiver xceiver = spy(DataXceiver.create(peer, dataNode, server)); + mockBlockReceiver(xceiver); + + DatanodeInfo target = DFSTestUtil.getDatanodeInfo( + "127.0.0.1", "localhost", 1); + xceiver.writeBlock( + new ExtendedBlock("bp-1", 1L), + StorageType.DISK, + createToken(), + "", + new DatanodeInfo[]{target}, + new StorageType[]{StorageType.DISK}, + target, + BlockConstructionStage.PIPELINE_SETUP_CREATE, + 0, 0, 0, 0, + createChecksum(), + CachingStrategy.newDefaultStrategy(), + false, + false, null, null, new String[0]); + + assertEquals(2, socketSendCount.get()); + assertEquals(1, keyFactory.clearCount); + } + + @Test + public void testReplaceBlockRetriesInvalidEncryptionKeyToProxy() + throws Exception { + Peer peer = createPeer(); + Configuration conf = new Configuration(); + CountingKeyFactory keyFactory = new CountingKeyFactory(); + RetryDataNode dataNode = new RetryDataNode(conf, keyFactory); + AtomicInteger socketSendCount = new AtomicInteger(); + org.mockito.Mockito.doAnswer(invocation -> { + if (socketSendCount.getAndIncrement() == 0) { + throw new InvalidEncryptionKeyException("test invalid key"); + } + return new IOStreamPair( + new ByteArrayInputStream(new byte[0]), new ByteArrayOutputStream()); + }).when(dataNode.saslClient).socketSend( + any(Socket.class), any(OutputStream.class), any(InputStream.class), + any(DataEncryptionKeyFactory.class), any(), + any(DatanodeID.class)); + + DataXceiverServer server = new DataXceiverServer( + mock(PeerServer.class), conf, dataNode); + DataXceiver xceiver = DataXceiver.create(peer, dataNode, server); + DatanodeInfo proxySource = DFSTestUtil.getDatanodeInfo( + "127.0.0.1", "localhost", 1); + try { + xceiver.replaceBlock(new ExtendedBlock("bp-1", 1L), + StorageType.DISK, createToken(), "delHint", proxySource, + "storage-id"); + } catch (Exception ignored) { + // The test only exercises the connection setup path; after the retry + // succeeds, the fake proxy has no copyBlock response to read. + } + + assertEquals(2, socketSendCount.get()); + assertEquals(1, keyFactory.clearCount); + } + + private static Peer createPeer() throws Exception { + Peer peer = mock(Peer.class); + doReturn("").when(peer).getRemoteAddressString(); + doReturn("").when(peer).getLocalAddressString(); + doReturn(new ByteArrayInputStream(new byte[0])).when(peer).getInputStream(); + doReturn(new ByteArrayOutputStream()).when(peer).getOutputStream(); + return peer; + } + + private static Token createToken() { + Token token = (Token) mock( + Token.class); + doReturn("".getBytes()).when(token).getIdentifier(); + doReturn("".getBytes()).when(token).getPassword(); + doReturn(new Text("")).when(token).getKind(); + doReturn(new Text("")).when(token).getService(); + return token; + } + + private static DataChecksum createChecksum() { + DataChecksum checksum = mock(DataChecksum.class); + doReturn(DataChecksum.Type.NULL).when(checksum).getChecksumType(); + return checksum; + } + + private static void mockBlockReceiver(DataXceiver xceiver) + throws Exception { + BlockReceiver mockBlockReceiver = mock(BlockReceiver.class); + Replica replica = mock(Replica.class); + doReturn(replica).when(mockBlockReceiver).getReplica(); + doReturn("storage-id").when(replica).getStorageUuid(); + doReturn(false).when(replica).isOnTransientStorage(); + doReturn(mock(FsVolumeSpi.class)).when(replica).getVolume(); + doReturn(mockBlockReceiver).when(xceiver).getBlockReceiver( + any(ExtendedBlock.class), any(StorageType.class), + any(), anyString(), any(), + any(BlockConstructionStage.class), anyLong(), anyLong(), anyLong(), + anyString(), any(DatanodeInfo.class), any(DataNode.class), + any(DataChecksum.class), any(CachingStrategy.class), + anyBoolean(), anyBoolean(), any()); + } + + private static final class RetryDataNode extends DataNode { + private final CountingKeyFactory keyFactory; + + private RetryDataNode(Configuration conf, CountingKeyFactory keyFactory) + throws Exception { + super(conf); + this.keyFactory = keyFactory; + data = (FsDatasetSpi) mock(FsDatasetSpi.class); + saslClient = mock(SaslDataTransferClient.class); + metrics = mock(DataNodeMetrics.class); + } + + @Override + public DatanodeRegistration getDNRegistrationForBP(String bpid) { + return null; + } + + @Override + public Socket newSocket() { + return new FakeSocket(); + } + + @Override + public DataEncryptionKeyFactory getDataEncryptionKeyFactoryForBlock( + ExtendedBlock block) { + return keyFactory; + } + + @Override + void closeBlock(ExtendedBlock block, String delHint, String storageUuid, + boolean isTransientStorage) { + } + + @Override + void incrDatanodeNetworkErrors(String host) { + } + } + + private static final class CountingKeyFactory + implements DataEncryptionKeyFactory { + private int clearCount; + + @Override + public DataEncryptionKey newDataEncryptionKey() { + return null; + } + + @Override + public void clearDataEncryptionKey() { + clearCount++; + } + } + + private static final class FakeSocket extends Socket { + private final ByteArrayOutputStream out = new ByteArrayOutputStream(); + private final ByteArrayInputStream in = + new ByteArrayInputStream(new byte[0]); + + @Override + public void connect(SocketAddress endpoint, int timeout) { + } + + @Override + public SocketChannel getChannel() { + return null; + } + + @Override + public OutputStream getOutputStream() { + return out; + } + + @Override + public InputStream getInputStream() { + return in; + } + + @Override + public void close() { + } + } +} diff --git a/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/erasurecode/TestStripedBlockWriterEncryptionKey.java b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/erasurecode/TestStripedBlockWriterEncryptionKey.java new file mode 100644 index 00000000000000..790a8488d46ab4 --- /dev/null +++ b/hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/server/datanode/erasurecode/TestStripedBlockWriterEncryptionKey.java @@ -0,0 +1,62 @@ +/** + * Licensed to 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 org.apache.hadoop.hdfs.server.datanode.erasurecode; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.apache.hadoop.hdfs.protocol.datatransfer.sasl.DataEncryptionKeyFactory; +import org.apache.hadoop.hdfs.security.token.block.DataEncryptionKey; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +/** + * Test StripedBlockWriter handling of InvalidEncryptionKeyException. + */ +@Timeout(60) +public class TestStripedBlockWriterEncryptionKey { + + @Test + public void testClearEncryptionKeyOnRetry() { + CountingKeyFactory keyFactory = new CountingKeyFactory(); + + assertTrue(StripedBlockWriter.prepareRetryAfterInvalidEncryptionKey( + keyFactory, 1)); + assertEquals(1, keyFactory.clearCount); + + assertFalse(StripedBlockWriter.prepareRetryAfterInvalidEncryptionKey( + keyFactory, 2)); + assertEquals(1, keyFactory.clearCount); + } + + private static final class CountingKeyFactory + implements DataEncryptionKeyFactory { + private int clearCount; + + @Override + public DataEncryptionKey newDataEncryptionKey() { + return null; + } + + @Override + public void clearDataEncryptionKey() { + clearCount++; + } + } +}