make handler use custom dialer

This commit is contained in:
Vladislav Yarmak 2023-02-07 20:26:12 +02:00
parent 907d52b7fa
commit daac1baa9a
2 changed files with 13 additions and 5 deletions

View File

@ -10,21 +10,29 @@ import (
"time" "time"
) )
type HandlerDialer interface {
DialContext(ctx context.Context, net, address string) (net.Conn, error)
}
type ProxyHandler struct { type ProxyHandler struct {
timeout time.Duration timeout time.Duration
auth Auth auth Auth
logger *CondLogger logger *CondLogger
dialer HandlerDialer
httptransport http.RoundTripper httptransport http.RoundTripper
outbound map[string]string outbound map[string]string
outboundMux sync.RWMutex outboundMux sync.RWMutex
} }
func NewProxyHandler(timeout time.Duration, auth Auth, logger *CondLogger) *ProxyHandler { func NewProxyHandler(timeout time.Duration, auth Auth, dialer HandlerDialer, logger *CondLogger) *ProxyHandler {
httptransport := &http.Transport{} httptransport := &http.Transport{
DialContext: dialer.DialContext,
}
return &ProxyHandler{ return &ProxyHandler{
timeout: timeout, timeout: timeout,
auth: auth, auth: auth,
logger: logger, logger: logger,
dialer: dialer,
httptransport: httptransport, httptransport: httptransport,
outbound: make(map[string]string), outbound: make(map[string]string),
} }
@ -32,8 +40,7 @@ func NewProxyHandler(timeout time.Duration, auth Auth, logger *CondLogger) *Prox
func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request) { func (s *ProxyHandler) HandleTunnel(wr http.ResponseWriter, req *http.Request) {
ctx, _ := context.WithTimeout(req.Context(), s.timeout) ctx, _ := context.WithTimeout(req.Context(), s.timeout)
dialer := net.Dialer{} conn, err := s.dialer.DialContext(ctx, "tcp", req.RequestURI)
conn, err := dialer.DialContext(ctx, "tcp", req.RequestURI)
if err != nil { if err != nil {
s.logger.Error("Can't satisfy CONNECT request: %v", err) s.logger.Error("Can't satisfy CONNECT request: %v", err)
http.Error(wr, "Can't satisfy CONNECT request", http.StatusBadGateway) http.Error(wr, "Can't satisfy CONNECT request", http.StatusBadGateway)

View File

@ -5,6 +5,7 @@ import (
"flag" "flag"
"fmt" "fmt"
"log" "log"
"net"
"net/http" "net/http"
"os" "os"
"path/filepath" "path/filepath"
@ -141,7 +142,7 @@ func run() int {
server := http.Server{ server := http.Server{
Addr: args.bind_address, Addr: args.bind_address,
Handler: NewProxyHandler(args.timeout, auth, proxyLogger), Handler: NewProxyHandler(args.timeout, auth, new(net.Dialer), proxyLogger),
ErrorLog: log.New(logWriter, "HTTPSRV : ", log.LstdFlags|log.Lshortfile), ErrorLog: log.New(logWriter, "HTTPSRV : ", log.LstdFlags|log.Lshortfile),
ReadTimeout: 0, ReadTimeout: 0,
ReadHeaderTimeout: 0, ReadHeaderTimeout: 0,