new dump_limit_size parameter: stop dumping records beyond this size. Bugfix: dumps now have the appropriate size and content.

Author: cyrilst

cmd/sshproxy/recorder.go

@@ -23,6 +23,13 @@ import (
 	"github.com/cea-hpc/sshproxy/pkg/utils"
 )
 
+// Dup duplicates a []byte slice.
+func Dup(a []byte, n int) []byte {
+	b := make([]byte, n)
+	copy(b, a)
+	return b
+}
+
 // A Splitter reads from and/or writes to a file descriptor and sends a
 // record.Record struct to a channel for each read/write operation.
 type Splitter struct {
@@ -48,23 +55,25 @@ func (s *Splitter) Close() error {
 // its internal channel.
 func (s *Splitter) Read(p []byte) (int, error) {
 	n, err := s.f.Read(p)
+	pp := Dup(p, n)
 	s.ch <- record.Record{
 		Time: time.Now(),
 		Fd:   s.fd,
 		Size: n,
-		Data: p,
+		Data: pp,
 	}
 	return n, err
 }
 
 // Write implements the Writer Write method. It sends a copy of the written
 // slice to its internal channel.
 func (s *Splitter) Write(p []byte) (int, error) {
+	pp := Dup(p, len(p))
 	s.ch <- record.Record{
 		Time: time.Now(),
 		Fd:   s.fd,
 		Size: len(p),
-		Data: p,
+		Data: pp,
 	}
 	return s.f.Write(p)
 }
@@ -78,7 +87,7 @@ func (s *Splitter) Write(p []byte) (int, error) {
 // The file is a succession of serialized record.Record structs. See the
 // record.Record documentation for details on the format.
 type Recorder struct {
-	Stdin, Stdout, Stderr io.ReadWriteCloser // standard input, output and error to be used instead of the standard file descriptors.
+	Stdin, Stdout, Stderr io.ReadWriteCloser // standard input, output and error to be used instead of the standard file descriptors
 	start                 time.Time          // when the Recorder was started
 	etcdStatsInterval     time.Duration      // interval at which bandwidth is updated in etcd
 	logStatsInterval      time.Duration      // interval at which basic statistics of transferred bytes are logged
@@ -87,16 +96,17 @@ type Recorder struct {
 	ch                    chan record.Record // channel to read record.Record structs
 	conninfo              *ConnInfo          // specific SSH connection information
 	command               string             // initial user command
-	dumpfile              string             // path to filename where the raw records are dumped.
-	writer                *record.Writer     // *record.Writer where the raw records are dumped.
+	dumpfile              string             // path to filename where the raw records are dumped
+	dumpLimitSize         uint64             // number of bytes beyond which records are no longer dumped
+	writer                *record.Writer     // *record.Writer where the raw records are dumped
 }
 
 // NewRecorder returns a new Recorder struct.
 //
 // If dumpfile is not empty, the intercepted raw data will be written in this
 // file. Logging of basic statistics will be done every logStatsInterval seconds. Bandwidth will be updated in etcd every etcdStatsInterval seconds.
 // It will stop recording when the context is cancelled.
-func NewRecorder(conninfo *ConnInfo, dumpfile, command string, etcdStatsInterval time.Duration, logStatsInterval time.Duration) *Recorder {
+func NewRecorder(conninfo *ConnInfo, dumpfile, command string, etcdStatsInterval time.Duration, logStatsInterval time.Duration, dumpLimitSize uint64) *Recorder {
 	ch := make(chan record.Record)
 
 	return &Recorder{
@@ -111,6 +121,7 @@ func NewRecorder(conninfo *ConnInfo, dumpfile, command string, etcdStatsInterval
 		conninfo:          conninfo,
 		command:           command,
 		dumpfile:          dumpfile,
+		dumpLimitSize:     dumpLimitSize,
 		writer:            nil,
 	}
 }
@@ -246,6 +257,11 @@ func (r *Recorder) Run(ctx context.Context, cli *utils.Client, etcdPath string)
 				r.totals[rec.Fd] += uint64(rec.Size)
 				if r.writer != nil {
 					r.dump(rec)
+					if r.dumpLimitSize != 0 && r.totals[rec.Fd] > r.dumpLimitSize {
+						fd.Close()
+						fd = nil
+						r.writer = nil
+					}
 				}
 			case <-ctx.Done():
 				return
@@ -258,6 +274,11 @@ func (r *Recorder) Run(ctx context.Context, cli *utils.Client, etcdPath string)
 				r.totals[rec.Fd] += uint64(rec.Size)
 				if r.writer != nil {
 					r.dump(rec)
+					if r.dumpLimitSize != 0 && r.totals[rec.Fd] > r.dumpLimitSize {
+						fd.Close()
+						fd = nil
+						r.writer = nil
+					}
 				}
 			case <-ctx.Done():
 				return

cmd/sshproxy/sshproxy.go

@@ -300,6 +300,7 @@ func mainExitCode() int {
 	log.Debugf("config.log = %s", config.Log)
 	log.Debugf("config.check_interval = %s", config.CheckInterval.Duration())
 	log.Debugf("config.dump = %s", config.Dump)
+	log.Debugf("config.dump_limit_size = %d", config.DumpLimitSize)
 	log.Debugf("config.etcd_stats_interval = %s", config.EtcdStatsInterval.Duration())
 	log.Debugf("config.log_stats_interval = %s", config.LogStatsInterval.Duration())
 	log.Debugf("config.etcd = %+v", config.Etcd)
@@ -437,11 +438,7 @@ func mainExitCode() int {
 
 	var recorder *Recorder
 	if config.Dump != "" {
-		if !interactiveCommand {
-			// dont't dump non-interactive commands, but still log and update etcd
-			config.Dump = "etcd"
-		}
-		recorder = NewRecorder(conninfo, config.Dump, originalCmd, config.EtcdStatsInterval.Duration(), config.LogStatsInterval.Duration())
+		recorder = NewRecorder(conninfo, config.Dump, originalCmd, config.EtcdStatsInterval.Duration(), config.LogStatsInterval.Duration(), config.DumpLimitSize)
 
 		wg.Add(1)
 		go func() {

config/sshproxy.yaml

@@ -28,8 +28,12 @@
 # 'TCP:host:port' (the TCP is case sensitive), e.g. 'TCP:collector:5555'.
 #dump: ""
 
+# Maximum size of a dump (bytes). Beyond the limit, records are no more dumped.
+# Defaults to 0 (no limit).
+#dump_limit_size: 0
+
 # Interval at which basic statistics of transferred bytes are logged.
-# 0 by default, the string can contain a unit suffix such as 'h', 'm' and
+# "0" by default, the string can contain a unit suffix such as 'h', 'm' and
 # 's' (e.g. "2m30s"). These statistics are only available when the 'dump'
 # option is set.
 #log_stats_interval: "0"

doc/sshproxy.yaml.txt

@@ -63,6 +63,10 @@ user needs to have the right to write in this directory. For example:
 It can also be a network address where to send dumps if specified as
 'TCP:host:port' (the TCP is case sensitive), e.g.  'TCP:collector:5555'.
 
+*dump_limit_size*::
+	an integer specifying the maximum size (bytes) of a dump. Beyond this
+	size, records are no longer dumped. Defaults to 0 (no limit).
+
 *log_stats_interval*::
 	a string specifying the interval at which basic statistics of
 	transferred bytes are logged. 0 by default. The string can contain a

pkg/utils/config.go

@@ -29,6 +29,7 @@ type Config struct {
 	Log               string
 	CheckInterval     Duration `yaml:"check_interval"` // Minimum interval between host checks
 	Dump              string
+	DumpLimitSize     uint64   `yaml:"dump_limit_size"`
 	Etcd              etcdConfig
 	EtcdStatsInterval Duration `yaml:"etcd_stats_interval"`
 	LogStatsInterval  Duration `yaml:"log_stats_interval"`
@@ -75,8 +76,9 @@ type subConfig struct {
 	Debug             interface{}
 	Log               interface{}
 	Dump              interface{}
-	EtcdStatsInterval interface{} `yaml:"stats_interval"`
-	LogStatsInterval  interface{} `yaml:"stats_interval"`
+	DumpLimitSize     interface{} `yaml:"dump_limit_size"`
+	EtcdStatsInterval interface{} `yaml:"etcd_stats_interval"`
+	LogStatsInterval  interface{} `yaml:"log_stats_interval"`
 	BgCommand         interface{} `yaml:"bg_command"`
 	Environment       map[string]string
 	Routes            map[string]*RouteConfig
@@ -96,6 +98,10 @@ func parseSubConfig(config *Config, subconfig *subConfig) error {
 		config.Dump = subconfig.Dump.(string)
 	}
 
+	if subconfig.DumpLimitSize != nil {
+		config.DumpLimitSize = subconfig.DumpLimitSize.(uint64)
+	}
+
 	if subconfig.EtcdStatsInterval != nil {
 		var err error
 		config.EtcdStatsInterval, err = ParseDuration(subconfig.EtcdStatsInterval.(string))