org.apache.hadoop.yarn.server.TestContainerTokenSecretManager.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.hadoop.yarn.server.TestContainerTokenSecretManager.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.yarn.server;

import static org.junit.Assert.fail;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileWriter;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.security.PrivilegedAction;
import java.util.ArrayList;
import java.util.List;

import junit.framework.Assert;

import org.apache.avro.AvroRuntimeException;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.CommonConfigurationKeysPublic;
import org.apache.hadoop.fs.FileContext;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.fs.UnsupportedFileSystemException;
import org.apache.hadoop.io.DataInputBuffer;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.net.NetUtils;
import org.apache.hadoop.security.AccessControlException;
import org.apache.hadoop.security.SecurityInfo;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.yarn.api.AMRMProtocol;
import org.apache.hadoop.yarn.api.ContainerManager;
import org.apache.hadoop.yarn.api.protocolrecords.AllocateRequest;
import org.apache.hadoop.yarn.api.protocolrecords.GetContainerStatusRequest;
import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterRequest;
import org.apache.hadoop.yarn.api.protocolrecords.SubmitApplicationRequest;
import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.api.records.ApplicationMaster;
import org.apache.hadoop.yarn.api.records.ApplicationSubmissionContext;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.api.records.ContainerToken;
import org.apache.hadoop.yarn.api.records.LocalResource;
import org.apache.hadoop.yarn.api.records.LocalResourceType;
import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
import org.apache.hadoop.yarn.api.records.Priority;
import org.apache.hadoop.yarn.api.records.Resource;
import org.apache.hadoop.yarn.api.records.ResourceRequest;
import org.apache.hadoop.yarn.api.records.URL;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.exceptions.YarnRemoteException;
import org.apache.hadoop.yarn.factories.RecordFactory;
import org.apache.hadoop.yarn.factory.providers.RecordFactoryProvider;
import org.apache.hadoop.yarn.ipc.YarnRPC;
import org.apache.hadoop.yarn.security.ApplicationTokenIdentifier;
import org.apache.hadoop.yarn.security.ApplicationTokenSecretManager;
import org.apache.hadoop.yarn.security.ContainerManagerSecurityInfo;
import org.apache.hadoop.yarn.security.ContainerTokenIdentifier;
import org.apache.hadoop.yarn.security.SchedulerSecurityInfo;
import org.apache.hadoop.yarn.server.resourcemanager.ResourceManager;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttempt;
import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptState;
import org.apache.hadoop.yarn.util.ConverterUtils;
import org.junit.BeforeClass;
import org.junit.Test;

public class TestContainerTokenSecretManager {

    private static Log LOG = LogFactory.getLog(TestContainerTokenSecretManager.class);
    private static final RecordFactory recordFactory = RecordFactoryProvider.getRecordFactory(null);
    private static FileContext localFS = null;
    private static final File localDir = new File("target",
            TestContainerTokenSecretManager.class.getName() + "-localDir").getAbsoluteFile();

    @BeforeClass
    public static void setup()
            throws AccessControlException, FileNotFoundException, UnsupportedFileSystemException, IOException {
        localFS = FileContext.getLocalFSFileContext();
        localFS.delete(new Path(localDir.getAbsolutePath()), true);
        localDir.mkdir();
    }

    @Test
    public void test() throws IOException, InterruptedException {

        final ApplicationId appID = recordFactory.newRecordInstance(ApplicationId.class);
        appID.setClusterTimestamp(1234);
        appID.setId(5);

        final Configuration conf = new Configuration();
        conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION, "kerberos");
        // Set AM expiry interval to be very long.
        conf.setLong(YarnConfiguration.AM_EXPIRY_INTERVAL, 100000L);
        UserGroupInformation.setConfiguration(conf);
        MiniYARNCluster yarnCluster = new MiniYARNCluster(TestContainerTokenSecretManager.class.getName());
        yarnCluster.init(conf);
        yarnCluster.start();

        ResourceManager resourceManager = yarnCluster.getResourceManager();

        final YarnRPC yarnRPC = YarnRPC.create(conf);

        // Submit an application
        ApplicationSubmissionContext appSubmissionContext = recordFactory
                .newRecordInstance(ApplicationSubmissionContext.class);
        appSubmissionContext.setApplicationId(appID);
        appSubmissionContext.setMasterCapability(recordFactory.newRecordInstance(Resource.class));
        appSubmissionContext.getMasterCapability().setMemory(1024);
        //    appSubmissionContext.resources = new HashMap<String, URL>();
        appSubmissionContext.setUser("testUser");
        //    appSubmissionContext.environment = new HashMap<String, String>();
        //    appSubmissionContext.command = new ArrayList<String>();
        appSubmissionContext.addCommand("sleep");
        appSubmissionContext.addCommand("100");

