Undertow SAML SP

gistfile1.txt

/*
 * JBoss, Home of Professional Open Source
 *
 * Copyright 2013 Red Hat, Inc. and/or its affiliates.
 *
 * Licensed 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.picketlink.test.identity.federation.bindings.wildfly;

import com.meterware.httpunit.GetMethodWebRequest;
import com.meterware.httpunit.SubmitButton;
import com.meterware.httpunit.WebConversation;
import com.meterware.httpunit.WebForm;
import com.meterware.httpunit.WebRequest;
import com.meterware.httpunit.WebResponse;
import io.undertow.Undertow;
import io.undertow.security.idm.Account;
import io.undertow.security.idm.Credential;
import io.undertow.security.idm.IdentityManager;
import io.undertow.security.idm.PasswordCredential;
import io.undertow.server.HttpHandler;
import io.undertow.server.handlers.PathHandler;
import io.undertow.server.handlers.resource.Resource;
import io.undertow.server.handlers.resource.ResourceChangeListener;
import io.undertow.server.handlers.resource.ResourceManager;
import io.undertow.server.handlers.resource.URLResource;
import io.undertow.servlet.api.DeploymentInfo;
import io.undertow.servlet.api.DeploymentManager;
import io.undertow.servlet.api.FilterInfo;
import io.undertow.servlet.api.InstanceFactory;
import io.undertow.servlet.api.InstanceHandle;
import io.undertow.servlet.api.ListenerInfo;
import io.undertow.servlet.api.LoginConfig;
import io.undertow.servlet.api.ServletContainer;
import io.undertow.servlet.api.ServletInfo;
import io.undertow.servlet.api.ServletSecurityInfo;
import org.apache.http.Header;
import org.apache.http.HttpEntity;
import org.apache.http.HttpResponse;
import org.apache.http.NameValuePair;
import org.apache.http.client.entity.UrlEncodedFormEntity;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.message.BasicNameValuePair;
import org.apache.http.util.EntityUtils;
import org.junit.After;
import org.junit.Before;
import org.junit.Ignore;
import org.junit.Test;
import org.picketlink.identity.federation.bindings.wildfly.sp.SPFormAuthenticationMechanism;
import org.picketlink.identity.federation.web.filters.IDPFilter;

import javax.servlet.DispatcherType;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Writer;
import java.net.URL;
import java.security.Principal;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.EventListener;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import static junit.framework.Assert.assertNotNull;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

/**
 * Simple Workflow for SAML SSO using Undertow
 * @author Anil Saldhana
 * @since November 14, 2013
 */
public class SPInitiatedSSOWorkflowTestCase extends UndertowTestCase {
    final PathHandler path = new PathHandler();

    @Override
    protected HttpHandler getHandler() {
        System.out.println("Inside SPInitiatedSSOWorkflowTestCase -> getHandler");
        return path;
    }

    @Before
    public void setup() throws Exception{
        super.setup();
        assertNotNull(server);
        deployIDP();
        deploySP();
    }

    @After
    public void tearDown() throws Exception{
        if(server != null){
            server.stop();
        }
    }

    public void deployIDP() throws Exception{
        final ServletContainer container = ServletContainer.Factory.newInstance();
        FilterInfo idpFilterInfo = new FilterInfo("IDPFilter", IDPFilter.class);

        ServletInfo regularServletInfo = new ServletInfo("servlet", SendUsernameServlet.class)
                .setServletSecurityInfo(new ServletSecurityInfo()
                        .addRoleAllowed("role1"))
                .addMapping("/*")
                ;

        ServletInfo formServletInfo = new ServletInfo("loginPage", FormLoginServlet.class)
                .setServletSecurityInfo(new ServletSecurityInfo()
                        .addRoleAllowed("group1"))
                .addMapping("/FormLoginServlet");

        TestIdentityManager identityManager = new TestIdentityManager();
        identityManager.addUser("user1", "password1", "role1");

        LoginConfig loginConfig = new LoginConfig("FORM", "Test Realm", "/FormLoginServlet","/error.html");

        DeploymentInfo deploymentInfo = new DeploymentInfo()
                .setClassLoader(SPInitiatedSSOWorkflowTestCase.class.getClassLoader())
                .setContextPath("/idp")
                .setDeploymentName("idp.war")
                .setClassIntrospecter(TestClassIntrospector.INSTANCE)
                .setIdentityManager(identityManager)
                .setLoginConfig(loginConfig)
                .setResourceManager(new TestResourceManager("idp"))
                .addServlets(regularServletInfo, formServletInfo)
                .addFilter(idpFilterInfo)
                .addFilterUrlMapping(idpFilterInfo.getName(), "/*", DispatcherType.REQUEST);

        DeploymentManager manager = container.addDeployment(deploymentInfo);
        manager.deploy();

        try{
            path.addPath(deploymentInfo.getContextPath(), manager.start());
        }catch(ServletException se){
            throw new RuntimeException(se);
        }
        System.out.println("Deployment success:" + deploymentInfo.getContextPath());
    }

