Ability to set custom socket factory

Walkyst · Mar 2, 2023 · eb3c473 · eb3c473
1 parent c86afa6
commit eb3c473
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/.../src/main/java/com/sedmelluq/discord/lavaplayer/tools/http/ExtendedHttpClientBuilder.java b/.../src/main/java/com/sedmelluq/discord/lavaplayer/tools/http/ExtendedHttpClientBuilder.java
@@ -47,6 +47,8 @@ public class ExtendedHttpClientBuilder extends HttpClientBuilder {
 
   private SSLContext sslContextOverride;
   private String[] sslSupportedProtocols;
+  private PlainConnectionSocketFactory plainSocketFactory;
+  private SSLConnectionSocketFactory sslSocketFactory;
   private ConnectionManagerFactory connectionManagerFactory = ExtendedHttpClientBuilder::createDefaultConnectionManager;
 
   @Override
@@ -70,6 +72,14 @@ public void setSslSupportedProtocols(String[] protocols) {
     this.sslSupportedProtocols = protocols;
   }
 
+  public void setPlainConnectionSocketFactory(PlainConnectionSocketFactory plainSocketFactory) {
+    this.plainSocketFactory = plainSocketFactory;
+  }
+
+  public void setSslConnectionSocketFactory(SSLConnectionSocketFactory sslSocketFactory) {
+    this.sslSocketFactory = sslSocketFactory;
+  }
+
   public void setConnectionManagerFactory(ConnectionManagerFactory factory) {
     this.connectionManagerFactory = factory;
   }
@@ -88,12 +98,12 @@ private HttpClientConnectionManager createConnectionManager() {
 
   private Registry<ConnectionSocketFactory> createConnectionSocketFactory() {
     HostnameVerifier hostnameVerifier = new DefaultHostnameVerifier(PublicSuffixMatcherLoader.getDefault());
-    ConnectionSocketFactory sslSocketFactory = new SSLConnectionSocketFactory(sslContextOverride != null ?
+    ConnectionSocketFactory defaultSslSocketFactory = new SSLConnectionSocketFactory(sslContextOverride != null ?
         sslContextOverride : defaultSslContext, sslSupportedProtocols, null, hostnameVerifier);
 
     return RegistryBuilder.<ConnectionSocketFactory>create()
-        .register("http", PlainConnectionSocketFactory.getSocketFactory())
-        .register("https", sslSocketFactory)
+        .register("http", plainSocketFactory != null ? plainSocketFactory : PlainConnectionSocketFactory.getSocketFactory())
+        .register("https", sslSocketFactory != null ? sslSocketFactory : defaultSslSocketFactory)
         .build();
   }