Merge pull request #514 from square/jw/okcurl

Initial implementation of an OkHttp-backed curl clone.
square · Feb 8, 2014 · de6d505 · de6d505
2 parents 1151c98 + acd45e1
commit de6d505
Show file tree

Hide file tree

Showing 7 changed files with 352 additions and 19 deletions.
diff --git a/okcurl/pom.xml b/okcurl/pom.xml
@@ -0,0 +1,94 @@
+<?xml version="1.0" encoding="UTF-8"?>
+
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+
+  <parent>
+    <groupId>com.squareup.okhttp</groupId>
+    <artifactId>parent</artifactId>
+    <version>2.0.0-SNAPSHOT</version>
+  </parent>
+
+  <artifactId>okcurl</artifactId>
+  <name>OkCurl</name>
+
+  <dependencies>
+    <dependency>
+      <groupId>com.squareup.okhttp</groupId>
+      <artifactId>okhttp</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>org.bouncycastle</groupId>
+      <artifactId>bcprov-jdk15on</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.mortbay.jetty.npn</groupId>
+      <artifactId>npn-boot</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.airlift</groupId>
+      <artifactId>airline</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>com.google.guava</groupId>
+      <artifactId>guava</artifactId>
+    </dependency>
+
+    <dependency>
+      <groupId>junit</groupId>
+      <artifactId>junit</artifactId>
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+
+  <build>
+    <resources>
+      <resource>
+        <directory>src/main/resources</directory>
+        <filtering>true</filtering>
+      </resource>
+    </resources>
+
+    <plugins>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-assembly-plugin</artifactId>
+        <configuration>
+          <descriptorRefs>
+            <descriptorRef>jar-with-dependencies</descriptorRef>
+          </descriptorRefs>
+          <archive>
+            <manifest>
+              <mainClass>com.squareup.okhttp.curl.Main</mainClass>
+            </manifest>
+          </archive>
+        </configuration>
+        <executions>
+          <execution>
+            <phase>package</phase>
+            <goals>
+              <goal>single</goal>
+            </goals>
+          </execution>
+        </executions>
+      </plugin>
+      <plugin>
+        <groupId>org.skife.maven</groupId>
+        <artifactId>really-executable-jar-maven-plugin</artifactId>
+        <version>1.1.0</version>
+        <executions>
+          <execution>
+            <phase>package</phase>
+            <goals>
+              <goal>really-executable-jar</goal>
+            </goals>
+          </execution>
+        </executions>
+        <configuration>
+          <flags>-Xbootclasspath/p:$0</flags>
+        </configuration>
+      </plugin>
+    </plugins>
+  </build>
+</project>
diff --git a/okcurl/src/main/java/com/squareup/okhttp/curl/Main.java b/okcurl/src/main/java/com/squareup/okhttp/curl/Main.java
@@ -0,0 +1,203 @@
+package com.squareup.okhttp.curl;
+
+import com.google.common.base.Function;
+import com.google.common.base.Joiner;
+import com.google.common.collect.Lists;
+import com.squareup.okhttp.ConnectionPool;
+import com.squareup.okhttp.Failure;
+import com.squareup.okhttp.Headers;
+import com.squareup.okhttp.OkHttpClient;
+import com.squareup.okhttp.Protocol;
+import com.squareup.okhttp.Request;
+import com.squareup.okhttp.Response;
+import io.airlift.command.Arguments;
+import io.airlift.command.Command;
+import io.airlift.command.HelpOption;
+import io.airlift.command.Option;
+import io.airlift.command.SingleCommand;
+import java.io.IOException;
+import java.io.InputStream;
+import java.security.cert.CertificateException;
+import java.security.cert.X509Certificate;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Properties;
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.SSLSocketFactory;
+import javax.net.ssl.TrustManager;
+import javax.net.ssl.X509TrustManager;
+
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+@Command(name = Main.NAME, description = "A curl for the next-generation web.")
+public class Main extends HelpOption implements Runnable, Response.Receiver {
+  static final String NAME = "okcurl";
+  static final int DEFAULT_TIMEOUT = -1;
+
+  public static void main(String... args) {
+    SingleCommand.singleCommand(Main.class).parse(args).run();
+  }
+
+  private static String versionString() {
+    try {
+      Properties prop = new Properties();
+      InputStream in = Main.class.getResourceAsStream("/okcurl-version.properties");
+      prop.load(in);
+      in.close();
+      return prop.getProperty("version");
+    } catch (IOException e) {
+      throw new AssertionError("Could not load okcurl-version.properties.");
+    }
+  }
+
+  private static String protocols() {
+    return Joiner.on(", ").join(Lists.transform(Arrays.asList(Protocol.values()),
+        new Function<Protocol, String>() {
+          @Override public String apply(Protocol protocol) {
+            return protocol.name.utf8();
+          }
+        }));
+  }
+
+  @Option(name = { "-H", "--header" }, description = "Custom header to pass to server")
+  public List<String> headers;
+
+  @Option(name = { "-A", "--user-agent" }, description = "User-Agent to send to server")
+  public String userAgent = NAME + "/" + versionString();
+
+  @Option(name = "--connect-timeout", description = "Maximum time allowed for connection (seconds)")
+  public int connectTimeout = DEFAULT_TIMEOUT;
+
+  @Option(name = "--read-timeout", description = "Maximum time allowed for reading data (seconds)")
+  public int readTimeout = DEFAULT_TIMEOUT;
+
+  @Option(name = { "-L", "--location" }, description = "Follow redirects")
+  public boolean followRedirects;
+
+  @Option(name = { "-k", "--insecure" },
+      description = "Allow connections to SSL sites without certs")
+  public boolean allowInsecure;
+
+  @Option(name = { "-i", "--include" }, description = "Include protocol headers in the output")
+  public boolean showHeaders;
+
+  @Option(name = { "-e", "--referer" }, description = "Referer URL")
+  public String referer;
+
+  @Option(name = { "-V", "--version" }, description = "Show version number and quit")
+  public boolean version;
+
+  @Arguments(title = "url", description = "Remote resource URL")
+  public String url;
+
+  private OkHttpClient client;
+
+  @Override public void run() {
+    if (showHelpIfRequested()) {
+      return;
+    }
+    if (version) {
+      System.out.println(NAME + " " + versionString());
+      System.out.println("Protocols: " + protocols());
+      return;
+    }
+
+    client = getConfiguredClient();
+    Request request = getConfiguredRequest();
+    client.enqueue(request, this);
+
+    // Immediately begin triggering an executor shutdown so that after execution of the above
+    // request the threads do not stick around until timeout.
+    client.getDispatcher().getExecutorService().shutdown();
+  }
+
+  private OkHttpClient getConfiguredClient() {
+    OkHttpClient client = new OkHttpClient();
+    client.setFollowProtocolRedirects(followRedirects);
+    if (connectTimeout != DEFAULT_TIMEOUT) {
+      client.setConnectTimeout(connectTimeout, SECONDS);
+    }
+    if (readTimeout != DEFAULT_TIMEOUT) {
+      client.setReadTimeout(readTimeout, SECONDS);
+    }
+    if (allowInsecure) {
+      client.setSslSocketFactory(createInsecureSslSocketFactory());
+    }
+    // If we don't set this reference, there's no way to clean shutdown persistent connections.
+    client.setConnectionPool(ConnectionPool.getDefault());
+    return client;
+  }
+
+  private Request getConfiguredRequest() {
+    Request.Builder request = new Request.Builder();
+    request.url(url);
+    if (headers != null) {
+      for (String header : headers) {
+        String[] parts = header.split(":", -1);
+        request.header(parts[0], parts[1]);
+      }
+    }
+    if (referer != null) {
+      request.header("Referer", referer);
+    }
+    request.header("User-Agent", userAgent);
+    return request.build();
+  }
+
+  @Override public void onFailure(Failure failure) {
+    failure.exception().printStackTrace();
+    close();
+  }
+
+  @Override public boolean onResponse(Response response) throws IOException {
+    if (showHeaders) {
+      System.out.println(response.statusLine());
+      Headers headers = response.headers();
+      for (int i = 0, count = headers.size(); i < count; i++) {
+        System.out.println(headers.name(i) + ": " + headers.value(i));
+      }
+      System.out.println();
+    }
+
+    Response.Body body = response.body();
+    byte[] buffer = new byte[1024];
+    while (body.ready()) {
+      int c = body.byteStream().read(buffer);
+      if (c == -1) {
+        close();
+        return true;
+      }
+
+      System.out.write(buffer, 0, c);
+    }
+    close();
+    return false;
+  }
+
+  private void close() {
+    client.getConnectionPool().evictAll(); // Close any persistent connections.
+  }
+
+  private static SSLSocketFactory createInsecureSslSocketFactory() {
+    try {
+      SSLContext context = SSLContext.getInstance("TLS");
+      TrustManager permissive = new X509TrustManager() {
+        @Override public void checkClientTrusted(X509Certificate[] chain, String authType)
+            throws CertificateException {
+        }
+
+        @Override public void checkServerTrusted(X509Certificate[] chain, String authType)
+            throws CertificateException {
+        }
+
+        @Override public X509Certificate[] getAcceptedIssuers() {
+          return null;
+        }
+      };
+      context.init(null, new TrustManager[] { permissive }, null);
+      return context.getSocketFactory();
+    } catch (Exception e) {
+      throw new AssertionError(e);
+    }
+  }
+}
diff --git a/okcurl/src/main/resources/okcurl-version.properties b/okcurl/src/main/resources/okcurl-version.properties
@@ -0,0 +1 @@
+version=${project.version}
diff --git a/okhttp/pom.xml b/okhttp/pom.xml
@@ -18,12 +18,6 @@
       <artifactId>okhttp-protocols</artifactId>
       <version>${project.version}</version>
     </dependency>
