Преглед изворни кода

Simplify and optimize genericProxyRewriter

- Reduce the amount of nested loops: it's preferable to search the whole page
  once and filter on it (even with filters that should always be false),
  than searching it again for every element we're looking for.
- Factorize the proxying conditions into a `shouldProxy` function to reduce the
  copy-pasta.
jvoisin пре 2 година
родитељ
комит
9fe99ce7fa
1 измењених фајлова са 12 додато и 33 уклоњено
  1. 12 33
      internal/proxy/media_proxy.go

+ 12 - 33
internal/proxy/media_proxy.go

@@ -43,9 +43,9 @@ func genericProxyRewriter(router *mux.Router, proxifyFunction urlProxyRewriter,
 	for _, mediaType := range config.Opts.ProxyMediaTypes() {
 		switch mediaType {
 		case "image":
-			doc.Find("img").Each(func(i int, img *goquery.Selection) {
+			doc.Find("img, picture source").Each(func(i int, img *goquery.Selection) {
 				if srcAttrValue, ok := img.Attr("src"); ok {
-					if !isDataURL(srcAttrValue) && (proxyOption == "all" || !urllib.IsHTTPS(srcAttrValue)) {
+					if shouldProxy(srcAttrValue, proxyOption) {
 						img.SetAttr("src", proxifyFunction(router, srcAttrValue))
 					}
 				}
@@ -55,59 +55,37 @@ func genericProxyRewriter(router *mux.Router, proxifyFunction urlProxyRewriter,
 				}
 			})
 
-			doc.Find("picture source").Each(func(i int, sourceElement *goquery.Selection) {
-				if srcsetAttrValue, ok := sourceElement.Attr("srcset"); ok {
-					proxifySourceSet(sourceElement, router, proxifyFunction, proxyOption, srcsetAttrValue)
-				}
-			})
-
 			doc.Find("video").Each(func(i int, video *goquery.Selection) {
 				if posterAttrValue, ok := video.Attr("poster"); ok {
-					if !isDataURL(posterAttrValue) && (proxyOption == "all" || !urllib.IsHTTPS(posterAttrValue)) {
+					if shouldProxy(posterAttrValue, proxyOption) {
 						video.SetAttr("poster", proxifyFunction(router, posterAttrValue))
 					}
 				}
 			})
 
 		case "audio":
-			doc.Find("audio").Each(func(i int, audio *goquery.Selection) {
+			doc.Find("audio, audio source").Each(func(i int, audio *goquery.Selection) {
 				if srcAttrValue, ok := audio.Attr("src"); ok {
-					if !isDataURL(srcAttrValue) && (proxyOption == "all" || !urllib.IsHTTPS(srcAttrValue)) {
+					if shouldProxy(srcAttrValue, proxyOption) {
 						audio.SetAttr("src", proxifyFunction(router, srcAttrValue))
 					}
 				}
 			})
 
-			doc.Find("audio source").Each(func(i int, sourceElement *goquery.Selection) {
-				if srcAttrValue, ok := sourceElement.Attr("src"); ok {
-					if !isDataURL(srcAttrValue) && (proxyOption == "all" || !urllib.IsHTTPS(srcAttrValue)) {
-						sourceElement.SetAttr("src", proxifyFunction(router, srcAttrValue))
-					}
-				}
-			})
-
 		case "video":
-			doc.Find("video").Each(func(i int, video *goquery.Selection) {
+			doc.Find("video, video source").Each(func(i int, video *goquery.Selection) {
 				if srcAttrValue, ok := video.Attr("src"); ok {
-					if !isDataURL(srcAttrValue) && (proxyOption == "all" || !urllib.IsHTTPS(srcAttrValue)) {
+					if shouldProxy(srcAttrValue, proxyOption) {
 						video.SetAttr("src", proxifyFunction(router, srcAttrValue))
 					}
 				}
 
 				if posterAttrValue, ok := video.Attr("poster"); ok {
-					if !isDataURL(posterAttrValue) && (proxyOption == "all" || !urllib.IsHTTPS(posterAttrValue)) {
+					if shouldProxy(posterAttrValue, proxyOption) {
 						video.SetAttr("poster", proxifyFunction(router, posterAttrValue))
 					}
 				}
 			})
-
-			doc.Find("video source").Each(func(i int, sourceElement *goquery.Selection) {
-				if srcAttrValue, ok := sourceElement.Attr("src"); ok {
-					if !isDataURL(srcAttrValue) && (proxyOption == "all" || !urllib.IsHTTPS(srcAttrValue)) {
-						sourceElement.SetAttr("src", proxifyFunction(router, srcAttrValue))
-					}
-				}
-			})
 		}
 	}
 
@@ -123,7 +101,7 @@ func proxifySourceSet(element *goquery.Selection, router *mux.Router, proxifyFun
 	imageCandidates := sanitizer.ParseSrcSetAttribute(srcsetAttrValue)
 
 	for _, imageCandidate := range imageCandidates {
-		if !isDataURL(imageCandidate.ImageURL) && (proxyOption == "all" || !urllib.IsHTTPS(imageCandidate.ImageURL)) {
+		if shouldProxy(imageCandidate.ImageURL, proxyOption) {
 			imageCandidate.ImageURL = proxifyFunction(router, imageCandidate.ImageURL)
 		}
 	}
@@ -131,6 +109,7 @@ func proxifySourceSet(element *goquery.Selection, router *mux.Router, proxifyFun
 	element.SetAttr("srcset", imageCandidates.String())
 }
 
-func isDataURL(s string) bool {
-	return strings.HasPrefix(s, "data:")
+func shouldProxy(attrValue, proxyOption string) bool {
+	return !strings.HasPrefix(attrValue, "data:") &&
+		(proxyOption == "all" || !urllib.IsHTTPS(attrValue))
 }