todo realUpstreamProxy

proxy/connection.go

@@ -20,10 +20,12 @@ import (
 // client connection
 type ClientConn struct {
  Id uuid.UUID
- Conn net.Conn
+ Conn net.Conn // rawClientConnContextKey is this
  Tls bool
 }
 
+var rawClientConnContextKey = new(struct{})
+
 func newClientConn(c net.Conn) *ClientConn {
  return &ClientConn{
  Id: uuid.NewV4(),
@@ -115,7 +117,7 @@ func (connCtx *ConnContext) initHttpServerConn() {
  serverConn := newServerConn()
  serverConn.client = &http.Client{
  Transport: &http.Transport{
- Proxy: clientProxy(connCtx.proxy.Opts.Upstream),
+ Proxy: connCtx.proxy.realUpstreamProxy(),
  DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
  c, err := (&net.Dialer{}).DialContext(ctx, network, addr)
  if err != nil {
@@ -150,13 +152,13 @@ func (connCtx *ConnContext) initHttpServerConn() {
  connCtx.ServerConn = serverConn
 }
 
-func (connCtx *ConnContext) initServerTcpConn(req *http.Request) error {
+func (connCtx *ConnContext) initServerTcpConn(req *http.Request, rawClientConn net.Conn) error {
  log.Debugln("in initServerTcpConn")
  ServerConn := newServerConn()
  connCtx.ServerConn = ServerConn
  ServerConn.Address = connCtx.pipeConn.host
 
- plainConn, err := getConnFrom(req.Host, connCtx.proxy.Opts.Upstream)
+ plainConn, err := getConnFrom(req.Host, connCtx.proxy, rawClientConn)
  if err != nil {
  return err
  }
@@ -369,9 +371,16 @@ func getProxyConn(proxyUrl *url.URL, address string) (net.Conn, error) {
  return conn, nil
 }
 
-func getConnFrom(address string, upstream string) (net.Conn, error) {
- clientReq := &http.Request{URL: &url.URL{Scheme: "https", Host: address}}
- proxyUrl, err := clientProxy(upstream)(clientReq)
+func getConnFrom(address string, proxy *Proxy, rawClientConn net.Conn) (net.Conn, error) {
+ clientReqCtx := context.WithValue(context.Background(), rawClientConnContextKey, rawClientConn)
+ clientReq, err := http.NewRequestWithContext(clientReqCtx, "CONNECT", "https://"+address, nil)
+ if err != nil {
+ return nil, err
+ }
+ clientReq.URL.Scheme = "https"
+ clientReq.URL.Host = address
+
+ proxyUrl, err := proxy.realUpstreamProxy()(clientReq)
  if err != nil {
  return nil, err
  }

proxy/interceptor.go

@@ -137,9 +137,9 @@ func (m *middle) close() error {
  return nil
 }
 
-func (m *middle) dial(req *http.Request) (net.Conn, error) {
+func (m *middle) dial(req *http.Request, rawClientConn net.Conn) (net.Conn, error) {
  pipeClientConn, pipeServerConn := newPipes(req)
- err := pipeServerConn.connContext.initServerTcpConn(req)
+ err := pipeServerConn.connContext.initServerTcpConn(req, rawClientConn)
  if err != nil {
  pipeClientConn.Close()
  pipeServerConn.Close()

proxy/proxy.go

@@ -8,6 +8,7 @@ import (
  "io"
  "net"
  "net/http"
+ "net/url"
 
  log "github.com/sirupsen/logrus"
 )
@@ -30,6 +31,7 @@ type Proxy struct {
  server *http.Server
  interceptor *middle
  shouldIntercept func(address string) bool
+ upstreamProxy func(*http.Request, net.Conn) (*url.URL, error)
 }
 
 func NewProxy(opts *Options) (*Proxy, error) {
@@ -45,7 +47,7 @@ func NewProxy(opts *Options) (*Proxy, error) {
 
  proxy.client = &http.Client{
  Transport: &http.Transport{
- Proxy: clientProxy(opts.Upstream),
+ Proxy: proxy.realUpstreamProxy(),
  ForceAttemptHTTP2: false, // disable http2
  DisableCompression: true, // To get the original response from the server, set Transport.DisableCompression to true.
  TLSClientConfig: &tls.Config{
@@ -229,7 +231,9 @@ func (proxy *Proxy) ServeHTTP(res http.ResponseWriter, req *http.Request) {
  for _, addon := range proxy.Addons {
  reqBody = addon.StreamRequestModifier(f, reqBody)
  }
- proxyReq, err := http.NewRequest(f.Request.Method, f.Request.URL.String(), reqBody)
+
+ proxyReqCtx := context.WithValue(context.Background(), rawClientConnContextKey, f.ConnContext.ClientConn.Conn)
+ proxyReq, err := http.NewRequestWithContext(proxyReqCtx, f.Request.Method, f.Request.URL.String(), reqBody)
  if err != nil {
  log.Error(err)
  res.WriteHeader(502)
@@ -327,10 +331,10 @@ func (proxy *Proxy) handleConnect(res http.ResponseWriter, req *http.Request) {
  var err error
  if shouldIntercept {
  log.Debugf("begin intercept %v", req.Host)
- conn, err = proxy.interceptor.dial(req)
+ conn, err = proxy.interceptor.dial(req, f.ConnContext.ClientConn.Conn)
  } else {
  log.Debugf("begin transpond %v", req.Host)
- conn, err = getConnFrom(req.Host, proxy.Opts.Upstream)
+ conn, err = getConnFrom(req.Host, proxy, f.ConnContext.ClientConn.Conn)
  }
  if err != nil {
  log.Error(err)
@@ -382,3 +386,17 @@ func (proxy *Proxy) GetCertificate() x509.Certificate {
 func (proxy *Proxy) SetShouldInterceptRule(rule func(address string) bool) {
  proxy.shouldIntercept = rule
 }
+
+func (proxy *Proxy) SetUpstreamProxy(fn func(*http.Request, net.Conn) (*url.URL, error)) {
+ proxy.upstreamProxy = fn
+}
+
+func (proxy *Proxy) realUpstreamProxy() func(*http.Request) (*url.URL, error) {
+ return func(req *http.Request) (*url.URL, error) {
+ if proxy.upstreamProxy != nil {
+ rawClientConn := req.Context().Value(rawClientConnContextKey).(net.Conn)
+ return proxy.upstreamProxy(req, rawClientConn)
+ }
+ return clientProxy(proxy.Opts.Upstream)(req)
+ }
+}