getlantern/lantern#2045 Added support for hooking into connections to…

… server (to facilitate stat reporting)

dialer.go

@@ -73,6 +73,11 @@ type Config struct {
 	// values indicating higher QOS.
 	QOS int
 
+	// OnDial: optional callback that gets invoked whenever we dial the server.
+	// The Conn and error returned from this callback will be used in lieu of
+	// the originals.
+	OnDial func(conn net.Conn, err error) (net.Conn, error)
+
 	// OnDialStats is an optional callback that will get called on every dial to
 	// the server to report stats on what was dialed and how long each step
 	// took.
@@ -105,7 +110,11 @@ func NewDialer(cfg *Config) *Dialer {
 		Dial:         d.dialServer,
 	}
 	d.enproxyConfig = d.enproxyConfigWith(func(addr string) (net.Conn, error) {
-		return d.connPool.Get()
+		conn, err := d.connPool.Get()
+		if d.OnDial != nil {
+			conn, err = d.OnDial(conn, err)
+		}
+		return conn, err
 	})
 	return d
 }

fronted_test.go

@@ -1,6 +1,7 @@
 package fronted
 
 import (
+	"fmt"
 	"io/ioutil"
 	"net"
 	"net/http"
@@ -34,6 +35,18 @@ func TestBadEnproxyConn(t *testing.T) {
 	assert.Error(t, err, "Dialing using a non-existent host should have failed")
 }
 
+func TestReplaceBadOnDial(t *testing.T) {
+	d := NewDialer(&Config{
+		Host: "fallbacks.getiantem.org",
+		Port: 443,
+		OnDial: func(conn net.Conn, err error) (net.Conn, error) {
+			return nil, fmt.Errorf("Gotcha!")
+		},
+	})
+	_, err := d.Dial("tcp", "www.google.com")
+	assert.Error(t, err, "Dialing using a bad OnDial should fail")
+}
+
 func TestHttpClientWithBadEnproxyConn(t *testing.T) {
 	d := NewDialer(&Config{
 		Host: "localhost",