1 // Copyright 2021 The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 package com.google.android.downloader;
16 
17 import static com.google.common.base.Preconditions.checkNotNull;
18 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
19 
20 import com.google.common.collect.ImmutableMultimap;
21 import com.google.common.collect.ImmutableSet;
22 import com.google.common.util.concurrent.Futures;
23 import com.google.common.util.concurrent.ListenableFuture;
24 import com.google.common.util.concurrent.ListeningExecutorService;
25 import com.google.common.util.concurrent.SettableFuture;
26 import com.squareup.okhttp.Call;
27 import com.squareup.okhttp.Callback;
28 import com.squareup.okhttp.OkHttpClient;
29 import com.squareup.okhttp.Request;
30 import com.squareup.okhttp.Response;
31 import com.squareup.okhttp.ResponseBody;
32 import java.io.IOException;
33 import java.io.OutputStream;
34 import java.nio.channels.Channels;
35 import java.nio.channels.WritableByteChannel;
36 import java.util.List;
37 import java.util.Map;
38 import java.util.Set;
39 import okio.Okio;
40 import okio.Sink;
41 
42 /** {@link UrlEngine} implementation that uses OkHttp3 for network connectivity. */
43 public class OkHttp2UrlEngine implements UrlEngine {
44   private static final ImmutableSet<String> HTTP_SCHEMES = ImmutableSet.of("http", "https");
45 
46   private final OkHttpClient client;
47   private final ListeningExecutorService transferExecutorService;
48 
49   /**
50    * Constructs an instance of the OkHttp URL engine, for the given OkHttpClient instance.
51    *
52    * <p>Note that due to how OkHttp is implemented, reads from the network are blocking operations,
53    * and thus threads in the provided {@link ListeningExecutorService} can be tied up for long
54    * periods of time waiting on network responses. To mitigate, set {@link
55    * OkHttpClient#setReadTimeout(long, java.util.concurrent.TimeUnit)} to a value that is reasonable
56    * for your use case.
57    *
58    * @param transferExecutorService Executor on which the requests are synchronously executed.
59    */
OkHttp2UrlEngine(OkHttpClient client, ListeningExecutorService transferExecutorService)60   public OkHttp2UrlEngine(OkHttpClient client, ListeningExecutorService transferExecutorService) {
61     checkNotNull(client.getDispatcher());
62     this.client = client;
63     this.transferExecutorService = transferExecutorService;
64   }
65 
66   @Override
createRequest(String url)67   public UrlRequest.Builder createRequest(String url) {
68     return new OkHttpUrlRequestBuilder(url);
69   }
70 
71   @Override
supportedSchemes()72   public Set<String> supportedSchemes() {
73     return HTTP_SCHEMES;
74   }
75 
76   class OkHttpUrlRequestBuilder implements UrlRequest.Builder {
77     private final String url;
78     private final ImmutableMultimap.Builder<String, String> headers = ImmutableMultimap.builder();
79 
OkHttpUrlRequestBuilder(String url)80     OkHttpUrlRequestBuilder(String url) {
81       this.url = url;
82     }
83 
84     @Override
addHeader(String key, String value)85     public UrlRequest.Builder addHeader(String key, String value) {
86       headers.put(key, value);
87       return this;
88     }
89 
90     @Override
build()91     public UrlRequest build() {
92       return new OkHttpUrlRequest(url, headers.build());
93     }
94   }
95 
96   /**
97    * Implementation of {@link UrlRequest} for OkHttp. Wraps OkHttp's {@link Call} to make network
98    * requests.
99    */
100   class OkHttpUrlRequest implements UrlRequest {
101     private final String url;
102     private final ImmutableMultimap<String, String> headers;
103 
OkHttpUrlRequest(String url, ImmutableMultimap<String, String> headers)104     OkHttpUrlRequest(String url, ImmutableMultimap<String, String> headers) {
105       this.url = url;
106       this.headers = headers;
107     }
108 
109     @Override
send()110     public ListenableFuture<UrlResponse> send() {
111       Request.Builder requestBuilder = new Request.Builder();
112 
113       try {
114         requestBuilder.url(url);
115       } catch (IllegalArgumentException e) {
116         return Futures.immediateFailedFuture(new RequestException(e));
117       }
118 
119       for (String key : headers.keys()) {
120         for (String value : headers.get(key)) {
121           requestBuilder.header(key, value);
122         }
123       }
124 
125       SettableFuture<UrlResponse> responseFuture = SettableFuture.create();
126       Call call = client.newCall(requestBuilder.build());
127       call.enqueue(
128           new Callback() {
129             @Override
130             public void onResponse(Response response) {
131               if (response.isSuccessful()) {
132                 responseFuture.set(new OkHttpUrlResponse(response));
133               } else {
134                 responseFuture.setException(
135                     new RequestException(
136                         ErrorDetails.createFromHttpErrorResponse(
137                             response.code(), response.headers().toMultimap(), response.message())));
138                 try {
139                   response.body().close();
140                 } catch (IOException e) {
141                   // Ignore, an exception on the future was already set.
142                 }
143               }
144             }
145 
146             @Override
147             public void onFailure(Request request, IOException exception) {
148               responseFuture.setException(new RequestException(exception));
149             }
150           });
151       responseFuture.addListener(
152           () -> {
153             if (responseFuture.isCancelled()) {
154               call.cancel();
155             }
156           },
157           directExecutor());
158       return responseFuture;
159     }
160   }
161 
162   /**
163    * Implementation of {@link UrlResponse} for OkHttp. Wraps OkHttp's {@link okhttp3.Response} to
164    * complete its operations.
165    */
166   class OkHttpUrlResponse implements UrlResponse {
167     private final Response response;
168 
OkHttpUrlResponse(Response response)169     OkHttpUrlResponse(Response response) {
170       this.response = response;
171     }
172 
173     @Override
getResponseCode()174     public int getResponseCode() {
175       return response.code();
176     }
177 
178     @Override
getResponseHeaders()179     public Map<String, List<String>> getResponseHeaders() {
180       return response.headers().toMultimap();
181     }
182 
183     @Override
readResponseBody(WritableByteChannel destinationChannel)184     public ListenableFuture<Long> readResponseBody(WritableByteChannel destinationChannel) {
185       IOUtil.validateChannel(destinationChannel);
186       return transferExecutorService.submit(
187           () -> {
188             try (ResponseBody body = checkNotNull(response.body())) {
189               // Transfer the response body to the destination channel via OkHttp's Okio API.
190               // Sadly this needs to operate on OutputStream instead of Channels, but at least
191               // Okio manages buffers efficiently internally.
192               OutputStream outputStream = Channels.newOutputStream(destinationChannel);
193               Sink sink = Okio.sink(outputStream);
194               return body.source().readAll(sink);
195             } catch (IllegalStateException e) {
196               // OkHttp throws an IllegalStateException if the stream is closed while
197               // trying to write. Catch and rethrow.
198               throw new RequestException(e);
199             } catch (IOException e) {
200               if (e instanceof RequestException) {
201                 throw e;
202               } else {
203                 throw new RequestException(e);
204               }
205             } finally {
206               response.body().close();
207             }
208           });
209     }
210 
211     @Override
close()212     public void close() throws IOException {
213       response.body().close();
214     }
215   }
216 }
217