org.apache.hadoop.hdfs.security.token.block.TestBlockToken.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.hadoop.hdfs.security.token.block.TestBlockToken.java

Source

/**
 * 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.security.token.block;

import java.io.ByteArrayInputStream;
import java.io.DataInputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.Arrays;
import java.util.EnumSet;
import java.util.Set;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.commons.logging.impl.Log4JLogger;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hdfs.protocol.ClientDatanodeProtocol;
import org.apache.hadoop.hdfs.protocol.Block;
import org.apache.hadoop.hdfs.protocol.DatanodeInfo;
import org.apache.hadoop.hdfs.protocol.LocatedBlock;
import org.apache.hadoop.io.TestWritable;
import org.apache.hadoop.ipc.Client;
import org.apache.hadoop.ipc.RPC;
import org.apache.hadoop.ipc.Server;
import org.apache.hadoop.net.NetUtils;
import org.apache.hadoop.security.SaslInputStream;
import org.apache.hadoop.security.SaslRpcClient;
import org.apache.hadoop.security.SaslRpcServer;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.security.token.TokenIdentifier;
import org.apache.log4j.Level;

import org.junit.Test;

import static org.apache.hadoop.fs.CommonConfigurationKeys.HADOOP_SECURITY_AUTHENTICATION;
import static org.junit.Assert.*;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyBoolean;
import static org.mockito.Matchers.anyLong;
import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;

/** Unit tests for block tokens */
public class TestBlockToken {
    public static final Log LOG = LogFactory.getLog(TestBlockToken.class);
    private static final String ADDRESS = "0.0.0.0";

    static final String SERVER_PRINCIPAL_KEY = "test.ipc.server.principal";
    private static Configuration conf;
    static {
        conf = new Configuration();
        conf.set(HADOOP_SECURITY_AUTHENTICATION, "kerberos");
        UserGroupInformation.setConfiguration(conf);
    }

    static {
        ((Log4JLogger) Client.LOG).getLogger().setLevel(Level.ALL);
        ((Log4JLogger) Server.LOG).getLogger().setLevel(Level.ALL);
        ((Log4JLogger) SaslRpcClient.LOG).getLogger().setLevel(Level.ALL);
        ((Log4JLogger) SaslRpcServer.LOG).getLogger().setLevel(Level.ALL);
        ((Log4JLogger) SaslInputStream.LOG).getLogger().setLevel(Level.ALL);
    }

    long blockKeyUpdateInterval = 10 * 60 * 1000; // 10 mins
    long blockTokenLifetime = 2 * 60 * 1000; // 2 mins
    Block block1 = new Block(0L);
    Block block2 = new Block(10L);
    Block block3 = new Block(-108L);

    private static class getLengthAnswer implements Answer<LocatedBlock> {
        BlockTokenSecretManager sm;
        BlockTokenIdentifier ident;

        public getLengthAnswer(BlockTokenSecretManager sm, BlockTokenIdentifier ident) {
            this.sm = sm;
            this.ident = ident;
        }

        @Override
        public LocatedBlock answer(InvocationOnMock invocation) throws IOException {
            Object args[] = invocation.getArguments();
            assertEquals(3, args.length);
            Block block = (Block) args[0];
            Set<TokenIdentifier> tokenIds = UserGroupInformation.getCurrentUser().getTokenIdentifiers();
            assertEquals("Only one BlockTokenIdentifier expected", 1, tokenIds.size());
            LocatedBlock result = null;
            for (TokenIdentifier tokenId : tokenIds) {
                BlockTokenIdentifier id = (BlockTokenIdentifier) tokenId;
                LOG.info("Got: " + id.toString());
                assertTrue("Received BlockTokenIdentifier is wrong", ident.equals(id));
                sm.checkAccess(id, null, block, BlockTokenSecretManager.AccessMode.WRITE);
                result = new LocatedBlock(new Block(id.getBlockIds()[0]), null);
            }
            return result;
        }
    }

    private BlockTokenIdentifier generateTokenId(BlockTokenSecretManager sm, Block block,
            EnumSet<BlockTokenSecretManager.AccessMode> accessModes) throws IOException {
        Token<BlockTokenIdentifier> token = sm.generateToken(block, accessModes);
        BlockTokenIdentifier id = sm.createIdentifier();
        id.readFields(new DataInputStream(new ByteArrayInputStream(token.getIdentifier())));
        return id;
    }

    @Test
    public void testWritable() throws Exception {
        TestWritable.testWritable(new BlockTokenIdentifier());
        BlockTokenSecretManager sm = new BlockTokenSecretManager(true, blockKeyUpdateInterval, blockTokenLifetime);
        TestWritable
                .testWritable(generateTokenId(sm, block1, EnumSet.allOf(BlockTokenSecretManager.AccessMode.class)));
        TestWritable
                .testWritable(generateTokenId(sm, block2, EnumSet.of(BlockTokenSecretManager.AccessMode.WRITE)));
        TestWritable.testWritable(
                generateTokenId(sm, block3, EnumSet.noneOf(BlockTokenSecretManager.AccessMode.class)));
    }