-    <dependency>
-      <groupId>com.squareup.okhttp</groupId>
-      <artifactId>mockwebserver</artifactId>
-      <version>${project.version}</version>
-      <scope>test</scope>
-    </dependency>
     <dependency>
       <groupId>org.mortbay.jetty.npn</groupId>
       <artifactId>npn-boot</artifactId>
@@ -35,6 +29,12 @@
       <artifactId>junit</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.squareup.okhttp</groupId>
+      <artifactId>mockwebserver</artifactId>
+      <version>${project.version}</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

diff --git a/okhttp/src/main/java/com/squareup/okhttp/Dispatcher.java b/okhttp/src/main/java/com/squareup/okhttp/Dispatcher.java
@@ -19,15 +19,15 @@
 import java.util.ArrayDeque;
 import java.util.Deque;
 import java.util.Iterator;
-import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 
 /**
  * Policy on when async requests are executed.
  *
- * <p>Each dispatcher uses an {@link Executor} to run jobs internally. If you
+ * <p>Each dispatcher uses an {@link ExecutorService} to run jobs internally. If you
  * supply your own executor, it should be able to run {@link #getMaxRequests the
  * configured maximum} number of jobs concurrently.
  */
@@ -36,27 +36,27 @@ public final class Dispatcher {
   private int maxRequestsPerHost = 5;
 
   /** Executes jobs. Created lazily. */
-  private Executor executor;
+  private ExecutorService executorService;
 
   /** Ready jobs in the order they'll be run. */
   private final Deque<Job> readyJobs = new ArrayDeque<Job>();
 
   /** Running jobs. Includes canceled jobs that haven't finished yet. */
   private final Deque<Job> runningJobs = new ArrayDeque<Job>();
 
-  public Dispatcher(Executor executor) {
-    this.executor = executor;
+  public Dispatcher(ExecutorService executorService) {
+    this.executorService = executorService;
   }
 
   public Dispatcher() {
   }
 
-  public synchronized Executor getExecutor() {
-    if (executor == null) {
-      executor = new ThreadPoolExecutor(0, Integer.MAX_VALUE, 60, TimeUnit.SECONDS,
+  public synchronized ExecutorService getExecutorService() {
+    if (executorService == null) {
+      executorService = new ThreadPoolExecutor(0, Integer.MAX_VALUE, 60, TimeUnit.SECONDS,
           new LinkedBlockingQueue<Runnable>(), Util.threadFactory("OkHttp Dispatcher", false));
     }
-    return executor;
+    return executorService;
   }
 
   /**
@@ -107,7 +107,7 @@ synchronized void enqueue(OkHttpClient client, Request request, Response.Receive
 
     if (runningJobs.size() < maxRequests && runningJobsForHost(job) < maxRequestsPerHost) {
       runningJobs.add(job);
-      getExecutor().execute(job);
+      getExecutorService().execute(job);
     } else {
       readyJobs.add(job);
     }
@@ -143,7 +143,7 @@ private void promoteJobs() {
       if (runningJobsForHost(job) < maxRequestsPerHost) {
         i.remove();
         runningJobs.add(job);
-        getExecutor().execute(job);
+        getExecutorService().execute(job);
       }
 
       if (runningJobs.size() >= maxRequests) return; // Reached max capacity.