/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2014 Red Hat, Inc., and individual contributors
 * as indicated by the @author tags.
 *
 * 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 io.undertow.server.handlers;

import io.undertow.server.ExchangeCompletionListener;
import io.undertow.server.HandlerWrapper;
import io.undertow.server.HttpHandler;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.handlers.builder.HandlerBuilder;
import io.undertow.server.handlers.cache.LRUCache;
import io.undertow.server.session.Session;
import io.undertow.server.session.SessionConfig;
import io.undertow.server.session.SessionManager;
import io.undertow.util.DateUtils;
import io.undertow.util.HeaderMap;
import io.undertow.util.Headers;
import io.undertow.util.Methods;

import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

Handler that builds up a cache of resources that a browsers requests, and uses server push to push them when supported.
Author:Stuart Douglas
/** * Handler that builds up a cache of resources that a browsers requests, and uses * server push to push them when supported. * * @author Stuart Douglas */
public class LearningPushHandler implements HttpHandler { private static final String SESSION_ATTRIBUTE = "io.undertow.PUSHED_RESOURCES"; private static final int DEFAULT_MAX_CACHE_ENTRIES = 1000; private static final int DEFAULT_MAX_CACHE_AGE = -1; private final LRUCache<String, Map<String, PushedRequest>> cache; private final HttpHandler next; public LearningPushHandler(final HttpHandler next) { this(DEFAULT_MAX_CACHE_ENTRIES, DEFAULT_MAX_CACHE_AGE, next); } public LearningPushHandler(int maxEntries, int maxAge, HttpHandler next) { this.next = next; cache = new LRUCache<>(maxEntries, maxAge); } @Override public void handleRequest(HttpServerExchange exchange) throws Exception { String fullPath; String requestPath; if(exchange.getQueryString().isEmpty()) { fullPath = exchange.getRequestURL(); requestPath = exchange.getRequestPath(); } else{ fullPath = exchange.getRequestURL() + "?" + exchange.getQueryString(); requestPath = exchange.getRequestPath() + "?" + exchange.getQueryString(); } doPush(exchange, fullPath); String referrer = exchange.getRequestHeaders().getFirst(Headers.REFERER); if (referrer != null) { String accept = exchange.getRequestHeaders().getFirst(Headers.ACCEPT); if (accept == null || !accept.contains("text/html")) { //if accept contains text/html it generally means the user has clicked //a link to move to a new page, and is not a resource load for the current page //we only care about resources for the current page exchange.addExchangeCompleteListener(new PushCompletionListener(fullPath, requestPath, referrer)); } } next.handleRequest(exchange); } private void doPush(HttpServerExchange exchange, String fullPath) { if (exchange.getConnection().isPushSupported()) { Map<String, PushedRequest> toPush = cache.get(fullPath); if (toPush != null) { Session session = getSession(exchange); if (session == null) { return; } Map<String, Object> pushed = (Map<String, Object>) session.getAttribute(SESSION_ATTRIBUTE); if (pushed == null) { pushed = Collections.synchronizedMap(new HashMap<String, Object>()); } for (Map.Entry<String, PushedRequest> entry : toPush.entrySet()) { PushedRequest request = entry.getValue(); Object pushedKey = pushed.get(request.getPath()); boolean doPush = pushedKey == null; if (!doPush) { if (pushedKey instanceof String && !pushedKey.equals(request.getEtag())) { doPush = true; } else if (pushedKey instanceof Long && ((Long) pushedKey) != request.getLastModified()) { doPush = true; } } if (doPush) { exchange.getConnection().pushResource(request.getPath(), Methods.GET, request.getRequestHeaders()); if(request.getEtag() != null) { pushed.put(request.getPath(), request.getEtag()); } else { pushed.put(request.getPath(), request.getLastModified()); } } } session.setAttribute(SESSION_ATTRIBUTE, pushed); } } } protected Session getSession(HttpServerExchange exchange) { SessionConfig sc = exchange.getAttachment(SessionConfig.ATTACHMENT_KEY); SessionManager sm = exchange.getAttachment(SessionManager.ATTACHMENT_KEY); if (sc == null || sm == null) { return null; } Session session = sm.getSession(exchange, sc); if (session == null) { return sm.createSession(exchange, sc); } return session; } private final class PushCompletionListener implements ExchangeCompletionListener { private final String fullPath; private final String requestPath; private final String referer; private PushCompletionListener(String fullPath, String requestPath, String referer) { this.fullPath = fullPath; this.requestPath = requestPath; this.referer = referer; } @Override public void exchangeEvent(HttpServerExchange exchange, NextListener nextListener) { if (exchange.getStatusCode() == 200 && referer != null) { //for now only cache 200 response codes String lmString = exchange.getResponseHeaders().getFirst(Headers.LAST_MODIFIED); String etag = exchange.getResponseHeaders().getFirst(Headers.ETAG); long lastModified = -1; if(lmString != null) { Date dt = DateUtils.parseDate(lmString); if(dt != null) { lastModified = dt.getTime(); } } Map<String, PushedRequest> pushes = cache.get(referer); if(pushes == null) { synchronized (cache) { pushes = cache.get(referer); if(pushes == null) { cache.add(referer, pushes = Collections.synchronizedMap(new HashMap<String, PushedRequest>())); } } } pushes.put(fullPath, new PushedRequest(new HeaderMap(), requestPath, etag, lastModified)); } nextListener.proceed(); } } private static class PushedRequest { private final HeaderMap requestHeaders; private final String path; private final String etag; private final long lastModified; private PushedRequest(HeaderMap requestHeaders, String path, String etag, long lastModified) { this.requestHeaders = requestHeaders; this.path = path; this.etag = etag; this.lastModified = lastModified; } public HeaderMap getRequestHeaders() { return requestHeaders; } public String getPath() { return path; } public String getEtag() { return etag; } public long getLastModified() { return lastModified; } } public static class Builder implements HandlerBuilder { @Override public String name() { return "learning-push"; } @Override public Map<String, Class<?>> parameters() { Map<String, Class<?>> params = new HashMap<>(); params.put("max-age", Integer.class); params.put("max-entries", Integer.class); return params; } @Override public Set<String> requiredParameters() { return null; } @Override public String defaultParameter() { return null; } @Override public HandlerWrapper build(Map<String, Object> config) { final int maxAge = config.containsKey("max-age") ? (Integer)config.get("max-age") : DEFAULT_MAX_CACHE_AGE; final int maxEntries = config.containsKey("max-entries") ? (Integer)config.get("max-entries") : DEFAULT_MAX_CACHE_ENTRIES; return new HandlerWrapper() { @Override public HttpHandler wrap(HttpHandler handler) { return new LearningPushHandler(maxEntries, maxAge, handler); } }; } } }