        // TODO: Use a resource to work around bugs. Today NM doesn't create local
        // app-dirs if there are no file to download!!
        File file = new File(localDir.getAbsolutePath(), "testFile");
        FileWriter tmpFile = new FileWriter(file);
        tmpFile.write("testing");
        tmpFile.close();
        URL testFileURL = ConverterUtils.getYarnUrlFromPath(
                FileContext.getFileContext().makeQualified(new Path(localDir.getAbsolutePath(), "testFile")));
        LocalResource rsrc = recordFactory.newRecordInstance(LocalResource.class);
        rsrc.setResource(testFileURL);
        rsrc.setSize(file.length());
        rsrc.setTimestamp(file.lastModified());
        rsrc.setType(LocalResourceType.FILE);
        rsrc.setVisibility(LocalResourceVisibility.PRIVATE);
        appSubmissionContext.setResourceTodo("testFile", rsrc);
        SubmitApplicationRequest submitRequest = recordFactory.newRecordInstance(SubmitApplicationRequest.class);
        submitRequest.setApplicationSubmissionContext(appSubmissionContext);
        resourceManager.getClientRMService().submitApplication(submitRequest);

        // Wait till container gets allocated for AM
        int waitCounter = 0;
        RMApp app = resourceManager.getRMContext().getRMApps().get(appID);
        RMAppAttempt appAttempt = app == null ? null : app.getCurrentAppAttempt();
        RMAppAttemptState state = appAttempt == null ? null : appAttempt.getAppAttemptState();
        while ((app == null || appAttempt == null || state == null || !state.equals(RMAppAttemptState.LAUNCHED))
                && waitCounter++ != 20) {
            LOG.info("Waiting for applicationAttempt to be created.. ");
            Thread.sleep(1000);
            app = resourceManager.getRMContext().getRMApps().get(appID);
            appAttempt = app == null ? null : app.getCurrentAppAttempt();
            state = appAttempt == null ? null : appAttempt.getAppAttemptState();
        }
        Assert.assertNotNull(app);
        Assert.assertNotNull(appAttempt);
        Assert.assertNotNull(state);
        Assert.assertEquals(RMAppAttemptState.LAUNCHED, state);

        UserGroupInformation currentUser = UserGroupInformation.getCurrentUser();

        // Ask for a container from the RM
        String schedulerAddressString = conf.get(YarnConfiguration.SCHEDULER_ADDRESS,
                YarnConfiguration.DEFAULT_SCHEDULER_BIND_ADDRESS);
        final InetSocketAddress schedulerAddr = NetUtils.createSocketAddr(schedulerAddressString);
        ApplicationTokenIdentifier appTokenIdentifier = new ApplicationTokenIdentifier(appID);
        ApplicationTokenSecretManager appTokenSecretManager = new ApplicationTokenSecretManager();
        appTokenSecretManager.setMasterKey(ApplicationTokenSecretManager.createSecretKey("Dummy".getBytes())); // TODO: FIX. Be in Sync with
                                                                                                               // ResourceManager.java
        Token<ApplicationTokenIdentifier> appToken = new Token<ApplicationTokenIdentifier>(appTokenIdentifier,
                appTokenSecretManager);
        appToken.setService(new Text(schedulerAddressString));
        currentUser.addToken(appToken);

        conf.setClass(YarnConfiguration.YARN_SECURITY_INFO, SchedulerSecurityInfo.class, SecurityInfo.class);
        AMRMProtocol scheduler = currentUser.doAs(new PrivilegedAction<AMRMProtocol>() {
            @Override
            public AMRMProtocol run() {
                return (AMRMProtocol) yarnRPC.getProxy(AMRMProtocol.class, schedulerAddr, conf);
            }
        });

        // Register the appMaster 
        RegisterApplicationMasterRequest request = recordFactory
                .newRecordInstance(RegisterApplicationMasterRequest.class);
        ApplicationMaster applicationMaster = recordFactory.newRecordInstance(ApplicationMaster.class);
        request.setApplicationAttemptId(
                resourceManager.getRMContext().getRMApps().get(appID).getCurrentAppAttempt().getAppAttemptId());
        scheduler.registerApplicationMaster(request);

        // Now request a container allocation.
        List<ResourceRequest> ask = new ArrayList<ResourceRequest>();
        ResourceRequest rr = recordFactory.newRecordInstance(ResourceRequest.class);
        rr.setCapability(recordFactory.newRecordInstance(Resource.class));
        rr.getCapability().setMemory(1024);
        rr.setHostName("*");
        rr.setNumContainers(1);
        rr.setPriority(recordFactory.newRecordInstance(Priority.class));
        rr.getPriority().setPriority(0);
        ask.add(rr);
        ArrayList<ContainerId> release = new ArrayList<ContainerId>();

        AllocateRequest allocateRequest = recordFactory.newRecordInstance(AllocateRequest.class);
        allocateRequest.setApplicationAttemptId(appAttempt.getAppAttemptId());
        allocateRequest.setResponseId(0);
        allocateRequest.addAllAsks(ask);
        allocateRequest.addAllReleases(release);
        List<Container> allocatedContainers = scheduler.allocate(allocateRequest).getAMResponse()
                .getNewContainerList();

