1 /* 2 * Copyright (C) 2023 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package com.android.server.wifi.entitlement.http; 18 19 import static java.nio.charset.StandardCharsets.UTF_8; 20 21 import com.google.auto.value.AutoValue; 22 import com.google.common.collect.ImmutableList; 23 import com.google.common.collect.ImmutableMap; 24 import com.google.common.net.HttpHeaders; 25 26 import java.io.ByteArrayInputStream; 27 import java.io.ByteArrayOutputStream; 28 import java.io.IOException; 29 import java.io.InputStream; 30 import java.io.OutputStream; 31 import java.net.URL; 32 import java.net.URLConnection; 33 import java.net.URLStreamHandler; 34 import java.net.URLStreamHandlerFactory; 35 import java.security.cert.Certificate; 36 import java.util.ArrayList; 37 import java.util.List; 38 import java.util.Map; 39 40 import javax.annotation.Nullable; 41 import javax.net.ssl.HttpsURLConnection; 42 43 /** 44 * A {@link URLStreamHandlerFactory} to return faked {@link URLConnection}, as {@link URL} is a 45 * final class and {@link URL#openConnection} cannot be mocked using mockito. 46 */ 47 public class FakeURLStreamHandler extends URLStreamHandler implements URLStreamHandlerFactory { 48 private Map<String, FakeResponse> mResponses = ImmutableMap.of(); 49 private List<FakeHttpsURLConnection> mConnections = new ArrayList<>(); 50 51 private static final String ACCESS_TOKEN = "8dGozfI6%2FEaSsE7LaTfJKwdy"; 52 private static final String LOCATION = "Location"; 53 private static final String CONTENT_TYPE = "Content-Type"; 54 private static final String RETRY_AFTER = "Retry-After"; 55 56 @Override openConnection(URL u)57 public URLConnection openConnection(URL u) { 58 FakeHttpsURLConnection connection = 59 new FakeHttpsURLConnection(u, mResponses.get(u.toString())); 60 mConnections.add(connection); 61 return connection; 62 } 63 64 @Override createURLStreamHandler(String protocol)65 public URLStreamHandler createURLStreamHandler(String protocol) { 66 return this; 67 } 68 69 /** 70 * Prepares canned responses. Must be called before using this handler to open any {@link 71 * URLConnection}. 72 */ stubResponse(Map<String, FakeResponse> response)73 public FakeURLStreamHandler stubResponse(Map<String, FakeResponse> response) { 74 mResponses = response; 75 mConnections = new ArrayList<>(); 76 return this; 77 } 78 79 /** Returns {@link URLConnection}s opened by this handler since last {@link #stubResponse}. */ getConnections()80 public ImmutableList<FakeHttpsURLConnection> getConnections() { 81 return ImmutableList.copyOf(mConnections); 82 } 83 84 /** Faked {@link HttpsURLConnection} to avoid making any network connection. */ 85 public static class FakeHttpsURLConnection extends HttpsURLConnection { 86 private final FakeResponse mResponse; 87 private final ByteArrayOutputStream mOutputStream = new ByteArrayOutputStream(); 88 FakeHttpsURLConnection(URL url, FakeResponse response)89 public FakeHttpsURLConnection(URL url, FakeResponse response) { 90 super(url); 91 mResponse = response; 92 } 93 94 @Override getInputStream()95 public InputStream getInputStream() throws IOException { 96 if (mResponse.hasException() && mResponse.responseBody().length == 0) { 97 throw new IOException("stub exception"); 98 } 99 return new ByteArrayInputStream(mResponse.responseBody()); 100 } 101 102 @Override getOutputStream()103 public OutputStream getOutputStream() { 104 return mOutputStream; 105 } 106 107 @Override getErrorStream()108 public InputStream getErrorStream() { 109 return new ByteArrayInputStream("stub error".getBytes(UTF_8)); 110 } 111 getBytesWrittenToOutputStream()112 public byte[] getBytesWrittenToOutputStream() { 113 return mOutputStream.toByteArray(); 114 } 115 116 @Override getResponseCode()117 public int getResponseCode() throws IOException { 118 if (mResponse.hasException() && mResponse.responseCode() == 0) { 119 throw new IOException("stub exception"); 120 } 121 return mResponse.responseCode(); 122 } 123 124 @Override getHeaderFields()125 public Map<String, List<String>> getHeaderFields() { 126 List<String> locationList = ImmutableList.of("access_token=" + ACCESS_TOKEN); 127 return ImmutableMap.of("Location", locationList); 128 } 129 130 @Override getHeaderField(String name)131 public String getHeaderField(String name) { 132 switch (name) { 133 case LOCATION: 134 return "Location: " + mResponse.responseLocation(); 135 case CONTENT_TYPE: 136 return mResponse.contentType(); 137 case RETRY_AFTER: 138 return mResponse.retryAfter(); 139 // Consistent with the string used by java.net.URLConnection.getContentEncoding() 140 case HttpHeaders.CONTENT_ENCODING: 141 return mResponse.contentEncoding(); 142 default: 143 return null; 144 } 145 } 146 147 @Override connect()148 public void connect() { 149 } 150 151 @Override disconnect()152 public void disconnect() { 153 } 154 155 @Override usingProxy()156 public boolean usingProxy() { 157 return false; 158 } 159 160 @Override getCipherSuite()161 public String getCipherSuite() { 162 return null; 163 } 164 165 @Override getLocalCertificates()166 public Certificate[] getLocalCertificates() { 167 return null; 168 } 169 170 @Override getServerCertificates()171 public Certificate[] getServerCertificates() { 172 return null; 173 } 174 } 175 176 @AutoValue 177 public abstract static class FakeResponse { responseCode()178 public abstract int responseCode(); 179 180 @Nullable responseLocation()181 public abstract String responseLocation(); 182 183 @SuppressWarnings("mutable") // For test only responseBody()184 public abstract byte[] responseBody(); 185 contentType()186 public abstract String contentType(); 187 contentEncoding()188 public abstract String contentEncoding(); 189 retryAfter()190 public abstract String retryAfter(); 191 hasException()192 abstract boolean hasException(); 193 builder()194 public static Builder builder() { 195 return new AutoValue_FakeURLStreamHandler_FakeResponse.Builder() 196 .setResponseBody(new byte[]{}) 197 .setContentType("") 198 .setContentEncoding("") 199 .setResponseCode(0) 200 .setResponseLocation("") 201 .setRetryAfter("") 202 .setHasException(false); 203 } 204 205 @AutoValue.Builder 206 public abstract static class Builder { setResponseCode(int value)207 public abstract Builder setResponseCode(int value); 208 setResponseLocation(String value)209 public abstract Builder setResponseLocation(String value); 210 setResponseBody(byte[] value)211 public abstract Builder setResponseBody(byte[] value); 212 setContentType(String contentType)213 public abstract Builder setContentType(String contentType); 214 setContentEncoding(String contentEncoding)215 public abstract Builder setContentEncoding(String contentEncoding); 216 setRetryAfter(String retryAfter)217 public abstract Builder setRetryAfter(String retryAfter); 218 setHasException(boolean hasException)219 public abstract Builder setHasException(boolean hasException); 220 build()221 public abstract FakeResponse build(); 222 } 223 } 224 } 225