fix: defer WebSocket logs until handshake delivery

proxy/proxy.go

@@ -209,7 +209,12 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		tracked := h.trackConn(conn)
 		if upgrade := st.upgrade.Load(); upgrade != nil {
 			cw.SetHijackedResponse(upgrade.status, upgrade.header)
-			submitOnce.Do(func() { submit(time.Now()) })
+			return &headerDeliveryConn{
+				Conn: tracked,
+				onDelivered: func() {
+					submitOnce.Do(func() { submit(time.Now()) })
+				},
+			}
 		}
 		return tracked
 	})

proxy/writer.go

@@ -97,6 +97,7 @@ func (c *captureWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 	c.hijacked = true
 	if c.onHijack != nil {
 		conn = c.onHijack(conn)
+		rw = bufio.NewReadWriter(rw.Reader, bufio.NewWriter(conn))
 	}
 	return conn, rw, nil
 }
@@ -131,3 +132,26 @@ func (c *trackedConn) Close() error {
 	c.closeOnce.Do(c.onClose)
 	return err
 }
+
+type headerDeliveryConn struct {
+	net.Conn
+	tail        []byte
+	onDelivered func()
+	delivered   bool
+}
+
+func (c *headerDeliveryConn) Write(p []byte) (int, error) {
+	n, err := c.Conn.Write(p)
+	if c.delivered || err != nil || n == 0 {
+		return n, err
+	}
+	c.tail = append(c.tail, p[:n]...)
+	if bytes.Contains(c.tail, []byte("\r\n\r\n")) {
+		c.delivered = true
+		c.tail = nil
+		c.onDelivered()
+	} else if len(c.tail) > 3 {
+		c.tail = append(c.tail[:0], c.tail[len(c.tail)-3:]...)
+	}
+	return n, err
+}