content: /*
   * Copyright (C) 2016 The Android Open Source Project
   *
   * 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 Lice`nse 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 androidx.test.rule;
  
  import static androidx.test.internal.util.Checks.checkArgument;
  import static androidx.test.internal.util.Checks.checkNotNull;
  
  import android.util.Log;
  import androidx.annotation.NonNull;
  import androidx.annotation.RestrictTo;
  import androidx.annotation.RestrictTo.Scope;
  import androidx.annotation.VisibleForTesting;
  import java.util.Properties;
  import org.junit.rules.TestRule;
  import org.junit.runner.Description;
  import org.junit.runners.model.Statement;
  
  /**
   * A {@code TestRule} to forward network traffic to a specific port. By default all traffic is
   * forwarded to default address {@code #DEFAULT_PROXY_HOST}:{@code #DEFAULT_PROXY_PORT} unless
   * otherwise specified
   *
   * <p>Note: Traffic forwarding will only apply to the current process under test.
   *
   * <p><b>This API is currently in beta.</b>
   *
   * @hide
   */
  @RestrictTo(Scope.LIBRARY_GROUP)
  public class PortForwardingRule implements TestRule {
  
    private static final String TAG = "PortForwardingRule";
  
    public static final int MIN_PORT = 1024;
    public static final int MAX_PORT = 65535;
  
    @VisibleForTesting static final int DEFAULT_PROXY_PORT = 8080;
    @VisibleForTesting static final String DEFAULT_PROXY_HOST = "127.0.0.1";
    @VisibleForTesting static final String HTTP_HOST_PROPERTY = "http.proxyHost";
    @VisibleForTesting static final String HTTPS_HOST_PROPERTY = "https.proxyHost";
    @VisibleForTesting static final String HTTP_PORT_PROPERTY = "http.proxyPort";
    @VisibleForTesting static final String HTTPS_PORT_PROPERTY = "https.proxyPort";
  
    @VisibleForTesting final String proxyHost;
    @VisibleForTesting final int proxyPort;
    @VisibleForTesting Properties prop;
  
    private Properties backUpProp;
  
    /** @hide */
    public static class Builder {
  
      private String proxyHost = DEFAULT_PROXY_HOST;
      private int proxyPort = DEFAULT_PROXY_PORT;
      private Properties prop = System.getProperties();
  
      /**
       * Builder to set a specific host address to forward the network traffic to.
       *
       * @param proxyHost The host address to which the network traffic is forwarded during tests.
       */
      public Builder withProxyHost(@NonNull String proxyHost) {
        this.proxyHost = checkNotNull(proxyHost);
        return this;
      }
  
      /**
       * Builder to set a specific port number to forward the network traffic to.
       *
       * @param proxyPort The port number to which the network traffic is forwarded during tests.
       */
      public Builder withProxyPort(int proxyPort) {
        checkArgument(
            proxyPort >= MIN_PORT && proxyPort <= MAX_PORT,
            "%d is used as a proxy port, must in range [%d, %d]",
            proxyPort,
            MIN_PORT,
            MAX_PORT);
        this.proxyPort = proxyPort;
        return this;
      }
  
      /**
       * Builder which allows to pass a {@link Properties} object for testing. This will help to avoid
       * the system properties being affected by tests.
       *
       * @param properties A pre-constructed properties object for testing.
       */
      public Builder withProperties(@NonNull Properties properties) {
        prop = checkNotNull(properties);
        return this;
      }
  
      public PortForwardingRule build() {
        return new PortForwardingRule(this);
      }
    }
  
    private PortForwardingRule(Builder builder) {
      this(builder.proxyHost, builder.proxyPort, builder.prop);
    }
  
    protected PortForwardingRule(int proxyPort) {
      this(DEFAULT_PROXY_HOST, proxyPort, System.getProperties());
    }
  
    @VisibleForTesting
    PortForwardingRule(String proxyHost, int proxyPort, @NonNull Properties properties) {
      this.proxyHost = proxyHost;
      this.proxyPort = proxyPort;
      prop = checkNotNull(properties);
      backUpProp = new Properties();
      backUpProperties();
    }
  
    protected static int getDefaultPort() {
      return DEFAULT_PROXY_PORT;
    }
  
    /**
     * Override this method to execute any code that should run before port forwarding. This method is
     * called before each test method, including any method annotated with <a
     * href="http://junit.sourceforge.net/javadoc/org/junit/Before.html"><code>Before</code></a>.
     */
    protected void beforePortForwarding() {
      // empty by default
    }
  
    /**
     * Override this method to execute any code that should run after port forwarding is set up, but
     * before any test code is run including any method annotated with <a
     * href="http://junit.sourceforge.net/javadoc/org/junit/Before.html"><code>Before</code></a>.
     */
    protected void afterPortForwarding() {
      // empty by default
    }
  
    /**
     * Override this method to execute any code that should run before port forwarding is restored.
     * This method is called after each test method, including any method annotated with <a
     * href="http://junit.sourceforge.net/javadoc/org/junit/After.html"><code>After</code></a>.
     */
    protected void beforeRestoreForwarding() {
      // empty by default
    }
  
    /**
     * Override this method to execute any code that should run after port forwarding is restored.
     * This method is called after each test method, including any method annotated with <a
     * href="http://junit.sourceforge.net/javadoc/org/junit/After.html"><code>After</code></a>.
     */
    protected void afterRestoreForwarding() {
      // empty by default
    }
  
    /**
     * Set the port forwarding system properties (Note: for the process under test only). This method
     * will also back up the existing system properties for later to restore.
     */
    private void setPortForwarding() {
      beforePortForwarding();
      prop.setProperty(HTTP_HOST_PROPERTY, proxyHost);
      prop.setProperty(HTTPS_HOST_PROPERTY, proxyHost);
      prop.setProperty(HTTP_PORT_PROPERTY, String.valueOf(proxyPort));
      prop.setProperty(HTTPS_PORT_PROPERTY, String.valueOf(proxyPort));
      afterPortForwarding();
    }
  
    /** Restore the system properties backed up (Note: for the process under test only) */
    private void restorePortForwarding() {
      try {
        beforeRestoreForwarding();
      } finally {
        restoreOneProperty(prop, backUpProp, HTTP_HOST_PROPERTY);
        restoreOneProperty(prop, backUpProp, HTTPS_HOST_PROPERTY);
        restoreOneProperty(prop, backUpProp, HTTP_PORT_PROPERTY);
        restoreOneProperty(prop, backUpProp, HTTPS_PORT_PROPERTY);
        afterRestoreForwarding();
      }
    }
  
    private void backUpProperties() {
      if (prop.getProperty(HTTP_HOST_PROPERTY) != null) {
        backUpProp.setProperty(HTTP_HOST_PROPERTY, prop.getProperty(HTTP_HOST_PROPERTY));
      }
      if (prop.getProperty(HTTPS_HOST_PROPERTY) != null) {
        backUpProp.setProperty(HTTPS_HOST_PROPERTY, prop.getProperty(HTTPS_HOST_PROPERTY));
      }
      if (prop.getProperty(HTTP_PORT_PROPERTY) != null) {
        backUpProp.setProperty(HTTP_PORT_PROPERTY, prop.getProperty(HTTP_PORT_PROPERTY));
      }
      if (prop.getProperty(HTTPS_PORT_PROPERTY) != null) {
        backUpProp.setProperty(HTTPS_PORT_PROPERTY, prop.getProperty(HTTPS_PORT_PROPERTY));
      }
    }
  
    private void restoreOneProperty(Properties prop, Properties backUpProp, String key) {
      if (backUpProp.getProperty(key) != null) {
        prop.setProperty(key, backUpProp.getProperty(key));
      } else {
        prop.remove(key);
      }
    }
  
    @Override
    public Statement apply(final Statement base, Description description) {
      return new PortForwardingStatement(base);
    }
  
    /** {@link Statement} that set/restore port forwarding before/after the execution of the test. */
    private class PortForwardingStatement extends Statement {
  
      private final Statement base;
  
      public PortForwardingStatement(Statement base) {
        this.base = base;
      }
  
      @Override
      public void evaluate() throws Throwable {
        try {
          setPortForwarding();
          Log.i(
              TAG,
              String.format(
                  "The current process traffic is forwarded to %s:%d", proxyHost, proxyPort));
          base.evaluate();
        } finally {
          restorePortForwarding();
          Log.i(TAG, "Current process traffic forwarding is cancelled");
        }
      }
    }
  }
  