Java tutorial
/** * 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; } }); } }