    public void deploySP() throws Exception{
        final ServletContainer container = ServletContainer.Factory.newInstance();

        ServletInfo welcomeServlet = new ServletInfo("/", WelcomeServlet.class)
                .addMapping("/WelcomeServlet");

        ServletInfo regularServletInfo = new ServletInfo("servlet", SendUsernameServlet.class)
                .setServletSecurityInfo(new ServletSecurityInfo()
                        .addRoleAllowed("role1"))
                .addMapping("/secured/*");

        ServletInfo formServletInfo = new ServletInfo("loginPage", FormLoginServlet.class)
                .setServletSecurityInfo(new ServletSecurityInfo()
                        .addRoleAllowed("group1"))
                .addMapping("/FormLoginServlet");

        TestIdentityManager identityManager = new TestIdentityManager();
        identityManager.addUser("user1", "password1", "role1");

        LoginConfig loginConfig = new LoginConfig("FORM", "Test Realm", "/FormLoginServlet","/error.html");
        final SPFormAuthenticationMechanism spFormAuthenticationMechanism = new SPFormAuthenticationMechanism("FORM", "/FormLoginServlet", "/error.html");
        DeploymentInfo deploymentInfo = new DeploymentInfo()
                .setClassLoader(SPInitiatedSSOWorkflowTestCase.class.getClassLoader())
                .setContextPath("/sp")
                .setDeploymentName("sp.war")
                .setClassIntrospecter(TestClassIntrospector.INSTANCE)
                .setIdentityManager(identityManager)
                .setLoginConfig(loginConfig)
                .setResourceManager(new TestResourceManager("sp"))
                .addServlets(regularServletInfo, formServletInfo)
                .addAuthenticationMechanism(spFormAuthenticationMechanism)
                .setIgnoreStandardAuthenticationMechanism(true);

        ListenerInfo listenerInfo = new ListenerInfo(SPFormAuthenticationMechanism.class, new InstanceFactory<EventListener>() {
            @Override
            public InstanceHandle<EventListener> createInstance() throws InstantiationException {
                return new InstanceHandle<EventListener>() {
                    @Override
                    public EventListener getInstance() {
                        return spFormAuthenticationMechanism;
                    }

                    @Override
                    public void release() {
                    }
                };
            }
        });
        deploymentInfo.addListener(listenerInfo);

        DeploymentManager manager = container.addDeployment(deploymentInfo);
        manager.deploy();

        try{
            path.addPath(deploymentInfo.getContextPath(), manager.start());
        }catch(ServletException se){
            throw new RuntimeException(se);
        }
        System.out.println("Deployment success:" + deploymentInfo.getContextPath());
    }

    /**
     * Customized implementation of {@link IdentityManager}
     */
    public class TestIdentityManager implements IdentityManager{
        private String userName,password,role;

        public void addUser(String userName, String password, String role){
            this.userName = userName;
            this.password = password;
            this.role = role;
        }
        @Override
        public Account verify(Account account) {
            return account;
        }

        @Override
        public Account verify(String userName, Credential credential) {
            PasswordCredential pwd = (PasswordCredential) credential;
            String pw = new String(pwd.getPassword());

            if(userName.equals("user1") && pw.equals("password1")){
                return new Account() {
                    @Override
                    public Principal getPrincipal() {
                        return new Principal() {
                            @Override
                            public String getName() {
                                return "user1";
                            }
                        };
                    }

                    @Override
                    public Set<String> getRoles() {
                        HashSet<String> set = new HashSet<String>();
                        set.add("role1");
                        return set;
                    }
                };
            }else {
                return null;
            }
        }

        @Override
        public Account verify(Credential credential) {
            throw new RuntimeException();
        }
    }
    public static class WelcomeServlet extends HttpServlet {
        @Override
        protected void doGet(final HttpServletRequest req, final HttpServletResponse resp) throws ServletException, IOException {
            OutputStream stream = resp.getOutputStream();
            stream.write("Welcome".getBytes());
        }
    }

    /**
     * @author Stuart Douglas
     */
    public static class SendUsernameServlet extends HttpServlet {
        @Override
        protected void doGet(final HttpServletRequest req, final HttpServletResponse resp) throws ServletException, IOException {
            OutputStream stream = resp.getOutputStream();
            Principal principal = req.getUserPrincipal();
            String name = principal.getName();
            stream.write(name.getBytes());
        }
        @Override
        protected void doPost(final HttpServletRequest req, final HttpServletResponse resp) throws ServletException, IOException {
            OutputStream stream = resp.getOutputStream();
            Principal principal = req.getUserPrincipal();
            String name = principal.getName();
            stream.write(name.getBytes());
        }
    }
    /*
     * @author Stuart Douglas
     */
    public static class FormLoginServlet extends HttpServlet {

