On client CloseWrite, do CloseWrite instead of Close for backend

Co-authored-by: Mathieu Lonjaret <mathieu.lonjaret@gmail.com>

pkg/server/server_entrypoint_tcp.go

@@ -37,7 +37,7 @@ func newHTTPForwarder(ln net.Listener) *httpForwarder {
 }
 
 // ServeTCP uses the connection to serve it later in "Accept"
-func (h *httpForwarder) ServeTCP(conn net.Conn) {
+func (h *httpForwarder) ServeTCP(conn tcp.WriteCloser) {
 	h.connChan <- conn
 }
 
@@ -99,7 +99,36 @@ func NewTCPEntryPoint(ctx context.Context, configuration *static.EntryPoint) (*T
 	}, nil
 }
 
+// writeCloserWrapper wraps together a connection, and the concrete underlying
+// connection type that was found to satisfy WriteCloser.
+type writeCloserWrapper struct {
+	net.Conn
+	writeCloser tcp.WriteCloser
+}
+
+func (c *writeCloserWrapper) CloseWrite() error {
+	return c.writeCloser.CloseWrite()
+}
+
+// writeCloser returns the given connection, augmented with the WriteCloser
+// implementation, if any was found within the underlying conn.
+func writeCloser(conn net.Conn) (tcp.WriteCloser, error) {
+	switch typedConn := conn.(type) {
+	case *proxyprotocol.Conn:
+		underlying, err := writeCloser(typedConn.Conn)
+		if err != nil {
+			return nil, err
+		}
+		return &writeCloserWrapper{writeCloser: underlying, Conn: typedConn}, nil
+	case *net.TCPConn:
+		return typedConn, nil
+	default:
+		return nil, fmt.Errorf("unknown connection type %T", typedConn)
+	}
+}
+
 func (e *TCPEntryPoint) startTCP(ctx context.Context) {
+
 	log.FromContext(ctx).Debugf("Start TCP Server")
 
 	for {
@@ -109,8 +138,13 @@ func (e *TCPEntryPoint) startTCP(ctx context.Context) {
 			return
 		}
 
+		writeCloser, err := writeCloser(conn)
+		if err != nil {
+			panic(err)
+		}
+
 		safe.Go(func() {
-			e.switcher.ServeTCP(newTrackedConnection(conn, e.tracker))
+			e.switcher.ServeTCP(newTrackedConnection(writeCloser, e.tracker))
 		})
 	}
 }
@@ -374,20 +408,20 @@ func createHTTPServer(ln net.Listener, configuration *static.EntryPoint, withH2c
 	}, nil
 }
 
-func newTrackedConnection(conn net.Conn, tracker *connectionTracker) *trackedConnection {
+func newTrackedConnection(conn tcp.WriteCloser, tracker *connectionTracker) *trackedConnection {
 	tracker.AddConnection(conn)
 	return &trackedConnection{
-		Conn:    conn,
-		tracker: tracker,
+		WriteCloser: conn,
+		tracker:     tracker,
 	}
 }
 
 type trackedConnection struct {
 	tracker *connectionTracker
-	net.Conn
+	tcp.WriteCloser
 }
 
 func (t *trackedConnection) Close() error {
-	t.tracker.RemoveConnection(t.Conn)
-	return t.Conn.Close()
+	t.tracker.RemoveConnection(t.WriteCloser)
+	return t.WriteCloser.Close()
 }