Fix connections leak: close connections after dispatching request

cycletls/index.go

@@ -35,22 +35,22 @@ type cycleTLSRequest struct {
 	Options   Options `json:"options"`
 }
 
-//rename to request+client+options
+// rename to request+client+options
 type fullRequest struct {
 	req     *http.Request
 	client  http.Client
 	options cycleTLSRequest
 }
 
-//Response contains Cycletls response data
+// Response contains Cycletls response data
 type Response struct {
 	RequestID string
 	Status    int
 	Body      string
 	Headers   map[string]string
 }
 
-//JSONBody converts response body to json
+// JSONBody converts response body to json
 func (re Response) JSONBody() map[string]interface{} {
 	var data map[string]interface{}
 	err := json.Unmarshal([]byte(re.Body), &data)
@@ -60,7 +60,7 @@ func (re Response) JSONBody() map[string]interface{} {
 	return data
 }
 
-//CycleTLS creates full request and response
+// CycleTLS creates full request and response
 type CycleTLS struct {
 	ReqChan  chan fullRequest
 	RespChan chan Response
@@ -168,6 +168,8 @@ func processRequest(request cycleTLSRequest) (result fullRequest) {
 }
 
 func dispatcher(res fullRequest) (response Response, err error) {
+	defer res.client.CloseIdleConnections()
+
 	resp, err := res.client.Do(res.req)
 	if err != nil {
 

cycletls/roundtripper.go

@@ -34,7 +34,7 @@ func (rt *roundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	// Fix this later for proper cookie parsing
 	for _, properties := range rt.Cookies {
 		req.AddCookie(&http.Cookie{
-			Name: properties.Name,
+			Name:       properties.Name,
 			Value:      properties.Value,
 			Path:       properties.Path,
 			Domain:     properties.Domain,
@@ -67,7 +67,7 @@ func (rt *roundTripper) getTransport(req *http.Request, addr string) error {
 		return fmt.Errorf("invalid URL scheme: [%v]", req.URL.Scheme)
 	}
 
-	_, err := rt.dialTLS(context.Background(), "tcp", addr)
+	_, err := rt.dialTLS(req.Context(), "tcp", addr)
 	switch err {
 	case errProtocolNegotiated:
 	case nil:
@@ -87,7 +87,6 @@ func (rt *roundTripper) dialTLS(ctx context.Context, network, addr string) (net.
 	// If we have the connection from when we determined the HTTPS
 	// cachedTransports to use, return that.
 	if conn := rt.cachedConnections[addr]; conn != nil {
-		delete(rt.cachedConnections, addr)
 		return conn, nil
 	}
 	rawConn, err := rt.dialer.DialContext(ctx, network, addr)
@@ -165,6 +164,13 @@ func (rt *roundTripper) getDialTLSAddr(req *http.Request) string {
 	return net.JoinHostPort(req.URL.Host, "443") // we can assume port is 443 at this point
 }
 
+func (rt *roundTripper) CloseIdleConnections() {
+	for addr, conn := range rt.cachedConnections {
+		_ = conn.Close()
+		delete(rt.cachedConnections, addr)
+	}
+}
+
 func newRoundTripper(browser browser, dialer ...proxy.ContextDialer) http.RoundTripper {
 	if len(dialer) > 0 {