From ca3b9e892de23f649f104efded67006c5ac3996e Mon Sep 17 00:00:00 2001
From: "Quentin McGaw (desktop)" <quentin.mcgaw@gmail.com>
Date: Wed, 25 Aug 2021 17:02:50 +0000
Subject: [PATCH] Maint: http proxy HTTPS handling simplifications

---
 internal/httpproxy/https.go | 37 +++++++++++++++++++------------------
 1 file changed, 19 insertions(+), 18 deletions(-)

diff --git a/internal/httpproxy/https.go b/internal/httpproxy/https.go
index 8ccf60b4..6f4fb244 100644
--- a/internal/httpproxy/https.go
+++ b/internal/httpproxy/https.go
@@ -1,11 +1,9 @@
 package httpproxy
 
 import (
-	"context"
 	"io"
 	"net"
 	"net/http"
-	"sync"
 )
 
 func (h *handler) handleHTTPS(responseWriter http.ResponseWriter, request *http.Request) {
@@ -38,27 +36,30 @@ func (h *handler) handleHTTPS(responseWriter http.ResponseWriter, request *http.
 	}
 
 	h.wg.Add(1)
-	ctx, cancel := context.WithCancel(h.ctx)
-	const transferGoroutines = 2
-	wg := &sync.WaitGroup{}
-	wg.Add(transferGoroutines)
-	go func() { // trigger cleanup when done
-		wg.Wait()
-		cancel()
-	}()
-	go func() { // cleanup
-		<-ctx.Done()
+
+	serverToClientDone := make(chan struct{})
+	clientToServerClientDone := make(chan struct{})
+	go transfer(destinationConn, clientConnection, clientToServerClientDone)
+	go transfer(clientConnection, destinationConn, serverToClientDone)
+
+	select {
+	case <-h.ctx.Done():
 		destinationConn.Close()
 		clientConnection.Close()
-		h.wg.Done()
-	}()
-	go transfer(destinationConn, clientConnection, wg)
-	go transfer(clientConnection, destinationConn, wg)
+		<-serverToClientDone
+		<-clientToServerClientDone
+	case <-serverToClientDone:
+		<-clientToServerClientDone
+	case <-clientToServerClientDone: // happens more rarely, when a connection is closed on the client side
+		<-serverToClientDone
+	}
+
+	h.wg.Done()
 }
 
-func transfer(destination io.WriteCloser, source io.ReadCloser, wg *sync.WaitGroup) {
+func transfer(destination io.WriteCloser, source io.ReadCloser, done chan<- struct{}) {
 	_, _ = io.Copy(destination, source)
 	_ = source.Close()
 	_ = destination.Close()
-	wg.Done()
+	close(done)
 }