all: handle revdial redirects on connect

In dev mode, the coordinator uses hostPathHandler, which redirects
/reverse and /revdial to /farmer.golang.org/reverse, etc.

Establishing a revdial connection chokes on this redirect, as there it
expects to complete the protocol switch in a single request.

Add rudimentary redirect support via a helper so this works in dev mode.

Note that linkRewriter must implement http.Hijacker as
revdial.ConnHandler type-asserts the http.ResponseWriter to a Hijacker.

For golang/go#48803

Change-Id: I191fa6ff17bbd334203430f3c1f2c5e03db407ff
Reviewed-on: https://go-review.googlesource.com/c/build/+/354630
Trust: Michael Pratt <[email protected]>
Run-TryBot: Michael Pratt <[email protected]>
TryBot-Result: Go Bot <[email protected]>
Reviewed-by: Alexander Rakoczy <[email protected]>

cmd/buildlet/reverse.go

@@ -10,6 +10,7 @@ import (
 	"crypto/hmac"
 	"crypto/md5"
 	"crypto/tls"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -118,28 +119,37 @@ func dialCoordinator() (net.Listener, error) {
 	bufw := bufio.NewWriter(conn)
 
 	log.Printf("Registering reverse mode with coordinator...")
-	req, err := http.NewRequest("GET", "/reverse", nil)
-	if err != nil {
-		log.Fatal(err)
-	}
-	req.Header.Set("X-Go-Host-Type", *reverseType)
-	req.Header.Set("X-Go-Builder-Key", key)
-	req.Header.Set("X-Go-Builder-Hostname", *hostname)
-	req.Header.Set("X-Go-Builder-Version", strconv.Itoa(buildletVersion))
-	req.Header.Set("X-Revdial-Version", "2")
-	if err := req.Write(bufw); err != nil {
-		return nil, fmt.Errorf("coordinator /reverse request failed: %v", err)
-	}
-	if err := bufw.Flush(); err != nil {
-		return nil, fmt.Errorf("coordinator /reverse request flush failed: %v", err)
-	}
-	resp, err := http.ReadResponse(bufr, req)
-	if err != nil {
-		return nil, fmt.Errorf("coordinator /reverse response failed: %v", err)
+
+	success := false
+	location := "/reverse"
+	const maxRedirects = 2
+	for i := 0; i < maxRedirects; i++ {
+		req, err := http.NewRequest("GET", location, nil)
+		if err != nil {
+			log.Fatal(err)
+		}
+		req.Header.Set("X-Go-Host-Type", *reverseType)
+		req.Header.Set("X-Go-Builder-Key", key)
+		req.Header.Set("X-Go-Builder-Hostname", *hostname)
+		req.Header.Set("X-Go-Builder-Version", strconv.Itoa(buildletVersion))
+		req.Header.Set("X-Revdial-Version", "2")
+		if err := req.Write(bufw); err != nil {
+			return nil, fmt.Errorf("coordinator /reverse request failed: %v", err)
+		}
+		if err := bufw.Flush(); err != nil {
+			return nil, fmt.Errorf("coordinator /reverse request flush failed: %v", err)
+		}
+		location, err = revdial.ReadProtoSwitchOrRedirect(bufr, req)
+		if err != nil {
+			return nil, fmt.Errorf("coordinator registration failed: %v", err)
+		}
+		if location == "" {
+			success = true
+			break
+		}
 	}
-	if resp.StatusCode != 101 {
-		msg, _ := ioutil.ReadAll(resp.Body)
-		return nil, fmt.Errorf("coordinator registration failed; want HTTP status 101; got %v:\n\t%s", resp.Status, msg)
+	if !success {
+		return nil, errors.New("coordinator /reverse: too many redirects")
 	}
 
 	log.Printf("Connected to coordinator; reverse dialing active")

cmd/buildlet/reverse_test.go

@@ -128,3 +128,33 @@ func TestReverseDial(t *testing.T) {
 	const hostType = "test-reverse-dial"
 	testReverseDial(t, srv.Addr, hostType)
 }
+
+// TestReverseDialRedirect verfies that a revdial connection works with a 307
+// redirect to the endpoints. The coordinator will do this in dev mode.
+func TestReverseDialRedirect(t *testing.T) {
+	pool.SetBuilderMasterKey([]byte(devMasterKey))
+
+	srv, ln, err := coordinatorServer()
+	if err != nil {
+		t.Fatalf("serveCoordinator got err %v want nil", err)
+	}
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/redirected/reverse", pool.HandleReverse)
+	mux.Handle("/redirected/revdial", revdial.ConnHandler())
+
+	redirect := func(w http.ResponseWriter, r *http.Request) {
+		u := *r.URL
+		u.Path = "/redirected/" + u.Path
+		http.Redirect(w, r, u.String(), http.StatusTemporaryRedirect)
+	}
+	mux.HandleFunc("/reverse", redirect)
+	mux.HandleFunc("/revdial", redirect)
+	srv.Handler = mux
+
+	go srv.Serve(ln)
+	defer srv.Close()
+
+	const hostType = "test-reverse-dial-redirect"
+	testReverseDial(t, srv.Addr, hostType)
+}

cmd/coordinator/coordinator.go

@@ -246,11 +246,21 @@ var validHosts = map[string]bool{
 	"build.golang.org":  true,
 }
 
+type hijackResponseWriter interface {
+	http.Hijacker
+	http.ResponseWriter
+}
+
 // hostPathHandler infers the host from the first element of the URL path,
 // and rewrites URLs in the output HTML accordingly. It disables response
 // compression to simplify the process of link rewriting.
 func hostPathHandler(h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		hw, ok := w.(hijackResponseWriter)
+		if !ok {
+			log.Fatalf("hostPathHandler ResponseWriter does not implement Hijack: %T", w)
+		}
+
 		elem, rest := strings.TrimPrefix(r.URL.Path, "/"), ""
 		if i := strings.Index(elem, "/"); i >= 0 {
 			elem, rest = elem[:i], elem[i+1:]
@@ -268,7 +278,7 @@ func hostPathHandler(h http.Handler) http.Handler {
 		r.URL.Host = elem
 		r.URL.Path = "/" + rest
 		r.Header.Set("Accept-Encoding", "identity") // Disable compression for link rewriting.
-		lw := &linkRewriter{ResponseWriter: w, host: r.Host}
+		lw := &linkRewriter{hijackResponseWriter: hw, host: r.Host}
 		h.ServeHTTP(lw, r)
 		lw.Flush()
 	})
@@ -279,7 +289,7 @@ func hostPathHandler(h http.Handler) http.Handler {
 // https://h/foo, where h is in validHosts, to be /h/foo. This corrects the
 // links to have the right form for the test server.
 type linkRewriter struct {
-	http.ResponseWriter
+	hijackResponseWriter
 	host string
 	buf  []byte
 	ct   string // content-type
@@ -295,7 +305,7 @@ func (r *linkRewriter) Write(data []byte) (int, error) {
 		r.ct = ct
 	}
 	if !strings.HasPrefix(r.ct, "text/html") {
-		return r.ResponseWriter.Write(data)
+		return r.hijackResponseWriter.Write(data)
 	}
 	r.buf = append(r.buf, data...)
 	return len(data), nil
@@ -308,7 +318,7 @@ func (r *linkRewriter) Flush() {
 	for host := range validHosts {
 		repl = append(repl, `href="https://`+host, `href="/`+host)
 	}
-	strings.NewReplacer(repl...).WriteString(r.ResponseWriter, string(r.buf))
+	strings.NewReplacer(repl...).WriteString(r.hijackResponseWriter, string(r.buf))
 	r.buf = nil
 }
 

revdial/v2/revdial.go

@@ -24,9 +24,11 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
+	"net/url"
 	"strings"
 	"sync"
 	"time"
@@ -314,21 +316,31 @@ func (ln *Listener) grabConn(path string) {
 		log.Printf("revdial.Listener: failed to pick up connection to %s: %v", path, err)
 		ln.sendMessage(controlMsg{Command: "pickup-failed", ConnPath: path, Err: err.Error()})
 	}
-	req, _ := http.NewRequest("GET", path, nil)
-	if err := req.Write(c); err != nil {
-		failPickup(err)
-		return
-	}
 	bufr := bufio.NewReader(c)
-	resp, err := http.ReadResponse(bufr, req)
-	if err != nil {
-		failPickup(err)
-		return
+
+	success := false
+	const maxRedirects = 2
+	for i := 0; i < maxRedirects; i++ {
+		req, _ := http.NewRequest("GET", path, nil)
+		if err := req.Write(c); err != nil {
+			failPickup(err)
+			return
+		}
+		path, err = ReadProtoSwitchOrRedirect(bufr, req)
+		if err != nil {
+			failPickup(fmt.Errorf("switch failed: %v", err))
+			return
+		}
+		if path == "" {
+			success = true
+			break
+		}
 	}
-	if resp.StatusCode != 101 {
-		failPickup(fmt.Errorf("non-101 response %v", resp.Status))
+	if !success {
+		failPickup(errors.New("too many redirects"))
 		return
 	}
+
 	select {
 	case ln.connc <- c:
 	case <-ln.donec:
@@ -420,3 +432,50 @@ func connHandler(w http.ResponseWriter, r *http.Request) {
 	(&http.Response{StatusCode: http.StatusSwitchingProtocols, Proto: "HTTP/1.1"}).Write(conn)
 	d.matchConn(conn)
 }
+
+// checkRelativeURL verifies that URL s does not change scheme or host.
+func checkRelativeURL(s string) error {
+	u, err := url.Parse(s)
+	if err != nil {
+		return err
+	}
+
+	// A relative URL should have no schema or host.
+	if u.Scheme != "" {
+		return fmt.Errorf("URL %q is not relative: contains scheme", s)
+	}
+	if u.Host != "" {
+		return fmt.Errorf("URL %q is not relative: contains host", s)
+	}
+	return nil
+}
+
+// ReadProtoSwitchOrRedirect is a helper for completing revdial protocol switch
+// requests. If the response indicates successful switch, nothing is returned.
+// If the response indicates a redirect, the new location is returned.
+func ReadProtoSwitchOrRedirect(r *bufio.Reader, req *http.Request) (location string, err error) {
+	resp, err := http.ReadResponse(r, req)
+	if err != nil {
+		return "", fmt.Errorf("error reading response: %v", err)
+	}
+	switch resp.StatusCode {
+	case http.StatusSwitchingProtocols:
+		// Success! Don't read body, as caller may want it.
+		return "", nil
+	case http.StatusTemporaryRedirect:
+		// Redirect. Discard body.
+		msg, _ := io.ReadAll(resp.Body)
+		location := resp.Header.Get("Location")
+		if location == "" {
+			return "", fmt.Errorf("redirect missing Location header; got %+v:\n\t%s", resp, msg)
+		}
+		if err := checkRelativeURL(location); err != nil {
+			return "", fmt.Errorf("redirect Location must be relative: %w", err)
+		}
+		// Retry at new location.
+		return location, nil
+	default:
+		msg, _ := io.ReadAll(resp.Body)
+		return "", fmt.Errorf("want HTTP status 101 or 307; got %v:\n\t%s", resp.Status, msg)
+	}
+}