libp2phttp: Return connection: close when doing http over streams

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))
 		}()
 	}
 
@@ -422,6 +422,9 @@ func (rt *streamRoundTripper) RoundTrip(r *http.Request) (*http.Response, error)
 		return nil, err
 	}
 
+	// Write connection: close header to ensure the stream is closed after the response
+	r.Header.Add("connection", "close")
+
 	go func() {
 		defer s.CloseWrite()
 		r.Write(s)
@@ -816,3 +819,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)
+	})
+}

p2p/http/libp2phttp_test.go

@@ -68,6 +68,84 @@ func TestHTTPOverStreams(t *testing.T) {
 	require.Equal(t, "hello", string(body))
 }
 
+func TestHTTPOverStreamsSendsConnectionClose(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}
+
+	connectionHeaderVal := make(chan string, 1)
+	httpHost.SetHTTPHandlerAtPath("/hello", "/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte("hello"))
+		connectionHeaderVal <- r.Header.Get("Connection")
+	}))
+
+	// Start server
+	go httpHost.Serve()
+	defer httpHost.Close()
+
+	// run client
+	clientHost, err := libp2p.New(libp2p.NoListenAddrs)
+	require.NoError(t, err)
+	clientHost.Connect(context.Background(), peer.AddrInfo{
+		ID:    serverHost.ID(),
+		Addrs: serverHost.Addrs(),
+	})
+	clientHttpHost := libp2phttp.Host{StreamHost: clientHost}
+	rt, err := clientHttpHost.NewConstrainedRoundTripper(peer.AddrInfo{ID: serverHost.ID()})
+	require.NoError(t, err)
+	client := &http.Client{Transport: rt}
+	_, err = client.Get("/")
+	require.NoError(t, err)
+
+	select {
+	case val := <-connectionHeaderVal:
+		require.Equal(t, "close", strings.ToLower(val))
+	case <-time.After(5 * time.Second):
+		t.Fatal("timeout waiting for connection header")
+	}
+}
+
+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)
+	if err != io.EOF {
+		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"),