chore: improve code

nxtreaming · Jul 5, 2024 · 169cd86 · 169cd86
1 parent 439d70a
commit 169cd86
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/internal/transporter/base.go b/internal/transporter/base.go
@@ -87,18 +87,18 @@ func (b *baseTransporter) RelayTCPConn(c net.Conn, handshakeF TCPHandShakeF) err
 		}
 	}
 
+	// rate limit
+	if b.cfg.MaxReadRateKbps > 0 {
+		c = conn.NewRateLimitedConn(c, b.cfg.MaxReadRateKbps)
+	}
+
 	clonedRemote := remote.Clone()
 	rc, err := handshakeF(clonedRemote)
 	if err != nil {
 		return err
 	}
 	defer rc.Close()
 
-	// rate limit
-	if b.cfg.MaxReadRateKbps > 0 {
-		c = conn.NewRateLimitedConn(c, b.cfg.MaxReadRateKbps)
-	}
-
 	b.l.Infof("RelayTCPConn from %s to %s", c.LocalAddr(), remote.Address)
 	relayConn := conn.NewRelayConn(
 		b.cfg.Label, c, rc, conn.WithHandshakeDuration(clonedRemote.HandShakeDuration))