Rewrite "Location" header in apiserver proxy.

pkg/apiserver/proxy.go

@@ -185,6 +185,10 @@ func (t *proxyTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 		return resp, nil
 	}
 
+	if redirect := resp.Header.Get("Location"); redirect != "" {
+		resp.Header.Set("Location", t.rewriteURL(redirect, req.URL))
+	}
+
 	cType := resp.Header.Get("Content-Type")
 	cType = strings.TrimSpace(strings.SplitN(cType, ";", 2)[0])
 	if cType != "text/html" {
@@ -195,6 +199,38 @@ func (t *proxyTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	return t.fixLinks(req, resp)
 }
 
+// rewriteURL rewrites a single URL to go through the proxy, if the URL refers
+// to the same host as sourceURL, which is the page on which the target URL
+// occurred. If any error occurs (e.g. parsing), it returns targetURL.
+func (t *proxyTransport) rewriteURL(targetURL string, sourceURL *url.URL) string {
+	url, err := url.Parse(targetURL)
+	if err != nil {
+		return targetURL
+	}
+	if url.Host != "" && url.Host != sourceURL.Host {
+		return targetURL
+	}
+
+	url.Scheme = t.proxyScheme
+	url.Host = t.proxyHost
+	origPath := url.Path
+
+	if strings.HasPrefix(url.Path, "/") {
+		// The path is rooted at the host. Just add proxy prepend.
+		url.Path = path.Join(t.proxyPathPrepend, url.Path)
+	} else {
+		// The path is relative to sourceURL.
+		url.Path = path.Join(t.proxyPathPrepend, path.Dir(sourceURL.Path), url.Path)
+	}
+
+	if strings.HasSuffix(origPath, "/") {
+		// Add back the trailing slash, which was stripped by path.Join().
+		url.Path += "/"
+	}
+
+	return url.String()
+}
+
 // updateURLs checks and updates any of n's attributes that are listed in tagsToAttrs.
 // Any URLs found are, if they're relative, updated with the necessary changes to make
 // a visit to that URL also go through the proxy.
@@ -212,32 +248,7 @@ func (t *proxyTransport) updateURLs(n *html.Node, sourceURL *url.URL) {
 		if !attrs.Has(attr.Key) {
 			continue
 		}
-		url, err := url.Parse(attr.Val)
-		if err != nil {
-			continue
-		}
-
-		// Is this URL referring to the same host as sourceURL?
-		if url.Host == "" || url.Host == sourceURL.Host {
-			url.Scheme = t.proxyScheme
-			url.Host = t.proxyHost
-			origPath := url.Path
-
-			if strings.HasPrefix(url.Path, "/") {
-				// The path is rooted at the host. Just add proxy prepend.
-				url.Path = path.Join(t.proxyPathPrepend, url.Path)
-			} else {
-				// The path is relative to sourceURL.
-				url.Path = path.Join(t.proxyPathPrepend, path.Dir(sourceURL.Path), url.Path)
-			}
-
-			if strings.HasSuffix(origPath, "/") {
-				// Add back the trailing slash, which was stripped by path.Join().
-				url.Path += "/"
-			}
-
-			n.Attr[i].Val = url.String()
-		}
+		n.Attr[i].Val = t.rewriteURL(attr.Val, sourceURL)
 	}
 }
 

pkg/apiserver/proxy_test.go

@@ -63,15 +63,18 @@ func TestProxyTransport(t *testing.T) {
 		proxyHost:        "foo.com",
 		proxyPathPrepend: "/proxy/minion/minion1:8080",
 	}
-
-	table := map[string]struct {
+	type Item struct {
 		input        string
 		sourceURL    string
 		transport    *proxyTransport
 		output       string
 		contentType  string
 		forwardedURI string
-	}{
+		redirect     string
+		redirectWant string
+	}
+
+	table := map[string]Item{
 		"normal": {
 			input:        `<pre><a href="kubelet.log">kubelet.log</a><a href="/google.log">google.log</a></pre>`,
 			sourceURL:    "http://myminion.com/logs/log.log",
@@ -136,9 +139,30 @@ func TestProxyTransport(t *testing.T) {
 			contentType:  "text/html",
 			forwardedURI: "/proxy/minion/minion1:10250/any/path/",
 		},
+		"redirect rel": {
+			sourceURL:    "http://myminion.com/redirect",
+			transport:    testTransport,
+			redirect:     "/redirected/target/",
+			redirectWant: "http://foo.com/proxy/minion/minion1:10250/redirected/target/",
+			forwardedURI: "/proxy/minion/minion1:10250/redirect",
+		},
+		"redirect abs same host": {
+			sourceURL:    "http://myminion.com/redirect",
+			transport:    testTransport,
+			redirect:     "http://myminion.com/redirected/target/",
+			redirectWant: "http://foo.com/proxy/minion/minion1:10250/redirected/target/",
+			forwardedURI: "/proxy/minion/minion1:10250/redirect",
+		},
+		"redirect abs other host": {
+			sourceURL:    "http://myminion.com/redirect",
+			transport:    testTransport,
+			redirect:     "http://example.com/redirected/target/",
+			redirectWant: "http://example.com/redirected/target/",
+			forwardedURI: "/proxy/minion/minion1:10250/redirect",
+		},
 	}
 
-	for name, item := range table {
+	testItem := func(name string, item *Item) {
 		// Canonicalize the html so we can diff.
 		item.input = fmtHTML(item.input)
 		item.output = fmtHTML(item.output)
@@ -156,34 +180,51 @@ func TestProxyTransport(t *testing.T) {
 			}
 
 			// Send response.
+			if item.redirect != "" {
+				http.Redirect(w, r, item.redirect, http.StatusMovedPermanently)
+				return
+			}
 			w.Header().Set("Content-Type", item.contentType)
 			fmt.Fprint(w, item.input)
 		}))
+		defer server.Close()
+
 		// Replace source URL with our test server address.
 		sourceURL := parseURLOrDie(item.sourceURL)
 		serverURL := parseURLOrDie(server.URL)
 		item.input = strings.Replace(item.input, sourceURL.Host, serverURL.Host, -1)
+		item.redirect = strings.Replace(item.redirect, sourceURL.Host, serverURL.Host, -1)
 		sourceURL.Host = serverURL.Host
 
 		req, err := http.NewRequest("GET", sourceURL.String(), nil)
 		if err != nil {
 			t.Errorf("%v: Unexpected error: %v", name, err)
-			continue
+			return
 		}
 		resp, err := item.transport.RoundTrip(req)
 		if err != nil {
 			t.Errorf("%v: Unexpected error: %v", name, err)
-			continue
+			return
+		}
+		if item.redirect != "" {
+			// Check that redirect URLs get rewritten properly.
+			if got, want := resp.Header.Get("Location"), item.redirectWant; got != want {
+				t.Errorf("%v: Location header = %q, want %q", name, got, want)
+			}
+			return
 		}
 		body, err := ioutil.ReadAll(resp.Body)
 		if err != nil {
 			t.Errorf("%v: Unexpected error: %v", name, err)
-			continue
+			return
 		}
 		if e, a := item.output, string(body); e != a {
 			t.Errorf("%v: expected %v, but got %v", name, e, a)
 		}
-		server.Close()
+	}
+
+	for name, item := range table {
+		testItem(name, &item)
 	}
 }