    private void tokenGenerationAndVerification(BlockTokenSecretManager master, BlockTokenSecretManager slave)
            throws Exception {
        // single-mode tokens
        for (BlockTokenSecretManager.AccessMode mode : BlockTokenSecretManager.AccessMode.values()) {
            // generated by master
            Token<BlockTokenIdentifier> token1 = master.generateToken(block1, EnumSet.of(mode));
            master.checkAccess(token1, null, block1, mode);
            slave.checkAccess(token1, null, block1, mode);
            // generated by slave
            Token<BlockTokenIdentifier> token2 = slave.generateToken(block2, EnumSet.of(mode));
            master.checkAccess(token2, null, block2, mode);
            slave.checkAccess(token2, null, block2, mode);
        }
        // multi-mode tokens
        Token<BlockTokenIdentifier> mtoken = master.generateToken(block3,
                EnumSet.allOf(BlockTokenSecretManager.AccessMode.class));
        for (BlockTokenSecretManager.AccessMode mode : BlockTokenSecretManager.AccessMode.values()) {
            master.checkAccess(mtoken, null, block3, mode);
            slave.checkAccess(mtoken, null, block3, mode);
        }
    }

    /** test block key and token handling */
    @Test
    public void testBlockTokenSecretManager() throws Exception {
        BlockTokenSecretManager masterHandler = new BlockTokenSecretManager(true, blockKeyUpdateInterval,
                blockTokenLifetime);
        BlockTokenSecretManager slaveHandler = new BlockTokenSecretManager(false, blockKeyUpdateInterval,
                blockTokenLifetime);
        ExportedBlockKeys keys = masterHandler.exportKeys();
        slaveHandler.setKeys(keys);
        tokenGenerationAndVerification(masterHandler, slaveHandler);
        // key updating
        masterHandler.updateKeys();
        tokenGenerationAndVerification(masterHandler, slaveHandler);
        keys = masterHandler.exportKeys();
        slaveHandler.setKeys(keys);
        tokenGenerationAndVerification(masterHandler, slaveHandler);
    }

    @Test
    public void testBlockTokenRpc() throws Exception {
        BlockTokenSecretManager sm = new BlockTokenSecretManager(true, blockKeyUpdateInterval, blockTokenLifetime);
        Token<BlockTokenIdentifier> token = sm.generateToken(block3,
                EnumSet.allOf(BlockTokenSecretManager.AccessMode.class));

        ClientDatanodeProtocol mockDN = mock(ClientDatanodeProtocol.class);
        when(mockDN.getProtocolVersion(anyString(), anyLong())).thenReturn(ClientDatanodeProtocol.versionID);
        BlockTokenIdentifier id = sm.createIdentifier();
        id.readFields(new DataInputStream(new ByteArrayInputStream(token.getIdentifier())));
        doAnswer(new getLengthAnswer(sm, id)).when(mockDN).recoverBlock(any(Block.class), anyBoolean(),
                any(DatanodeInfo[].class));

        final Server server = RPC.getServer(mockDN, ADDRESS, 0, 5, true, conf, sm);

        server.start();

        final InetSocketAddress addr = NetUtils.getConnectAddress(server);
        final UserGroupInformation ticket = UserGroupInformation.createRemoteUser(block3.toString());
        ticket.addToken(token);

        ClientDatanodeProtocol proxy = null;
        try {
            proxy = (ClientDatanodeProtocol) RPC.getProxy(ClientDatanodeProtocol.class,
                    ClientDatanodeProtocol.versionID, addr, ticket, conf, NetUtils.getDefaultSocketFactory(conf));
            LocatedBlock lb = proxy.recoverBlock(block3, true, null);
            assertEquals(block3.getBlockId(), lb.getBlock().getBlockId());
        } finally {
            server.stop();
            if (proxy != null) {
                RPC.stopProxy(proxy);
            }
        }
    }

    @Test
    public void collectionOfBlocksActsSanely() {
        final long[][] testBlockIds = new long[][] { { 99l, 7l, -32l, 0l }, {}, { 42l }, { -5235l, 2352 } };
        final long[] notBlockIds = new long[] { 32l, 1l, -23423423l };

        for (long[] bids : testBlockIds) {
            BlockTokenIdentifier bti = new BlockTokenIdentifier("Madame Butterfly", bids,
                    EnumSet.noneOf(BlockTokenSecretManager.AccessMode.class));

            for (long bid : bids)
                assertTrue(bti.isBlockIncluded(bid));

            for (long nbid : notBlockIds)
                assertFalse(bti.isBlockIncluded(nbid));

            // BlockTokenIdentifiers maintain a sorted array of the block Ids.
            long[] sorted = Arrays.copyOf(bids, bids.length);
            Arrays.sort(sorted);

            assertTrue(Arrays.toString(bids) + " doesn't equal " + Arrays.toString(sorted),
                    Arrays.equals(bti.getBlockIds(), sorted));
        }
    }
}