Add configurable buffer size

2 files changed, 10 insertions, 8 deletions

diff --git a/cmd/main.go b/cmd/main.go
index 92ec8bd..db3f185 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -19,6 +19,7 @@ var opts struct {
 	Port          int           `short:"p" long:"port" description:"TODO" required:"true"`
 	Verbose       []bool        `short:"v" long:"verbose" description:"TODO"`
 	Latency       time.Duration `short:"l" long:"latency" default:"0ms" description:"TODO"`
+	BufferSize    int           `short:"b" long:"buffer" default:"1500" description:"TODO"`
 }
 
 func init() {
@@ -48,8 +49,8 @@ func configureLogger() {
 }
 
 func main() {
-	log.Infof("Proxying from :%d to %s with latency of %s", opts.Port, opts.TargetAddress, opts.Latency)
-
+	log.Infof("Proxying between :%d and %s with %s of latency", opts.Port, opts.TargetAddress, opts.Latency)
+	log.Debugf("Transfer buffer size set to %d bytes", opts.BufferSize)
 	listener := getListener(opts.Port)
 	waitForClients(listener)
 }
@@ -74,7 +75,7 @@ func waitForClients(listener net.Listener) {
 		}
 		log.Infof("Accepted connection from client %v\n", client.RemoteAddr())
 
-		s := slowpoke.New(client, opts.TargetAddress, opts.Latency, log)
+		s := slowpoke.New(client, opts.TargetAddress, opts.Latency, opts.BufferSize, log)
 
 		go s.StartTransfer()
 	}
diff --git a/slowpoke.go b/slowpoke.go
index 3178196..98182f6 100644
--- a/slowpoke.go
+++ b/slowpoke.go
@@ -12,16 +12,18 @@ type Slowpoke struct {
 	conn       net.Conn
 	targetAddr string
 	latency    time.Duration
+	bufferSize int
 	isClosed   bool
 	close      chan bool
 	logger     *logging.Logger
 }
 
-func New(conn net.Conn, targetAddr string, latency time.Duration, logger *logging.Logger) *Slowpoke {
+func New(conn net.Conn, targetAddr string, latency time.Duration, bufferSize int, logger *logging.Logger) *Slowpoke {
 	return &Slowpoke{
 		conn:       conn,
 		targetAddr: targetAddr,
 		latency:    latency,
+		bufferSize: bufferSize,
 		isClosed:   false,
 		close:      make(chan bool),
 		logger:     logger,
@@ -47,13 +49,12 @@ func (s *Slowpoke) StartTransfer() {
 	s.logger.Infof("Connection between client %s and target %s closed", s.conn.RemoteAddr(), target.RemoteAddr())
 }
 
-func createBuffer() []byte {
-	// TODO configurable
-	return make([]byte, 1500)
+func (s *Slowpoke) createBuffer() []byte {
+	return make([]byte, s.bufferSize)
 }
 
 func (s *Slowpoke) transferWithLatency(source net.Conn, target net.Conn) {
-	byteBuffer := createBuffer()
+	byteBuffer := s.createBuffer()
 
 	var transferDirection string
 	// If the data source is the client then we are sending