Changeset 131 in code for trunk/morty.go

Timestamp:

Jul 24, 2020, 2:56:05 PM (5 years ago)

Author:

alex

Message:

Merge remote-tracking branch 'origin/master' into redirect

File:

• trunk/morty.go (modified) (8 diffs)

trunk/morty.go

@@ -39 +39 @@
 
 const VERSION = "v0.2.0"
+
+const MAX_REDIRECT_COUNT = 5
 
 var CLIENT *fasthttp.Client = &fasthttp.Client{
@@ -182 +184 @@
         Key            []byte
         RequestTimeout time.Duration
+        FollowRedirect bool
 }
 
@@ -313 +316 @@
         }
 
-        parsedURI, err := url.Parse(string(requestURI))
+        p.ProcessUri(ctx, string(requestURI), 0)
+}
+
+func (p *Proxy) ProcessUri(ctx *fasthttp.RequestCtx, requestURIStr string, redirectCount int) {
+        parsedURI, err := url.Parse(requestURIStr)
 
         if err != nil {
@@ -322 +329 @@
 
         if parsedURI.Scheme == "" {
-                requestURI = append([]byte("https://"), requestURI...)
-                parsedURI, err = url.Parse(string(requestURI))
+                requestURIStr = "https://" + requestURIStr
+                parsedURI, err = url.Parse(requestURIStr)
                 if err != nil {
                         p.serveMainPage(ctx, 500, err)
@@ -339 +346 @@
         defer fasthttp.ReleaseRequest(req)
         req.SetConnectionClose()
-
-        requestURIStr := string(requestURI)
 
         if cfg.Debug {
@@ -375 +380 @@
                         loc := resp.Header.Peek("Location")
                         if loc != nil {
-                                rc := &RequestConfig{Key: p.Key, BaseURL: parsedURI}
-                                url, err := rc.ProxifyURI(loc)
-                                if err == nil {
-                                        ctx.SetStatusCode(resp.StatusCode())
-                                        ctx.Response.Header.Add("Location", url)
-                                        if cfg.Debug {
-                                                log.Println("redirect to", string(loc))
+                                if p.FollowRedirect && ctx.IsGet() {
+                                        // GET method: Morty follows the redirect
+                                        if redirectCount < MAX_REDIRECT_COUNT {
+                                                if cfg.Debug {
+                                                        log.Println("follow redirect to", string(loc))
+                                                }
+                                                p.ProcessUri(ctx, string(loc), redirectCount+1)
+                                        } else {
+                                                p.serveMainPage(ctx, 310, errors.New("Too many redirects"))
                                         }
                                         return
+                                } else {
+                                        // Other HTTP methods: Morty does NOT follow the redirect
+                                        rc := &RequestConfig{Key: p.Key, BaseURL: parsedURI}
+                                        url, err := rc.ProxifyURI(loc)
+                                        if err == nil {
+                                                ctx.SetStatusCode(resp.StatusCode())
+                                                ctx.Response.Header.Add("Location", url)
+                                                if cfg.Debug {
+                                                        log.Println("redirect to", string(loc))
+                                                }
+                                                return
+                                        }
                                 }
                         }
@@ -1041 +1060 @@
         cfg.Debug = *flag.Bool("debug", cfg.Debug, "Debug mode")
         cfg.RequestTimeout = *flag.Uint("timeout", cfg.RequestTimeout, "Request timeout")
+        cfg.FollowRedirect = *flag.Bool("followredirect", cfg.FollowRedirect, "Follow HTTP GET redirect")
         version := flag.Bool("version", false, "Show version")
         socks5 := flag.String("socks5", "", "SOCKS5 proxy")
@@ -1058 +1078 @@
         }
 
-        p := &Proxy{RequestTimeout: time.Duration(cfg.RequestTimeout) * time.Second}
+        p := &Proxy{RequestTimeout: time.Duration(cfg.RequestTimeout) * time.Second,
+                FollowRedirect: cfg.FollowRedirect}
 
         if cfg.Key != "" {