target based upstream options

internal/http/server.go

@@ -3,6 +3,7 @@ package http
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"io"
 	"log"
@@ -12,6 +13,7 @@ import (
 	"time"
 
 	"edgaru089.ink/go/regolith/internal/perm"
+	"golang.org/x/net/proxy"
 )
 
 const (
@@ -21,6 +23,8 @@ const (
 var (
 	dialer      = net.Dialer{Timeout: outgoing_client_timeout}
 	http_client = http.Client{Timeout: outgoing_client_timeout}
+
+	dialer_func func(ctx context.Context, network, addr string) (net.Conn, error) = dialer.DialContext
 )
 
 type Server struct {
@@ -58,6 +62,22 @@ func (s *Server) Serve(listener net.Listener) (err error) {
 	}
 }
 
+// SetDialer sets the global dialer every Server uses to dial outgoing HTTP/CONNECT connections.
+//
+// When set to nil, net.Dial is used instead.
+func (s *Server) SetDialer(dialer proxy.ContextDialer) {
+	if http_client.Transport == nil {
+		http_client.Transport = http.DefaultTransport
+	}
+
+	http_client.Transport.(*http.Transport).DialContext = dialer.DialContext
+	if dialer != nil {
+		dialer_func = dialer.DialContext
+	} else {
+		dialer_func = dialer.DialContext
+	}
+}
+
 func (s *Server) dispatch(conn net.Conn) {
 	buf := bufio.NewReader(conn)
 	for {
@@ -156,7 +176,9 @@ func (s *Server) handle_connect(conn net.Conn, req *http.Request) {
 	}()
 
 	// dial
-	remote_conn, err := dialer.Dial("tcp", req_addr)
+	ctx, cancel := context.WithTimeout(context.Background(), outgoing_client_timeout)
+	remote_conn, err := dialer_func(ctx, "tcp", req_addr)
+	cancel()
 	if err != nil {
 		simple_respond(conn, req, http.StatusBadGateway)
 		close_err = err