1 /* 2 * Copyright (C) 2021 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.libraries.entitlement.testing; 18 19 import com.google.auto.value.AutoValue; 20 import com.google.common.collect.ImmutableList; 21 import com.google.common.collect.ImmutableMap; 22 23 import java.io.ByteArrayInputStream; 24 import java.io.ByteArrayOutputStream; 25 import java.io.IOException; 26 import java.io.InputStream; 27 import java.io.OutputStream; 28 import java.net.URL; 29 import java.net.URLConnection; 30 import java.net.URLStreamHandler; 31 import java.net.URLStreamHandlerFactory; 32 import java.security.cert.Certificate; 33 import java.util.ArrayList; 34 import java.util.List; 35 import java.util.Map; 36 37 import javax.annotation.Nullable; 38 import javax.net.ssl.HttpsURLConnection; 39 40 /** 41 * A {@link URLStreamHandlerFactory} to return faked {@link URLConnection}, as {@link URL} is a 42 * final class and {@link URL#openConnection} cannot be mocked using mockito. 43 */ 44 public class FakeURLStreamHandler extends URLStreamHandler implements URLStreamHandlerFactory { 45 private Map<String, FakeResponse> mResponses = ImmutableMap.of(); 46 private List<FakeHttpsURLConnection> mConnections = new ArrayList<>(); 47 48 private static final String ACCESS_TOKEN = "8dGozfI6%2FEaSsE7LaTfJKwdy"; 49 private static final String LOCATION = "Location"; 50 private static final String CONTENT_TYPE = "Content-Type"; 51 private static final String RETRY_AFTER = "Retry-After"; 52 53 @Override openConnection(URL u)54 public URLConnection openConnection(URL u) { 55 FakeHttpsURLConnection connection = 56 new FakeHttpsURLConnection(u, mResponses.get(u.toString())); 57 mConnections.add(connection); 58 return connection; 59 } 60 61 @Override createURLStreamHandler(String protocol)62 public URLStreamHandler createURLStreamHandler(String protocol) { 63 return this; 64 } 65 66 /** 67 * Prepares canned responses. Must be called before using this handler to open any {@link 68 * URLConnection}. 69 */ stubResponse(Map<String, FakeResponse> response)70 public FakeURLStreamHandler stubResponse(Map<String, FakeResponse> response) { 71 mResponses = response; 72 mConnections = new ArrayList<>(); 73 return this; 74 } 75 76 /** Returns {@link URLConnection}s opened by this handler since last {@link #stubResponse}. */ getConnections()77 public ImmutableList<FakeHttpsURLConnection> getConnections() { 78 return ImmutableList.copyOf(mConnections); 79 } 80 81 /** Faked {@link HttpsURLConnection} to avoid making any network connection. */ 82 public static class FakeHttpsURLConnection extends HttpsURLConnection { 83 private final FakeResponse mResponse; 84 private final ByteArrayOutputStream mOutputStream = new ByteArrayOutputStream(); 85 FakeHttpsURLConnection(URL url, FakeResponse response)86 public FakeHttpsURLConnection(URL url, FakeResponse response) { 87 super(url); 88 mResponse = response; 89 } 90 91 @Override getInputStream()92 public InputStream getInputStream() throws IOException { 93 return new ByteArrayInputStream(mResponse.responseBody()); 94 } 95 96 @Override getOutputStream()97 public OutputStream getOutputStream() { 98 return mOutputStream; 99 } 100 getBytesWrittenToOutputStream()101 public byte[] getBytesWrittenToOutputStream() { 102 return mOutputStream.toByteArray(); 103 } 104 105 @Override getResponseCode()106 public int getResponseCode() { 107 return mResponse.responseCode(); 108 } 109 110 @Override getHeaderFields()111 public Map<String, List<String>> getHeaderFields() { 112 List<String> locationList = ImmutableList.of("access_token=" + ACCESS_TOKEN); 113 return ImmutableMap.of("Location", locationList); 114 } 115 116 @Override getHeaderField(String name)117 public String getHeaderField(String name) { 118 switch (name) { 119 case LOCATION: 120 return "Location: " + mResponse.responseLocation(); 121 case CONTENT_TYPE: 122 return mResponse.contentType(); 123 case RETRY_AFTER: 124 return mResponse.retryAfter(); 125 default: 126 return null; 127 } 128 } 129 130 @Override connect()131 public void connect() { 132 } 133 134 @Override disconnect()135 public void disconnect() { 136 } 137 138 @Override usingProxy()139 public boolean usingProxy() { 140 return false; 141 } 142 143 @Override getCipherSuite()144 public String getCipherSuite() { 145 return null; 146 } 147 148 @Override getLocalCertificates()149 public Certificate[] getLocalCertificates() { 150 return null; 151 } 152 153 @Override getServerCertificates()154 public Certificate[] getServerCertificates() { 155 return null; 156 } 157 } 158 159 @AutoValue 160 public abstract static class FakeResponse { responseCode()161 public abstract int responseCode(); 162 163 @Nullable responseLocation()164 public abstract String responseLocation(); 165 166 @SuppressWarnings("mutable") // For test only responseBody()167 public abstract byte[] responseBody(); 168 contentType()169 public abstract String contentType(); 170 retryAfter()171 public abstract String retryAfter(); 172 builder()173 public static Builder builder() { 174 return new AutoValue_FakeURLStreamHandler_FakeResponse.Builder() 175 .setResponseBody(new byte[]{}) 176 .setContentType("") 177 .setResponseCode(0) 178 .setResponseLocation("") 179 .setRetryAfter(""); 180 } 181 182 @AutoValue.Builder 183 public abstract static class Builder { setResponseCode(int value)184 public abstract Builder setResponseCode(int value); 185 setResponseLocation(String value)186 public abstract Builder setResponseLocation(String value); 187 setResponseBody(byte[] value)188 public abstract Builder setResponseBody(byte[] value); 189 setContentType(String contentType)190 public abstract Builder setContentType(String contentType); 191 setRetryAfter(String retryAfter)192 public abstract Builder setRetryAfter(String retryAfter); 193 build()194 public abstract FakeResponse build(); 195 } 196 } 197 } 198