add read stub to tcp connections

src/core/packetgen/connection.go

@@ -23,6 +23,7 @@
 package packetgen
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"net"
@@ -40,7 +41,7 @@ type ConnectionConfig struct {
 	Args map[string]any
 }
 
-func OpenConnection(c ConnectionConfig) (Connection, error) {
+func OpenConnection(ctx context.Context, c ConnectionConfig) (Connection, error) {
 	switch c.Type {
 	case "raw":
 		var cfg rawConnConfig
@@ -55,7 +56,7 @@ func OpenConnection(c ConnectionConfig) (Connection, error) {
 			return nil, fmt.Errorf("error decoding connection config: %w", err)
 		}
 
-		return openNetConn(cfg)
+		return openNetConn(ctx, cfg)
 	default:
 		return nil, fmt.Errorf("unknown connection type: %v", c.Type)
 	}
@@ -118,13 +119,32 @@ type netConn struct {
 	target string
 }
 
-func openNetConn(c netConnConfig) (*netConn, error) {
+func readStub(ctx context.Context, conn net.Conn) {
+	const bufSize = 1024
+	buf := make([]byte, bufSize)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+			_, err := conn.Read(buf)
+			if err != nil {
+				return
+			}
+		}
+	}
+}
+
+func openNetConn(ctx context.Context, c netConnConfig) (*netConn, error) {
 	conn, err := utils.GetProxyFunc(c.ProxyURLs, c.Timeout, false)(c.Protocol, c.Address)
 
 	switch {
 	case err != nil:
 		return nil, err
 	case c.TLSClientConfig == nil:
+		go readStub(ctx, conn)
+
 		return &netConn{Conn: conn, target: c.Protocol + "://" + c.Address}, nil
 	}
 
@@ -135,6 +155,8 @@ func openNetConn(c netConnConfig) (*netConn, error) {
 		return nil, err
 	}
 
+	go readStub(ctx, tlsConn)
+
 	return &netConn{Conn: tlsConn, target: c.Protocol + "://" + c.Address}, nil
 }
 

src/job/packetgen.go

@@ -65,7 +65,10 @@ func packetgenJob(ctx context.Context, args config.Args, globalConfig *GlobalCon
 }
 
 func sendPacket(ctx context.Context, logger *zap.Logger, jobConfig *packetgenJobConfig, a *metrics.Accumulator) error {
-	conn, err := packetgen.OpenConnection(jobConfig.Connection)
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	conn, err := packetgen.OpenConnection(ctx, jobConfig.Connection)
 	if err != nil {
 		return err
 	}