package io.undertow.servlet.compat.rewrite;
import io.undertow.UndertowOptions;
import io.undertow.server.HttpHandler;
import io.undertow.server.HttpServerExchange;
import io.undertow.servlet.UndertowServletLogger;
import io.undertow.servlet.handlers.ServletRequestContext;
import io.undertow.servlet.spec.HttpServletRequestImpl;
import io.undertow.servlet.spec.HttpServletResponseImpl;
import io.undertow.util.Headers;
import io.undertow.util.QueryParameterUtils;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import javax.servlet.http.Cookie;
import javax.servlet.http.HttpServletResponse;
public class RewriteHandler implements HttpHandler {
private final RewriteConfig config;
private final HttpHandler next;
protected ThreadLocal<Boolean> invoked = new ThreadLocal<>();
public RewriteHandler(RewriteConfig config, HttpHandler next) {
this.config = config;
this.next = next;
}
public void handleRequest(HttpServerExchange exchange) throws Exception {
RewriteRule[] rules = config.getRules();
if (rules == null || rules.length == 0) {
next.handleRequest(exchange);
return;
}
if (Boolean.TRUE.equals(invoked.get())) {
next.handleRequest(exchange);
invoked.set(null);
return;
}
ServletRequestContext src = exchange.getAttachment(ServletRequestContext.ATTACHMENT_KEY);
HttpServletRequestImpl request = src.getOriginalRequest();
HttpServletResponseImpl response = src.getOriginalResponse();
UndertowResolver resolver = new UndertowResolver(src, src.getOriginalRequest());
invoked.set(Boolean.TRUE);
CharSequence url = exchange.getRelativePath();
CharSequence host = request.getServerName();
boolean rewritten = false;
boolean done = false;
for (int i = 0; i < rules.length; i++) {
CharSequence test = (rules[i].isHost()) ? host : url;
CharSequence newtest = rules[i].evaluate(test, resolver);
if (newtest != null && !test.equals(newtest.toString())) {
if (UndertowServletLogger.REQUEST_LOGGER.isDebugEnabled()) {
UndertowServletLogger.REQUEST_LOGGER.debug("Rewrote " + test + " as " + newtest
+ " with rule pattern " + rules[i].getPatternString());
}
if (rules[i].isHost()) {
host = newtest;
} else {
url = newtest;
}
rewritten = true;
}
if (rules[i].isForbidden() && newtest != null) {
response.sendError(HttpServletResponse.SC_FORBIDDEN);
done = true;
break;
}
if (rules[i].isGone() && newtest != null) {
response.sendError(HttpServletResponse.SC_GONE);
done = true;
break;
}
if (rules[i].isRedirect() && newtest != null) {
String queryString = request.getQueryString();
StringBuffer urlString = new StringBuffer(url);
if (queryString != null && queryString.length() > 0) {
int index = urlString.indexOf("?");
if (index != -1) {
if (rules[i].isQsappend()) {
urlString.append('&');
urlString.append(queryString);
}
else if (index == urlString.length() - 1) {
urlString.deleteCharAt(index);
}
} else {
urlString.append('?');
urlString.append(queryString);
}
}
if (urlString.charAt(0) == '/' && !hasScheme(urlString)) {
urlString.insert(0, request.getContextPath());
}
response.sendRedirect(urlString.toString());
response.setStatus(rules[i].getRedirectCode());
done = true;
break;
}
if (rules[i].isCookie() && newtest != null) {
Cookie cookie = new Cookie(rules[i].getCookieName(),
rules[i].getCookieResult());
cookie.setDomain(rules[i].getCookieDomain());
cookie.setMaxAge(rules[i].getCookieLifetime());
cookie.setPath(rules[i].getCookiePath());
cookie.setSecure(rules[i].isCookieSecure());
cookie.setHttpOnly(rules[i].isCookieHttpOnly());
response.addCookie(cookie);
}
if (rules[i].isEnv() && newtest != null) {
Map<String, String> attrs = exchange.getAttachment(HttpServerExchange.REQUEST_ATTRIBUTES);
if (attrs == null) {
attrs = new HashMap<>();
exchange.putAttachment(HttpServerExchange.REQUEST_ATTRIBUTES, attrs);
}
for (int j = 0; j < rules[i].getEnvSize(); j++) {
final String envName = rules[i].getEnvName(j);
final String envResult = rules[i].getEnvResult(j);
attrs.put(envName, envResult);
request.setAttribute(envName, envResult);
}
}
if (rules[i].isType() && newtest != null) {
exchange.getRequestHeaders().put(Headers.CONTENT_TYPE, rules[i].getTypeValue());
}
if (rules[i].isQsappend() && newtest != null) {
String queryString = request.getQueryString();
String urlString = url.toString();
if (urlString.indexOf('?') != -1 && queryString != null) {
url = urlString + "&" + queryString;
}
}
if (rules[i].isChain() && newtest == null) {
for (int j = i; j < rules.length; j++) {
if (!rules[j].isChain()) {
i = j;
break;
}
}
continue;
}
if (rules[i].isLast() && newtest != null) {
break;
}
if (rules[i].isNext() && newtest != null) {
i = 0;
continue;
}
if (newtest != null) {
i += rules[i].getSkip();
}
}
if (rewritten) {
if (!done) {
String urlString = url.toString();
String queryString = null;
int queryIndex = urlString.indexOf('?');
if (queryIndex != -1) {
queryString = urlString.substring(queryIndex + 1);
urlString = urlString.substring(0, queryIndex);
}
StringBuilder chunk = new StringBuilder();
chunk.append(request.getContextPath());
chunk.append(urlString);
String requestPath = chunk.toString();
exchange.setRequestURI(requestPath);
exchange.setRequestPath(requestPath);
exchange.setRelativePath(urlString);
if (queryString != null) {
exchange.setQueryString(queryString);
exchange.getQueryParameters().clear();
exchange.getQueryParameters().putAll(QueryParameterUtils.parseQueryString(queryString, exchange.getConnection().getUndertowOptions().get(UndertowOptions.URL_CHARSET, StandardCharsets.UTF_8.name())));
}
if (!host.equals(request.getServerName())) {
exchange.getRequestHeaders().put(Headers.HOST, host + ":" + exchange.getHostPort());
}
src.getDeployment().getHandler().handleRequest(exchange);
}
} else {
next.handleRequest(exchange);
}
invoked.set(null);
}
protected static boolean hasScheme(StringBuffer uri) {
int len = uri.length();
for (int i = 0; i < len; i++) {
char c = uri.charAt(i);
if (c == ':') {
return i > 0;
} else if (!isSchemeChar(c)) {
return false;
}
}
return false;
}
private static boolean isSchemeChar(char c) {
return Character.isLetterOrDigit(c) ||
c == '+' || c == '-' || c == '.';
}
}