From 6f548135e040a48b5c78e9d504d8e3607466e5fe Mon Sep 17 00:00:00 2001
From: Guilherme Salazar <gsz@acm.org>
Date: Tue, 17 Aug 2021 13:07:13 -0300
Subject: [PATCH] fix(grpc) allow setting :authority on balancer phase

---
 README.md                  |  5 +++-
 lualib/resty/kong/grpc.lua |  1 +
 t/003-grpc.t               | 59 ++++++++++++++++++++++++++++++++++++--
 3 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 99a7d6e0..a64adafd 100644
--- a/README.md
+++ b/README.md
@@ -248,7 +248,7 @@ resty.kong.grpc.set\_authority
 ------------------------------
 **syntax:** *ok, err = resty.kong.grpc.set_authority(new_authority)*
 
-**context:** *rewrite_by_lua&#42;, access_by_lua&#42;
+**context:** *rewrite_by_lua&#42;, access_by_lua&#42;, balancer_by_lua&#42;*
 
 **subsystems:** *http*
 
@@ -274,6 +274,9 @@ describing the error will be returned.
 This function can be called multiple times in the same request. Later calls override
 previous ones.
 
+If called in the `balancer_by_lua` context, the request needs to be recreated
+(see [balancer.recreate](https://github.com/openresty/lua-resty-core/blob/master/lib/ngx/balancer.md#recreate_request).
+
 [Back to TOC](#table-of-contents)
 
 resty.kong.tls.disable\_proxy\_ssl
diff --git a/lualib/resty/kong/grpc.lua b/lualib/resty/kong/grpc.lua
index b40e79bc..84b0cdad 100644
--- a/lualib/resty/kong/grpc.lua
+++ b/lualib/resty/kong/grpc.lua
@@ -42,6 +42,7 @@ do
     local ALLOWED_PHASES = {
         ['rewrite'] = true,
         ['access'] = true,
+        ['balancer'] = true,
     }
 
     function _M.set_authority(authority)
diff --git a/t/003-grpc.t b/t/003-grpc.t
index 842f4b6a..1d0b4a0f 100644
--- a/t/003-grpc.t
+++ b/t/003-grpc.t
@@ -3,8 +3,6 @@
 use Test::Nginx::Socket::Lua;
 use Cwd qw(cwd);
 
-repeat_each(2);
-
 plan tests => repeat_each() * (blocks() * 5);
 
 my $pwd = cwd();
@@ -225,3 +223,60 @@ falseincorrect argument, expects a string, got number
 [error]
 [crit]
 [alert]
+
+
+
+=== TEST 6: recreate upstream module requests with authority header change
+--- http_config
+    lua_package_path "../lua-resty-core/lib/?.lua;lualib/?.lua;;";
+
+    server {
+        listen 127.0.0.1:12346 http2;
+        server_name   example.com;
+
+        server_tokens off;
+
+        location / {
+            default_type 'text/plain';
+            more_clear_headers Date;
+            echo ':authority: $http_host';
+        }
+    }
+
+    upstream backend {
+        server 0.0.0.1;
+
+        balancer_by_lua_block {
+            print("here")
+            local b = require "ngx.balancer"
+            local grpc = require("resty.kong.grpc")
+
+            if ngx.ctx.balancer_run then
+                print"IF"
+                grpc.set_authority("try2")
+                assert(b.set_current_peer("127.0.0.1", 12346))
+                assert(b.recreate_request())
+
+            else
+                print"ELSE"
+                grpc.set_authority("try1")
+                ngx.ctx.balancer_run = true
+                assert(b.set_current_peer("127.0.0.3", 12345))
+                assert(b.set_more_tries(1))
+            end
+        }
+    }
+--- config
+    location = /t {
+         grpc_pass grpc://backend;
+    }
+
+--- request
+GET /t
+--- response_body
+:authority: try2
+--- error_log
+connect() failed (111: Connection refused) while connecting to upstream, client: 127.0.0.1
+--- no_error_log
+[warn]
+[crit]