Check host after redirect

Fixes #4

Author: jorinvo

httpsyet/crawler.go

@@ -122,6 +122,10 @@ func toURLs(links []string, parse func(string) (*url.URL, error)) (urls []*url.U
 			invalids = append(invalids, fmt.Sprintf("%s (%v)", s, e))
 			continue
 		}
+		// Default to https
+		if u.Scheme == "" {
+			u.Scheme = "https"
+		}
 		// Ignore invalid protocols
 		if u.Scheme == "http" || u.Scheme == "https" {
 			urls = append(urls, u)
@@ -240,6 +244,11 @@ func crawlSite(s site, get func(string) (*http.Response, error)) ([]string, bool
 		return nil, false, fmt.Errorf("%d %v", r.StatusCode, u)
 	}
 
+	// Stop when redirecting to external page
+	if r.Request.URL.Host != u.Host {
+		isExternal = true
+	}
+
 	// Stop when site is external.
 	// Also stop if depth one is reached, ignored when depth is set to 0.
 	if isExternal || s.Depth == 1 {

httpsyet/crawler_test.go

@@ -55,6 +55,10 @@ const (
 	This is an relative internal link to a page without children.
 </a>
 
+<a href="/redirect">
+	This is an internal link but it redirects to an external page.
+</a>
+
 <a href="{{ .Self }}/sub">
 	This is an internal link.
 </a>
@@ -93,6 +97,12 @@ const (
 
 	basic = `
 <h1>Welcome to the basic page</h1>
+`
+
+	redirectTarget = `
+<a href="/no-follow">
+	This link after redirecting should not be followed.
+</a>
 `
 )
 
@@ -113,20 +123,31 @@ func TestRun(t *testing.T) {
 		}
 	}
 
+	redirect := func(name, u string) http.HandlerFunc {
+		visited[name] = 0
+		return func(w http.ResponseWriter, r *http.Request) {
+			visited[name]++
+			t.Logf("visited page: %s", name)
+			http.Redirect(w, r, u, http.StatusMovedPermanently)
+		}
+	}
+
+	httpMux := http.NewServeMux()
+	httpMux.HandleFunc("/page-a", serve("http/page-a", basic))
+	httpMux.HandleFunc("/page-b", serve("http/page-b", basic))
+	httpMux.HandleFunc("/redirect-target", serve("http/redirect-target", redirectTarget))
+	httpServer := httptest.NewServer(httpMux)
+	defer httpServer.Close()
+
 	pageMux := http.NewServeMux()
 	pageMux.HandleFunc("/base", serve("page/base", basePage))
 	pageMux.HandleFunc("/sub", serve("page/sub", subPage))
 	pageMux.HandleFunc("/empty-sub", serve("page/empty-sub", basic))
 	pageMux.HandleFunc("/sub/sub", serve("page/sub/sub", basic))
+	pageMux.HandleFunc("/redirect", redirect("page/redirect", httpServer.URL+"/redirect-target"))
 	pageServer := httptest.NewServer(pageMux)
 	defer pageServer.Close()
 
-	httpMux := http.NewServeMux()
-	httpMux.HandleFunc("/page-a", serve("http/page-a", basic))
-	httpMux.HandleFunc("/page-b", serve("http/page-b", basic))
-	httpServer := httptest.NewServer(httpMux)
-	defer httpServer.Close()
-
 	tlsMux := http.NewServeMux()
 	tlsMux.HandleFunc("/base", serve("tls/base", basic))
 	tlsMux.HandleFunc("/base2", serve("tls/base2", basic))
@@ -244,7 +265,13 @@ func TestRunSingle(t *testing.T) {
 	noErr(t, err)
 
 	expect := fmt.Sprintf(
-		"404 %s/404 on page %s/base\n404 %s/404 on page %s/sub\n404 %s/404 on page %s/sub\n",
+		`404 %s/404 on page %s/base
+404 %s/redirect on page %s/base
+404 %s/404 on page %s/sub
+404 %s/404 on page %s/sub
+`,
+		pageServer.URL,
+		pageServer.URL,
 		pageServer.URL,
 		pageServer.URL,
 		httpServer.URL,
@@ -343,6 +370,8 @@ verbose: GET %s/page-a
 verbose: GET %s/404
 404 %s/404 on page %s/base
 verbose: GET %s/empty-sub
+verbose: GET %s/redirect
+404 %s/redirect on page %s/base
 verbose: GET %s/sub
 `,
 		pageServer.URL,
@@ -358,6 +387,9 @@ verbose: GET %s/sub
 		pageServer.URL,
 		pageServer.URL,
 		pageServer.URL,
+		pageServer.URL,
+		pageServer.URL,
+		pageServer.URL,
 	)
 	eqLines(t, expect, errs.String(), "unexpected errors")