From 34c9b87538e7f0c4dcb08d1940ccb507205b7102 Mon Sep 17 00:00:00 2001 From: Paul Kulchenko Date: Sat, 13 May 2023 14:18:46 -0700 Subject: [PATCH] Extend redbean Fetch to add option to keeping connection open. --- tool/net/fetch.inc | 95 +++++++++++++++++++++++++++++++--------------- 1 file changed, 65 insertions(+), 30 deletions(-) diff --git a/tool/net/fetch.inc b/tool/net/fetch.inc index 2428f2ede..69bad4190 100644 --- a/tool/net/fetch.inc +++ b/tool/net/fetch.inc @@ -4,6 +4,11 @@ #define FetchHeaderEqualCase(H, S) \ SlicesEqualCase(S, strlen(S), FetchHeaderData(H), FetchHeaderLength(H)) +#define kaNONE 0 +#define kaOPEN 1 +#define kaKEEP 2 +#define kaCLOSE 3 + static int LuaFetch(lua_State *L) { #define ssl nope // TODO(jart): make this file less huge char *p; @@ -11,7 +16,7 @@ static int LuaFetch(lua_State *L) { bool usingssl; uint32_t ip; struct Url url; - int t, ret, sock, methodidx, hdridx; + int t, ret, sock = -1, methodidx, hdridx; char *host, *port, *request; struct TlsBio *bio; struct addrinfo *addr; @@ -22,11 +27,13 @@ static int LuaFetch(lua_State *L) { char *conlenhdr = ""; char *headers = 0; char *hosthdr = 0; + char *connhdr = 0; char *agenthdr = brand; char *key, *val, *hdr; size_t keylen, vallen; size_t urlarglen, requestlen, paylen, bodylen; size_t i, g, n, hdrsize; + int keepalive = kaNONE; int imethod, numredirects = 0, maxredirects = 5; bool followredirect = true; struct addrinfo hints = {.ai_family = AF_INET, @@ -56,6 +63,16 @@ static int LuaFetch(lua_State *L) { maxredirects = luaL_optinteger(L, -1, maxredirects); lua_getfield(L, 2, "numredirects"); numredirects = luaL_optinteger(L, -1, numredirects); + lua_getfield(L, 2, "keepalive"); // accept keepalive=true|num + if (lua_isboolean(L, -1)) keepalive = lua_toboolean(L, -1) ? kaOPEN : kaNONE; + if (lua_isinteger(L, -1)) { + keepalive = kaKEEP; + sock = lua_tointeger(L, -1); + if (sock < 0) { + sock = -sock; + keepalive = kaCLOSE; + } + } lua_getfield(L, 2, "headers"); if (!lua_isnil(L, -1)) { if (!lua_istable(L, -1)) @@ -72,15 +89,16 @@ static int LuaFetch(lua_State *L) { if (!(hdr = _gc(EncodeHttpHeaderValue(val, vallen, 0)))) return LuaNilError(L, "invalid header %s value encoding", key); - // Content-Length and Connection will be overwritten; - // skip them to avoid duplicates; + // Content-Length will be overwritten; skip it to avoid duplicates; // also allow unknown headers if ((hdridx = GetHttpHeader(key, keylen)) == -1 || - hdridx != kHttpContentLength && hdridx != kHttpConnection) { + hdridx != kHttpContentLength) { if (hdridx == kHttpUserAgent) { agenthdr = hdr; } else if (hdridx == kHttpHost) { hosthdr = hdr; + } else if (hdridx == kHttpConnection) { + connhdr = hdr; } else { appendd(&headers, key, keylen); appendw(&headers, READ16LE(": ")); @@ -128,6 +146,7 @@ static int LuaFetch(lua_State *L) { } #ifndef UNSECURE + if (usingssl) keepalive = kaNONE; if (usingssl && !sslinitialized) TlsInit(); #endif @@ -173,39 +192,43 @@ static int LuaFetch(lua_State *L) { appendf(&request, "%s %s HTTP/1.1\r\n" "Host: %s\r\n" - "Connection: close\r\n" + "Connection: %s\r\n" "User-Agent: %s\r\n" "%s%s" "\r\n", - method, _gc(EncodeUrl(&url, 0)), hosthdr, agenthdr, conlenhdr, - headers ? headers : ""); + method, _gc(EncodeUrl(&url, 0)), hosthdr, + keepalive == kaNONE || keepalive == kaCLOSE ? "close" + : (connhdr ? connhdr : "keep-alive"), + agenthdr, conlenhdr, headers ? headers : ""); appendd(&request, body, bodylen); requestlen = appendz(request).i; _gc(request); - /* - * Perform DNS lookup. - */ - DEBUGF("(ftch) client resolving %s", host); - if ((rc = getaddrinfo(host, port, &hints, &addr)) != EAI_SUCCESS) { - return LuaNilError(L, "getaddrinfo(%s:%s) error: EAI_%s %s", host, port, - gai_strerror(rc), strerror(errno)); - } + if (keepalive != kaKEEP && keepalive != kaCLOSE) { + /* + * Perform DNS lookup. + */ + DEBUGF("(ftch) client resolving %s", host); + if ((rc = getaddrinfo(host, port, &hints, &addr)) != EAI_SUCCESS) { + return LuaNilError(L, "getaddrinfo(%s:%s) error: EAI_%s %s", host, port, + gai_strerror(rc), strerror(errno)); + } - /* - * Connect to server. - */ - ip = ntohl(((struct sockaddr_in *)addr->ai_addr)->sin_addr.s_addr); - DEBUGF("(ftch) client connecting %hhu.%hhu.%hhu.%hhu:%d", ip >> 24, ip >> 16, - ip >> 8, ip, ntohs(((struct sockaddr_in *)addr->ai_addr)->sin_port)); - CHECK_NE(-1, (sock = GoodSocket(addr->ai_family, addr->ai_socktype, - addr->ai_protocol, false, &timeout))); - rc = connect(sock, addr->ai_addr, addr->ai_addrlen); - freeaddrinfo(addr), addr = 0; - if (rc == -1) { - close(sock); - return LuaNilError(L, "connect(%s:%s) error: %s", host, port, - strerror(errno)); + /* + * Connect to server. + */ + ip = ntohl(((struct sockaddr_in *)addr->ai_addr)->sin_addr.s_addr); + DEBUGF("(ftch) client connecting %hhu.%hhu.%hhu.%hhu:%d", ip >> 24, ip >> 16, + ip >> 8, ip, ntohs(((struct sockaddr_in *)addr->ai_addr)->sin_port)); + CHECK_NE(-1, (sock = GoodSocket(addr->ai_family, addr->ai_socktype, + addr->ai_protocol, false, &timeout))); + rc = connect(sock, addr->ai_addr, addr->ai_addrlen); + freeaddrinfo(addr), addr = 0; + if (rc == -1) { + close(sock); + return LuaNilError(L, "connect(%s:%s) error: %s", host, port, + strerror(errno)); + } } #ifndef UNSECURE @@ -424,6 +447,13 @@ Finished: lua_pushstring(L, method); lua_setfield(L, -2, "method"); + // since the connection is going to be closed on redirect, + // make sure that keepalive socket is reset + if (keepalive) { + lua_pushboolean(L, true); + lua_setfield(L, -2, "keepalive"); + } + lua_pushinteger(L, numredirects + 1); lua_setfield(L, -2, "numredirects"); // replace URL with Location header @@ -436,12 +466,17 @@ Finished: close(sock); return LuaFetch(L); } else { + if (keepalive && keepalive != kaCLOSE && lua_istable(L, 2)) { + lua_pushinteger(L, sock); + lua_setfield(L, -2, "keepalive"); + } + lua_pushinteger(L, msg.status); LuaPushHeaders(L, &msg, inbuf.p); lua_pushlstring(L, inbuf.p + hdrsize, paylen); DestroyHttpMessage(&msg); free(inbuf.p); - close(sock); + if (!keepalive || keepalive == kaCLOSE) close(sock); return 3; } TransportError: