LoginUnifier.java

package extensions;

import java.io.IOException;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.LinkedList;
import java.util.List;
import java.util.Collections;
import java.util.TreeMap;

import dareka.common.Logger;
import dareka.extensions.Extension2;
import dareka.extensions.ExtensionManager;
import dareka.extensions.RequestFilter;
import dareka.processor.HttpHeader;
import dareka.processor.HttpRequestHeader;
import dareka.processor.HttpResponseHeader;
import dareka.processor.Processor;
import dareka.processor.URLResource;

public class LoginUnifier
		implements Extension2, RequestFilter {
		
	// Extensionのバージョン
	public String getVersionString() {
		return "LoginUnifier ver.2.2";
	}
	
	// ExtensionをNicoCache_nl本体に登録する
	public void registerExtensions(ExtensionManager mgr) {
		// 必要な行のみコメントを外す
		//mgr.registerProcessor(this);      // Processorとして登録
		//mgr.registerRewriter(this);       // Rewriterとして登録
		mgr.registerRequestFilter(this);  // RequestFilterとして登録
	}
	
	private static final String SESSION_COOKIE = "user_session";
	
	// user_session cookie known to be valid
	private static String validSession = null;
	// user_session cookie known to be expired
	private static List<String> expiredSessions = 
		Collections.synchronizedList(new LinkedList<String>());
	// number of expired user_session to store
	private static final int expiredSessionsMax = 100;

	private void addExpiredSession(String session) {
		synchronized(expiredSessions) {
			expiredSessions.add(session); // push
			while (expiredSessions.size() >= expiredSessionsMax) {
				expiredSessions.remove(0); // shift
			}
		}
	}

	private String sniffAuthFlag(HttpRequestHeader header)
		throws IOException {
		URLResource urlResource = new URLResource("http://www.nicovideo.jp/");
		HttpRequestHeader tmpRequestHeader = new HttpRequestHeader(header.toString());
		tmpRequestHeader.setMethod(HttpHeader.HEAD);
		tmpRequestHeader.setURI("http://www.nicovideo.jp/");
		tmpRequestHeader.setContentLength(0);
		tmpRequestHeader.setMessageHeader("host", "www.nicovideo.jp");
		HttpResponseHeader tmpResponseHeader = urlResource.getResponseHeader(null, tmpRequestHeader);
		return tmpResponseHeader.getMessageHeader("x-niconico-authflag");
	}
	
	// Processor Interface
	private static final Pattern PROCESSOR_SUPPORTED_PATTERN = Pattern.compile(
			"^http://(?:[^/]+\\.)?nicovideo\\.jp/\\S*$");

	private static final Pattern PROCESSOR_UNSUPPORTED_PATTERN = Pattern.compile(
			"^http://(?:ext)\\.nicovideo\\.jp/\\S*$");

	// RequestFilter interface
	// stores and rewrites Cookie header
	public int onRequest(HttpRequestHeader requestHeader)
			throws IOException {
		if (requestHeader == null) { // failed to receive header
			return RequestFilter.DROP;
		}
		Matcher m = PROCESSOR_SUPPORTED_PATTERN.matcher(requestHeader.getURI());
		if (!m.matches()) {
			return RequestFilter.OK; // just pass
		}
		m = PROCESSOR_UNSUPPORTED_PATTERN.matcher(requestHeader.getURI());
		if (m.matches()) {
			return RequestFilter.OK; // just pass
		}
		String url = requestHeader.getURI();
		//Logger.info("LoginUnifier: " + url);
		
		// copy validSession to local variable (will be written back at the end)
		String storedSession = validSession; // could be null
		
		TreeMap<String, String> cookie = CookieToMap(requestHeader.getMessageHeader("cookie"));
		String uaSession = cookie.get(SESSION_COOKIE);
		String newSession = null;
		if (uaSession == null || // ログアウトした、または別のブラウザに移った
				expiredSessions.contains(uaSession) // 別のブラウザを使って戻ってきた
				) {
			if (storedSession == null) {
				// LoginUnifier スタート直後、または前回のセッションが切れた後
				newSession = null;
			} else {
				// LoginUnifier にセッションが残ってた
				newSession = storedSession;
			}
		} else if (uaSession.equals(storedSession)) {
			// 普通の場合。LoginUnifier を使ってログインし、そのまま同じブラウザを使ってる
			newSession = uaSession;
		} else {
			// 再ログイン直後、またはとても古いセッション
			if (storedSession == null) {
				// LoginUnifier スタート直後、または前回のセッションが切れた後
				newSession = uaSession;
			} else {
				// LoginUnifier にセッションが残ってた
				// とりあえず UA のセッションでリクエストし、
				// 切れてたら保存してたセッションを使って再度リクエスト
				Logger.info("LoginUnifier: HEAD request to check validity of " + uaSession);
				String authflag = sniffAuthFlag(requestHeader);
				Logger.info("LoginUnifier: authflag = " + (authflag == null ? "null" : authflag));
				if (authflag == null || // something went wrong (should not happen)
						authflag.equals("0")) { // expired
					addExpiredSession(uaSession);
					newSession = storedSession;
				} else { // 1 = normal user, 3 = premium user
					newSession = uaSession;
				}
			}
		}

		Logger.info(
				"LoginUnifier: " + url + "\n" +
				"\tUA session     = " + (uaSession == null ? "null" : uaSession) + "\n" +
				"\tstored session = " + (storedSession == null ? "null" : storedSession) + "\n" +
				"\tNew session    = " + (newSession == null ? "null" : newSession));

		validSession = newSession;
		if (newSession == null) {
			cookie.remove(SESSION_COOKIE);
		} else {
			cookie.put(SESSION_COOKIE, newSession);
		}
		if (cookie.isEmpty()) {
			requestHeader.removeMessageHeader("cookie");
		} else {
			requestHeader.setMessageHeader("cookie", MapToCookie(cookie));
		}
		return RequestFilter.OK;
	}

	private static final Pattern COOKIE_PATTERN = Pattern.compile(
			"([^=; ]*)=([^; ]*)");
	
	// Cookie header filed "hoge=fuga; foo=bar;" to TreeMap
	public static TreeMap<String, String> CookieToMap(String cookie) {
		TreeMap<String, String> map = new TreeMap<String, String>();
		if (cookie == null) {
			return map;
		}
		Matcher m = COOKIE_PATTERN.matcher(cookie);
		while(m.find()) {
			map.put(m.group(1), m.group(2));
		}
		return map;
	}

	// TreeMap back to Cookie field String
	public static String MapToCookie(TreeMap<String, String> map) {
		StringBuffer cookie = new StringBuffer();
		Boolean first = true;
		for (Object key: map.keySet()) {
			if (first) {
				first = false;
			} else {
				cookie.append("; ");
			}
			cookie.append(key).append("=").append(map.get(key));
		}
		return cookie.toString();
	}
}