Merge branch 'main' into 39-request-header-fields-too-large

cmd/main.go

@@ -38,6 +38,11 @@ func main() {
 		Help:     "This will spawn multiple processes listening",
 	})
 
+	ruleset := parser.String("r", "ruleset", &argparse.Options{
+		Required: false,
+		Help:     "File, Directory or URL to a ruleset.yml. Overrides RULESET environment variable.",
+	})
+
 	err := parser.Parse(os.Args)
 	if err != nil {
 		fmt.Print(parser.Usage(err))
@@ -89,7 +94,6 @@ func main() {
 
 	app.Get("raw/*", handlers.Raw)
 	app.Get("api/*", handlers.Api)
-	app.Get("/*", handlers.ProxySite)
-
+	app.Get("/*", handlers.ProxySite(*ruleset))
 	log.Fatal(app.Listen(":" + *port))
 }

handlers/proxy.go

@@ -10,18 +10,26 @@ import (
 	"regexp"
 	"strings"
 
+	"ladder/pkg/ruleset"
+
 	"github.com/PuerkitoBio/goquery"
 	"github.com/gofiber/fiber/v2"
-	"gopkg.in/yaml.v3"
 )
 
 var (
 	UserAgent      = getenv("USER_AGENT", "Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)")
 	ForwardedFor   = getenv("X_FORWARDED_FOR", "66.249.66.1")
-	rulesSet       = loadRules()
-	allowedDomains = strings.Split(os.Getenv("ALLOWED_DOMAINS"), ",")
+	rulesSet       = ruleset.NewRulesetFromEnv()
+	allowedDomains = []string{}
 )
 
+func init() {
+	allowedDomains = strings.Split(os.Getenv("ALLOWED_DOMAINS"), ",")
+	if os.Getenv("ALLOWED_DOMAINS_RULESET") == "true" {
+		allowedDomains = append(allowedDomains, rulesSet.Domains()...)
+	}
+}
+
 // extracts a URL from the request ctx. If the URL in the request
 // is a relative path, it reconstructs the full URL using the referer header.
 func extractUrl(c *fiber.Ctx) (string, error) {
@@ -75,29 +83,39 @@ func extractUrl(c *fiber.Ctx) (string, error) {
 
 }
 
-func ProxySite(c *fiber.Ctx) error {
-	// Get the url from the URL
-	url, err := extractUrl(c)
-	if err != nil {
-		log.Println("ERROR In URL extraction:", err)
+func ProxySite(rulesetPath string) fiber.Handler {
+	if rulesetPath != "" {
+		rs, err := ruleset.NewRuleset(rulesetPath)
+		if err != nil {
+			panic(err)
+		}
+		rulesSet = rs
 	}
 
-	queries := c.Queries()
-	body, _, resp, err := fetchSite(url, queries)
-	if err != nil {
-		log.Println("ERROR:", err)
-		c.SendStatus(fiber.StatusInternalServerError)
-		return c.SendString(err.Error())
-	}
+	return func(c *fiber.Ctx) error {
+		// Get the url from the URL
+		url, err := extractUrl(c)
+		if err != nil {
+			log.Println("ERROR In URL extraction:", err)
+		}
+
+		queries := c.Queries()
+		body, _, resp, err := fetchSite(url, queries)
+		if err != nil {
+			log.Println("ERROR:", err)
+			c.SendStatus(fiber.StatusInternalServerError)
+			return c.SendString(err.Error())
+		}
 
 	c.Cookie(&fiber.Cookie{})
 	c.Set("Content-Type", resp.Header.Get("Content-Type"))
 	c.Set("Content-Security-Policy", resp.Header.Get("Content-Security-Policy"))
 
-	return c.SendString(body)
+		return c.SendString(body)
+	}
 }
 
-func modifyURL(uri string, rule Rule) (string, error) {
+func modifyURL(uri string, rule ruleset.Rule) (string, error) {
 	newUrl, err := url.Parse(uri)
 	if err != nil {
 		return "", err
@@ -205,7 +223,7 @@ func fetchSite(urlpath string, queries map[string]string) (string, *http.Request
 	}
 
 	if rule.Headers.CSP != "" {
-		log.Println(rule.Headers.CSP)
+		//log.Println(rule.Headers.CSP)
 		resp.Header.Set("Content-Security-Policy", rule.Headers.CSP)
 	}
 
@@ -214,7 +232,7 @@ func fetchSite(urlpath string, queries map[string]string) (string, *http.Request
 	return body, req, resp, nil
 }
 
-func rewriteHtml(bodyB []byte, u *url.URL, rule Rule) string {
+func rewriteHtml(bodyB []byte, u *url.URL, rule ruleset.Rule) string {
 	// Rewrite the HTML
 	body := string(bodyB)
 
@@ -248,63 +266,11 @@ func getenv(key, fallback string) string {
 	return value
 }
 
-func loadRules() RuleSet {
-	rulesUrl := os.Getenv("RULESET")
-	if rulesUrl == "" {
-		RulesList := RuleSet{}
-		return RulesList
-	}
-	log.Println("Loading rules")
-
-	var ruleSet RuleSet
-	if strings.HasPrefix(rulesUrl, "http") {
-
-		resp, err := http.Get(rulesUrl)
-		if err != nil {
-			log.Println("ERROR:", err)
-		}
-		defer resp.Body.Close()
-
-		if resp.StatusCode >= 400 {
-			log.Println("ERROR:", resp.StatusCode, rulesUrl)
-		}
-
-		body, err := io.ReadAll(resp.Body)
-		if err != nil {
-			log.Println("ERROR:", err)
-		}
-		yaml.Unmarshal(body, &ruleSet)
-
-		if err != nil {
-			log.Println("ERROR:", err)
-		}
-	} else {
-		yamlFile, err := os.ReadFile(rulesUrl)
-		if err != nil {
-			log.Println("ERROR:", err)
-		}
-		yaml.Unmarshal(yamlFile, &ruleSet)
-	}
-
-	domains := []string{}
-	for _, rule := range ruleSet {
-
-		domains = append(domains, rule.Domain)
-		domains = append(domains, rule.Domains...)
-		if os.Getenv("ALLOWED_DOMAINS_RULESET") == "true" {
-			allowedDomains = append(allowedDomains, domains...)
-		}
-	}
-
-	log.Println("Loaded ", len(ruleSet), " rules for", len(domains), "Domains")
-	return ruleSet
-}
-
-func fetchRule(domain string, path string) Rule {
+func fetchRule(domain string, path string) ruleset.Rule {
 	if len(rulesSet) == 0 {
-		return Rule{}
+		return ruleset.Rule{}
 	}
-	rule := Rule{}
+	rule := ruleset.Rule{}
 	for _, rule := range rulesSet {
 		domains := rule.Domains
 		if rule.Domain != "" {
@@ -323,7 +289,7 @@ func fetchRule(domain string, path string) Rule {
 	return rule
 }
 
-func applyRules(body string, rule Rule) string {
+func applyRules(body string, rule ruleset.Rule) string {
 	if len(rulesSet) == 0 {
 		return body
 	}

handlers/proxy.test.go

@@ -2,6 +2,7 @@
 package handlers
 
 import (
+	"ladder/pkg/ruleset"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
@@ -13,7 +14,7 @@ import (
 
 func TestProxySite(t *testing.T) {
 	app := fiber.New()
-	app.Get("/:url", ProxySite)
+	app.Get("/:url", ProxySite(""))
 
 	req := httptest.NewRequest("GET", "/https://example.com", nil)
 	resp, err := app.Test(req)
@@ -51,7 +52,7 @@ func TestRewriteHtml(t *testing.T) {
 		</html>
 	`
 
-	actual := rewriteHtml(bodyB, u, Rule{})
+	actual := rewriteHtml(bodyB, u, ruleset.Rule{})
 	assert.Equal(t, expected, actual)
 }