Merge pull request #104 from yashsinghcodes/fix-proxy-moca

Fixing SHUFFLE_INTERNAL_PROXY
Shuffle · Oct 1, 2024 · eb82d99 · eb82d99
2 parents a2756b7 + 5b74a57
commit eb82d99
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 14 deletions.
diff --git a/shared.go b/shared.go
@@ -23821,23 +23821,46 @@ func DecideExecution(ctx context.Context, workflowExecution WorkflowExecution, e
 	return workflowExecution, relevantActions
 }
 
-func GetExternalClient(baseUrl string) *http.Client {
-	httpProxy := os.Getenv("HTTP_PROXY")
-	httpsProxy := os.Getenv("HTTPS_PROXY")
+func HandleInternalProxy(handler *http.Client) *http.Client {
+	httpProxy := os.Getenv("SHUFFLE_INTERNAL_HTTP_PROXY")
+	httpsProxy := os.Getenv("SHUFFLE_INTERNAL_HTTPS_PROXY")
 
-	// Look for internal proxy instead
-	// in case apps need a different one: https://jamboard.google.com/d/1KNr4JJXmTcH44r5j_5goQYinIe52lWzW-12Ii_joi-w/viewer?mtt=9r8nrqpnbz6z&f=0
+	transport := &http.Transport{}
 
-	overrideHttpProxy := os.Getenv("SHUFFLE_INTERNAL_HTTP_PROXY")
-	overrideHttpsProxy := os.Getenv("SHUFFLE_INTERNAL_HTTPS_PROXY")
-	if len(overrideHttpProxy) > 0 && strings.ToLower(overrideHttpProxy) != "noproxy" {
-		httpProxy = overrideHttpProxy
-	}
+	if (len(httpProxy) > 0 || len(httpsProxy) > 0) && (strings.ToLower(httpProxy) != "noproxy" || strings.ToLower(httpsProxy) != "noproxy") {
+		if len(httpProxy) > 0 && strings.ToLower(httpProxy) != "noproxy" {
+			log.Printf("[INFO] Running with HTTP proxy %s (env: HTTP_PROXY)", httpProxy)
+
+			url_i := url.URL{}
+			url_proxy, err := url_i.Parse(httpProxy)
+			if err == nil {
+				transport.Proxy = http.ProxyURL(url_proxy)
+			}
+		}
+
+		if len(httpsProxy) > 0 && strings.ToLower(httpsProxy) != "noproxy" {
+			log.Printf("[INFO] Running with HTTPS proxy %s (env: HTTPS_PROXY)", httpsProxy)
 
-	if len(overrideHttpsProxy) > 0 && strings.ToLower(overrideHttpsProxy) != "noproxy" {
-		httpsProxy = overrideHttpsProxy
+			url_i := url.URL{}
+			url_proxy, err := url_i.Parse(httpsProxy)
+			if err == nil {
+				transport.Proxy = http.ProxyURL(url_proxy)
+			}
+		}
 	}
 
+	handler.Transport = transport
+
+	return handler
+}
+
+func GetExternalClient(baseUrl string) *http.Client {
+	// Look for internal proxy instead
+	// in case apps need a different one: https://jamboard.google.com/d/1KNr4JJXmTcH44r5j_5goQYinIe52lWzW-12Ii_joi-w/viewer?mtt=9r8nrqpnbz6z&f=0
+	httpProxy := os.Getenv("SHUFFLE_INTERNAL_HTTP_PROXY")
+	httpsProxy := os.Getenv("SHUFFLE_INTERNAL_HTTPS_PROXY")
+
+
 	transport := http.DefaultTransport.(*http.Transport)
 	transport.MaxIdleConnsPerHost = 100
 	transport.ResponseHeaderTimeout = time.Second * 60
@@ -23894,8 +23917,27 @@ func GetExternalClient(baseUrl string) *http.Client {
 
 	if (len(httpProxy) > 0 || len(httpsProxy) > 0) && baseUrl != "http://shuffle-backend:5001" {
 		//client = &http.Client{}
+        if len(httpProxy) > 0 && httpProxy != "noproxy"{
+			log.Printf("[INFO] Running with HTTP proxy %s (env: HTTP_PROXY)", httpProxy)
+
+			url_i := url.URL{}
+			url_proxy, err := url_i.Parse(httpProxy)
+			if err == nil {
+				transport.Proxy = http.ProxyURL(url_proxy)
+			}
+		}
+		if len(httpsProxy) > 0 && httpsProxy != "noproxy"{
+			log.Printf("[INFO] Running with HTTPS proxy %s (env: HTTPS_PROXY)", httpsProxy)
+
+			url_i := url.URL{}
+			url_proxy, err := url_i.Parse(httpsProxy)
+			if err == nil {
+				transport.Proxy = http.ProxyURL(url_proxy)
+			}
+		}
 	} else {
-		if len(httpProxy) > 0 {
+        // keeping this here for now
+		if len(httpProxy) > 0 && httpProxy != "noproxy" {
 			log.Printf("[INFO] Running with HTTP proxy %s (env: HTTP_PROXY)", httpProxy)
 
 			url_i := url.URL{}
@@ -23904,7 +23946,7 @@ func GetExternalClient(baseUrl string) *http.Client {
 				transport.Proxy = http.ProxyURL(url_proxy)
 			}
 		}
-		if len(httpsProxy) > 0 {
+		if len(httpsProxy) > 0 && httpsProxy != "noproxy" {
 			log.Printf("[INFO] Running with HTTPS proxy %s (env: HTTPS_PROXY)", httpsProxy)
 
 			url_i := url.URL{}

diff --git a/shared_test.go b/shared_test.go
@@ -2,6 +2,7 @@ package shuffle
 
 import (
     "testing"
+    "net/http"
 )
 
 func TestIsLoop(t *testing.T) {
@@ -26,3 +27,17 @@ func TestIsLoop(t *testing.T) {
         }
     }
 }
+
+// Simple Test for HandleInternalProxy(client)
+// set env SHUFFLE_INTERNAL_HTTP_PROXY to test the function.
+func TestHandleInternalProxy(t *testing.T) {
+	client := &http.Client{}
+	result := HandleInternalProxy(client)
+
+	if result.Transport.(*http.Transport).Proxy != nil {
+		proxyURL, _ := result.Transport.(*http.Transport).Proxy(nil)
+		t.Logf("Proxy URL set: %v", proxyURL)
+	} else {
+		t.Log("No proxy set")
+	}
+}