        waitCounter = 0;
        while ((allocatedContainers == null || allocatedContainers.size() == 0) && waitCounter++ != 20) {
            LOG.info("Waiting for container to be allocated..");
            Thread.sleep(1000);
            allocateRequest.setResponseId(allocateRequest.getResponseId() + 1);
            allocatedContainers = scheduler.allocate(allocateRequest).getAMResponse().getNewContainerList();
        }

        Assert.assertNotNull("Container is not allocted!", allocatedContainers);
        Assert.assertEquals("Didn't get one container!", 1, allocatedContainers.size());

        // Now talk to the NM for launching the container.
        final Container allocatedContainer = allocatedContainers.get(0);
        ContainerToken containerToken = allocatedContainer.getContainerToken();
        Token<ContainerTokenIdentifier> token = new Token<ContainerTokenIdentifier>(
                containerToken.getIdentifier().array(), containerToken.getPassword().array(),
                new Text(containerToken.getKind()), new Text(containerToken.getService()));
        currentUser.addToken(token);
        conf.setClass(YarnConfiguration.YARN_SECURITY_INFO, ContainerManagerSecurityInfo.class, SecurityInfo.class);
        currentUser.doAs(new PrivilegedAction<Void>() {
            @Override
            public Void run() {
                ContainerManager client = (ContainerManager) yarnRPC.getProxy(ContainerManager.class,
                        NetUtils.createSocketAddr(allocatedContainer.getNodeId().toString()), conf);
                try {
                    LOG.info("Going to make a getContainerStatus() legal request");
                    GetContainerStatusRequest request = recordFactory
                            .newRecordInstance(GetContainerStatusRequest.class);
                    ContainerId containerID = recordFactory.newRecordInstance(ContainerId.class);
                    ApplicationAttemptId appAttemptId = recordFactory.newRecordInstance(ApplicationAttemptId.class);
                    appAttemptId.setApplicationId(appID);
                    appAttemptId.setAttemptId(1);
                    containerID.setAppId(appID);
                    containerID.setId(1);
                    containerID.setAppAttemptId(appAttemptId);
                    request.setContainerId(containerID);
                    client.getContainerStatus(request);
                } catch (YarnRemoteException e) {
                    LOG.info("Error", e);
                } catch (AvroRuntimeException e) {
                    LOG.info("Got the expected exception");
                }
                return null;
            }
        });

        UserGroupInformation maliceUser = UserGroupInformation.createRemoteUser(currentUser.getShortUserName());
        byte[] identifierBytes = containerToken.getIdentifier().array();
        DataInputBuffer di = new DataInputBuffer();
        di.reset(identifierBytes, identifierBytes.length);
        ContainerTokenIdentifier dummyIdentifier = new ContainerTokenIdentifier();
        dummyIdentifier.readFields(di);
        Resource modifiedResource = recordFactory.newRecordInstance(Resource.class);
        modifiedResource.setMemory(2048);
        ContainerTokenIdentifier modifiedIdentifier = new ContainerTokenIdentifier(dummyIdentifier.getContainerID(),
                dummyIdentifier.getNmHostName(), modifiedResource);
        // Malice user modifies the resource amount
        Token<ContainerTokenIdentifier> modifiedToken = new Token<ContainerTokenIdentifier>(
                modifiedIdentifier.getBytes(), containerToken.getPassword().array(),
                new Text(containerToken.getKind()), new Text(containerToken.getService()));
        maliceUser.addToken(modifiedToken);
        maliceUser.doAs(new PrivilegedAction<Void>() {
            @Override
            public Void run() {
                ContainerManager client = (ContainerManager) yarnRPC.getProxy(ContainerManager.class,
                        NetUtils.createSocketAddr(allocatedContainer.getNodeId().toString()), conf);
                ContainerId containerID;

                LOG.info("Going to contact NM:  ilLegal request");
                GetContainerStatusRequest request = recordFactory
                        .newRecordInstance(GetContainerStatusRequest.class);
                containerID = recordFactory.newRecordInstance(ContainerId.class);
                ApplicationAttemptId appAttemptId = recordFactory.newRecordInstance(ApplicationAttemptId.class);
                appAttemptId.setApplicationId(appID);
                appAttemptId.setAttemptId(1);
                containerID.setAppId(appID);
                containerID.setId(1);
                containerID.setAppAttemptId(appAttemptId);
                request.setContainerId(containerID);
                try {
                    client.getContainerStatus(request);
                    fail("Connection initiation with illegally modified " + "tokens is expected to fail.");
                } catch (YarnRemoteException e) {
                    LOG.error("Got exception", e);
                    fail("Cannot get a YARN remote exception as " + "it will indicate RPC success");
                } catch (Exception e) {
                    Assert.assertEquals(java.lang.reflect.UndeclaredThrowableException.class.getCanonicalName(),
                            e.getClass().getCanonicalName());
                    Assert.assertEquals("DIGEST-MD5: digest response format violation. Mismatched response.",
                            e.getCause().getCause().getMessage());
                }
                return null;
            }
        });
    }
}