From 5c0f551a6d70deb90e18c7993787b648b1aa2726 Mon Sep 17 00:00:00 2001
From: Marco Munizaga <git@marcopolo.io>
Date: Tue, 2 Apr 2024 20:09:01 -0700
Subject: [PATCH] Return connection: close when doing http over streams

---
 p2p/http/libp2phttp.go      | 10 +++++++++-
 p2p/http/libp2phttp_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/p2p/http/libp2phttp.go b/p2p/http/libp2phttp.go
index fbd8a64f4e..9b6de1ac05 100644
--- a/p2p/http/libp2phttp.go
+++ b/p2p/http/libp2phttp.go
@@ -304,7 +304,7 @@ func (h *Host) Serve() error {
 		h.httpTransport.listenAddrs = append(h.httpTransport.listenAddrs, h.StreamHost.Addrs()...)
 
 		go func() {
-			errCh <- http.Serve(listener, h.ServeMux)
+			errCh <- http.Serve(listener, connectionCloseHeaderMiddleware(h.ServeMux))
 		}()
 	}
 
@@ -816,3 +816,11 @@ func (h *Host) RemovePeerMetadata(server peer.ID) {
 	}
 	h.peerMetadata.Remove(server)
 }
+
+func connectionCloseHeaderMiddleware(next http.Handler) http.Handler {
+	// Sets connection: close. It's preferable to not reuse streams for HTTP.
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Connection", "close")
+		next.ServeHTTP(w, r)
+	})
+}
diff --git a/p2p/http/libp2phttp_test.go b/p2p/http/libp2phttp_test.go
index 57dd0cc8e2..cc5788500c 100644
--- a/p2p/http/libp2phttp_test.go
+++ b/p2p/http/libp2phttp_test.go
@@ -68,6 +68,42 @@ func TestHTTPOverStreams(t *testing.T) {
 	require.Equal(t, "hello", string(body))
 }
 
+func TestHTTPOverStreamsReturnsConnectionClose(t *testing.T) {
+	serverHost, err := libp2p.New(
+		libp2p.ListenAddrStrings("/ip4/127.0.0.1/udp/0/quic-v1"),
+	)
+	require.NoError(t, err)
+
+	httpHost := libp2phttp.Host{StreamHost: serverHost}
+
+	httpHost.SetHTTPHandlerAtPath("/hello", "/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("hello"))
+	}))
+
+	// Start server
+	go httpHost.Serve()
+	defer httpHost.Close()
+
+	// Start client
+	clientHost, err := libp2p.New(libp2p.NoListenAddrs)
+	require.NoError(t, err)
+	clientHost.Connect(context.Background(), peer.AddrInfo{
+		ID:    serverHost.ID(),
+		Addrs: serverHost.Addrs(),
+	})
+
+	s, err := clientHost.NewStream(context.Background(), serverHost.ID(), libp2phttp.ProtocolIDForMultistreamSelect)
+	require.NoError(t, err)
+	_, err = s.Write([]byte("GET / HTTP/1.1\r\nHost: \r\n\r\n"))
+	require.NoError(t, err)
+
+	out := make([]byte, 1024)
+	n, err := s.Read(out)
+	require.NoError(t, err)
+
+	require.Contains(t, strings.ToLower(string(out[:n])), "connection: close")
+}
+
 func TestRoundTrippers(t *testing.T) {
 	serverHost, err := libp2p.New(
 		libp2p.ListenAddrStrings("/ip4/127.0.0.1/udp/0/quic-v1"),