add fix on customClieent

Author: lupin012

internal/perf/vegeta.go

@@ -278,6 +278,43 @@ func (pt *PerfTest) loadTargets(filepath string) ([]vegeta.Target, error) {
 	return targets, nil
 }
 
+// limitBodyTransport wraps an http.RoundTripper and pre-limits response bodies.
+//
+// Vegeta v12 library reads maxBody bytes then drains the remainder with
+// io.Copy(io.Discard, r.Body). If the server resets the connection during
+// that drain, Vegeta records the result as failed (Code=0) even though a 200
+// OK was already received. The Python/vegeta-CLI binary (older) skips the
+// drain step and counts 200 as success immediately.
+//
+// By capping resp.Body at maxBytes here, the subsequent drain reads 0 bytes
+// and succeeds, so res.Code is set to the actual HTTP status, matching the
+// Python/CLI behavior: success = received 2xx status.
+type limitBodyTransport struct {
+	base     http.RoundTripper
+	maxBytes int64
+}
+
+func (t *limitBodyTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	resp, err := t.base.RoundTrip(req)
+	if err != nil || resp == nil {
+		return resp, err
+	}
+	if resp.Body != nil {
+		resp.Body = &limitedReadCloser{
+			Reader: io.LimitReader(resp.Body, t.maxBytes),
+			Closer: resp.Body,
+		}
+	}
+	return resp, nil
+}
+
+// limitedReadCloser combines a limited Reader with the original Closer so
+// resp.Body.Close() still reaches the underlying TCP connection.
+type limitedReadCloser struct {
+	io.Reader
+	io.Closer
+}
+
 // runVegetaAttack executes a Vegeta attack using the library.
 func (pt *PerfTest) runVegetaAttack(ctx context.Context, targets []vegeta.Target, qps int, duration time.Duration, outputFile string) (*vegeta.Metrics, error) {
 	rate := vegeta.Rate{Freq: qps, Per: time.Second}
@@ -296,7 +333,7 @@ func (pt *PerfTest) runVegetaAttack(ctx context.Context, targets []vegeta.Target
 	}
 
 	customClient := &http.Client{
-		Transport: tr,
+		Transport: &limitBodyTransport{base: tr, maxBytes: int64(maxBodyInt)},
 	}
 
         //