org.apache.hadoop.mapreduce.security.TestJHSSecurity.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.hadoop.mapreduce.security.TestJHSSecurity.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.mapreduce.security;

import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.security.PrivilegedAction;
import java.security.PrivilegedExceptionAction;

import org.junit.Assert;

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.mapred.JobConf;
import org.apache.hadoop.mapreduce.v2.api.HSClientProtocol;
import org.apache.hadoop.mapreduce.v2.api.MRClientProtocol;
import org.apache.hadoop.mapreduce.v2.api.protocolrecords.CancelDelegationTokenRequest;
import org.apache.hadoop.mapreduce.v2.api.protocolrecords.GetDelegationTokenRequest;
import org.apache.hadoop.mapreduce.v2.api.protocolrecords.GetJobReportRequest;
import org.apache.hadoop.mapreduce.v2.api.protocolrecords.RenewDelegationTokenRequest;
import org.apache.hadoop.mapreduce.v2.hs.HistoryClientService;
import org.apache.hadoop.mapreduce.v2.hs.HistoryServerStateStoreService;
import org.apache.hadoop.mapreduce.v2.hs.JHSDelegationTokenSecretManager;
import org.apache.hadoop.mapreduce.v2.hs.JobHistoryServer;
import org.apache.hadoop.mapreduce.v2.jobhistory.JHAdminConfig;
import org.apache.hadoop.mapreduce.v2.util.MRBuilderUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.UserGroupInformation.AuthenticationMethod;
import org.apache.hadoop.yarn.api.records.Token;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.ipc.YarnRPC;
import org.apache.hadoop.yarn.util.ConverterUtils;
import org.apache.hadoop.yarn.util.Records;
import org.apache.log4j.Level;
import org.apache.log4j.LogManager;
import org.apache.log4j.Logger;
import org.junit.Test;

public class TestJHSSecurity {

    private static final Log LOG = LogFactory.getLog(TestJHSSecurity.class);

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

        Logger rootLogger = LogManager.getRootLogger();
        rootLogger.setLevel(Level.DEBUG);

        final YarnConfiguration conf = new YarnConfiguration(new JobConf());
        // Just a random principle
        conf.set(JHAdminConfig.MR_HISTORY_PRINCIPAL, "RandomOrc/localhost@apache.org");

        conf.set(CommonConfigurationKeysPublic.HADOOP_SECURITY_AUTHENTICATION, "kerberos");
        UserGroupInformation.setConfiguration(conf);

        final long initialInterval = 10000l;
        final long maxLifetime = 20000l;
        final long renewInterval = 10000l;

