h2 cl conn pool prep fn

Author: Snawoot

dialer/h2.go

@@ -61,6 +61,12 @@ func H2ProxyDialerFromURL(u *url.URL, next xproxy.Dialer) (xproxy.Dialer, error)
 		AllowHTTP:       h2c,
 		TLSClientConfig: tlsConfig,
 	}
+	t.ConnPool = &clientConnPool{
+		t: t,
+		prepare: func(ctx context.Context, c *http2.ClientConn) (*http2.ClientConn, error) {
+			return c, nil
+		},
+	}
 	nextDialer := MaybeWrapWithContextDialer(next)
 	if h2c {
 		t.DialTLSContext = func(ctx context.Context, network, _ string, _ *tls.Config) (net.Conn, error) {

dialer/h2_client_conn_pool.go

@@ -29,6 +29,7 @@ type clientConnPool struct {
 	dialing      map[string]*dialCall           // currently in-flight dials
 	keys         map[*http2.ClientConn][]string
 	addConnCalls map[string]*addConnCall // in-flight addConnIfNeeded calls
+	prepare      func(context.Context, *http2.ClientConn) (*http2.ClientConn, error)
 }
 
 func (p *clientConnPool) GetClientConn(req *http.Request, addr string) (*http2.ClientConn, error) {
@@ -78,7 +79,14 @@ func (p *clientConnPool) dialClientConn(ctx context.Context, addr string) (*http
 	if err != nil {
 		return nil, err
 	}
-	return p.t.NewClientConn(tconn)
+	cc, err := p.t.NewClientConn(tconn)
+	if err != nil {
+		return nil, err
+	}
+	if p.prepare != nil {
+		return p.prepare(ctx, cc)
+	}
+	return cc, nil
 }
 
 func (p *clientConnPool) dialTLS(ctx context.Context, network, addr string, tlsCfg *tls.Config) (net.Conn, error) {