/*
 * 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.calcite.avatica.server;

import org.apache.calcite.avatica.ConnectionSpec;
import org.apache.calcite.avatica.jdbc.JdbcMeta;
import org.apache.calcite.avatica.remote.AuthenticationType;
import org.apache.calcite.avatica.remote.Driver;
import org.apache.calcite.avatica.remote.LocalService;

import org.eclipse.jetty.security.UserAuthentication;
import org.eclipse.jetty.server.Authentication;
import org.eclipse.jetty.server.Connector;
import org.eclipse.jetty.server.Handler;
import org.eclipse.jetty.server.Request;
import org.eclipse.jetty.server.Server;
import org.eclipse.jetty.server.ServerConnector;
import org.eclipse.jetty.server.UserIdentity;
import org.eclipse.jetty.server.handler.DefaultHandler;
import org.eclipse.jetty.server.handler.HandlerList;
import org.eclipse.jetty.server.handler.HandlerWrapper;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;

import java.sql.SQLException;
import java.util.Arrays;
import java.util.Properties;
import java.util.concurrent.Callable;
import javax.servlet.http.HttpServletRequest;

import static org.hamcrest.core.StringContains.containsString;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.fail;

/**
 * Test class for providing CustomAvaticaServerConfiguration to the HTTP Server
 */
public class CustomAuthHttpServerTest extends HttpAuthBase {
  private static final ConnectionSpec CONNECTION_SPEC = ConnectionSpec.HSQLDB;
  private static HttpServer server;
  private static String url;

  // Counters to keep track of number of function calls
  private static int methodCallCounter1 = 0;
  private static int methodCallCounter2 = 0;
  private static int methodCallCounter3 = 0;

  @Before
  public void before() {
    methodCallCounter1 = 0;
    methodCallCounter2 = 0;
    methodCallCounter3 = 0;
  }

  @After
  public void stopServer() {
    if (null != server) {
      server.stop();
    }
  }

  @Test
  public void testCustomImpersonationConfig() throws Exception {
    AvaticaServerConfiguration configuration = new CustomImpersonationConfig();
    createServer(configuration, false);

    readWriteData(url, "CUSTOM_CONFIG_1_TABLE", new Properties());
    Assert.assertEquals("supportsImpersonation should be called same number of "
            + "times as doAsRemoteUser method", methodCallCounter1, methodCallCounter2);
    Assert.assertEquals("supportsImpersonation should be called same number of "
            + "times as getRemoteUserExtractor method", methodCallCounter1, methodCallCounter3);
  }

  @Test
  public void testCustomBasicImpersonationConfigWithAllowedUser() throws Exception {
    AvaticaServerConfiguration configuration = new CustomBasicImpersonationConfig();
    createServer(configuration, true);

    final Properties props = new Properties();
    props.put("avatica_user", "USER2");
    props.put("avatica_password", "password2");
    props.put("user", "USER2");
    props.put("password", "password2");

    readWriteData(url, "CUSTOM_CONFIG_2_ALLOWED_TABLE", props);
    Assert.assertEquals("supportsImpersonation should be called same number of "
            + "times as doAsRemoteUser method", methodCallCounter1, methodCallCounter2);
    Assert.assertEquals("supportsImpersonation should be called same number of "
            + "times as getRemoteUserExtractor method", methodCallCounter1, methodCallCounter3);
  }

  @Test
  public void testCustomBasicImpersonationConfigWithDisallowedUser() throws Exception {
    AvaticaServerConfiguration configuration = new CustomBasicImpersonationConfig();
    createServer(configuration, true);

    final Properties props = new Properties();
    props.put("avatica_user", "USER1");
    props.put("avatica_password", "password1");
    props.put("user", "USER1");
    props.put("password", "password1");

    try {
      readWriteData(url, "CUSTOM_CONFIG_2_DISALLOWED_TABLE", props);
      fail("Expected an exception");
    } catch (RuntimeException e) {
      assertThat(e.getMessage(), containsString("Failed to execute HTTP Request, got HTTP/403"));
    }
  }

  @Test(expected = IllegalStateException.class)
  public void testCustomConfigDisallowsWithHandlerMethod() {
    AvaticaServerConfiguration configuration = new CustomBasicImpersonationConfig();
    server = new HttpServer.Builder()
            .withCustomAuthentication(configuration)
            .withHandler(Mockito.mock(AvaticaHandler.class))
            .withPort(0)
            .build();
    server.start();
  }

  public static HttpServer getAvaticaServer() {
    return server;
  }

  @SuppressWarnings("unchecked") // needed for the mocked customizers, not the builder
  protected void createServer(AvaticaServerConfiguration config, boolean isBasicAuth)
      throws SQLException {
    final JdbcMeta jdbcMeta = new JdbcMeta(CONNECTION_SPEC.url,
        CONNECTION_SPEC.username, CONNECTION_SPEC.password);
    LocalService service = new LocalService(jdbcMeta);

    ConnectorCustomizer connectorCustomizer = new ConnectorCustomizer();
    BasicAuthHandlerCustomizer basicAuthCustomizer =
            new BasicAuthHandlerCustomizer(config, service, isBasicAuth);

    server = new HttpServer.Builder()
            .withCustomAuthentication(config)
            .withPort(0)
            .withServerCustomizers(
                    Arrays.asList(connectorCustomizer, basicAuthCustomizer), Server.class)
            .build();
    server.start();

    // Create and grant permissions to our users
    createHsqldbUsers();
    url = "jdbc:avatica:remote:url=http://localhost:" + connectorCustomizer.getLocalPort()
            + ";authentication=BASIC;serialization=PROTOBUF";
  }