        JobHistoryServer jobHistoryServer = null;
        MRClientProtocol clientUsingDT = null;
        long tokenFetchTime;
        try {
            jobHistoryServer = new JobHistoryServer() {
                protected void doSecureLogin(Configuration conf) throws IOException {
                    // no keytab based login
                };

                @Override
                protected JHSDelegationTokenSecretManager createJHSSecretManager(Configuration conf,
                        HistoryServerStateStoreService store) {
                    return new JHSDelegationTokenSecretManager(initialInterval, maxLifetime, renewInterval, 3600000,
                            store);
                }

                @Override
                protected HistoryClientService createHistoryClientService() {
                    return new HistoryClientService(historyContext, this.jhsDTSecretManager) {
                        @Override
                        protected void initializeWebApp(Configuration conf) {
                            // Don't need it, skip.;
                        }
                    };
                }
            };
            //      final JobHistoryServer jobHistoryServer = jhServer;
            jobHistoryServer.init(conf);
            jobHistoryServer.start();
            final MRClientProtocol hsService = jobHistoryServer.getClientService().getClientHandler();

            // Fake the authentication-method
            UserGroupInformation loggedInUser = UserGroupInformation.createRemoteUser("testrenewer@APACHE.ORG");
            Assert.assertEquals("testrenewer", loggedInUser.getShortUserName());
            // Default realm is APACHE.ORG
            loggedInUser.setAuthenticationMethod(AuthenticationMethod.KERBEROS);

            Token token = getDelegationToken(loggedInUser, hsService, loggedInUser.getShortUserName());
            tokenFetchTime = System.currentTimeMillis();
            LOG.info("Got delegation token at: " + tokenFetchTime);

            // Now try talking to JHS using the delegation token
            clientUsingDT = getMRClientProtocol(token, jobHistoryServer.getClientService().getBindAddress(),
                    "TheDarkLord", conf);

            GetJobReportRequest jobReportRequest = Records.newRecord(GetJobReportRequest.class);
            jobReportRequest.setJobId(MRBuilderUtils.newJobId(123456, 1, 1));
            try {
                clientUsingDT.getJobReport(jobReportRequest);
            } catch (IOException e) {
                Assert.assertEquals("Unknown job job_123456_0001", e.getMessage());
            }

            // Renew after 50% of token age.
            while (System.currentTimeMillis() < tokenFetchTime + initialInterval / 2) {
                Thread.sleep(500l);
            }
            long nextExpTime = renewDelegationToken(loggedInUser, hsService, token);
            long renewalTime = System.currentTimeMillis();
            LOG.info("Renewed token at: " + renewalTime + ", NextExpiryTime: " + nextExpTime);

            // Wait for first expiry, but before renewed expiry.
            while (System.currentTimeMillis() > tokenFetchTime + initialInterval
                    && System.currentTimeMillis() < nextExpTime) {
                Thread.sleep(500l);
            }
            Thread.sleep(50l);

            // Valid token because of renewal.
            try {
                clientUsingDT.getJobReport(jobReportRequest);
            } catch (IOException e) {
                Assert.assertEquals("Unknown job job_123456_0001", e.getMessage());
            }

            // Wait for expiry.
            while (System.currentTimeMillis() < renewalTime + renewInterval) {
                Thread.sleep(500l);
            }
            Thread.sleep(50l);
            LOG.info("At time: " + System.currentTimeMillis() + ", token should be invalid");
            // Token should have expired.      
            try {
                clientUsingDT.getJobReport(jobReportRequest);
                fail("Should not have succeeded with an expired token");
            } catch (IOException e) {
                assertTrue(e.getCause().getMessage().contains("is expired"));
            }

            // Test cancellation
            // Stop the existing proxy, start another.
            if (clientUsingDT != null) {
                //        RPC.stopProxy(clientUsingDT);
                clientUsingDT = null;
            }
            token = getDelegationToken(loggedInUser, hsService, loggedInUser.getShortUserName());
            tokenFetchTime = System.currentTimeMillis();
            LOG.info("Got delegation token at: " + tokenFetchTime);

            // Now try talking to HSService using the delegation token
            clientUsingDT = getMRClientProtocol(token, jobHistoryServer.getClientService().getBindAddress(),
                    "loginuser2", conf);

            try {
                clientUsingDT.getJobReport(jobReportRequest);
            } catch (IOException e) {
                fail("Unexpected exception" + e);
            }
            cancelDelegationToken(loggedInUser, hsService, token);

            // Testing the token with different renewer to cancel the token
            Token tokenWithDifferentRenewer = getDelegationToken(loggedInUser, hsService, "yarn");
            cancelDelegationToken(loggedInUser, hsService, tokenWithDifferentRenewer);
            if (clientUsingDT != null) {
                //        RPC.stopProxy(clientUsingDT);
                clientUsingDT = null;
            }

            // Creating a new connection.
            clientUsingDT = getMRClientProtocol(token, jobHistoryServer.getClientService().getBindAddress(),
                    "loginuser2", conf);
            LOG.info("Cancelled delegation token at: " + System.currentTimeMillis());
            // Verify cancellation worked.
            try {
                clientUsingDT.getJobReport(jobReportRequest);
                fail("Should not have succeeded with a cancelled delegation token");
            } catch (IOException e) {
            }

        } finally {
            jobHistoryServer.stop();
        }
    }

    private Token getDelegationToken(final UserGroupInformation loggedInUser, final MRClientProtocol hsService,
            final String renewerString) throws IOException, InterruptedException {
        // Get the delegation token directly as it is a little difficult to setup
        // the kerberos based rpc.
        Token token = loggedInUser.doAs(new PrivilegedExceptionAction<Token>() {
            @Override
            public Token run() throws IOException {
                GetDelegationTokenRequest request = Records.newRecord(GetDelegationTokenRequest.class);
                request.setRenewer(renewerString);
                return hsService.getDelegationToken(request).getDelegationToken();
            }

        });
        return token;
    }

    private long renewDelegationToken(final UserGroupInformation loggedInUser, final MRClientProtocol hsService,
            final Token dToken) throws IOException, InterruptedException {
        long nextExpTime = loggedInUser.doAs(new PrivilegedExceptionAction<Long>() {

            @Override
            public Long run() throws IOException {
                RenewDelegationTokenRequest request = Records.newRecord(RenewDelegationTokenRequest.class);
                request.setDelegationToken(dToken);
                return hsService.renewDelegationToken(request).getNextExpirationTime();
            }
        });
        return nextExpTime;
    }

    private void cancelDelegationToken(final UserGroupInformation loggedInUser, final MRClientProtocol hsService,
            final Token dToken) throws IOException, InterruptedException {

        loggedInUser.doAs(new PrivilegedExceptionAction<Void>() {
            @Override
            public Void run() throws IOException {
                CancelDelegationTokenRequest request = Records.newRecord(CancelDelegationTokenRequest.class);
                request.setDelegationToken(dToken);
                hsService.cancelDelegationToken(request);
                return null;
            }
        });
    }

    private MRClientProtocol getMRClientProtocol(Token token, final InetSocketAddress hsAddress, String user,
            final Configuration conf) {
        UserGroupInformation ugi = UserGroupInformation.createRemoteUser(user);
        ugi.addToken(ConverterUtils.convertFromYarn(token, hsAddress));

        final YarnRPC rpc = YarnRPC.create(conf);
        MRClientProtocol hsWithDT = ugi.doAs(new PrivilegedAction<MRClientProtocol>() {

            @Override
            public MRClientProtocol run() {
                return (MRClientProtocol) rpc.getProxy(HSClientProtocol.class, hsAddress, conf);
            }
        });
        return hsWithDT;
    }

}