LoginAbstractAzkabanServlet.java

package azkaban.webapp.servlet;

import java.io.IOException;
import java.io.Writer;
import java.util.UUID;

import javax.servlet.ServletException;
import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.apache.log4j.Logger;

import azkaban.user.User;
import azkaban.user.UserManager;
import azkaban.user.UserManagerException;
import azkaban.webapp.session.Session;

/**
 * Abstract Servlet that handles auto login when the session hasn't been
 * verified.
 */
public abstract class LoginAbstractAzkabanServlet extends
		AbstractAzkabanServlet {

	private static final long serialVersionUID = 1L;

	private static final Logger logger = Logger
			.getLogger(LoginAbstractAzkabanServlet.class.getName());
	private static final String SESSION_ID_NAME = "azkaban.session.id";

	@Override
	protected void doGet(HttpServletRequest req, HttpServletResponse resp)
			throws ServletException, IOException {

		// Set session id
		Session session = getSessionFromRequest(req);
		if (hasParam(req, "logout")) {
			resp.sendRedirect(req.getContextPath());
			if (session != null) {
				getApplication().getSessionCache().removeSession(
						session.getSessionId());
			}
			return;
		}
		else if (hasParam(req, "action")) {
			String action = getParam(req, "action");
			if (action.equals("login")) {
				handleLoginAction(req, resp);
			}
			
			String referer = req.getHeader("Referer");
			resp.sendRedirect(referer);			
			return;
		}

		
		if (session != null) {
			logger.info("Found session " + session.getUser());
			handleGet(req, resp, session);
		} else {
			handleLogin(req, resp);
		}
	}

	private Session getSessionFromRequest(HttpServletRequest req) {
		Cookie cookie = getCookieByName(req, SESSION_ID_NAME);
		String sessionId = null;

		if (cookie != null) {
			sessionId = cookie.getValue();
			logger.info("Session id " + sessionId);
		}
		if (sessionId == null) {
			return null;
		} else {
			return getApplication().getSessionCache().getSession(sessionId);
		}
	}

	private void handleLogin(HttpServletRequest req, HttpServletResponse resp)
			throws ServletException, IOException {
		handleLogin(req, resp, null);
	}

	private void handleLogin(HttpServletRequest req, HttpServletResponse resp,
			String errorMsg) throws ServletException, IOException {

		Page page = newPage(req, resp,
				"azkaban/webapp/servlet/velocity/login.vm");
		if (errorMsg != null) {
			page.add("errorMsg", errorMsg);
		}

		page.render();
	}

	@Override
	protected void doPost(HttpServletRequest req, HttpServletResponse resp)
			throws ServletException, IOException {
		if (hasParam(req, "action")) {
			String action = getParam(req, "action");
			if (action.equals("login")) {
				handleLoginAction(req, resp);
			} else {
				Session session = getSessionFromRequest(req);
				if (session == null) {
					if (isAjaxCall(req)) {
						String response = createJsonResponse("error",
								"Invalid Session. Need to re-login", "login",
								null);
						writeResponse(resp, response);
					} else {
						handleLogin(req, resp, "Enter username and password");
					}
				} else {
					handlePost(req, resp, session);
				}
			}
		} else {
			Session session = getSessionFromRequest(req);
			if (session == null) {
				if (isAjaxCall(req)) {
					String response = createJsonResponse("error",
							"Invalid Session. Need to re-login", "login", null);
					writeResponse(resp, response);
				} else {
					handleLogin(req, resp, "Enter username and password");
				}
			} else {
				handlePost(req, resp, session);
			}
		}
	}

	protected void handleLoginAction(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
		if (hasParam(req, "username") && hasParam(req, "password")) {
			String username = getParam(req, "username");
			String password = getParam(req, "password");

			UserManager manager = getApplication().getUserManager();

			User user = null;
			try {
				user = manager.getUser(username, password);
			} catch (UserManagerException e) {
				handleLogin(req, resp, e.getMessage());
				return;
			}

			String randomUID = UUID.randomUUID().toString();
			Session session = new Session(randomUID, user);
			resp.addCookie(new Cookie(SESSION_ID_NAME, randomUID));
			getApplication().getSessionCache().addSession(session);
			handleGet(req, resp, session);
		} else {
			if (isAjaxCall(req)) {
				String response = createJsonResponse("error",
						"Incorrect Login.", "login", null);
				writeResponse(resp, response);
			} else {
				handleLogin(req, resp, "Enter username and password");
			}
		}
	}
	
	protected void writeResponse(HttpServletResponse resp, String response)
			throws IOException {
		Writer writer = resp.getWriter();
		writer.append(response);
		writer.flush();
	}

	protected boolean isAjaxCall(HttpServletRequest req)
			throws ServletException {
		String value = req.getHeader("X-Requested-With");
		if (value != null) {
			logger.info("has X-Requested-With " + value);
			return value.equals("XMLHttpRequest");
		}

		return false;
	}

	/**
	 * The get request is handed off to the implementor after the user is logged
	 * in.
	 * 
	 * @param req
	 * @param resp
	 * @param session
	 * @throws ServletException
	 * @throws IOException
	 */
	protected abstract void handleGet(HttpServletRequest req,
			HttpServletResponse resp, Session session) throws ServletException,
			IOException;

	/**
	 * The post request is handed off to the implementor after the user is
	 * logged in.
	 * 
	 * @param req
	 * @param resp
	 * @param session
	 * @throws ServletException
	 * @throws IOException
	 */
	protected abstract void handlePost(HttpServletRequest req,
			HttpServletResponse resp, Session session) throws ServletException,
			IOException;
}