        @Override
        protected void doGet(final HttpServletRequest req, final HttpServletResponse resp) throws ServletException, IOException {
            Writer writer = resp.getWriter();
            writer.write("Login Page");
            writer.write("<form id=\"login_form\" name=\"login_form\" method=\"post\"\n" +
                    "                        action=\"j_security_check\" enctype=\"application/x-www-form-urlencoded\">\n" +
                    "                        <div style=\"margin-left: 15px;\">\n" +
                    "                                <p>\n" +
                    "                                        <label for=\"username\"> Username</label><br /> <input id=\"username\"\n" +
                    "                                                type=\"text\" name=\"j_username\" size=\"20\" />\n" +
                    "                                </p>\n" +
                    "                                <p>\n" +
                    "                                        <label for=\"password\"> Password</label><br /> <input id=\"password\"\n" +
                    "                                                type=\"password\" name=\"j_password\" value=\"\" size=\"20\" />\n" +
                    "                                </p>\n" +
                    "                                <center>\n" +
                    "                                        <input id=\"submit\" type=\"submit\" name=\"submit\" value=\"Login\"\n" +
                    "                                                class=\"buttonmed\" />\n" +
                    "                                </center>\n" +
                    "                        </div>\n" +
                    "                </form>");

        }
    }


    @Test
    public void testServerUp() throws Exception{
    }

    public class TestResourceManager implements ResourceManager{

        private final String basePath;

        public TestResourceManager(String basePath){
            this.basePath = basePath;
        }

        @Override
        public Resource getResource(String path) throws IOException {
            String temp = path;
            //Remove WEB-INF
            temp = temp.replace("/WEB-INF","");

            URL url = getClass().getClassLoader().getResource(basePath+temp);
            return new URLResource(url, url.openConnection(), path);
        }

        @Override
        public boolean isResourceChangeListenerSupported() {
            throw new RuntimeException();
        }

        @Override
        public void registerResourceChangeListener(ResourceChangeListener listener) {
            throw new RuntimeException();
        }

        @Override
        public void removeResourceChangeListener(ResourceChangeListener listener) {
            throw new RuntimeException();
        }

        @Override
        public void close() throws IOException {
            throw new RuntimeException();
        }
    }

    @Test
    public void testSSO() throws Exception{
        String spURI = "http://localhost:8080/sp/secured/test";

        WebRequest serviceRequest1 = new GetMethodWebRequest(spURI);
        WebConversation webConversation = new WebConversation();

        WebResponse webResponse = webConversation.getResponse(serviceRequest1);
        WebForm loginForm = webResponse.getForms()[0];
        loginForm.setParameter("j_username", "user1");
        loginForm.setParameter("j_password", "password1");
        SubmitButton submitButton = loginForm.getSubmitButtons()[0];
        submitButton.click();


        webResponse = webConversation.getCurrentPage();
        assertTrue(" Reached the sales index page ", webResponse.getText().contains("SalesTool"));

        /*String uri = "http://localhost:8080/sp/secured/test";
        HttpGet get = new HttpGet(uri);

        HttpResponse response = httpClient.execute(get);
        assertEquals(200, response.getStatusLine().getStatusCode());

        Header[] headers = response.getHeaders("Location");

        Header header = headers[0];

        String idpURL = header.getValue();
        System.out.println("IDP URL=" + idpURL);
        assertTrue(idpURL.contains("/idp"));

        HttpEntity entity = response.getEntity();
        EntityUtils.consume(entity);

        get = new HttpGet(idpURL);
        response = httpClient.execute(get);

        String responseAsString = readResponse(response);
        assertEquals("Login Page", responseAsString);

        BasicNameValuePair jUserName = new BasicNameValuePair("j_username", "user1");
        BasicNameValuePair jPassword =  new BasicNameValuePair("j_password", "password1");
        BasicNameValuePair[] pairs = new BasicNameValuePair[]{jUserName,jPassword};
        final List<NameValuePair> data = new ArrayList<NameValuePair>();
        data.addAll(Arrays.asList(pairs));
        HttpPost post = new HttpPost("http://localhost:8080/idp/j_security_check");
        post.setEntity(new UrlEncodedFormEntity(data));

        response = httpClient.execute(post);
        assertEquals(200, response.getStatusLine().getStatusCode());

        responseAsString = readResponse(response);
        assertEquals("user1", response);*/
    }

    private String readResponse(final HttpResponse response) throws IOException {
        HttpEntity entity = response.getEntity();
        if(entity == null) {
            return "";
        }
        return readResponse(entity.getContent());
    }
    private String readResponse(InputStream stream) throws IOException {
        final StringBuilder builder = new StringBuilder();
        byte[] data = new byte[100];
        int read;
        while ((read = stream.read(data)) != -1) {
            builder.append(new String(data,0,read,"UTF-8"));
        }
        return builder.toString();
    }
}