// Copyright 2018 The Bazel Authors. All rights reserved. // // 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 License 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 com.google.devtools.build.lib.remote.http; import static com.google.common.truth.Truth.assertThat; import static com.google.devtools.build.lib.remote.util.Utils.getFromFuture; import static java.util.Collections.singletonList; import static org.junit.Assert.assertThrows; import static org.junit.Assert.fail; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.reset; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; import build.bazel.remote.execution.v2.ActionResult; import build.bazel.remote.execution.v2.Digest; import com.google.auth.Credentials; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.ListeningScheduledExecutorService; import com.google.common.util.concurrent.MoreExecutors; import com.google.devtools.build.lib.authandtls.AuthAndTLSOptions; import com.google.devtools.build.lib.remote.RemoteRetrier; import com.google.devtools.build.lib.remote.Retrier; import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext; import com.google.devtools.build.lib.remote.common.RemoteCacheClient; import com.google.devtools.build.lib.remote.util.DigestUtil; import com.google.devtools.build.lib.remote.util.TracingMetadataUtils; import com.google.devtools.build.lib.vfs.DigestHashFunction; import com.google.devtools.build.lib.vfs.SyscallCache; import com.google.devtools.build.remote.worker.http.HttpCacheServerHandler; import com.google.devtools.common.options.Options; import com.google.protobuf.ByteString; import io.netty.bootstrap.ServerBootstrap; import io.netty.buffer.ByteBuf; import io.netty.buffer.ByteBufUtil; import io.netty.buffer.Unpooled; import io.netty.channel.Channel; import io.netty.channel.ChannelFutureListener; import io.netty.channel.ChannelHandler; import io.netty.channel.ChannelHandler.Sharable; import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelInboundHandler; import io.netty.channel.ChannelInboundHandlerAdapter; import io.netty.channel.ChannelInitializer; import io.netty.channel.EventLoopGroup; import io.netty.channel.ServerChannel; import io.netty.channel.SimpleChannelInboundHandler; import io.netty.channel.epoll.Epoll; import io.netty.channel.epoll.EpollEventLoopGroup; import io.netty.channel.epoll.EpollServerDomainSocketChannel; import io.netty.channel.kqueue.KQueue; import io.netty.channel.kqueue.KQueueEventLoopGroup; import io.netty.channel.kqueue.KQueueServerDomainSocketChannel; import io.netty.channel.nio.NioEventLoopGroup; import io.netty.channel.socket.nio.NioServerSocketChannel; import io.netty.channel.unix.DomainSocketAddress; import io.netty.handler.codec.TooLongFrameException; import io.netty.handler.codec.http.DefaultFullHttpResponse; import io.netty.handler.codec.http.DefaultHttpContent; import io.netty.handler.codec.http.DefaultHttpResponse; import io.netty.handler.codec.http.DefaultLastHttpContent; import io.netty.handler.codec.http.FullHttpRequest; import io.netty.handler.codec.http.FullHttpResponse; import io.netty.handler.codec.http.HttpHeaderNames; import io.netty.handler.codec.http.HttpHeaderValues; import io.netty.handler.codec.http.HttpObjectAggregator; import io.netty.handler.codec.http.HttpResponseStatus; import io.netty.handler.codec.http.HttpServerCodec; import io.netty.handler.codec.http.HttpUtil; import io.netty.handler.codec.http.HttpVersion; import java.io.ByteArrayOutputStream; import java.io.File; import java.io.IOException; import java.io.OutputStream; import java.net.ConnectException; import java.net.InetSocketAddress; import java.net.SocketAddress; import java.net.URI; import java.nio.channels.ClosedChannelException; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executors; import java.util.function.IntFunction; import javax.annotation.Nullable; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.Parameterized; import org.junit.runners.Parameterized.Parameters; import org.mockito.AdditionalAnswers; import org.mockito.Mockito; /** Tests for {@link HttpCacheClient}. */ @RunWith(Parameterized.class) @SuppressWarnings("FutureReturnValueIgnored") public class HttpCacheClientTest { private static final DigestUtil DIGEST_UTIL = new DigestUtil(SyscallCache.NO_CACHE, DigestHashFunction.SHA256); private static final Digest DIGEST = DIGEST_UTIL.computeAsUtf8("File Contents"); private RemoteActionExecutionContext remoteActionExecutionContext; private static ServerChannel createServer( Class serverChannelClass, IntFunction newEventLoopGroup, SocketAddress socketAddress, ChannelHandler handler) { EventLoopGroup eventLoop = newEventLoopGroup.apply(1); ServerBootstrap sb = new ServerBootstrap() .group(eventLoop) .channel(serverChannelClass) .childHandler( new ChannelInitializer() { @Override protected void initChannel(Channel ch) { ch.pipeline().addLast(new HttpServerCodec()); ch.pipeline().addLast(new HttpObjectAggregator(1000)); ch.pipeline().addLast(handler); } }); try { return ((ServerChannel) sb.bind(socketAddress).sync().channel()); } catch (Exception e) { throw new IllegalStateException(e); } } private static DomainSocketAddress newDomainSocketAddress() { try { File file = File.createTempFile("bazel", ".sock", new File("/tmp")); file.delete(); return new DomainSocketAddress(file.getAbsoluteFile()); } catch (Exception e) { throw new IllegalStateException(e); } } interface TestServer { ServerChannel start(ChannelInboundHandler handler); void stop(ServerChannel serverChannel); } private static final class InetTestServer implements TestServer { @Override public ServerChannel start(ChannelInboundHandler handler) { return createServer( NioServerSocketChannel.class, NioEventLoopGroup::new, new InetSocketAddress("localhost", 0), handler); } @Override public void stop(ServerChannel serverChannel) { try { serverChannel.close(); serverChannel.closeFuture().sync(); serverChannel.eventLoop().shutdownGracefully().sync(); } catch (Exception e) { throw new IllegalStateException(e); } } } private static final class UnixDomainServer implements TestServer { // Note: this odd implementation is a workaround because we're unable to shut down and restart // KQueue backed implementations. See https://github.com/netty/netty/issues/7047. private final ServerChannel serverChannel; private ChannelInboundHandler handler = null; public UnixDomainServer( Class serverChannelClass, IntFunction newEventLoopGroup) { EventLoopGroup eventLoop = newEventLoopGroup.apply(1); ServerBootstrap sb = new ServerBootstrap() .group(eventLoop) .channel(serverChannelClass) .childHandler( new ChannelInitializer() { @Override protected void initChannel(Channel ch) { ch.pipeline().addLast(new HttpServerCodec()); ch.pipeline().addLast(new HttpObjectAggregator(1000)); ch.pipeline().addLast(Preconditions.checkNotNull(handler)); } }); try { ServerChannel actual = ((ServerChannel) sb.bind(newDomainSocketAddress()).sync().channel()); this.serverChannel = mock(ServerChannel.class, AdditionalAnswers.delegatesTo(actual)); } catch (Exception e) { throw new IllegalStateException(e); } } @Override public ServerChannel start(ChannelInboundHandler handler) { reset(this.serverChannel); this.handler = handler; return this.serverChannel; } @Override public void stop(ServerChannel serverChannel) { // Note: In the tests, we expect that connecting to a closed server channel results // in a channel connection error. Netty doesn't seem to handle closing domain socket // addresses very well-- often connecting to a closed domain socket will result in a // read timeout instead of a connection timeout. // // This is a hack to ensure connection timeouts are "received" by the tests for this // dummy domain socket server. In particular, this lets the timeoutShouldWork_connect // test work for both inet and domain sockets. // // This is also part of the workaround for https://github.com/netty/netty/issues/7047. when(this.serverChannel.localAddress()).thenReturn(new DomainSocketAddress("")); this.handler = null; } } @Parameters public static List createInputValues() { ArrayList parameters = new ArrayList(Arrays.asList(new Object[][] {{new InetTestServer()}})); if (Epoll.isAvailable()) { parameters.add( new Object[] { new UnixDomainServer(EpollServerDomainSocketChannel.class, EpollEventLoopGroup::new) }); } if (KQueue.isAvailable()) { parameters.add( new Object[] { new UnixDomainServer(KQueueServerDomainSocketChannel.class, KQueueEventLoopGroup::new) }); } return parameters; } private final TestServer testServer; public HttpCacheClientTest(TestServer testServer) { this.testServer = testServer; } private HttpCacheClient createHttpBlobStore( ServerChannel serverChannel, int timeoutSeconds, boolean remoteVerifyDownloads, @Nullable final Credentials creds, AuthAndTLSOptions authAndTlsOptions, Optional optRetrier) throws Exception { SocketAddress socketAddress = serverChannel.localAddress(); RemoteRetrier retrier = optRetrier.orElseGet( () -> { ListeningScheduledExecutorService retryScheduler = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(1)); return new RemoteRetrier( () -> RemoteRetrier.RETRIES_DISABLED, (e) -> false, retryScheduler, Retrier.ALLOW_ALL_CALLS); }); if (socketAddress instanceof DomainSocketAddress) { DomainSocketAddress domainSocketAddress = (DomainSocketAddress) socketAddress; URI uri = new URI("http://localhost"); return HttpCacheClient.create( domainSocketAddress, uri, timeoutSeconds, /* remoteMaxConnections= */ 0, remoteVerifyDownloads, ImmutableList.of(), DIGEST_UTIL, retrier, creds, authAndTlsOptions); } else if (socketAddress instanceof InetSocketAddress) { InetSocketAddress inetSocketAddress = (InetSocketAddress) socketAddress; URI uri = new URI("http://localhost:" + inetSocketAddress.getPort()); return HttpCacheClient.create( uri, timeoutSeconds, /* remoteMaxConnections= */ 0, remoteVerifyDownloads, ImmutableList.of(), DIGEST_UTIL, retrier, creds, authAndTlsOptions); } else { throw new IllegalStateException( "unsupported socket address class " + socketAddress.getClass()); } } private HttpCacheClient createHttpBlobStore( ServerChannel serverChannel, int timeoutSeconds, @Nullable final Credentials creds, AuthAndTLSOptions authAndTlsOptions) throws Exception { return createHttpBlobStore( serverChannel, timeoutSeconds, /* remoteVerifyDownloads= */ true, creds, authAndTlsOptions, Optional.empty()); } @Before public void setUp() throws Exception { remoteActionExecutionContext = RemoteActionExecutionContext.create( TracingMetadataUtils.buildMetadata( "none", "none", Digest.getDefaultInstance().getHash(), null)); } @Test public void testUpload() throws Exception { ServerChannel server = null; try { ConcurrentHashMap cacheContents = new ConcurrentHashMap<>(); server = testServer.start(new HttpCacheServerHandler(cacheContents)); HttpCacheClient blobStore = createHttpBlobStore( server, /* timeoutSeconds= */ 1, /* creds= */ null, new AuthAndTLSOptions()); ByteString data = ByteString.copyFrom("foo bar", StandardCharsets.UTF_8); Digest digest = DIGEST_UTIL.compute(data.toByteArray()); blobStore.uploadBlob(remoteActionExecutionContext, digest, data).get(); assertThat(cacheContents).hasSize(1); String cacheKey = "/cas/" + digest.getHash(); assertThat(cacheContents).containsKey(cacheKey); assertThat(cacheContents.get(cacheKey)).isEqualTo(data.toByteArray()); } finally { testServer.stop(server); } } @Test(timeout = 30000) public void connectTimeout() throws Exception { ServerChannel server = testServer.start(new ChannelInboundHandlerAdapter() {}); testServer.stop(server); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore(server, /* timeoutSeconds= */ 1, credentials, authAndTlsOptions); assertThrows( ConnectException.class, () -> getFromFuture( blobStore.downloadBlob( remoteActionExecutionContext, DIGEST, new ByteArrayOutputStream()))); } @Test(timeout = 30000) public void uploadTimeout() throws Exception { ServerChannel server = null; try { server = testServer.start( new SimpleChannelInboundHandler() { @Override protected void channelRead0( ChannelHandlerContext channelHandlerContext, FullHttpRequest fullHttpRequest) { // Don't respond and force a client timeout. } }); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore(server, /* timeoutSeconds= */ 1, credentials, authAndTlsOptions); byte[] data = "File Contents".getBytes(StandardCharsets.US_ASCII); assertThrows( UploadTimeoutException.class, () -> getFromFuture( blobStore.uploadBlob( remoteActionExecutionContext, DIGEST_UTIL.compute(data), ByteString.copyFrom(data)))); } finally { testServer.stop(server); } } @Test(timeout = 30000) public void downloadTimeout() throws Exception { ServerChannel server = null; try { server = testServer.start( new SimpleChannelInboundHandler() { @Override protected void channelRead0( ChannelHandlerContext channelHandlerContext, FullHttpRequest fullHttpRequest) { // Don't respond and force a client timeout. } }); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore(server, /* timeoutSeconds= */ 1, credentials, authAndTlsOptions); assertThrows( DownloadTimeoutException.class, () -> getFromFuture( blobStore.downloadBlob( remoteActionExecutionContext, DIGEST, new ByteArrayOutputStream()))); } finally { testServer.stop(server); } } @Test public void uploadResponseTooLarge() throws Exception { ServerChannel server = null; try { server = testServer.start( new SimpleChannelInboundHandler() { @Override protected void channelRead0( ChannelHandlerContext channelHandlerContext, FullHttpRequest request) { ByteBuf longMessage = channelHandlerContext.alloc().buffer(50000).writerIndex(50000); DefaultFullHttpResponse response = new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.INTERNAL_SERVER_ERROR, longMessage); channelHandlerContext .writeAndFlush(response) .addListener(ChannelFutureListener.CLOSE); } }); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore(server, /* timeoutSeconds= */ 1, credentials, authAndTlsOptions); ByteString data = ByteString.copyFrom("File Contents", StandardCharsets.US_ASCII); IOException e = assertThrows( IOException.class, () -> getFromFuture( blobStore.uploadBlob( remoteActionExecutionContext, DIGEST_UTIL.compute(data.toByteArray()), data))); assertThat(e.getCause()).isInstanceOf(TooLongFrameException.class); } finally { testServer.stop(server); } } @Test public void testDownloadFailsOnDigestMismatch() throws Exception { // Test that the download fails when a blob/file has a different content hash than expected. ServerChannel server = null; try { server = testServer.start( new SimpleChannelInboundHandler() { @Override protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest request) { ByteBuf data = ctx.alloc().buffer(); ByteBufUtil.writeUtf8(data, "bar"); DefaultFullHttpResponse response = new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.OK, data); HttpUtil.setContentLength(response, data.readableBytes()); ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); } }); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore( server, /* timeoutSeconds= */ 1, /* remoteVerifyDownloads= */ true, credentials, authAndTlsOptions, Optional.empty()); Digest fooDigest = DIGEST_UTIL.compute("foo".getBytes(StandardCharsets.UTF_8)); try (OutputStream out = new ByteArrayOutputStream()) { IOException e = assertThrows( IOException.class, () -> getFromFuture( blobStore.downloadBlob(remoteActionExecutionContext, fooDigest, out))); assertThat(e).hasMessageThat().contains(fooDigest.getHash()); assertThat(e).hasMessageThat().contains(DIGEST_UTIL.computeAsUtf8("bar").getHash()); } } finally { testServer.stop(server); } } @Test public void testDisablingDigestVerification() throws Exception { // Test that when digest verification is disabled a corrupted download works. ServerChannel server = null; try { server = testServer.start( new SimpleChannelInboundHandler() { @Override protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest request) { ByteBuf data = ctx.alloc().buffer(); ByteBufUtil.writeUtf8(data, "bar"); DefaultFullHttpResponse response = new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.OK, data); HttpUtil.setContentLength(response, data.readableBytes()); ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); } }); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore( server, /* timeoutSeconds= */ 1, /* remoteVerifyDownloads= */ false, credentials, authAndTlsOptions, Optional.empty()); Digest fooDigest = DIGEST_UTIL.compute("foo".getBytes(StandardCharsets.UTF_8)); try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { getFromFuture(blobStore.downloadBlob(remoteActionExecutionContext, fooDigest, out)); assertThat(out.toByteArray()).isEqualTo("bar".getBytes(StandardCharsets.UTF_8)); } } finally { testServer.stop(server); } } @Test public void partialDownloadFailsWithoutRetry() throws Exception { ServerChannel server = null; try { ByteBuf chunk1 = Unpooled.wrappedBuffer("File ".getBytes(StandardCharsets.US_ASCII)); ByteBuf chunk2 = Unpooled.wrappedBuffer("Contents".getBytes(StandardCharsets.US_ASCII)); server = testServer.start(new IntermittentFailureHandler(chunk1, chunk2)); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore(server, /* timeoutSeconds= */ 1, credentials, authAndTlsOptions); assertThrows( ClosedChannelException.class, () -> getFromFuture( blobStore.downloadBlob( remoteActionExecutionContext, DIGEST, new ByteArrayOutputStream()))); } finally { testServer.stop(server); } } @Test public void partialDownloadSucceedsWithRetry() throws Exception { ServerChannel server = null; try { ByteBuf chunk1 = Unpooled.wrappedBuffer("File ".getBytes(StandardCharsets.US_ASCII)); // Replace first chunk to test that the client skips the redundant prefix on retry. ByteBuf chunk1Attempt2 = Unpooled.wrappedBuffer("abcde".getBytes(StandardCharsets.US_ASCII)); ByteBuf chunk2 = Unpooled.wrappedBuffer("Contents".getBytes(StandardCharsets.US_ASCII)); server = testServer.start(new IntermittentFailureHandler(chunk1, chunk1Attempt2, chunk2)); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); ListeningScheduledExecutorService retryScheduler = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(1)); RemoteRetrier retrier = new RemoteRetrier( () -> new Retrier.ZeroBackoff(1), (e) -> { return e instanceof ClosedChannelException; }, retryScheduler, Retrier.ALLOW_ALL_CALLS); HttpCacheClient blobStore = createHttpBlobStore( server, /* timeoutSeconds= */ 1, /* remoteVerifyDownloads= */ false, credentials, authAndTlsOptions, Optional.of(retrier)); ByteArrayOutputStream download = new ByteArrayOutputStream(); getFromFuture(blobStore.downloadBlob(remoteActionExecutionContext, DIGEST, download)); assertThat(download.toByteArray()) .isEqualTo("File Contents".getBytes(StandardCharsets.US_ASCII)); } finally { testServer.stop(server); } } @Test public void actionResultRetryReadsFromStart() throws Exception { ServerChannel server = null; try { ActionResult.Builder builder1 = ActionResult.newBuilder(); builder1 .addOutputFilesBuilder() .setPath("attempt1/filename") .setDigest(DIGEST_UTIL.computeAsUtf8("digest1")) .setIsExecutable(true); ActionResult action1 = builder1.build(); ByteArrayOutputStream buffer1 = new ByteArrayOutputStream(); action1.writeTo(buffer1); int splitAt = buffer1.size() / 2; ByteBuf chunk1 = Unpooled.copiedBuffer(buffer1.toByteArray(), 0, splitAt); // Replace first chunk to test that the client starts a fresh ActionResult download on retry. ActionResult.Builder builder2 = ActionResult.newBuilder(); builder2 .addOutputFilesBuilder() .setPath("attempt2/filename") .setDigest(DIGEST_UTIL.computeAsUtf8("digest2")) .setIsExecutable(false); ActionResult action2 = builder2.build(); ByteArrayOutputStream buffer2 = new ByteArrayOutputStream(); action2.writeTo(buffer2); ByteBuf chunk1Attempt2 = Unpooled.copiedBuffer(buffer2.toByteArray(), 0, splitAt); ByteBuf chunk2 = Unpooled.copiedBuffer(buffer2.toByteArray(), splitAt, buffer2.size() - splitAt); server = testServer.start(new IntermittentFailureHandler(chunk1, chunk1Attempt2, chunk2)); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); ListeningScheduledExecutorService retryScheduler = MoreExecutors.listeningDecorator(Executors.newScheduledThreadPool(1)); RemoteRetrier retrier = new RemoteRetrier( () -> new Retrier.ZeroBackoff(1), (e) -> { return e instanceof ClosedChannelException; }, retryScheduler, Retrier.ALLOW_ALL_CALLS); HttpCacheClient blobStore = createHttpBlobStore( server, /* timeoutSeconds= */ 1, /* remoteVerifyDownloads= */ false, credentials, authAndTlsOptions, Optional.of(retrier)); RemoteCacheClient.CachedActionResult download = getFromFuture( blobStore.downloadActionResult( remoteActionExecutionContext, new RemoteCacheClient.ActionKey(DIGEST), false)); assertThat(download.actionResult()).isEqualTo(action2); } finally { testServer.stop(server); } } @Test public void expiredAuthTokensShouldBeRetried_get() throws Exception { expiredAuthTokensShouldBeRetried_get( HttpCacheClientTest.NotAuthorizedHandler.ErrorType.UNAUTHORIZED); expiredAuthTokensShouldBeRetried_get( HttpCacheClientTest.NotAuthorizedHandler.ErrorType.INVALID_TOKEN); } private void expiredAuthTokensShouldBeRetried_get( HttpCacheClientTest.NotAuthorizedHandler.ErrorType errorType) throws Exception { ServerChannel server = null; try { server = testServer.start(new NotAuthorizedHandler(errorType)); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore(server, /* timeoutSeconds= */ 1, credentials, authAndTlsOptions); ByteArrayOutputStream out = Mockito.spy(new ByteArrayOutputStream()); getFromFuture(blobStore.downloadBlob(remoteActionExecutionContext, DIGEST, out)); assertThat(out.toString(StandardCharsets.US_ASCII.name())).isEqualTo("File Contents"); verify(credentials, times(1)).refresh(); verify(credentials, times(2)).getRequestMetadata(any(URI.class)); verify(credentials, times(2)).hasRequestMetadata(); // The caller is responsible to the close the stream. verify(out, never()).close(); verifyNoMoreInteractions(credentials); } finally { testServer.stop(server); } } @Test public void expiredAuthTokensShouldBeRetried_put() throws Exception { expiredAuthTokensShouldBeRetried_put( HttpCacheClientTest.NotAuthorizedHandler.ErrorType.UNAUTHORIZED); expiredAuthTokensShouldBeRetried_put( HttpCacheClientTest.NotAuthorizedHandler.ErrorType.INVALID_TOKEN); } private void expiredAuthTokensShouldBeRetried_put( HttpCacheClientTest.NotAuthorizedHandler.ErrorType errorType) throws Exception { ServerChannel server = null; try { server = testServer.start(new NotAuthorizedHandler(errorType)); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore(server, /* timeoutSeconds= */ 1, credentials, authAndTlsOptions); byte[] data = "File Contents".getBytes(StandardCharsets.US_ASCII); blobStore .uploadBlob( remoteActionExecutionContext, DIGEST_UTIL.compute(data), ByteString.copyFrom(data)) .get(); verify(credentials, times(1)).refresh(); verify(credentials, times(2)).getRequestMetadata(any(URI.class)); verify(credentials, times(2)).hasRequestMetadata(); verifyNoMoreInteractions(credentials); } finally { testServer.stop(server); } } @Test public void errorCodesThatShouldNotBeRetried_get() { errorCodeThatShouldNotBeRetried_get( HttpCacheClientTest.NotAuthorizedHandler.ErrorType.INSUFFICIENT_SCOPE); errorCodeThatShouldNotBeRetried_get( HttpCacheClientTest.NotAuthorizedHandler.ErrorType.INVALID_REQUEST); } private void errorCodeThatShouldNotBeRetried_get( HttpCacheClientTest.NotAuthorizedHandler.ErrorType errorType) { ServerChannel server = null; try { server = testServer.start(new NotAuthorizedHandler(errorType)); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore(server, /* timeoutSeconds= */ 1, credentials, authAndTlsOptions); getFromFuture( blobStore.downloadBlob( remoteActionExecutionContext, DIGEST, new ByteArrayOutputStream())); fail("Exception expected."); } catch (Exception e) { assertThat(e).isInstanceOf(HttpException.class); assertThat(((HttpException) e).response().status()) .isEqualTo(HttpResponseStatus.UNAUTHORIZED); } finally { testServer.stop(server); } } @Test public void errorCodesThatShouldNotBeRetried_put() { errorCodeThatShouldNotBeRetried_put( HttpCacheClientTest.NotAuthorizedHandler.ErrorType.INSUFFICIENT_SCOPE); errorCodeThatShouldNotBeRetried_put( HttpCacheClientTest.NotAuthorizedHandler.ErrorType.INVALID_REQUEST); } private void errorCodeThatShouldNotBeRetried_put( HttpCacheClientTest.NotAuthorizedHandler.ErrorType errorType) { ServerChannel server = null; try { server = testServer.start(new NotAuthorizedHandler(errorType)); Credentials credentials = newCredentials(); AuthAndTLSOptions authAndTlsOptions = Options.getDefaults(AuthAndTLSOptions.class); HttpCacheClient blobStore = createHttpBlobStore(server, /* timeoutSeconds= */ 1, credentials, authAndTlsOptions); byte[] oneByte = new byte[] {0}; getFromFuture( blobStore.uploadBlob( remoteActionExecutionContext, DIGEST_UTIL.compute(oneByte), ByteString.copyFrom(oneByte))); fail("Exception expected."); } catch (Exception e) { assertThat(e).isInstanceOf(HttpException.class); assertThat(((HttpException) e).response().status()) .isEqualTo(HttpResponseStatus.UNAUTHORIZED); } finally { testServer.stop(server); } } private Credentials newCredentials() throws Exception { Credentials credentials = mock(Credentials.class); when(credentials.hasRequestMetadata()).thenReturn(true); Map> headers = new HashMap<>(); headers.put("Authorization", singletonList("Bearer invalidToken")); when(credentials.getRequestMetadata(any(URI.class))).thenReturn(headers); Mockito.doAnswer( (mock) -> { Map> headers2 = new HashMap<>(); headers2.put("Authorization", singletonList("Bearer validToken")); when(credentials.getRequestMetadata(any(URI.class))).thenReturn(headers2); return null; }) .when(credentials) .refresh(); return credentials; } /** * {@link ChannelHandler} that on the first request responds with a 401 UNAUTHORIZED status code, * which the client is expected to retry once with a new authentication token. */ @Sharable static class NotAuthorizedHandler extends SimpleChannelInboundHandler { enum ErrorType { UNAUTHORIZED, INVALID_TOKEN, INSUFFICIENT_SCOPE, INVALID_REQUEST } private final ErrorType errorType; private int messageCount; NotAuthorizedHandler(ErrorType errorType) { this.errorType = errorType; } @Override protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest request) { if (messageCount == 0) { if (!"Bearer invalidToken".equals(request.headers().get(HttpHeaderNames.AUTHORIZATION))) { ctx.writeAndFlush( new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.INTERNAL_SERVER_ERROR)) .addListener(ChannelFutureListener.CLOSE); return; } final FullHttpResponse response; if (errorType == ErrorType.UNAUTHORIZED) { response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.UNAUTHORIZED); } else { response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.UNAUTHORIZED); response .headers() .set( HttpHeaderNames.WWW_AUTHENTICATE, "Bearer realm=\"localhost\"," + "error=\"" + errorType.name().toLowerCase() + "\"," + "error_description=\"The access token expired\""); } ctx.writeAndFlush(response).addListener(ChannelFutureListener.CLOSE); messageCount++; } else if (messageCount == 1) { if (!"Bearer validToken".equals(request.headers().get(HttpHeaderNames.AUTHORIZATION))) { ctx.writeAndFlush( new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.INTERNAL_SERVER_ERROR)) .addListener(ChannelFutureListener.CLOSE); return; } ByteBuf content = ctx.alloc().buffer(); content.writeCharSequence("File Contents", StandardCharsets.US_ASCII); FullHttpResponse response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK, content); HttpUtil.setKeepAlive(response, true); HttpUtil.setContentLength(response, content.readableBytes()); ctx.writeAndFlush(response); messageCount++; } else { // No third message expected. ctx.writeAndFlush( new DefaultFullHttpResponse( HttpVersion.HTTP_1_1, HttpResponseStatus.INTERNAL_SERVER_ERROR)) .addListener(ChannelFutureListener.CLOSE); } } } /** * {@link ChannelHandler} that on the first request returns a partial response and then closes the * stream, and on any further requests returns a full response. */ @Sharable static class IntermittentFailureHandler extends SimpleChannelInboundHandler { private final ByteBuf attempt1Chunk1; private final ByteBuf attempt2Chunk1; private final ByteBuf attempt2Chunk2; private int messageCount; public IntermittentFailureHandler( ByteBuf attempt1Chunk1, ByteBuf attempt2Chunk1, ByteBuf attempt2Chunk2) { this.attempt1Chunk1 = attempt1Chunk1; this.attempt2Chunk1 = attempt2Chunk1; this.attempt2Chunk2 = attempt2Chunk2; } public IntermittentFailureHandler(ByteBuf chunk1, ByteBuf chunk2) { this(chunk1.copy(), chunk1, chunk2); } @Override protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest request) { DefaultHttpResponse response = new DefaultHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK); response.headers().set(HttpHeaderNames.TRANSFER_ENCODING, HttpHeaderValues.CHUNKED); ctx.write(response); if (messageCount == 0) { ctx.writeAndFlush(new DefaultHttpContent(attempt1Chunk1)) .addListener(ChannelFutureListener.CLOSE); } else { ctx.writeAndFlush(new DefaultHttpContent(attempt2Chunk1)); ctx.writeAndFlush(new DefaultLastHttpContent(attempt2Chunk2)) .addListener(ChannelFutureListener.CLOSE); } ++messageCount; } } }