From 0fb2da902b8fe7ced5d86730ae1b9a540ba21b00 Mon Sep 17 00:00:00 2001
From: Minhan Xia <mixia@google.com>
Date: Tue, 6 Nov 2018 13:27:22 -0800
Subject: [PATCH] add ability to close connection for serve-host

---
 test/images/serve-hostname/VERSION           |  2 +-
 test/images/serve-hostname/serve_hostname.go | 15 +++++++++++----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/test/images/serve-hostname/VERSION b/test/images/serve-hostname/VERSION
index 9459d4ba2a..5625e59da8 100644
--- a/test/images/serve-hostname/VERSION
+++ b/test/images/serve-hostname/VERSION
@@ -1 +1 @@
-1.1
+1.2
diff --git a/test/images/serve-hostname/serve_hostname.go b/test/images/serve-hostname/serve_hostname.go
index 21793fbbf9..f6d7507e1f 100644
--- a/test/images/serve-hostname/serve_hostname.go
+++ b/test/images/serve-hostname/serve_hostname.go
@@ -30,10 +30,11 @@ import (
 )
 
 var (
-	doTCP  = flag.Bool("tcp", false, "Serve raw over TCP.")
-	doUDP  = flag.Bool("udp", false, "Serve raw over UDP.")
-	doHTTP = flag.Bool("http", true, "Serve HTTP.")
-	port   = flag.Int("port", 9376, "Port number.")
+	doTCP   = flag.Bool("tcp", false, "Serve raw over TCP.")
+	doUDP   = flag.Bool("udp", false, "Serve raw over UDP.")
+	doHTTP  = flag.Bool("http", true, "Serve HTTP.")
+	doClose = flag.Bool("close", false, "Close connection per each HTTP request")
+	port    = flag.Int("port", 9376, "Port number.")
 )
 
 func main() {
@@ -88,6 +89,12 @@ func main() {
 	if *doHTTP {
 		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 			log.Printf("HTTP request from %s", r.RemoteAddr)
+
+			if *doClose {
+				// Add this header to force to close the connection after serving the request.
+				w.Header().Add("Connection", "close")
+			}
+
 			fmt.Fprintf(w, "%s", hostname)
 		})
 		go func() {