Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -894,8 +894,17 @@ public Map<String, List<String>> getSegments(URI controllerBaseUri, String rawTa
public Map<String, List<String>> getSegments(URI controllerBaseUri, String rawTableName,
@Nullable TableType tableType, boolean excludeReplacedSegments, @Nullable AuthProvider authProvider)
throws Exception {
return getSegments(controllerBaseUri, rawTableName, tableType, excludeReplacedSegments, authProvider,
HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
}

/// Returns table segments using an explicit socket timeout for the controller request.
public Map<String, List<String>> getSegments(URI controllerBaseUri, String rawTableName,
@Nullable TableType tableType, boolean excludeReplacedSegments, @Nullable AuthProvider authProvider,
int socketTimeoutMs)
throws Exception {
return getSegments(controllerBaseUri, rawTableName, tableType, excludeReplacedSegments, Long.MIN_VALUE,
Long.MAX_VALUE, false, authProvider);
Long.MAX_VALUE, false, authProvider, socketTimeoutMs);
}

/**
Expand All @@ -916,6 +925,15 @@ public Map<String, List<String>> getSegments(URI controllerBaseUri, String rawTa
@Nullable TableType tableType, boolean excludeReplacedSegments, long startTimestamp, long endTimestamp,
boolean excludeOverlapping, @Nullable AuthProvider authProvider)
throws Exception {
return getSegments(controllerBaseUri, rawTableName, tableType, excludeReplacedSegments, startTimestamp,
endTimestamp, excludeOverlapping, authProvider, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
}

/// Returns table segments using an explicit socket timeout for the controller request.
public Map<String, List<String>> getSegments(URI controllerBaseUri, String rawTableName,
@Nullable TableType tableType, boolean excludeReplacedSegments, long startTimestamp, long endTimestamp,
boolean excludeOverlapping, @Nullable AuthProvider authProvider, int socketTimeoutMs)
throws Exception {
List<String> tableTypes;
if (tableType == null) {
tableTypes = Arrays.asList(TableType.OFFLINE.toString(), TableType.REALTIME.toString());
Expand All @@ -936,7 +954,7 @@ public Map<String, List<String>> getSegments(URI controllerBaseUri, String rawTa
RetryPolicies.exponentialBackoffRetryPolicy(5, 10_000L, 2.0).attempt(() -> {
try {
SimpleHttpResponse response = HttpClient.wrapAndThrowHttpException(
_httpClient.sendRequest(requestBuilder.build(), HttpClient.DEFAULT_SOCKET_TIMEOUT_MS));
_httpClient.sendRequest(requestBuilder.build(), socketTimeoutMs));
LOGGER.info("Response {}: {} received for GET request to URI: {}", response.getStatusCode(),
response.getResponse(), uri);
tableTypeToSegments.put(tableTypeToFilter,
Expand Down Expand Up @@ -1150,9 +1168,16 @@ public SimpleHttpResponse sendSegmentJson(URI uri, String jsonString)
public SimpleHttpResponse startReplaceSegments(URI uri, StartReplaceSegmentsRequest startReplaceSegmentsRequest,
@Nullable AuthProvider authProvider)
throws IOException, HttpErrorStatusException {
return startReplaceSegments(uri, startReplaceSegmentsRequest, authProvider, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
}

/// Starts the consistent-push replace protocol using an explicit socket timeout.
public SimpleHttpResponse startReplaceSegments(URI uri, StartReplaceSegmentsRequest startReplaceSegmentsRequest,
@Nullable AuthProvider authProvider, int socketTimeoutMs)
throws IOException, HttpErrorStatusException {
return HttpClient.wrapAndThrowHttpException(_httpClient.sendRequest(
getStartReplaceSegmentsRequest(uri, JsonUtils.objectToString(startReplaceSegmentsRequest), authProvider),
HttpClient.DEFAULT_SOCKET_TIMEOUT_MS));
socketTimeoutMs));
}

/**
Expand Down Expand Up @@ -1197,8 +1222,14 @@ public SimpleHttpResponse revertReplaceSegments(URI uri)
*/
public SimpleHttpResponse revertReplaceSegments(URI uri, @Nullable AuthProvider authProvider)
throws IOException, HttpErrorStatusException {
return revertReplaceSegments(uri, authProvider, HttpClient.DEFAULT_SOCKET_TIMEOUT_MS);
}

/// Reverts the consistent-push replace protocol using an explicit socket timeout.
public SimpleHttpResponse revertReplaceSegments(URI uri, @Nullable AuthProvider authProvider, int socketTimeoutMs)
throws IOException, HttpErrorStatusException {
return HttpClient.wrapAndThrowHttpException(_httpClient.sendRequest(
getRevertReplaceSegmentRequest(uri, authProvider)));
getRevertReplaceSegmentRequest(uri, authProvider), socketTimeoutMs));
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
import javax.annotation.Nullable;
import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
Expand Down Expand Up @@ -250,30 +251,59 @@ public static void installDefaultSSLSocketFactory(TlsConfig tlsConfig) {
*/
public static void installDefaultSSLSocketFactory(String keyStoreType, String keyStorePath, String keyStorePassword,
String trustStoreType, String trustStorePath, String trustStorePassword) {
SSLContext sc = createSslContext(keyStoreType, keyStorePath, keyStorePassword, trustStoreType, trustStorePath,
trustStorePassword);
// HttpsURLConnection
HttpsURLConnection.setDefaultSSLSocketFactory(sc.getSocketFactory());
setSslContext(sc);
logTlsDiagnosticsOnce("https.default", sc, null, false);
}

/// Creates a client-side SSL context from key/trust store settings without mutating JVM defaults.
///
/// If both store paths are null, the returned context uses default key and trust managers. If either store path is
/// provided, its type and password must also be provided. File-backed stores are auto-renewed.
public static SSLContext createSslContext(@Nullable String keyStoreType, @Nullable String keyStorePath,
@Nullable String keyStorePassword, @Nullable String trustStoreType, @Nullable String trustStorePath,
@Nullable String trustStorePassword) {
return createSslContext(keyStoreType, keyStorePath, keyStorePassword, trustStoreType, trustStorePath,
trustStorePassword, true);
}

/// Creates a client-side SSL context without enabling file-store auto-renewal.
///
/// This is intended for short-lived clients whose owners can close the HTTP client but do not own the renewal
/// executors/watch services created by auto-renewal.
public static SSLContext createSslContextWithoutAutoRenewal(@Nullable String keyStoreType,
@Nullable String keyStorePath, @Nullable String keyStorePassword, @Nullable String trustStoreType,
@Nullable String trustStorePath, @Nullable String trustStorePassword) {
return createSslContext(keyStoreType, keyStorePath, keyStorePassword, trustStoreType, trustStorePath,
trustStorePassword, false);
}

private static SSLContext createSslContext(@Nullable String keyStoreType, @Nullable String keyStorePath,
@Nullable String keyStorePassword, @Nullable String trustStoreType, @Nullable String trustStorePath,
@Nullable String trustStorePassword, boolean enableAutoRenewal) {
try {
SecureRandom secureRandom = new SecureRandom();
SSLContext sc;
if (keyStorePath == null && trustStorePath == null) {
// When neither keyStorePath nor trustStorePath is provided, a SSLFactory cannot be created. create SSLContext
// directly and use the default key manager and trust manager.
sc = SSLContext.getInstance(SSL_CONTEXT_PROTOCOL);
sc.init(null, null, secureRandom);
} else {
SSLFactory sslFactory =
RenewableTlsUtils.createSSLFactory(keyStoreType, keyStorePath, keyStorePassword, trustStoreType,
trustStorePath, trustStorePassword, SSL_CONTEXT_PROTOCOL, secureRandom, true, false);
if (isKeyOrTrustStorePathNullOrHasFileScheme(keyStorePath) && isKeyOrTrustStorePathNullOrHasFileScheme(
trustStorePath)) {
RenewableTlsUtils.enableAutoRenewalFromFileStoreForSSLFactory(sslFactory, keyStoreType, keyStorePath,
keyStorePassword, trustStoreType, trustStorePath, trustStorePassword, SSL_CONTEXT_PROTOCOL, secureRandom,
PinotInsecureMode::isPinotInInsecureMode);
}
sc = sslFactory.getSslContext();
SSLContext sslContext = SSLContext.getInstance(SSL_CONTEXT_PROTOCOL);
sslContext.init(null, null, secureRandom);
return sslContext;
}

SSLFactory sslFactory =
RenewableTlsUtils.createSSLFactory(keyStoreType, keyStorePath, keyStorePassword, trustStoreType,
trustStorePath, trustStorePassword, SSL_CONTEXT_PROTOCOL, secureRandom, true, false);
if (enableAutoRenewal && isKeyOrTrustStorePathNullOrHasFileScheme(keyStorePath)
&& isKeyOrTrustStorePathNullOrHasFileScheme(trustStorePath)) {
RenewableTlsUtils.enableAutoRenewalFromFileStoreForSSLFactory(sslFactory, keyStoreType, keyStorePath,
keyStorePassword, trustStoreType, trustStorePath, trustStorePassword, SSL_CONTEXT_PROTOCOL, secureRandom,
PinotInsecureMode::isPinotInInsecureMode);
}
// HttpsURLConnection
HttpsURLConnection.setDefaultSSLSocketFactory(sc.getSocketFactory());
setSslContext(sc);
logTlsDiagnosticsOnce("https.default", sc, null, false);
return sslFactory.getSslContext();
} catch (GenericSSLContextException | GeneralSecurityException e) {
throw new IllegalStateException("Could not initialize SSL support", e);
}
Expand Down
Loading
Loading