Merge pull request #4050 from getlantern/balancecheck

Unified request modification for chained proxies

src/github.com/getlantern/balancer/balancer.go

@@ -6,6 +6,7 @@ import (
 	"container/heap"
 	"fmt"
 	"net"
+	"net/http"
 	"sync"
 
 	"github.com/getlantern/golog"
@@ -21,7 +22,7 @@ var (
 
 // Balancer balances connections established by one or more Dialers.
 type Balancer struct {
-	mu      sync.Mutex
+	mu      sync.RWMutex
 	dialers dialerHeap
 	trusted dialerHeap
 }
@@ -47,10 +48,11 @@ func New(st Strategy, dialers ...*Dialer) *Balancer {
 	return bal
 }
 
-// AllAuthTokens() returns a list of all auth tokens for all dialers on this
-// balancer.
-func (b *Balancer) AllAuthTokens() []string {
-	return b.dialers.AuthTokens()
+// OnRequest calls Dialer.OnRequest for every dialer in this balancer.
+func (b *Balancer) OnRequest(req *http.Request) {
+	b.mu.RLock()
+	b.dialers.onRequest(req)
+	b.mu.RUnlock()
 }
 
 // Dial dials (network, addr) using one of the currently active configured

src/github.com/getlantern/balancer/dialer.go

@@ -34,7 +34,8 @@ type Dialer struct {
 	// Determines whether a dialer can be trusted with unencrypted traffic.
 	Trusted bool
 
-	AuthToken string
+	// Modifies any HTTP requests made using connections from this dialer.
+	OnRequest func(req *http.Request)
 }
 
 var (
@@ -158,7 +159,9 @@ func (d *dialer) defaultCheck() bool {
 			log.Errorf("Could not create HTTP request?")
 			return false, nil
 		}
-		req.Header.Set("X-LANTERN-AUTH-TOKEN", d.AuthToken)
+		if d.OnRequest != nil {
+			d.OnRequest(req)
+		}
 		resp, err := client.Do(req)
 		if err != nil {
 			log.Debugf("Error testing dialer %s to humans.txt: %s", d.Label, err)

src/github.com/getlantern/balancer/strategy.go

@@ -1,6 +1,9 @@
 package balancer
 
-import "math/rand"
+import (
+	"math/rand"
+	"net/http"
+)
 
 // Random strategy gives even chance to each dialer, act as a baseline to other
 // strategies.
@@ -97,9 +100,11 @@ func (s *dialerHeap) Pop() interface{} {
 	return x
 }
 
-func (s *dialerHeap) AuthTokens() (tokens []string) {
+func (s *dialerHeap) onRequest(req *http.Request) {
 	for _, d := range s.dialers {
-		tokens = append(tokens, d.AuthToken)
+		if d.OnRequest != nil {
+			d.OnRequest(req)
+		}
 	}
 	return
 }

src/github.com/getlantern/flashlight/client/chained.go

@@ -79,7 +79,7 @@ func (s *ChainedServerInfo) Dialer(deviceID string) (*balancer.Dialer, error) {
 
 	ccfg.OnRequest = func(req *http.Request) {
 		if authToken != "" {
-			req.Header.Set("X-LANTERN-AUTH-TOKEN", authToken)
+			req.Header.Add("X-LANTERN-AUTH-TOKEN", authToken)
 		}
 		req.Header.Set("X-LANTERN-DEVICE-ID", deviceID)
 	}
@@ -103,6 +103,6 @@ func (s *ChainedServerInfo) Dialer(deviceID string) (*balancer.Dialer, error) {
 
 			return conn, nil
 		},
-		AuthToken: authToken,
+		OnRequest: ccfg.OnRequest,
 	}, nil
 }

src/github.com/getlantern/flashlight/client/reverseproxy.go

@@ -26,7 +26,7 @@ func (client *Client) newReverseProxy(bal *balancer.Balancer) *httputil.ReverseP
 	// ReverseProxies for different QOS's or something like that.
 	transport.Dial = client.proxiedDialer(bal.Dial)
 
-	allAuthTokens := bal.AllAuthTokens()
+	onRequest := bal.OnRequest
 	return &httputil.ReverseProxy{
 		// We need to set the authentication tokens for all servers that we might
 		// connect to because we don't know which one the dialer will actually
@@ -38,10 +38,7 @@ func (client *Client) newReverseProxy(bal *balancer.Balancer) *httputil.ReverseP
 			// Add back the Host header which was stripped by the ReverseProxy. This
 			// is needed for sites that do virtual hosting.
 			req.Header.Set("Host", req.Host)
-			req.Header.Set("X-LANTERN-DEVICE-ID", client.cfg().DeviceID)
-			for _, authToken := range allAuthTokens {
-				req.Header.Add("X-LANTERN-AUTH-TOKEN", authToken)
-			}
+			onRequest(req)
 		},
 		Transport: &errorRewritingRoundTripper{
 			&noForwardedForRoundTripper{withDumpHeaders(false, transport)},