From 902ebe89ac7bab76488f283bb45d015dd52bd789 Mon Sep 17 00:00:00 2001
From: Vladislav Yarmak <vladislav@vm-0.com>
Date: Tue, 19 May 2020 23:13:51 +0300
Subject: [PATCH] implement timeout for CONNECT method

---
 handler.go | 10 ++++++++--
 main.go    |  2 +-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/handler.go b/handler.go
index d643868..8ff18e2 100644
--- a/handler.go
+++ b/handler.go
@@ -4,25 +4,31 @@ import (
     "io"
     "net"
     "fmt"
+    "time"
     "net/http"
     "strings"
+    "context"
 )
 
 type ProxyHandler struct {
+    timeout time.Duration
     logger *CondLogger
     httptransport http.RoundTripper
 }
 
-func NewProxyHandler(logger *CondLogger) *ProxyHandler {
+func NewProxyHandler(timeout time.Duration, logger *CondLogger) *ProxyHandler {
 	httptransport := &http.Transport{}
     return &ProxyHandler{
+        timeout: timeout,
         logger: logger,
         httptransport: httptransport,
     }
 }
 
 func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request) {
-    conn, err := net.Dial("tcp", req.RequestURI)
+    ctx, _ := context.WithTimeout(req.Context(), s.timeout)
+    dialer := net.Dialer{}
+    conn, err := dialer.DialContext(ctx, "tcp", req.RequestURI)
     if err != nil {
         s.logger.Error("Can't satisfy CONNECT request: %v", err)
         http.Error(wr, "Can't satisfy CONNECT request", http.StatusBadGateway)
diff --git a/main.go b/main.go
index b51e8f8..7c289a4 100644
--- a/main.go
+++ b/main.go
@@ -51,7 +51,7 @@ func run() int {
                                 log.LstdFlags | log.Lshortfile),
                                 args.verbosity)
     mainLogger.Info("Starting proxy server...")
-    handler := NewProxyHandler(proxyLogger)
+    handler := NewProxyHandler(args.timeout, proxyLogger)
     err := http.ListenAndServe(args.bind_address, handler)
     mainLogger.Critical("Server terminated with a reason: %v", err)
     mainLogger.Info("Shutting down...")