Merge pull request #1036 from nghttp2/nghttpx-affinity-cookie

nghttpx: Cookie based session affinity

integration-tests/nghttpx_http1_test.go

@@ -9,6 +9,7 @@ import (
 	"golang.org/x/net/websocket"
 	"io"
 	"net/http"
+	"regexp"
 	"syscall"
 	"testing"
 	"time"
@@ -125,6 +126,54 @@ Content-Length: 0
 // 	}
 // }
 
+// TestH1H1AffinityCookie tests that affinity cookie is sent back in
+// cleartext http.
+func TestH1H1AffinityCookie(t *testing.T) {
+	st := newServerTester([]string{"--affinity-cookie"}, t, noopHandler)
+	defer st.Close()
+
+	res, err := st.http1(requestParam{
+		name: "TestH1H1AffinityCookie",
+	})
+	if err != nil {
+		t.Fatalf("Error st.http1() = %v", err)
+	}
+
+	if got, want := res.status, 200; got != want {
+		t.Errorf("status = %v; want %v", got, want)
+	}
+
+	const pattern = `affinity=[0-9a-f]{8}; Path=/foo/bar`
+	validCookie := regexp.MustCompile(pattern)
+	if got := res.header.Get("Set-Cookie"); !validCookie.MatchString(got) {
+		t.Errorf("Set-Cookie: %v; want pattern %v", got, pattern)
+	}
+}
+
+// TestH1H1AffinityCookieTLS tests that affinity cookie is sent back
+// in https.
+func TestH1H1AffinityCookieTLS(t *testing.T) {
+	st := newServerTesterTLS([]string{"--alpn-h1", "--affinity-cookie"}, t, noopHandler)
+	defer st.Close()
+
+	res, err := st.http1(requestParam{
+		name: "TestH1H1AffinityCookieTLS",
+	})
+	if err != nil {
+		t.Fatalf("Error st.http1() = %v", err)
+	}
+
+	if got, want := res.status, 200; got != want {
+		t.Errorf("status = %v; want %v", got, want)
+	}
+
+	const pattern = `affinity=[0-9a-f]{8}; Path=/foo/bar; Secure`
+	validCookie := regexp.MustCompile(pattern)
+	if got := res.header.Get("Set-Cookie"); !validCookie.MatchString(got) {
+		t.Errorf("Set-Cookie: %v; want pattern %v", got, pattern)
+	}
+}
+
 // TestH1H1GracefulShutdown tests graceful shutdown.
 func TestH1H1GracefulShutdown(t *testing.T) {
 	st := newServerTester(nil, t, noopHandler)

integration-tests/nghttpx_http2_test.go

@@ -1705,6 +1705,55 @@ func TestH2H1Code204TE(t *testing.T) {
 	}
 }
 
+// TestH2H1AffinityCookie tests that affinity cookie is sent back in
+// cleartext http.
+func TestH2H1AffinityCookie(t *testing.T) {
+	st := newServerTester([]string{"--affinity-cookie"}, t, noopHandler)
+	defer st.Close()
+
+	res, err := st.http2(requestParam{
+		name: "TestH2H1AffinityCookie",
+	})
+	if err != nil {
+		t.Fatalf("Error st.http2() = %v", err)
+	}
+
+	if got, want := res.status, 200; got != want {
+		t.Errorf("status = %v; want %v", got, want)
+	}
+
+	const pattern = `affinity=[0-9a-f]{8}; Path=/foo/bar`
+	validCookie := regexp.MustCompile(pattern)
+	if got := res.header.Get("Set-Cookie"); !validCookie.MatchString(got) {
+		t.Errorf("Set-Cookie: %v; want pattern %v", got, pattern)
+	}
+}
+
+// TestH2H1AffinityCookieTLS tests that affinity cookie is sent back
+// in https.
+func TestH2H1AffinityCookieTLS(t *testing.T) {
+	st := newServerTesterTLS([]string{"--affinity-cookie"}, t, noopHandler)
+	defer st.Close()
+
+	res, err := st.http2(requestParam{
+		name:   "TestH2H1AffinityCookieTLS",
+		scheme: "https",
+	})
+	if err != nil {
+		t.Fatalf("Error st.http2() = %v", err)
+	}
+
+	if got, want := res.status, 200; got != want {
+		t.Errorf("status = %v; want %v", got, want)
+	}
+
+	const pattern = `affinity=[0-9a-f]{8}; Path=/foo/bar; Secure`
+	validCookie := regexp.MustCompile(pattern)
+	if got := res.header.Get("Set-Cookie"); !validCookie.MatchString(got) {
+		t.Errorf("Set-Cookie: %v; want pattern %v", got, pattern)
+	}
+}
+
 // TestH2H1GracefulShutdown tests graceful shutdown.
 func TestH2H1GracefulShutdown(t *testing.T) {
 	st := newServerTester(nil, t, noopHandler)

integration-tests/server_tester.go

@@ -101,7 +101,7 @@ func newServerTesterInternal(src_args []string, t *testing.T, handler http.Handl
 
 	args := []string{}
 
-	var backendTLS, dns, externalDNS, acceptProxyProtocol, redirectIfNotTLS bool
+	var backendTLS, dns, externalDNS, acceptProxyProtocol, redirectIfNotTLS, affinityCookie, alpnH1 bool
 
 	for _, k := range src_args {
 		switch k {
@@ -116,6 +116,10 @@ func newServerTesterInternal(src_args []string, t *testing.T, handler http.Handl
 			acceptProxyProtocol = true
 		case "--redirect-if-not-tls":
 			redirectIfNotTLS = true
+		case "--affinity-cookie":
+			affinityCookie = true
+		case "--alpn-h1":
+			alpnH1 = true
 		default:
 			args = append(args, k)
 		}
@@ -168,6 +172,10 @@ func newServerTesterInternal(src_args []string, t *testing.T, handler http.Handl
 		b += ";redirect-if-not-tls"
 	}
 
+	if affinityCookie {
+		b += ";affinity=cookie;affinity-cookie-name=affinity;affinity-cookie-path=/foo/bar"
+	}
+
 	noTLS := ";no-tls"
 	if frontendTLS {
 		noTLS = ""
@@ -218,7 +226,11 @@ func newServerTesterInternal(src_args []string, t *testing.T, handler http.Handl
 				tlsConfig = clientConfig
 			}
 			tlsConfig.InsecureSkipVerify = true
-			tlsConfig.NextProtos = []string{"h2", "spdy/3.1"}
+			if alpnH1 {
+				tlsConfig.NextProtos = []string{"http/1.1"}
+			} else {
+				tlsConfig.NextProtos = []string{"h2", "spdy/3.1"}
+			}
 			conn, err = tls.Dial("tcp", authority, tlsConfig)
 		} else {
 			conn, err = net.Dial("tcp", authority)

src/shrpx-unittest.cc

@@ -117,6 +117,8 @@ int main(int argc, char *argv[]) {
                    shrpx::test_downstream_rewrite_location_response_header) ||
       !CU_add_test(pSuite, "downstream_supports_non_final_response",
                    shrpx::test_downstream_supports_non_final_response) ||
+      !CU_add_test(pSuite, "downstream_find_affinity_cookie",
+                   shrpx::test_downstream_find_affinity_cookie) ||
       !CU_add_test(pSuite, "config_parse_header",
                    shrpx::test_shrpx_config_parse_header) ||
       !CU_add_test(pSuite, "config_parse_log_format",
@@ -131,6 +133,8 @@ int main(int argc, char *argv[]) {
                    shrpx::test_shrpx_http_create_forwarded) ||
       !CU_add_test(pSuite, "http_create_via_header_value",
                    shrpx::test_shrpx_http_create_via_header_value) ||
+      !CU_add_test(pSuite, "http_create_affinity_cookie",
+                   shrpx::test_shrpx_http_create_affinity_cookie) ||
       !CU_add_test(pSuite, "router_match", shrpx::test_shrpx_router_match) ||
       !CU_add_test(pSuite, "router_match_wildcard",
                    shrpx::test_shrpx_router_match_wildcard) ||

src/shrpx.cc

@@ -1772,16 +1772,28 @@ Connections:
               The     session     affinity    is     enabled     using
               "affinity=<METHOD>"  parameter.   If  "ip" is  given  in
               <METHOD>, client  IP based session affinity  is enabled.
-              If  "none" is  given  in <METHOD>,  session affinity  is
+              If "cookie"  is given in <METHOD>,  cookie based session
-              disabled, and this is the default.  The session affinity
+              affinity is  enabled.  If  "none" is given  in <METHOD>,
-              is enabled per  <PATTERN>.  If at least  one backend has
+              session affinity  is disabled, and this  is the default.
-              "affinity" parameter,  and its  <METHOD> is  not "none",
+              The session  affinity is  enabled per <PATTERN>.   If at
-              session  affinity is  enabled  for  all backend  servers
+              least  one backend  has  "affinity"  parameter, and  its
-              sharing  the  same  <PATTERN>.   It is  advised  to  set
+              <METHOD> is not "none",  session affinity is enabled for
-              "affinity"  parameter  to   all  backend  explicitly  if
+              all backend  servers sharing the same  <PATTERN>.  It is
-              session affinity  is desired.  The session  affinity may
+              advised  to  set  "affinity" parameter  to  all  backend
-              break if one of the backend gets unreachable, or backend
+              explicitly if session affinity  is desired.  The session
-              settings are reloaded or replaced by API.
+              affinity  may   break  if   one  of  the   backend  gets
+              unreachable,  or   backend  settings  are   reloaded  or
+              replaced by API.
+
+              If   "affinity=cookie"    is   used,    the   additional
+              configuration                is                required.
+              "affinity-cookie-name=<NAME>" must be  used to specify a
+              name     of     cookie      to     use.      Optionally,
+              "affinity-cookie-path=<PATH>" can  be used to  specify a
+              path which cookie is applied.  The Secure attribute of a
+              cookie is determined by a  request scheme.  If a request
+              scheme  is  "https",  then Secure  attribute  is  added.
+              Otherwise, it is not added.
 
               By default, name resolution of backend host name is done
               at  start  up,  or reloading  configuration.   If  "dns"

src/shrpx_client_handler.cc

@@ -699,7 +699,7 @@ void ClientHandler::pool_downstream_connection(
 
   auto &shared_addr = group->shared_addr;
 
-  if (shared_addr->affinity == AFFINITY_NONE) {
+  if (shared_addr->affinity.type == AFFINITY_NONE) {
     auto &dconn_pool = group->shared_addr->dconn_pool;
     dconn_pool.add_downstream_connection(std::move(dconn));
 
@@ -947,6 +947,24 @@ uint32_t next_cycle(const WeightedPri &pri) {
 }
 } // namespace
 
+uint32_t ClientHandler::get_affinity_cookie(Downstream *downstream,
+                                            const StringRef &cookie_name) {
+  auto h = downstream->find_affinity_cookie(cookie_name);
+  if (h) {
+    return h;
+  }
+
+  auto d = std::uniform_int_distribution<uint32_t>(
+      1, std::numeric_limits<uint32_t>::max());
+  auto rh = d(worker_->get_randgen());
+  h = util::hash32(StringRef{reinterpret_cast<uint8_t *>(&rh),
+                             reinterpret_cast<uint8_t *>(&rh) + sizeof(rh)});
+
+  downstream->renew_affinity_cookie(h);
+
+  return h;
+}
+
 std::unique_ptr<DownstreamConnection>
 ClientHandler::get_downstream_connection(int &err, Downstream *downstream) {
   size_t group_idx;
@@ -1012,16 +1030,27 @@ ClientHandler::get_downstream_connection(int &err, Downstream *downstream) {
   auto &group = groups[group_idx];
   auto &shared_addr = group->shared_addr;
 
-  if (shared_addr->affinity == AFFINITY_IP) {
+  if (shared_addr->affinity.type != AFFINITY_NONE) {
-    if (!affinity_hash_computed_) {
+    uint32_t hash;
-      affinity_hash_ = compute_affinity_from_ip(ipaddr_);
+    switch (shared_addr->affinity.type) {
-      affinity_hash_computed_ = true;
+    case AFFINITY_IP:
+      if (!affinity_hash_computed_) {
+        affinity_hash_ = compute_affinity_from_ip(ipaddr_);
+        affinity_hash_computed_ = true;
+      }
+      hash = affinity_hash_;
+      break;
+    case AFFINITY_COOKIE:
+      hash = get_affinity_cookie(downstream, shared_addr->affinity.cookie.name);
+      break;
+    default:
+      assert(0);
     }
 
     const auto &affinity_hash = shared_addr->affinity_hash;
 
     auto it = std::lower_bound(
-        std::begin(affinity_hash), std::end(affinity_hash), affinity_hash_,
+        std::begin(affinity_hash), std::end(affinity_hash), hash,
         [](const AffinityHash &lhs, uint32_t rhs) { return lhs.hash < rhs; });
 
     if (it == std::end(affinity_hash)) {

src/shrpx_client_handler.h

@@ -153,6 +153,11 @@ public:
   Http2Session *select_http2_session_with_affinity(
       const std::shared_ptr<DownstreamAddrGroup> &group, DownstreamAddr *addr);
 
+  // Returns an affinity cookie value for |downstream|.  |cookie_name|
+  // is used to inspect cookie header field in request header fields.
+  uint32_t get_affinity_cookie(Downstream *downstream,
+                               const StringRef &cookie_name);
+
   const UpstreamAddr *get_upstream_addr() const;
 
   void repeat_read_timer();

src/shrpx_config.cc

@@ -789,10 +789,10 @@ int parse_upstream_params(UpstreamParams &out, const StringRef &src_params) {
 
 struct DownstreamParams {
   StringRef sni;
+  AffinityConfig affinity;
   size_t fall;
   size_t rise;
   shrpx_proto proto;
-  shrpx_session_affinity affinity;
   bool tls;
   bool dns;
   bool redirect_if_not_tls;
@@ -862,13 +862,27 @@ int parse_downstream_params(DownstreamParams &out,
     } else if (util::istarts_with_l(param, "affinity=")) {
       auto valstr = StringRef{first + str_size("affinity="), end};
       if (util::strieq_l("none", valstr)) {
-        out.affinity = AFFINITY_NONE;
+        out.affinity.type = AFFINITY_NONE;
       } else if (util::strieq_l("ip", valstr)) {
-        out.affinity = AFFINITY_IP;
+        out.affinity.type = AFFINITY_IP;
+      } else if (util::strieq_l("cookie", valstr)) {
+        out.affinity.type = AFFINITY_COOKIE;
       } else {
-        LOG(ERROR) << "backend: affinity: value must be either none or ip";
+        LOG(ERROR)
+            << "backend: affinity: value must be one of none, ip, and cookie";
         return -1;
       }
+    } else if (util::istarts_with_l(param, "affinity-cookie-name=")) {
+      auto val = StringRef{first + str_size("affinity-cookie-name="), end};
+      if (val.empty()) {
+        LOG(ERROR)
+            << "backend: affinity-cookie-name: non empty string is expected";
+        return -1;
+      }
+      out.affinity.cookie.name = val;
+    } else if (util::istarts_with_l(param, "affinity-cookie-path=")) {
+      out.affinity.cookie.path =
+          StringRef{first + str_size("affinity-cookie-path="), end};
     } else if (util::strieq_l("dns", param)) {
       out.dns = true;
     } else if (util::strieq_l("redirect-if-not-tls", param)) {
@@ -918,6 +932,13 @@ int parse_mapping(Config *config, DownstreamAddrConfig &addr,
     return -1;
   }
 
+  if (params.affinity.type == AFFINITY_COOKIE &&
+      params.affinity.cookie.name.empty()) {
+    LOG(ERROR) << "backend: affinity-cookie-name is mandatory if "
+                  "affinity=cookie is specified";
+    return -1;
+  }
+
   addr.fall = params.fall;
   addr.rise = params.rise;
   addr.proto = params.proto;
@@ -962,8 +983,24 @@ int parse_mapping(Config *config, DownstreamAddrConfig &addr,
       if (g.pattern == pattern) {
         // Last value wins if we have multiple different affinity
         // value under one group.
-        if (params.affinity != AFFINITY_NONE) {
+        if (params.affinity.type != AFFINITY_NONE) {
-          g.affinity = params.affinity;
+          if (g.affinity.type == AFFINITY_NONE) {
+            g.affinity.type = params.affinity.type;
+            if (params.affinity.type == AFFINITY_COOKIE) {
+              g.affinity.cookie.name = make_string_ref(
+                  downstreamconf.balloc, params.affinity.cookie.name);
+              if (!params.affinity.cookie.path.empty()) {
+                g.affinity.cookie.path = make_string_ref(
+                    downstreamconf.balloc, params.affinity.cookie.path);
+              }
+            }
+          } else if (g.affinity.type != params.affinity.type ||
+                     g.affinity.cookie.name != params.affinity.cookie.name ||
+                     g.affinity.cookie.path != params.affinity.cookie.path) {
+            LOG(ERROR) << "backend: affinity: multiple different affinity "
+                          "configurations found in a single group";
+            return -1;
+          }
         }
         // If at least one backend requires frontend TLS connection,
         // enable it for all backends sharing the same pattern.
@@ -983,7 +1020,15 @@ int parse_mapping(Config *config, DownstreamAddrConfig &addr,
     addr_groups.emplace_back(pattern);
     auto &g = addr_groups.back();
     g.addrs.push_back(addr);
-    g.affinity = params.affinity;
+    g.affinity.type = params.affinity.type;
+    if (params.affinity.type == AFFINITY_COOKIE) {
+      g.affinity.cookie.name =
+          make_string_ref(downstreamconf.balloc, params.affinity.cookie.name);
+      if (!params.affinity.cookie.path.empty()) {
+        g.affinity.cookie.path =
+            make_string_ref(downstreamconf.balloc, params.affinity.cookie.path);
+      }
+    }
     g.redirect_if_not_tls = params.redirect_if_not_tls;
 
     if (pattern[0] == '*') {
@@ -3823,7 +3868,7 @@ int configure_downstream_group(Config *config, bool http2_proxy,
       }
     }
 
-    if (g.affinity == AFFINITY_IP) {
+    if (g.affinity.type != AFFINITY_NONE) {
       size_t idx = 0;
       for (auto &addr : g.addrs) {
         StringRef key;

src/shrpx_config.h

@@ -356,6 +356,19 @@ enum shrpx_session_affinity {
   AFFINITY_NONE,
   // Client IP affinity
   AFFINITY_IP,
+  // Cookie based affinity
+  AFFINITY_COOKIE,
+};
+
+struct AffinityConfig {
+  // Type of session affinity.
+  shrpx_session_affinity type;
+  struct {
+    // Name of a cookie to use.
+    StringRef name;
+    // Path which a cookie is applied to.
+    StringRef path;
+  } cookie;
 };
 
 enum shrpx_forwarded_param {
@@ -449,15 +462,15 @@ struct AffinityHash {
 
 struct DownstreamAddrGroupConfig {
   DownstreamAddrGroupConfig(const StringRef &pattern)
-      : pattern(pattern), affinity(AFFINITY_NONE), redirect_if_not_tls(false) {}
+      : pattern(pattern), affinity{AFFINITY_NONE}, redirect_if_not_tls(false) {}
 
   StringRef pattern;
   std::vector<DownstreamAddrConfig> addrs;
   // Bunch of session affinity hash.  Only used if affinity ==
   // AFFINITY_IP.
   std::vector<AffinityHash> affinity_hash;
-  // Session affinity
+  // Cookie based session affinity configuration.
-  shrpx_session_affinity affinity;
+  AffinityConfig affinity;
   // true if this group requires that client connection must be TLS,
   // and the request must be redirected to https URI.
   bool redirect_if_not_tls;

src/shrpx_downstream.cc

@@ -131,6 +131,7 @@ Downstream::Downstream(Upstream *upstream, MemchunkPool *mcpool,
       assoc_stream_id_(-1),
       downstream_stream_id_(-1),
       response_rst_stream_error_code_(NGHTTP2_NO_ERROR),
+      affinity_cookie_(0),
       request_state_(INITIAL),
       response_state_(INITIAL),
       dispatch_state_(DISPATCH_NONE),
@@ -140,7 +141,8 @@ Downstream::Downstream(Upstream *upstream, MemchunkPool *mcpool,
       expect_final_response_(false),
       request_pending_(false),
       request_header_sent_(false),
-      accesslog_written_(false) {
+      accesslog_written_(false),
+      new_affinity_cookie_(false) {
 
   auto &timeoutconf = get_config()->http2.timeout;
 
@@ -305,6 +307,49 @@ StringRef Downstream::assemble_request_cookie() {
   return StringRef{iov.base, p};
 }
 
+uint32_t Downstream::find_affinity_cookie(const StringRef &name) {
+  for (auto &kv : req_.fs.headers()) {
+    if (kv.token != http2::HD_COOKIE) {
+      continue;
+    }
+
+    for (auto it = std::begin(kv.value); it != std::end(kv.value);) {
+      if (*it == '\t' || *it == ' ' || *it == ';') {
+        ++it;
+        continue;
+      }
+
+      auto end = std::find(it, std::end(kv.value), '=');
+      if (end == std::end(kv.value)) {
+        return 0;
+      }
+
+      if (!util::streq(name, StringRef{it, end})) {
+        it = std::find(it, std::end(kv.value), ';');
+        continue;
+      }
+
+      it = std::find(end + 1, std::end(kv.value), ';');
+      auto val = StringRef{end + 1, it};
+      if (val.size() != 8) {
+        return 0;
+      }
+      uint32_t h = 0;
+      for (auto c : val) {
+        auto n = util::hex_to_uint(c);
+        if (n == 256) {
+          return 0;
+        }
+        h <<= 4;
+        h += n;
+      }
+      affinity_cookie_ = h;
+      return h;
+    }
+  }
+  return 0;
+}
+
 size_t Downstream::count_crumble_request_cookie() {
   size_t n = 0;
   for (auto &kv : req_.fs.headers()) {
@@ -995,4 +1040,16 @@ const DownstreamAddr *Downstream::get_addr() const { return addr_; }
 
 void Downstream::set_accesslog_written(bool f) { accesslog_written_ = f; }
 
+void Downstream::renew_affinity_cookie(uint32_t h) {
+  affinity_cookie_ = h;
+  new_affinity_cookie_ = true;
+}
+
+uint32_t Downstream::get_affinity_cookie_to_send() const {
+  if (new_affinity_cookie_) {
+    return affinity_cookie_;
+  }
+  return 0;
+}
+
 } // namespace shrpx

src/shrpx_downstream.h

@@ -412,6 +412,18 @@ public:
 
   void set_accesslog_written(bool f);
 
+  // Finds affinity cookie from request header fields.  The name of
+  // cookie is given in |name|.  If an affinity cookie is found, it is
+  // assigned to a member function, and is returned.  If it is not
+  // found, or is malformed, returns 0.
+  uint32_t find_affinity_cookie(const StringRef &name);
+  // Set |h| as affinity cookie.
+  void renew_affinity_cookie(uint32_t h);
+  // Returns affinity cookie to send.  If it does not need to be sent,
+  // for example, because the value is retrieved from a request header
+  // field, returns 0.
+  uint32_t get_affinity_cookie_to_send() const;
+
   enum {
     EVENT_ERROR = 0x1,
     EVENT_TIMEOUT = 0x2,
@@ -474,6 +486,8 @@ private:
   int32_t downstream_stream_id_;
   // RST_STREAM error_code from downstream HTTP2 connection
   uint32_t response_rst_stream_error_code_;
+  // An affinity cookie value.
+  uint32_t affinity_cookie_;
   // request state
   int request_state_;
   // response state
@@ -497,6 +511,8 @@ private:
   bool request_header_sent_;
   // true if access.log has been written.
   bool accesslog_written_;
+  // true if affinity cookie is generated for this request.
+  bool new_affinity_cookie_;
 };
 
 } // namespace shrpx

src/shrpx_downstream_test.cc

@@ -189,4 +189,38 @@ void test_downstream_supports_non_final_response(void) {
   CU_ASSERT(!d.supports_non_final_response());
 }
 
+void test_downstream_find_affinity_cookie(void) {
+  Downstream d(nullptr, nullptr, 0);
+
+  auto &req = d.request();
+  req.fs.add_header_token(StringRef::from_lit("cookie"), StringRef{}, false,
+                          http2::HD_COOKIE);
+  req.fs.add_header_token(StringRef::from_lit("cookie"),
+                          StringRef::from_lit("a=b;;c=d"), false,
+                          http2::HD_COOKIE);
+  req.fs.add_header_token(StringRef::from_lit("content-length"),
+                          StringRef::from_lit("599"), false,
+                          http2::HD_CONTENT_LENGTH);
+  req.fs.add_header_token(StringRef::from_lit("cookie"),
+                          StringRef::from_lit("lb=deadbeef;LB=f1f2f3f4"), false,
+                          http2::HD_COOKIE);
+  req.fs.add_header_token(StringRef::from_lit("cookie"),
+                          StringRef::from_lit("short=e1e2e3e"), false,
+                          http2::HD_COOKIE);
+
+  uint32_t aff;
+
+  aff = d.find_affinity_cookie(StringRef::from_lit("lb"));
+
+  CU_ASSERT(0xdeadbeef == aff);
+
+  aff = d.find_affinity_cookie(StringRef::from_lit("LB"));
+
+  CU_ASSERT(0xf1f2f3f4 == aff);
+
+  aff = d.find_affinity_cookie(StringRef::from_lit("short"));
+
+  CU_ASSERT(0 == aff);
+}
+
 } // namespace shrpx

src/shrpx_downstream_test.h

@@ -37,6 +37,7 @@ void test_downstream_crumble_request_cookie(void);
 void test_downstream_assemble_request_cookie(void);
 void test_downstream_rewrite_location_response_header(void);
 void test_downstream_supports_non_final_response(void);
+void test_downstream_find_affinity_cookie(void);
 
 } // namespace shrpx
 

src/shrpx_http.cc

@@ -164,6 +164,41 @@ ssize_t select_padding_callback(nghttp2_session *session,
   return std::min(max_payload, frame->hd.length + get_config()->padding);
 }
 
+StringRef create_affinity_cookie(BlockAllocator &balloc, const StringRef &name,
+                                 uint32_t affinity_cookie,
+                                 const StringRef &path, bool secure) {
+  static constexpr auto PATH_PREFIX = StringRef::from_lit("; Path=");
+  static constexpr auto SECURE = StringRef::from_lit("; Secure");
+  // <name>=<value>[; Path=<path>][; Secure]
+  size_t len = name.size() + 1 + 8;
+
+  if (!path.empty()) {
+    len += PATH_PREFIX.size() + path.size();
+  }
+  if (secure) {
+    len += SECURE.size();
+  }
+
+  auto iov = make_byte_ref(balloc, len + 1);
+  auto p = iov.base;
+  p = std::copy(std::begin(name), std::end(name), p);
+  *p++ = '=';
+  affinity_cookie = htonl(affinity_cookie);
+  p = util::format_hex(p,
+                       StringRef{reinterpret_cast<uint8_t *>(&affinity_cookie),
+                                 reinterpret_cast<uint8_t *>(&affinity_cookie) +
+                                     sizeof(affinity_cookie)});
+  if (!path.empty()) {
+    p = std::copy(std::begin(PATH_PREFIX), std::end(PATH_PREFIX), p);
+    p = std::copy(std::begin(path), std::end(path), p);
+  }
+  if (secure) {
+    p = std::copy(std::begin(SECURE), std::end(SECURE), p);
+  }
+  *p = '\0';
+  return StringRef{iov.base, p};
+}
+
 } // namespace http
 
 } // namespace shrpx

src/shrpx_http.h

@@ -66,6 +66,13 @@ ssize_t select_padding_callback(nghttp2_session *session,
                                 const nghttp2_frame *frame, size_t max_payload,
                                 void *user_data);
 
+// Creates set-cookie-string for cookie based affinity.  If |path| is
+// not empty, "; <path>" is added.  If |secure| is true, "; Secure" is
+// added.
+StringRef create_affinity_cookie(BlockAllocator &balloc, const StringRef &name,
+                                 uint32_t affinity_cookie,
+                                 const StringRef &path, bool secure);
+
 } // namespace http
 
 } // namespace shrpx

src/shrpx_http2_upstream.cc

@@ -1657,9 +1657,9 @@ int Http2Upstream::on_downstream_header_complete(Downstream *downstream) {
   }
 
   auto nva = std::vector<nghttp2_nv>();
-  // 4 means :status and possible server, via and x-http2-push header
+  // 5 means :status and possible server, via, x-http2-push, and
-  // field.
+  // set-cookie (for affinity cookie) header field.
-  nva.reserve(resp.fs.headers().size() + 4 +
+  nva.reserve(resp.fs.headers().size() + 5 +
               httpconf.add_response_headers.size());
 
   auto response_status = http2::stringify_status(balloc, resp.http_status);
@@ -1700,6 +1700,21 @@ int Http2Upstream::on_downstream_header_complete(Downstream *downstream) {
     }
   }
 
+  if (req.method != HTTP_CONNECT || !downstream->get_upgraded()) {
+    auto affinity_cookie = downstream->get_affinity_cookie_to_send();
+    if (affinity_cookie) {
+      auto dconn = downstream->get_downstream_connection();
+      assert(dconn);
+      auto &group = dconn->get_downstream_addr_group();
+      auto &shared_addr = group->shared_addr;
+      auto &cookieconf = shared_addr->affinity.cookie;
+      auto cookie_str =
+          http::create_affinity_cookie(balloc, cookieconf.name, affinity_cookie,
+                                       cookieconf.path, req.scheme == "https");
+      nva.push_back(http2::make_nv_ls_nocopy("set-cookie", cookie_str));
+    }
+  }
+
   auto via = resp.fs.header(http2::HD_VIA);
   if (httpconf.no_via) {
     if (via) {

src/shrpx_http_downstream_connection.cc

@@ -269,8 +269,9 @@ int HttpDownstreamConnection::initiate_connection() {
     // initial_addr_idx_.
     size_t temp_idx = initial_addr_idx_;
 
-    auto &next_downstream =
+    auto &next_downstream = shared_addr->affinity.type == AFFINITY_NONE
-        shared_addr->affinity == AFFINITY_NONE ? shared_addr->next : temp_idx;
+                                ? shared_addr->next
+                                : temp_idx;
     auto end = next_downstream;
     for (;;) {
       auto check_dns_result = dns_query_.get() != nullptr;
@@ -757,7 +758,7 @@ void remove_from_pool(HttpDownstreamConnection *dconn) {
   auto &group = dconn->get_downstream_addr_group();
   auto &shared_addr = group->shared_addr;
 
-  if (shared_addr->affinity == AFFINITY_NONE) {
+  if (shared_addr->affinity.type == AFFINITY_NONE) {
     auto &dconn_pool =
         dconn->get_downstream_addr_group()->shared_addr->dconn_pool;
     dconn_pool.remove_downstream_connection(dconn);

src/shrpx_http_test.cc

@@ -91,4 +91,31 @@ void test_shrpx_http_create_via_header_value(void) {
   CU_ASSERT(("2 nghttpx" == StringRef{std::begin(buf), end}));
 }
 
+void test_shrpx_http_create_affinity_cookie(void) {
+  BlockAllocator balloc(1024, 1024);
+  StringRef c;
+
+  c = http::create_affinity_cookie(balloc, StringRef::from_lit("cookie-val"),
+                                   0xf1e2d3c4u, StringRef{}, false);
+
+  CU_ASSERT("cookie-val=f1e2d3c4" == c);
+
+  c = http::create_affinity_cookie(balloc, StringRef::from_lit("alpha"),
+                                   0x00000000u, StringRef{}, true);
+
+  CU_ASSERT("alpha=00000000; Secure" == c);
+
+  c = http::create_affinity_cookie(balloc, StringRef::from_lit("bravo"),
+                                   0x01111111u, StringRef::from_lit("bar"),
+                                   false);
+
+  CU_ASSERT("bravo=01111111; Path=bar" == c);
+
+  c = http::create_affinity_cookie(balloc, StringRef::from_lit("charlie"),
+                                   0x01111111u, StringRef::from_lit("bar"),
+                                   true);
+
+  CU_ASSERT("charlie=01111111; Path=bar; Secure" == c);
+}
+
 } // namespace shrpx

src/shrpx_http_test.h

@@ -33,6 +33,7 @@ namespace shrpx {
 
 void test_shrpx_http_create_forwarded(void);
 void test_shrpx_http_create_via_header_value(void);
+void test_shrpx_http_create_affinity_cookie(void);
 
 } // namespace shrpx
 

src/shrpx_https_upstream.cc

@@ -1147,6 +1147,23 @@ int HttpsUpstream::on_downstream_header_complete(Downstream *downstream) {
     }
   }
 
+  if (req.method != HTTP_CONNECT || !downstream->get_upgraded()) {
+    auto affinity_cookie = downstream->get_affinity_cookie_to_send();
+    if (affinity_cookie) {
+      auto dconn = downstream->get_downstream_connection();
+      assert(dconn);
+      auto &group = dconn->get_downstream_addr_group();
+      auto &shared_addr = group->shared_addr;
+      auto &cookieconf = shared_addr->affinity.cookie;
+      auto cookie_str =
+          http::create_affinity_cookie(balloc, cookieconf.name, affinity_cookie,
+                                       cookieconf.path, req.scheme == "https");
+      buf->append("Set-Cookie: ");
+      buf->append(cookie_str);
+      buf->append("\r\n");
+    }
+  }
+
   auto via = resp.fs.header(http2::HD_VIA);
   if (httpconf.no_via) {
     if (via) {

src/shrpx_worker.cc

@@ -76,11 +76,17 @@ bool match_shared_downstream_addr(
     return false;
   }
 
-  if (lhs->affinity != rhs->affinity ||
+  if (lhs->affinity.type != rhs->affinity.type ||
       lhs->redirect_if_not_tls != rhs->redirect_if_not_tls) {
     return false;
   }
 
+  if (lhs->affinity.type == AFFINITY_COOKIE &&
+      (lhs->affinity.cookie.name != rhs->affinity.cookie.name ||
+       lhs->affinity.cookie.path != rhs->affinity.cookie.path)) {
+    return false;
+  }
+
   auto used = std::vector<bool>(lhs->addrs.size());
 
   for (auto &a : lhs->addrs) {
@@ -156,7 +162,7 @@ void Worker::replace_downstream_config(
 
     auto &shared_addr = g->shared_addr;
 
-    if (shared_addr->affinity == AFFINITY_NONE) {
+    if (shared_addr->affinity.type == AFFINITY_NONE) {
       shared_addr->dconn_pool.remove_all();
       continue;
     }
@@ -186,7 +192,15 @@ void Worker::replace_downstream_config(
     auto shared_addr = std::make_shared<SharedDownstreamAddr>();
 
     shared_addr->addrs.resize(src.addrs.size());
-    shared_addr->affinity = src.affinity;
+    shared_addr->affinity.type = src.affinity.type;
+    if (src.affinity.type == AFFINITY_COOKIE) {
+      shared_addr->affinity.cookie.name =
+          make_string_ref(shared_addr->balloc, src.affinity.cookie.name);
+      if (!src.affinity.cookie.path.empty()) {
+        shared_addr->affinity.cookie.path =
+            make_string_ref(shared_addr->balloc, src.affinity.cookie.path);
+      }
+    }
     shared_addr->affinity_hash = src.affinity_hash;
     shared_addr->redirect_if_not_tls = src.redirect_if_not_tls;
 
@@ -268,7 +282,7 @@ void Worker::replace_downstream_config(
       shared_addr->http1_pri.weight = num_http1;
       shared_addr->http2_pri.weight = num_http2;
 
-      if (shared_addr->affinity != AFFINITY_NONE) {
+      if (shared_addr->affinity.type != AFFINITY_NONE) {
         for (auto &addr : shared_addr->addrs) {
           addr.dconn_pool = make_unique<DownstreamConnectionPool>();
         }

src/shrpx_worker.h

@@ -133,10 +133,10 @@ struct WeightedPri {
 struct SharedDownstreamAddr {
   SharedDownstreamAddr()
       : balloc(1024, 1024),
+        affinity{AFFINITY_NONE},
         next{0},
         http1_pri{},
         http2_pri{},
-        affinity{AFFINITY_NONE},
         redirect_if_not_tls{false} {}
 
   SharedDownstreamAddr(const SharedDownstreamAddr &) = delete;
@@ -161,6 +161,8 @@ struct SharedDownstreamAddr {
   // wise.
   DList<Http2Session> http2_avail_freelist;
   DownstreamConnectionPool dconn_pool;
+  // Configuration for session affinity
+  AffinityConfig affinity;
   // Next http/1.1 downstream address index in addrs.
   size_t next;
   // http1_pri and http2_pri are used to which protocols are used
@@ -171,7 +173,6 @@ struct SharedDownstreamAddr {
   WeightedPri http1_pri;
   WeightedPri http2_pri;
   // Session affinity
-  shrpx_session_affinity affinity;
   // true if this group requires that client connection must be TLS,
   // and the request must be redirected to https URI.
   bool redirect_if_not_tls;

src/util.cc

@@ -189,7 +189,7 @@ uint32_t hex_to_uint(char c) {
   if (c <= 'z') {
     return c - 'a' + 10;
   }
-  return c;
+  return 256;
 }
 
 StringRef quote_string(BlockAllocator &balloc, const StringRef &target) {
@@ -415,10 +415,6 @@ char upcase(char c) {
   }
 }
 
-namespace {
-constexpr char LOWER_XDIGITS[] = "0123456789abcdef";
-} // namespace
-
 std::string format_hex(const unsigned char *s, size_t len) {
   std::string res;
   res.resize(len * 2);

src/util.h

@@ -96,8 +96,8 @@ bool in_token(char c);
 
 bool in_attr_char(char c);
 
-// Returns integer corresponding to hex notation |c|.  It is undefined
+// Returns integer corresponding to hex notation |c|.  If
-// if is_hex_digit(c) is false.
+// is_hex_digit(c) is false, it returns 256.
 uint32_t hex_to_uint(char c);
 
 std::string percent_encode(const unsigned char *target, size_t len);
@@ -152,6 +152,19 @@ template <size_t N> std::string format_hex(const std::array<uint8_t, N> &s) {
 
 StringRef format_hex(BlockAllocator &balloc, const StringRef &s);
 
+static constexpr char LOWER_XDIGITS[] = "0123456789abcdef";
+
+template <typename OutputIt>
+OutputIt format_hex(OutputIt it, const StringRef &s) {
+  for (auto cc : s) {
+    uint8_t c = cc;
+    *it++ = LOWER_XDIGITS[c >> 4];
+    *it++ = LOWER_XDIGITS[c & 0xf];
+  }
+
+  return it;
+}
+
 // decode_hex decodes hex string |s|, returns the decoded byte string.
 // This function assumes |s| is hex string, that is is_hex_string(s)
 // == true.