/*

  * Licensed to the Apache Software Foundation (ASF) under one

  * or more contributor license agreements.  See the NOTICE file

  * distributed with this work for additional information

  * regarding copyright ownership.  The ASF licenses this file

  * to you 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.apache.shiro.web.servlet;

 import javax.servlet.ServletContext;

 import javax.servlet.http.HttpServletRequest;

 import javax.servlet.http.HttpServletResponse;

 import javax.servlet.http.HttpServletResponseWrapper;

 import javax.servlet.http.HttpSession;

 import java.io.IOException;

 import java.net.MalformedURLException;

 import java.net.URL;

 import java.net.URLEncoder;

 /**

  * HttpServletResponse implementation to support URL Encoding of Shiro Session IDs.

  * <p/>

  * It is only used when using Shiro's native Session Management configuration (and not when using the Servlet

  * Container session configuration, which is Shiro's default in a web environment).  Because the servlet container

  * already performs url encoding of its own session ids, instances of this class are only needed when using Shiro

  * native sessions.

  * <p/>

  * Note that this implementation relies in part on source code from the Tomcat 6.x distribution for

  * encoding URLs for session ID URL Rewriting (we didn't want to re-invent the wheel).  Since Shiro is also

  * Apache 2.0 license, all regular licenses and conditions have remained in tact.

  *

  * @since 0.2

  */

 public class ShiroHttpServletResponse extends HttpServletResponseWrapper {

     //TODO - complete JavaDoc

     private static final String DEFAULT_SESSION_ID_PARAMETER_NAME = ShiroHttpSession.DEFAULT_SESSION_ID_NAME;

     private ServletContext context = null;

     //the associated request

     private ShiroHttpServletRequest request = null;

     public ShiroHttpServletResponse(HttpServletResponse wrapped, ServletContext context, ShiroHttpServletRequest request) {

         super(wrapped);

         this.context = context;

         this.request = request;

     }

     @SuppressWarnings({"UnusedDeclaration"})

     public ServletContext getContext() {

         return context;

     }

     @SuppressWarnings({"UnusedDeclaration"})

     public void setContext(ServletContext context) {

         this.context = context;

     }

     public ShiroHttpServletRequest getRequest() {

         return request;

     }

     @SuppressWarnings({"UnusedDeclaration"})

     public void setRequest(ShiroHttpServletRequest request) {

         this.request = request;

     }

     /**

      * Encode the session identifier associated with this response

      * into the specified redirect URL, if necessary.

      *

      * @param url URL to be encoded

      */

     public String encodeRedirectURL(String url) {

         if (isEncodeable(toAbsolute(url))) {

             return toEncoded(url, request.getSession().getId());

         } else {

             return url;

         }

     }

     public String encodeRedirectUrl(String s) {

         return encodeRedirectURL(s);

     }

     /**

      * Encode the session identifier associated with this response

      * into the specified URL, if necessary.

      *

      * @param url URL to be encoded

      */

     public String encodeURL(String url) {

         String absolute = toAbsolute(url);

         if (isEncodeable(absolute)) {

             // W3c spec clearly said

             if (url.equalsIgnoreCase("")) {

                 url = absolute;

             }

             return toEncoded(url, request.getSession().getId());

         } else {

             return url;

         }

     }

     public String encodeUrl(String s) {

         return encodeURL(s);

     }

     /**

      * Return <code>true</code> if the specified URL should be encoded with

      * a session identifier.  This will be true if all of the following

      * conditions are met:

      * <ul>

      * <li>The request we are responding to asked for a valid session

      * <li>The requested session ID was not received via a cookie

      * <li>The specified URL points back to somewhere within the web

      * application that is responding to this request

      * </ul>

      *

      * @param location Absolute URL to be validated

      * @return {@code true} if the specified URL should be encoded with a session identifier, {@code false} otherwise.

      */

     protected boolean isEncodeable(final String location) {

         if (location == null)

             return (false);

         // Is this an intra-document reference?

         if (location.startsWith("#"))

             return (false);

         // Are we in a valid session that is not using cookies?

         final HttpServletRequest hreq = request;

         final HttpSession session = hreq.getSession(false);

         if (session == null)

             return (false);

         if (hreq.isRequestedSessionIdFromCookie())

             return (false);

         return doIsEncodeable(hreq, session, location);

     }

     private boolean doIsEncodeable(HttpServletRequest hreq, HttpSession session, String location) {

         // Is this a valid absolute URL?

         URL url;

         try {

             url = new URL(location);

         } catch (MalformedURLException e) {

             return (false);

         }

         // Does this URL match down to (and including) the context path?

         if (!hreq.getScheme().equalsIgnoreCase(url.getProtocol()))

             return (false);

         if (!hreq.getServerName().equalsIgnoreCase(url.getHost()))

             return (false);

         int serverPort = hreq.getServerPort();

         if (serverPort == -1) {

             if ("https".equals(hreq.getScheme()))

                 serverPort = 443;

             else

                 serverPort = 80;

         }

         int urlPort = url.getPort();

         if (urlPort == -1) {

             if ("https".equals(url.getProtocol()))

                 urlPort = 443;

             else

                 urlPort = 80;

         }

         if (serverPort != urlPort)

             return (false);

         String contextPath = getRequest().getContextPath();

         if (contextPath != null) {

             String file = url.getFile();

             if ((file == null) || !file.startsWith(contextPath))

                 return (false);

             String tok = ";" + DEFAULT_SESSION_ID_PARAMETER_NAME + "=" + session.getId();

             if (file.indexOf(tok, contextPath.length()) >= 0)

                 return (false);

         }

         // This URL belongs to our web application, so it is encodeable

         return (true);

     }

     /**

      * Convert (if necessary) and return the absolute URL that represents the

      * resource referenced by this possibly relative URL.  If this URL is

      * already absolute, return it unchanged.

      *

      * @param location URL to be (possibly) converted and then returned

      * @return resource location as an absolute url

      * @throws IllegalArgumentException if a MalformedURLException is

      *                                  thrown when converting the relative URL to an absolute one

      */

     private String toAbsolute(String location) {

         if (location == null)

             return (location);

         boolean leadingSlash = location.startsWith("/");

         if (leadingSlash || !hasScheme(location)) {

             StringBuilder buf = new StringBuilder();

             String scheme = request.getScheme();

             String name = request.getServerName();

             int port = request.getServerPort();

             try {

                 buf.append(scheme).append("://").append(name);

                 if ((scheme.equals("http") && port != 80)

                         || (scheme.equals("https") && port != 443)) {

                     buf.append(':').append(port);

                 }

                 if (!leadingSlash) {

                     String relativePath = request.getRequestURI();

                     int pos = relativePath.lastIndexOf('/');

                     relativePath = relativePath.substring(0, pos);

                     String encodedURI = URLEncoder.encode(relativePath, getCharacterEncoding());

                     buf.append(encodedURI).append('/');

                 }

                 buf.append(location);

             } catch (IOException e) {

                 IllegalArgumentException iae = new IllegalArgumentException(location);

                 iae.initCause(e);

                 throw iae;

             }

             return buf.toString();

         } else {

             return location;

         }

     }

     /**

      * Determine if the character is allowed in the scheme of a URI.

      * See RFC 2396, Section 3.1

      *

      * @param c the character to check

      * @return {@code true} if the character is allowed in a URI scheme, {@code false} otherwise.

      */

     public static boolean isSchemeChar(char c) {

         return Character.isLetterOrDigit(c) ||

                 c == '+' || c == '-' || c == '.';

     }

     /**

      * Returns {@code true} if the URI string has a {@code scheme} component, {@code false} otherwise.

      *

      * @param uri the URI string to check for a scheme component

      * @return {@code true} if the URI string has a {@code scheme} component, {@code false} otherwise.

      */

     private boolean hasScheme(String 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;

     }

     /**

      * Return the specified URL with the specified session identifier suitably encoded.

      *

      * @param url       URL to be encoded with the session id

      * @param sessionId Session id to be included in the encoded URL

      * @return the url with the session identifer properly encoded.

      */

     protected String toEncoded(String url, String sessionId) {

         if ((url == null) || (sessionId == null))

             return (url);

         String path = url;

         String query = "";

         String anchor = "";

         int question = url.indexOf('?');

         if (question >= 0) {

             path = url.substring(0, question);

             query = url.substring(question);

         }

         int pound = path.indexOf('#');

         if (pound >= 0) {

             anchor = path.substring(pound);

             path = path.substring(0, pound);

         }

         StringBuilder sb = new StringBuilder(path);

         if (sb.length() > 0) { // session id param can't be first.

             sb.append(";");

             sb.append(DEFAULT_SESSION_ID_PARAMETER_NAME);

             sb.append("=");

             sb.append(sessionId);

         }

         sb.append(anchor);

         sb.append(query);

         return (sb.toString());

     }

 }