Add more tests and fix various issues

extensions/omni_httpc/omni_httpc.c

@@ -209,8 +209,6 @@ struct request {
   h2o_iovec_t request_body;
   // Response body
   StringInfoData body;
-  // Pointer to a counter of completed requests
-  int *done;
   // Request URL
   h2o_url_t url;
   // Error (NULL if none)
@@ -246,7 +244,6 @@ static int on_body(h2o_httpclient_t *client, const char *errstr, h2o_header_t *t
   if (errstr != NULL) {
     if (errstr != h2o_httpclient_error_is_eos) {
       req->errstr = errstr;
-      (*req->done)++;
       req->complete = true;
       return -1;
     }
@@ -264,7 +261,6 @@ static int on_body(h2o_httpclient_t *client, const char *errstr, h2o_header_t *t
   // End of stream, complete the request
   if (errstr == h2o_httpclient_error_is_eos) {
     SET_VARSIZE(req->body.data, req->body.len);
-    (*req->done)++;
     req->complete = true;
   }
 
@@ -292,7 +288,6 @@ static h2o_httpclient_body_cb on_head(h2o_httpclient_t *client, const char *errs
   // If there's an error, report it
   if (errstr != NULL && errstr != h2o_httpclient_error_is_eos) {
     req->errstr = errstr;
-    (*req->done)++;
     req->complete = true;
     return NULL;
   }
@@ -311,10 +306,9 @@ static h2o_httpclient_body_cb on_head(h2o_httpclient_t *client, const char *errs
     if ((location_index = h2o_find_header(headers_vec, H2O_TOKEN_LOCATION, -1)) > -1) {
       // create a new request based on the original request
       struct request *new_req = palloc0(sizeof(struct request));
-      new_req->done = req->done;
       new_req->complete = false;
       new_req->connected = false;
-      new_req->headers = headers_vec->entries;
+      new_req->headers = req->headers;
 
       initStringInfo(&new_req->body);
       appendStringInfoSpaces(&new_req->body, sizeof(struct varlena));
@@ -330,7 +324,9 @@ static h2o_httpclient_body_cb on_head(h2o_httpclient_t *client, const char *errs
         new_req->request_body = req->request_body;
       } else {
         // omit the request body
-//        new_req->request_body = h2o_iovec_init(NULL, 0);
+        // we don't need to remove the header, as it will be automatically omitted if the body
+        // is empty
+        new_req->request_body = h2o_iovec_init(NULL, 0);
         // if the method is POST, override the method as GET
         if (h2o_memis(req->method.base, req->method.len, H2O_STRLIT("POST"))) {
           new_req->method = h2o_iovec_init(H2O_STRLIT("GET"));
@@ -340,6 +336,9 @@ static h2o_httpclient_body_cb on_head(h2o_httpclient_t *client, const char *errs
       }
 
       h2o_url_t url;
+      // re-evaluate location index after the headers have been modified
+      location_index = h2o_find_header(headers_vec, H2O_TOKEN_LOCATION, -1);
+
       h2o_iovec_t location_header = headers_vec->entries[location_index].value;
       // try to parse the URL as an absolute URL
       int result = h2o_url_parse(client->pool, location_header.base, location_header.len, &url);
@@ -368,7 +367,6 @@ static h2o_httpclient_body_cb on_head(h2o_httpclient_t *client, const char *errs
       // copy the response from the new request to the original request
       copy_request_fields(new_req, req);
       req->complete = true;
-      (*req->done)++;
       return NULL;
     } else {
       ereport(
@@ -402,7 +400,6 @@ static h2o_httpclient_body_cb on_head(h2o_httpclient_t *client, const char *errs
 
   // End of stream, complete the request
   if (errstr == h2o_httpclient_error_is_eos) {
-    (*req->done)++;
     req->complete = true;
     return NULL;
   }
@@ -422,7 +419,7 @@ static h2o_httpclient_head_cb on_connect(h2o_httpclient_t *client, const char *e
   // If there's an error, report it
   if (errstr != NULL) {
     req->errstr = errstr;
-    (*req->done)++;
+    req->complete = true;
     return NULL;
   }
   // Provide H2O client with necessary request payload
@@ -456,8 +453,6 @@ PG_FUNCTION_INFO_V1(http_execute);
 Datum http_execute(PG_FUNCTION_ARGS) {
   init();
 
-  int done = 0;
-
   // Prepare to return the set
   ReturnSetInfo *rsinfo = (ReturnSetInfo *)fcinfo->resultinfo;
   rsinfo->returnMode = SFRM_Materialize;
@@ -647,12 +642,12 @@ Datum http_execute(PG_FUNCTION_ARGS) {
       h2o_add_header(pool, headers_vec, H2O_TOKEN_CONTENT_LENGTH, NULL, clbuf, clbuf_len);
     }
 
-    request->done = &done;
     request->redirects_left = MAX_REDIRECTS;
     request->errstr = NULL;
     request->num_headers = headers_vec->size;
     request->headers = headers_vec->entries;
     request->connected = false;
+    request->complete = false;
     initStringInfo(&request->body);
     // Prepend the response with varlena's header so that we can populate it
     // in the end, when the ull length is known.
@@ -699,9 +694,12 @@ Datum http_execute(PG_FUNCTION_ARGS) {
   }
 
   // Run the event loop until all requests have been completed
-  while (done < num_requests) {
-    CHECK_FOR_INTERRUPTS();
-    h2o_evloop_run(ctx.loop, INT32_MAX);
+  for (int i = 0; i < num_requests; i++) {
+    struct request *request = &requests[i];
+    while (!request->complete) {
+      CHECK_FOR_INTERRUPTS();
+      h2o_evloop_run(ctx.loop, INT32_MAX);
+    }
   }
 
   // Start populating the result set

extensions/omni_httpc/tests/redirect.yml

@@ -4,12 +4,13 @@ instance:
     shared_preload_libraries: */env/OMNI_EXT_SO
     max_worker_processes: 64
   init:
-    - set session omni_httpd.init_port = 13337
+    - set session omni_httpd.init_port = 0
     - create extension omni_httpd cascade
     - create extension omni_httpc cascade
     - call omni_httpd.wait_for_configuration_reloads(1)
     - |
-      update omni_httpd.handlers set query = $$
+      update omni_httpd.handlers
+      set query = $$
         select omni_httpd.http_response(json_build_object(
         'method', request.method,
         'path', request.path,
@@ -25,31 +26,49 @@ instance:
         headers =>
           case
                   when
-                      request.query_string = 'q=absolute_url'
+                      request.query_string LIKE '%q=absolute_url%'
                   then
                       ARRAY[omni_http.http_header(
                           'Location',
                           'http://127.0.0.1:' || (select effective_port from omni_httpd.listeners) || '/test?q=destination')]
                   when
-                      request.query_string = 'q=relative_url_absolute_path'
+                      request.query_string LIKE '%q=relative_url_absolute_path%'
                   then
                       ARRAY[omni_http.http_header(
                           'Location',
                           '/test?q=destination')]
       
                   when
-                      request.query_string = 'q=relative_url_relative_path'
+                      request.query_string LIKE '%q=relative_url_relative_path%'
                   then
                       ARRAY[omni_http.http_header(
                           'Location',
                           'path1/path2?q=destination')]
+                    
+                  when
+                      request.query_string LIKE '%counter=%'
+                  then
+                      ARRAY[
+                          omni_http.http_header(
+                              'Location', '/test?status=302&counter=' || (
+                                      (regexp_match(request.query_string, 'counter=([^&]+)'))::int[]
+                                  )[1]+ 1
+                          )
+                      ]
                   else NULL
           end,
         status => case
           when 
-              request.query_string = 'q=absolute_url' or 
-              request.query_string = 'q=relative_url_absolute_path' or 
-              request.query_string = 'q=relative_url_relative_path' then 302 else 200
+              request.query_string LIKE '%status=%'
+              OR request.query_string LIKE 'counter=%'
+          then (
+              ((regexp_match(
+                  request.query_string,
+                  'status=([^&]+)'
+              ))::int[])[1]::int
+          )
+              
+          else 200
           end
         ) from request
       $$
@@ -60,36 +79,42 @@ tests:
     query: |
       with response as (select * from omni_httpc.http_execute_with_options(
              omni_httpc.http_execute_options(follow_redirects => false),
-             omni_httpc.http_request('http://127.0.0.1:' || (select effective_port from omni_httpd.listeners) || '/test?q=absolute_url'))
-            )
+             omni_httpc.http_request(
+              'http://127.0.0.1:' || 
+              (select effective_port from omni_httpd.listeners) ||
+              '/test?q=absolute_url&status=302')
+              )
+      )
       select
         response.status,
         (select
-             json_agg(json_build_object(h.name, h.value))
+             h.value ~ 'http://127.0.0.1:[0-9]+/test\?q=destination'
          from
              unnest(response.headers) h
          where
-             h.name not in ('server', 'content-length', 'connection')) as headers,
+             h.name = 'location') as location_matches,
         convert_from(response.body, 'utf-8')::json as body
         from response
     results:
     - status: 302
-      headers:
-      - location: 'http://127.0.0.1:13337/test?q=destination'
-      - content-type: application/json
+      location_matches: true
       body:
         method: GET
-        qs: q=absolute_url
+        qs: q=absolute_url&status=302
         path: /test
         headers: null
         body: ""
 
-  - name: redirect should be followed if follow_redirects is true
+  - name: redirect should be followed if follow_redirects is true for an absolute URL
     query: |
       with response as (select * from omni_httpc.http_execute_with_options(
           omni_httpc.http_execute_options(follow_redirects => true),
-          omni_httpc.http_request('http://127.0.0.1:' || (select effective_port from omni_httpd.listeners) || '/test?q=absolute_url'))
-      )
+          omni_httpc.http_request(
+              'http://127.0.0.1:' || 
+              (select effective_port from omni_httpd.listeners) || 
+              '/test?q=absolute_url&status=302'
+          )
+      ))
       select
           response.status,
           (select
@@ -106,17 +131,20 @@ tests:
           - content-type: application/json
         body:
           method: GET
-          qs: q=destination
           path: /test
-          headers:
-            - connection: keep-alive
+          qs: q=destination
+          headers: null
           body: ""
 
-  - name: redirect should be followed if follow_redirects is true
+  - name: redirect should be followed if follow_redirects is true for a relative URL and absolute path
     query: |
       with response as (select * from omni_httpc.http_execute_with_options(
               omni_httpc.http_execute_options(follow_redirects => true),
-              omni_httpc.http_request('http://127.0.0.1:' || (select effective_port from omni_httpd.listeners) || '/test?q=relative_url_absolute_path'))
+              omni_httpc.http_request(
+                  'http://127.0.0.1:' ||
+                  (select effective_port from omni_httpd.listeners) ||
+                  '/test?q=relative_url_absolute_path&status=302')
+              )
       )
       select
           response.status,
@@ -136,15 +164,17 @@ tests:
           method: GET
           qs: q=destination
           path: /test
-          headers:
-            - connection: keep-alive
+          headers: null
           body: ""
 
-  - name:
+  - name: redirect should be followed if follow_redirects is true for a relative URL and relative path
     query: |
       with response as (select * from omni_httpc.http_execute_with_options(
               omni_httpc.http_execute_options(follow_redirects => true),
-              omni_httpc.http_request('http://127.0.0.1:' || (select effective_port from omni_httpd.listeners) || '/test?q=relative_url_relative_path'))
+              omni_httpc.http_request(
+                  'http://127.0.0.1:' || 
+                  (select effective_port from omni_httpd.listeners) || 
+                  '/test?q=relative_url_relative_path&status=302'))
       )
       select
           response.status,
@@ -164,6 +194,71 @@ tests:
           method: GET
           qs: q=destination
           path: /path1/path2
-          headers:
-            - connection: keep-alive
+          headers: null
           body: ""
+
+  - name: up to 5 redirects should be followed if follow_redirects is true
+    query: |
+      with response as (select *
+          from omni_httpc.http_execute_with_options(
+                omni_httpc.http_execute_options(follow_redirects => true),
+                omni_httpc.http_request(
+                    'http://127.0.0.1:' || 
+                    (select effective_port from omni_httpd.listeners) ||
+                    '/test?counter=1&status=302'
+                )
+          )
+      )
+      select response.status,
+             (select json_agg(json_build_object(h.name, h.value))
+              from
+                  unnest(response.headers) h
+              where
+                  h.name not in ('server', 'content-length', 'connection')) as headers,
+             convert_from(response.body, 'utf-8')::json as body
+      from response
+    results:
+      - status: 302
+        headers:
+          - location: /test?status=302&counter=7
+          - content-type: application/json
+        body:
+          method: GET
+          path: /test
+          qs: status=302&counter=6
+          headers: null
+          body: ""
+
+  - name: 307/308 redirects should be redirected with the same method and body
+    query: |
+      with response as (select *
+                        from omni_httpc.http_execute_with_options(
+                                omni_httpc.http_execute_options(follow_redirects => true),
+                                omni_httpc.http_request(
+                                        'http://127.0.0.1:' ||
+                                        (select effective_port from omni_httpd.listeners) ||
+                                        '/test?q=absolute_url&status=308',
+                                         method => 'POST',
+                                         body => 'test'
+                                )
+                             )
+      )
+      select response.status,
+             (select json_agg(json_build_object(h.name, h.value))
+              from
+                  unnest(response.headers) h
+              where
+                  h.name not in ('server', 'content-length', 'connection')) as headers,
+             convert_from(response.body, 'utf-8')::json as body
+      from response
+    results:
+      - status: 200
+        headers:
+          - content-type: application/json
+        body:
+          method: POST
+          path: /test
+          qs: q=destination
+          headers: null
+          body: "test"
+