  /**
   * Customizer to add ServerConnectors to the server
   */
  static class ConnectorCustomizer implements ServerCustomizer<Server> {

    ServerConnector connector;

    @Override public void customize(Server server) {
      HttpServer avaticaServer = getAvaticaServer();
      connector = avaticaServer.configureConnector(avaticaServer.getServerConnector(), 0);
      server.setConnectors(new Connector[] { connector });
    }

    public int getLocalPort() {
      return connector.getLocalPort();
    }

  }

  /**
   * Customizer to add handlers to the server (with or without BasicAuth)
   */
  static class BasicAuthHandlerCustomizer implements ServerCustomizer<Server> {

    AvaticaServerConfiguration configuration;
    LocalService service;
    boolean isBasicAuth;

    BasicAuthHandlerCustomizer(AvaticaServerConfiguration configuration, LocalService service,
        boolean isBasicAuth) {
      this.configuration = configuration;
      this.service = service;
      this.isBasicAuth = isBasicAuth;
    }

    @Override public void customize(Server server) {
      HttpServer avaticaServer = getAvaticaServer();

      HandlerFactory factory = new HandlerFactory();
      Handler avaticaHandler = factory.getHandler(service,
              Driver.Serialization.PROTOBUF, null, configuration);

      if (isBasicAuth) {
        HandlerWrapper securityHandler =
                avaticaServer.configureBasicAuthentication(server, configuration);
        securityHandler.setHandler(avaticaHandler);
        avaticaHandler = securityHandler;
      }

      HandlerList handlerList = new HandlerList();
      handlerList.setHandlers(new Handler[] { avaticaHandler, new DefaultHandler()});
      server.setHandler(handlerList);
    }
  }

  /**
   * CustomImpersonationConfig doesn't authenticates the user but supports user impersonation
   */
  static class CustomImpersonationConfig implements AvaticaServerConfiguration {


    @Override public AuthenticationType getAuthenticationType() {
      return AuthenticationType.CUSTOM;
    }

    @Override public String getKerberosRealm() {
      return null;
    }

    @Override public String getKerberosPrincipal() {
      return null;
    }

    @Override public String getKerberosServiceName() {
      return null;
    }

    @Override public String getKerberosHostName() {
      return null;
    }

    @Override public String[] getAllowedRoles() {
      return new String[0];
    }

    @Override public String getHashLoginServiceRealm() {
      return null;
    }

    @Override public String getHashLoginServiceProperties() {
      return null;
    }

    @Override public boolean supportsImpersonation() {
      methodCallCounter1++;
      return true;
    }

    @Override public <T> T doAsRemoteUser(String remoteUserName,
              String remoteAddress, Callable<T> action) throws Exception {
      methodCallCounter2++;
      return action.call();
    }
    @Override public RemoteUserExtractor getRemoteUserExtractor() {
      return new RemoteUserExtractor() {
        @Override public String extract(HttpServletRequest request) {
          methodCallCounter3++;
          return "randomUser";
        }
      };
    }

  }

  /**
   * CustomBasicImpersonationConfig supports BasicAuthentication with user impersonation
   */
  static class CustomBasicImpersonationConfig implements AvaticaServerConfiguration {


    @Override public AuthenticationType getAuthenticationType() {
      return AuthenticationType.CUSTOM;
    }

    @Override public String getKerberosRealm() {
      return null;
    }

    @Override public String getKerberosPrincipal() {
      return null;
    }

    @Override public String getKerberosServiceName() {
      return null;
    }

    @Override public String getKerberosHostName() {
      return null;
    }

    @Override public String[] getAllowedRoles() {
      return new String[] { "users" };
    }

    @Override public String getHashLoginServiceRealm() {
      return "Avatica";
    }

    @Override public String getHashLoginServiceProperties() {
      return HttpAuthBase.getHashLoginServicePropertiesString();
    }

    @Override public boolean supportsImpersonation() {
      methodCallCounter1++;
      return true;
    }

    @Override public <T> T doAsRemoteUser(String remoteUserName,
      String remoteAddress, Callable<T> action) throws Exception {
      methodCallCounter2++;
      if (remoteUserName.equals("USER1")) {
        throw new RemoteUserDisallowedException("USER1 is a disallowed user!");
      }
      return action.call();
    }
    @Override public RemoteUserExtractor getRemoteUserExtractor() {
      return new RemoteUserExtractor() {
        @Override public String extract(HttpServletRequest request)
            throws RemoteUserExtractionException {
          methodCallCounter3++;
          if (request instanceof Request) {
            Authentication authentication = ((Request) request).getAuthentication();
            if (authentication instanceof UserAuthentication) {
              UserIdentity userIdentity = ((UserAuthentication) authentication).getUserIdentity();
              return userIdentity.getUserPrincipal().getName();
            }
          }
          throw new RemoteUserExtractionException("Request doesn't contain user credentials.");
        }
      };
    }
  }

}

// End CustomAuthHttpServerTest.java
