From b142ea71767b0575ac906f6b0890118e3408704d Mon Sep 17 00:00:00 2001
From: Paul Kulchenko <paul@kulchenko.com>
Date: Fri, 6 Aug 2021 04:48:46 -0700
Subject: [PATCH] Add following redirects to redbean Fetch (#226)

---
 tool/net/help.txt  | 14 ++++++++++-
 tool/net/redbean.c | 58 +++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 63 insertions(+), 9 deletions(-)

diff --git a/tool/net/help.txt b/tool/net/help.txt
index 73dc284b6..4b2bce2e3 100644
--- a/tool/net/help.txt
+++ b/tool/net/help.txt
@@ -545,7 +545,19 @@ FUNCTIONS
            provided, then a GET request is sent. If both URL and body parameters
            are specified, then a POST request is sent. If any other method needs
            to be specified (for example, PUT or DELETE), then passing a table as
-           the second value allows setting method and body values.
+           the second value allows setting method and body values as well other
+           options:
+             - followredirect (default = true): forces temporary and permanent
+               redirects to be followed. This behavior can be disabled by
+               passing `false`.
+             - maxredirects (default = 5): sets the number of allowed redirects
+               to minimize looping due to misconfigured servers. When the number
+               is exceeded, the result of the last redirect is returned.
+           When the redirect is being followed, the same method and body values
+           are being sent in all cases except when 303 status is returned. In
+           that case the method is set to GET and the body is removed before the
+           redirect is followed. Note that if these (method/body) values are
+           provided as table fields, they will be modified in place.
 
    FormatHttpDateTime(seconds:int) → rfc1123:str
            Converts UNIX timestamp to an RFC1123 string that looks like this:
diff --git a/tool/net/redbean.c b/tool/net/redbean.c
index a30611694..fdad733d9 100644
--- a/tool/net/redbean.c
+++ b/tool/net/redbean.c
@@ -3703,6 +3703,8 @@ static int LuaFetch(lua_State *L) {
   char *conlenhdr = "";
   size_t urlarglen, requestlen, paylen, bodylen;
   size_t g, i, n, hdrsize;
+  int numredirects = 0, maxredirects = 5;
+  bool followredirect = true;
   struct addrinfo hints = {.ai_family = AF_INET,
                            .ai_socktype = SOCK_STREAM,
                            .ai_protocol = IPPROTO_TCP,
@@ -3719,6 +3721,12 @@ static int LuaFetch(lua_State *L) {
     lua_getfield(L, 2, "method");
     // use GET by default if no method is provided
     method = strtoupper(luaL_optstring(L, -1, kHttpMethod[kHttpGet]));
+    lua_getfield(L, 2, "followredirect");
+    if (lua_isboolean(L, -1)) followredirect = lua_toboolean(L, -1);
+    lua_getfield(L, 2, "maxredirects");
+    maxredirects = luaL_optinteger(L, -1, maxredirects);
+    lua_getfield(L, 2, "numredirects");
+    numredirects = luaL_optinteger(L, -1, numredirects);
     lua_settop(L, 2);  // drop all added elements to keep the stack balanced
   } else if (lua_isnoneornil(L, 2)) {
     body = "";
@@ -3995,15 +4003,49 @@ static int LuaFetch(lua_State *L) {
 
 Finished:
   if (paylen && logbodies) LogBody("received", inbuf.p + hdrsize, paylen);
-  LOGF("FETCH HTTP%02d %d %s %`'.*s", msg.version, msg.status, urlarg,
+  LOGF("FETCH %s HTTP%02d %d %s %`'.*s", method, msg.version, msg.status, urlarg,
        HeaderLength(kHttpServer), HeaderData(kHttpServer));
-  lua_pushinteger(L, msg.status);
-  LuaPushHeaders(L, &msg, inbuf.p);
-  lua_pushlstring(L, inbuf.p + hdrsize, paylen);
-  DestroyHttpMessage(&msg);
-  free(inbuf.p);
-  close(sock);
-  return 3;
+  if (followredirect && HasHeader(kHttpLocation) &&
+    (msg.status == 301 || msg.status == 308 || // permanent redirects
+     msg.status == 302 || msg.status == 307 || // temporary redirects
+     msg.status == 303 // see other; non-GET changes to GET, body lost
+    ) && numredirects < maxredirects) {
+    // if 303, then remove body and set method to GET
+    if (msg.status == 303) {
+      body = "";
+      bodylen = 0;
+      method = kHttpMethod[kHttpGet];
+    }
+    // create table if needed
+    if (!lua_istable(L, 2)) {
+      lua_settop(L, 1); // pop body if present
+      lua_createtable(L, 0, 3); // body, method, numredirects
+    }
+    lua_pushlstring(L, body, bodylen);
+    lua_setfield(L, -2, "body");
+
+    lua_pushstring(L, method);
+    lua_setfield(L, -2, "method");
+
+    lua_pushinteger(L, numredirects + 1);
+    lua_setfield(L, -2, "numredirects");
+    // replace URL with Location header
+    lua_pushlstring(L, HeaderData(kHttpLocation), HeaderLength(kHttpLocation));
+    lua_replace(L, -3);
+
+    DestroyHttpMessage(&msg);
+    free(inbuf.p);
+    close(sock);
+    return LuaFetch(L);
+  } else {
+    lua_pushinteger(L, msg.status);
+    LuaPushHeaders(L, &msg, inbuf.p);
+    lua_pushlstring(L, inbuf.p + hdrsize, paylen);
+    DestroyHttpMessage(&msg);
+    free(inbuf.p);
+    close(sock);
+    return 3;
+  }
 TransportError:
   close(sock);
   free(inbuf.p);