From 8aa1ff70df7a7536f8e5ddb86ff40daacaa53be0 Mon Sep 17 00:00:00 2001 From: Keerat Singh Date: Tue, 3 Nov 2020 10:33:54 -0800 Subject: [PATCH 1/3] Address code review comments - Remove COOKIE2_HEADER - Store cookies in the concurrent hashmap in the factory for persistance. - Added end to end test for cookie persistence --- .../flight/client/ClientCookieMiddleware.java | 42 +++--- .../flight/client/TestCookieHandling.java | 142 +++++++++++++----- 2 files changed, 133 insertions(+), 51 deletions(-) diff --git a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java index a49851cfde0..19fdcfe4c3c 100644 --- a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java +++ b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java @@ -18,9 +18,9 @@ package org.apache.arrow.flight.client; import java.net.HttpCookie; -import java.util.HashMap; import java.util.List; -import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import java.util.function.Consumer; import java.util.stream.Collectors; @@ -37,29 +37,30 @@ */ public class ClientCookieMiddleware implements FlightClientMiddleware { private static final String SET_COOKIE_HEADER = "Set-Cookie"; - private static final String SET_COOKIE2_HEADER = "Set-Cookie2"; private static final String COOKIE_HEADER = "Cookie"; + private final Factory factory; - // Use a map to track the most recent version of a cookie from the server. - // Note that cookie names are case-sensitive (but header names aren't). - private Map cookies = new HashMap<>(); - - public ClientCookieMiddleware() { + public ClientCookieMiddleware(Factory factory) { + this.factory = factory; } /** * Factory used within FlightClient. */ public static class Factory implements FlightClientMiddleware.Factory { + // Use a map to track the most recent version of a cookie from the server. + // Note that cookie names are case-sensitive (but header names aren't). + private ConcurrentMap cookies = new ConcurrentHashMap<>(); + @Override public ClientCookieMiddleware onCallStarted(CallInfo info) { - return new ClientCookieMiddleware(); + return new ClientCookieMiddleware(this); } } @Override public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { - final String cookieValue = calculateCookieString(); + final String cookieValue = getValidCookiesAsString(); if (!cookieValue.isEmpty()) { outgoingHeaders.insert(COOKIE_HEADER, cookieValue); } @@ -74,10 +75,14 @@ public void onHeadersReceived(CallHeaders incomingHeaders) { // to signal that the client should stop using the cookie immediately. final Consumer handleSetCookieHeader = (headerValue) -> { final List parsedCookies = HttpCookie.parse(headerValue); - parsedCookies.forEach(parsedCookie -> cookies.put(parsedCookie.getName(), parsedCookie)); + parsedCookies.forEach(parsedCookie -> factory.cookies.put(parsedCookie.getName(), parsedCookie)); }; - incomingHeaders.getAll(SET_COOKIE_HEADER).forEach(handleSetCookieHeader); - incomingHeaders.getAll(SET_COOKIE2_HEADER).forEach(handleSetCookieHeader); + final Iterable setCookieHeaders = incomingHeaders.getAll(SET_COOKIE_HEADER); + if (setCookieHeaders != null) { + for (String cookieHeader : setCookieHeaders) { + handleSetCookieHeader.accept(cookieHeader); + } + } } @Override @@ -85,14 +90,17 @@ public void onCallCompleted(CallStatus status) { } + /** + * Discards expired cookies and returns the valid cookies as a String delimited by ';'. + */ @VisibleForTesting - String calculateCookieString() { + String getValidCookiesAsString() { // Discard expired cookies. - cookies.entrySet().removeIf(cookieEntry -> cookieEntry.getValue().hasExpired()); + factory.cookies.entrySet().removeIf(cookieEntry -> cookieEntry.getValue().hasExpired()); // Cookie header value format: - // =[; ==; = cookie.getValue().toString()) .collect(Collectors.joining("; ")); } diff --git a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java index 1a39b9bd138..b0d1afa743a 100644 --- a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java +++ b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java @@ -17,10 +17,26 @@ package org.apache.arrow.flight.client; +import java.io.IOException; + import org.apache.arrow.flight.CallHeaders; +import org.apache.arrow.flight.CallInfo; +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.flight.Criteria; import org.apache.arrow.flight.ErrorFlightMetadata; +import org.apache.arrow.flight.FlightClient; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightProducer; +import org.apache.arrow.flight.FlightServer; +import org.apache.arrow.flight.FlightServerMiddleware; +import org.apache.arrow.flight.FlightTestUtil; +import org.apache.arrow.flight.NoOpFlightProducer; +import org.apache.arrow.flight.RequestContext; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; import org.junit.After; import org.junit.Assert; +import org.junit.Before; import org.junit.Ignore; import org.junit.Test; @@ -29,13 +45,23 @@ */ public class TestCookieHandling { private static final String SET_COOKIE_HEADER = "Set-Cookie"; - private static final String SET_COOKIE2_HEADER = "Set-Cookie2"; - - private ClientCookieMiddleware cookieMiddleware = new ClientCookieMiddleware(); + private static final String COOKIE_HEADER = "Cookie"; + private BufferAllocator allocator; + private FlightServer server; + private FlightClient client; + + private ClientCookieMiddleware.Factory factory = new ClientCookieMiddleware.Factory(); + private ClientCookieMiddleware cookieMiddleware = new ClientCookieMiddleware(factory); + + @Before + public void setup() throws Exception { + allocator = new RootAllocator(Long.MAX_VALUE); + startServerAndClient(); + } @After public void cleanup() { - cookieMiddleware = new ClientCookieMiddleware(); + cookieMiddleware = new ClientCookieMiddleware(factory); } @Test @@ -43,7 +69,7 @@ public void basicCookie() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v"); cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("k=v", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); } @Test @@ -51,15 +77,15 @@ public void cookieStaysAfterMultipleRequests() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v"); cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("k=v", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("k=v", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("k=v", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); } @Test @@ -68,11 +94,11 @@ public void cookieAutoExpires() { headersToSend.insert(SET_COOKIE_HEADER, "k=v; Max-Age=2"); cookieMiddleware.onHeadersReceived(headersToSend); // Note: using max-age changes cookie version from 0->1, which quotes values. - Assert.assertEquals("k=\"v\"", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=\"v\"", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("k=\"v\"", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=\"v\"", cookieMiddleware.getValidCookiesAsString()); try { Thread.sleep(5000); @@ -80,7 +106,7 @@ public void cookieAutoExpires() { } // Verify that the k cookie was discarded because it expired. - Assert.assertTrue(cookieMiddleware.calculateCookieString().isEmpty()); + Assert.assertTrue(cookieMiddleware.getValidCookiesAsString().isEmpty()); } @Test @@ -89,14 +115,14 @@ public void cookieExplicitlyExpires() { headersToSend.insert(SET_COOKIE_HEADER, "k=v; Max-Age=2"); cookieMiddleware.onHeadersReceived(headersToSend); // Note: using max-age changes cookie version from 0->1, which quotes values. - Assert.assertEquals("k=\"v\"", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=\"v\"", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v; Max-Age=-2"); cookieMiddleware.onHeadersReceived(headersToSend); // Verify that the k cookie was discarded because the server told the client it is expired. - Assert.assertTrue(cookieMiddleware.calculateCookieString().isEmpty()); + Assert.assertTrue(cookieMiddleware.getValidCookiesAsString().isEmpty()); } @Ignore @@ -106,7 +132,7 @@ public void cookieExplicitlyExpiresWithMaxAgeMinusOne() { headersToSend.insert(SET_COOKIE_HEADER, "k=v; Max-Age=2"); cookieMiddleware.onHeadersReceived(headersToSend); // Note: using max-age changes cookie version from 0->1, which quotes values. - Assert.assertEquals("k=\"v\"", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=\"v\"", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); @@ -116,7 +142,7 @@ public void cookieExplicitlyExpiresWithMaxAgeMinusOne() { cookieMiddleware.onHeadersReceived(headersToSend); // Verify that the k cookie was discarded because the server told the client it is expired. - Assert.assertTrue(cookieMiddleware.calculateCookieString().isEmpty()); + Assert.assertTrue(cookieMiddleware.getValidCookiesAsString().isEmpty()); } @Test @@ -124,12 +150,12 @@ public void changeCookieValue() { CallHeaders headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v"); cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("k=v", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); headersToSend = new ErrorFlightMetadata(); headersToSend.insert(SET_COOKIE_HEADER, "k=v2"); cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("k=v2", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("k=v2", cookieMiddleware.getValidCookiesAsString()); } @Test @@ -138,27 +164,75 @@ public void multipleCookiesWithSetCookie() { headersToSend.insert(SET_COOKIE_HEADER, "firstKey=firstVal"); headersToSend.insert(SET_COOKIE_HEADER, "secondKey=secondVal"); cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("firstKey=firstVal; secondKey=secondVal", cookieMiddleware.calculateCookieString()); + Assert.assertEquals("firstKey=firstVal; secondKey=secondVal", cookieMiddleware.getValidCookiesAsString()); } @Test - public void basicCookiesWithSetCookie2() { - CallHeaders headersToSend = new ErrorFlightMetadata(); - headersToSend.insert(SET_COOKIE2_HEADER, "firstKey=firstVal"); - cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("firstKey=firstVal", cookieMiddleware.calculateCookieString()); + public void cookieStaysAfterMultipleRequestsEndToEnd() { + client.handshake(); + Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); + client.listFlights(Criteria.ALL); + Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); + client.listFlights(Criteria.ALL); + Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); } - @Ignore - @Test - public void multipleCookiesWithSetCookie2() { - // There seems to be a JDK bug with HttpCookie.parse() with multiple cookies - // in a Set-Cookie2 header. This is odd, because that method explictly returns a list - // of cookies because of Set-Cookie2. - // Set-Cookie2 itself is deprecated. - CallHeaders headersToSend = new ErrorFlightMetadata(); - headersToSend.insert(SET_COOKIE2_HEADER, "firstKey=firstVal, secondKey=secondVal"); - cookieMiddleware.onHeadersReceived(headersToSend); - Assert.assertEquals("firstKey=firstVal; secondKey=secondVal", cookieMiddleware.calculateCookieString()); + /** + * A server middleware component that injects SET_COOKIE_HEADER into the outgoing headers. + */ + static class SetCookieHeaderInjector implements FlightServerMiddleware { + private final Factory factory; + + public SetCookieHeaderInjector(Factory factory) { + this.factory = factory; + } + + @Override + public void onBeforeSendingHeaders(CallHeaders outgoingHeaders) { + if (!factory.receivedCookieHeader) { + outgoingHeaders.insert(SET_COOKIE_HEADER, "k=v"); + } + } + + @Override + public void onCallCompleted(CallStatus status) { + + } + + @Override + public void onCallErrored(Throwable err) { + + } + + static class Factory implements FlightServerMiddleware.Factory { + private boolean receivedCookieHeader = false; + + @Override + public SetCookieHeaderInjector onCallStarted(CallInfo info, CallHeaders incomingHeaders, + RequestContext context) { + if (null != incomingHeaders.get(COOKIE_HEADER)) { + receivedCookieHeader = true; + } + return new SetCookieHeaderInjector(this); + } + } + } + + private void startServerAndClient() throws IOException { + final FlightProducer flightProducer = new NoOpFlightProducer() { + public void listFlights(CallContext context, Criteria criteria, + StreamListener listener) { + listener.onCompleted(); + } + }; + + this.server = FlightTestUtil.getStartedServer((location) -> FlightServer + .builder(allocator, location, flightProducer) + .middleware(FlightServerMiddleware.Key.of("test"), new SetCookieHeaderInjector.Factory()) + .build()); + + this.client = FlightClient.builder(allocator, server.getLocation()) + .intercept(factory) + .build(); } } From 520a5e40fd6f38d7ccf82f94b66f97cbbe0dcabb Mon Sep 17 00:00:00 2001 From: Keerat Singh Date: Tue, 3 Nov 2020 17:01:00 -0800 Subject: [PATCH 2/3] Fixed test cookieStaysAfterMultipleRequestsEndToEnd - Fixed test cookieStaysAfterMultipleRequestsEndToEnd. - Cleanup TestCookieHandling resources. --- .../flight/client/ClientCookieMiddleware.java | 12 ++++++++---- .../arrow/flight/client/TestCookieHandling.java | 15 ++++++++++----- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java index 19fdcfe4c3c..47af4973f5c 100644 --- a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java +++ b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java @@ -48,13 +48,19 @@ public ClientCookieMiddleware(Factory factory) { * Factory used within FlightClient. */ public static class Factory implements FlightClientMiddleware.Factory { + public ClientCookieMiddleware getClientCookieMiddleware() { + return clientCookieMiddleware; + } + + private ClientCookieMiddleware clientCookieMiddleware; // Use a map to track the most recent version of a cookie from the server. // Note that cookie names are case-sensitive (but header names aren't). private ConcurrentMap cookies = new ConcurrentHashMap<>(); @Override public ClientCookieMiddleware onCallStarted(CallInfo info) { - return new ClientCookieMiddleware(this); + this.clientCookieMiddleware = new ClientCookieMiddleware(this); + return this.clientCookieMiddleware; } } @@ -79,9 +85,7 @@ public void onHeadersReceived(CallHeaders incomingHeaders) { }; final Iterable setCookieHeaders = incomingHeaders.getAll(SET_COOKIE_HEADER); if (setCookieHeaders != null) { - for (String cookieHeader : setCookieHeaders) { - handleSetCookieHeader.accept(cookieHeader); - } + setCookieHeaders.forEach(handleSetCookieHeader); } } diff --git a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java index b0d1afa743a..1c8bb41fd71 100644 --- a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java +++ b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java @@ -34,6 +34,7 @@ import org.apache.arrow.flight.RequestContext; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.util.AutoCloseables; import org.junit.After; import org.junit.Assert; import org.junit.Before; @@ -60,8 +61,12 @@ public void setup() throws Exception { } @After - public void cleanup() { + public void cleanup() throws Exception { cookieMiddleware = new ClientCookieMiddleware(factory); + AutoCloseables.close(client, server, allocator); + client = null; + server = null; + allocator = null; } @Test @@ -170,11 +175,11 @@ public void multipleCookiesWithSetCookie() { @Test public void cookieStaysAfterMultipleRequestsEndToEnd() { client.handshake(); - Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); - client.listFlights(Criteria.ALL); - Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); + Assert.assertEquals("k=v", factory.getClientCookieMiddleware().getValidCookiesAsString()); + client.handshake(); + Assert.assertEquals("k=v", factory.getClientCookieMiddleware().getValidCookiesAsString()); client.listFlights(Criteria.ALL); - Assert.assertEquals("k=v", cookieMiddleware.getValidCookiesAsString()); + Assert.assertEquals("k=v", factory.getClientCookieMiddleware().getValidCookiesAsString()); } /** From 110883004194f49fc46c76ef34ce84a5531977f1 Mon Sep 17 00:00:00 2001 From: Keerat Singh Date: Wed, 4 Nov 2020 11:59:07 -0800 Subject: [PATCH 3/3] Address code review feedback - Update TestCookieHandling to use the test Factory with the ability to intercept the returned ClientCookieMiddleware. --- .../flight/client/ClientCookieMiddleware.java | 11 ++----- .../flight/client/TestCookieHandling.java | 29 ++++++++++++------- 2 files changed, 22 insertions(+), 18 deletions(-) diff --git a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java index 47af4973f5c..474c02008ed 100644 --- a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java +++ b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/client/ClientCookieMiddleware.java @@ -40,7 +40,8 @@ public class ClientCookieMiddleware implements FlightClientMiddleware { private static final String COOKIE_HEADER = "Cookie"; private final Factory factory; - public ClientCookieMiddleware(Factory factory) { + @VisibleForTesting + ClientCookieMiddleware(Factory factory) { this.factory = factory; } @@ -48,19 +49,13 @@ public ClientCookieMiddleware(Factory factory) { * Factory used within FlightClient. */ public static class Factory implements FlightClientMiddleware.Factory { - public ClientCookieMiddleware getClientCookieMiddleware() { - return clientCookieMiddleware; - } - - private ClientCookieMiddleware clientCookieMiddleware; // Use a map to track the most recent version of a cookie from the server. // Note that cookie names are case-sensitive (but header names aren't). private ConcurrentMap cookies = new ConcurrentHashMap<>(); @Override public ClientCookieMiddleware onCallStarted(CallInfo info) { - this.clientCookieMiddleware = new ClientCookieMiddleware(this); - return this.clientCookieMiddleware; + return new ClientCookieMiddleware(this); } } diff --git a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java index 1c8bb41fd71..eea05d26876 100644 --- a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java +++ b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/client/TestCookieHandling.java @@ -51,8 +51,8 @@ public class TestCookieHandling { private FlightServer server; private FlightClient client; - private ClientCookieMiddleware.Factory factory = new ClientCookieMiddleware.Factory(); - private ClientCookieMiddleware cookieMiddleware = new ClientCookieMiddleware(factory); + private ClientCookieMiddlewareTestFactory testFactory = new ClientCookieMiddlewareTestFactory(); + private ClientCookieMiddleware cookieMiddleware = new ClientCookieMiddleware(testFactory); @Before public void setup() throws Exception { @@ -62,7 +62,7 @@ public void setup() throws Exception { @After public void cleanup() throws Exception { - cookieMiddleware = new ClientCookieMiddleware(factory); + cookieMiddleware = new ClientCookieMiddleware(testFactory); AutoCloseables.close(client, server, allocator); client = null; server = null; @@ -175,11 +175,11 @@ public void multipleCookiesWithSetCookie() { @Test public void cookieStaysAfterMultipleRequestsEndToEnd() { client.handshake(); - Assert.assertEquals("k=v", factory.getClientCookieMiddleware().getValidCookiesAsString()); + Assert.assertEquals("k=v", testFactory.clientCookieMiddleware.getValidCookiesAsString()); client.handshake(); - Assert.assertEquals("k=v", factory.getClientCookieMiddleware().getValidCookiesAsString()); + Assert.assertEquals("k=v", testFactory.clientCookieMiddleware.getValidCookiesAsString()); client.listFlights(Criteria.ALL); - Assert.assertEquals("k=v", factory.getClientCookieMiddleware().getValidCookiesAsString()); + Assert.assertEquals("k=v", testFactory.clientCookieMiddleware.getValidCookiesAsString()); } /** @@ -215,14 +215,23 @@ static class Factory implements FlightServerMiddleware.Factory