use a buffer to fix the header rewriting

mirror/header.go

@@ -2,10 +2,13 @@ package mirror
 
 import (
 	"bufio"
+	"bytes"
 	"crypto/tls"
+	"io"
 	"log"
 	"net"
 	"net/http"
+	"time"
 )
 
 // AddHeaders adds headers to the connection.
@@ -13,30 +16,56 @@ import (
 // It only adds headers if the connection is an HTTP connection.
 // It returns a net.Conn with the headers added.
 func AddHeaders(conn net.Conn, headers map[string]string) net.Conn {
-	// read a request from the connection
-	// if the request is an HTTP request, add the headers
-	// if the request is not an HTTP request, return the connection as is
-	req, err := http.ReadRequest(bufio.NewReader(conn))
+	// Create a buffer to store the original request
+	var buf bytes.Buffer
+	teeReader := io.TeeReader(conn, &buf)
+
+	// Try to read the request, but also save it to our buffer
+	req, err := http.ReadRequest(bufio.NewReader(teeReader))
 	if err != nil {
-		log.Println("Error reading request:", err)
-		// if the request is not an HTTP request, return the connection as is
+		// Not an HTTP request or couldn't parse, return original connection
 		return conn
 	}
-	log.Println("Adding headers to connection:", req.Method, req.URL)
+
+	// Add our headers
 	for key, value := range headers {
 		req.Header.Add(key, value)
-		log.Println("Added header:", key, value)
 	}
-	// write the request back to the connection
-	if err := req.Write(conn); err != nil {
-		log.Println("Error writing request:", err)
-		// if there is an error writing the request, return the connection as is
-		return conn
+
+	// Create a pipe to connect our modified request with the output
+	pr, pw := io.Pipe()
+
+	// Write the modified request to one end of the pipe
+	go func() {
+		req.Write(pw)
+		// Then copy the rest of the original connection
+		io.Copy(pw, conn)
+		pw.Close()
+	}()
+
+	// Return a ReadWriter that reads from our pipe and writes to the original connection
+	return &readWriteConn{
+		Reader: pr,
+		Writer: conn,
+		conn:   conn,
 	}
-	// If all goes well, return the connection with the headers added
-	return conn
 }
 
+// readWriteConn implements net.Conn
+type readWriteConn struct {
+	io.Reader
+	io.Writer
+	conn net.Conn
+}
+
+// Implement the rest of net.Conn interface by delegating to the original connection
+func (rwc *readWriteConn) Close() error                       { return rwc.conn.Close() }
+func (rwc *readWriteConn) LocalAddr() net.Addr                { return rwc.conn.LocalAddr() }
+func (rwc *readWriteConn) RemoteAddr() net.Addr               { return rwc.conn.RemoteAddr() }
+func (rwc *readWriteConn) SetDeadline(t time.Time) error      { return rwc.conn.SetDeadline(t) }
+func (rwc *readWriteConn) SetReadDeadline(t time.Time) error  { return rwc.conn.SetReadDeadline(t) }
+func (rwc *readWriteConn) SetWriteDeadline(t time.Time) error { return rwc.conn.SetWriteDeadline(t) }
+
 // Accept accepts a connection from the listener.
 // It takes a net.Listener as input and returns a net.Conn with the headers added.
 // It is used to accept connections from the meta listener and add headers to them.