version 1.3.3: reconnect to etcd after a connection lost. Avoid heavy

CPU usage in case of loss of etcd connection.

Author: cyrilst

Makefile

@@ -1,4 +1,4 @@
-SSHPROXY_VERSION ?= 1.3.2
+SSHPROXY_VERSION ?= 1.3.3
 SSHPROXY_GIT_URL ?= github.com/cea-hpc/sshproxy
 
 prefix		?= /usr

cmd/sshproxy/recorder.go

@@ -23,8 +23,6 @@ import (
 
 	"github.com/cea-hpc/sshproxy/pkg/record"
 	"github.com/cea-hpc/sshproxy/pkg/utils"
-
-	"go.etcd.io/etcd/clientv3"
 )
 
 // Dup duplicates a []byte slice.
@@ -103,7 +101,6 @@ type Recorder struct {
 	dumpfile              string             // path to filename where the raw records are dumped
 	dumpLimitSize         uint64             // number of bytes beyond which records are no longer dumped
 	dumpLimitWindow       time.Duration      // time window in which dump size is accounted
-	leaseID               clientv3.LeaseID   // etcd lease ID used for updating stats
 	lock                  sync.RWMutex       // mutex to avoid concurrent reads and writes in bandwidth and totals maps
 	writer                *record.Writer     // *record.Writer where the raw records are dumped
 }
@@ -113,7 +110,7 @@ type 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, dumpLimitSize uint64, dumpLimitWindow time.Duration, leaseID clientv3.LeaseID) *Recorder {
+func NewRecorder(conninfo *ConnInfo, dumpfile, command string, etcdStatsInterval time.Duration, logStatsInterval time.Duration, dumpLimitSize uint64, dumpLimitWindow time.Duration) *Recorder {
 	ch := make(chan record.Record)
 
 	return &Recorder{
@@ -130,23 +127,21 @@ func NewRecorder(conninfo *ConnInfo, dumpfile, command string, etcdStatsInterval
 		dumpfile:          dumpfile,
 		dumpLimitSize:     dumpLimitSize,
 		dumpLimitWindow:   dumpLimitWindow,
-		leaseID:           leaseID,
 		lock:              sync.RWMutex{},
 		writer:            nil,
 	}
 }
 
 // updateStats writes the bandwidth to etcd
 func (r *Recorder) updateStats(cli *utils.Client, etcdPath string) {
-	if cli != nil {
-		if cli.IsAlive() {
-			r.lock.RLock()
-			stats := r.bandwidth
-			r.lock.RUnlock()
-			err := cli.UpdateStats(etcdPath, stats, r.leaseID)
-			if err != nil {
-				log.Errorf("updating stats: %v", err)
-			}
+	if cli != nil && cli.IsAlive() {
+		r.lock.RLock()
+		stats := r.bandwidth
+		r.lock.RUnlock()
+		err := cli.UpdateStats(etcdPath, stats)
+		if err != nil {
+			log.Errorf("updating stats: %v", err)
+			cli.Disable()
 		}
 	}
 }

cmd/sshproxy/sshproxy.go

@@ -352,13 +352,12 @@ func mainExitCode() int {
 	}()
 
 	var etcdPath string
-	var leaseID clientv3.LeaseID
+	var tmpKeepAliveChan <-chan *clientv3.LeaseKeepAliveResponse
 	// Register destination in etcd and keep it alive while running.
 	if cli != nil && cli.IsAlive() {
 		key := fmt.Sprintf("%s@%s", username, service)
-		keepAliveChan, eP, lID, err := cli.SetDestination(ctx, key, sshInfos.Dst(), hostport)
+		keepAliveChan, eP, err := cli.SetDestination(ctx, key, sshInfos.Dst(), hostport)
 		etcdPath = eP
-		leaseID = lID
 		if err != nil {
 			log.Warningf("setting destination in etcd: %v", err)
 		}
@@ -368,6 +367,18 @@ func mainExitCode() int {
 			for {
 				select {
 				case <-keepAliveChan:
+					if !cli.IsAlive() {
+						tmpKeepAliveChan, err = cli.NewLease(ctx)
+						if err != nil {
+							log.Warningf("getting a new lease in etcd: %v", err)
+						} else {
+							keepAliveChan = tmpKeepAliveChan
+							cli.Enable()
+						}
+					} else {
+						// Avoid looping too fast in case of a loss of etcd
+						time.Sleep(time.Second)
+					}
 				case <-ctx.Done():
 					return
 				}
@@ -447,7 +458,7 @@ func mainExitCode() int {
 
 	var recorder *Recorder
 	if config.Dump != "" {
-		recorder = NewRecorder(conninfo, config.Dump, originalCmd, config.EtcdStatsInterval.Duration(), config.LogStatsInterval.Duration(), config.DumpLimitSize, config.DumpLimitWindow.Duration(), leaseID)
+		recorder = NewRecorder(conninfo, config.Dump, originalCmd, config.EtcdStatsInterval.Duration(), config.LogStatsInterval.Duration(), config.DumpLimitSize, config.DumpLimitWindow.Duration())
 
 		wg.Add(1)
 		go func() {

misc/sshproxy.spec

@@ -3,7 +3,7 @@
 %global debug_package   %{nil}
 
 Name:           sshproxy
-Version:        1.3.2
+Version:        1.3.3
 Release:        1%{?dist}
 Summary:        SSH proxy
 License:        CeCILL-B
@@ -51,6 +51,9 @@ install -p -m 0644 config/sshproxy.yaml %{buildroot}%{_sysconfdir}/sshproxy
 %{_mandir}/man8/sshproxy-replay.8*
 
 %changelog
+* Fri Oct 02 2020 Cyril Servant <[email protected]> - 1.3.3-1
+- sshproxy 1.3.3
+
 * Mon Sep 28 2020 Cyril Servant <[email protected]> - 1.3.2-1
 - sshproxy 1.3.2
 

pkg/utils/etcd.go

@@ -107,6 +107,8 @@ type Client struct {
 	log            *logging.Logger
 	requestTimeout time.Duration
 	keyTTL         int64
+	active         bool
+	leaseID        clientv3.LeaseID
 }
 
 // Host represents the state of a host.
@@ -172,6 +174,7 @@ func NewEtcdClient(config *Config, log *logging.Logger) (*Client, error) {
 		log:            log,
 		requestTimeout: 2 * time.Second,
 		keyTTL:         keyTTL,
+		active:         true,
 	}, nil
 }
 
@@ -205,35 +208,50 @@ func (c *Client) GetDestination(key string) (string, error) {
 }
 
 // SetDestination set current destination in etcd.
-func (c *Client) SetDestination(rootctx context.Context, key, sshdHostport string, dst string) (<-chan *clientv3.LeaseKeepAliveResponse, string, clientv3.LeaseID, error) {
+func (c *Client) SetDestination(rootctx context.Context, key, sshdHostport string, dst string) (<-chan *clientv3.LeaseKeepAliveResponse, string, error) {
 	path := fmt.Sprintf("%s/%s/%s/%s", toConnectionKey(key), dst, sshdHostport, time.Now().Format(time.RFC3339Nano))
 	ctx, cancel := context.WithTimeout(context.Background(), c.requestTimeout)
 	resp, err := c.cli.Grant(ctx, c.keyTTL)
 	cancel()
 	if err != nil {
-		return nil, "", 0, err
+		return nil, "", err
 	}
 
 	bytes, err := json.Marshal(&Bandwidth{
 		In:  0,
 		Out: 0,
 	})
 	if err != nil {
-		return nil, "", 0, err
+		return nil, "", err
 	}
 	ctx, cancel = context.WithTimeout(context.Background(), c.requestTimeout)
 	_, err = c.cli.Put(ctx, path, string(bytes), clientv3.WithLease(resp.ID))
 	cancel()
 	if err != nil {
-		return nil, "", 0, err
+		return nil, "", err
 	}
 
 	k, e := c.cli.KeepAlive(rootctx, resp.ID)
-	return k, path, resp.ID, e
+	c.leaseID = resp.ID
+	return k, path, e
+}
+
+// NewLease creates a new lease in etcd.
+func (c *Client) NewLease(rootctx context.Context) (<-chan *clientv3.LeaseKeepAliveResponse, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), c.requestTimeout)
+	resp, err := c.cli.Grant(ctx, c.keyTTL)
+	cancel()
+	if err != nil {
+		return nil, err
+	}
+
+	k, e := c.cli.KeepAlive(rootctx, resp.ID)
+	c.leaseID = resp.ID
+	return k, e
 }
 
 // UpdateStats updates the stats (bandwidth in and out in kB/s) of a connection.
-func (c *Client) UpdateStats(etcdPath string, stats map[int]uint64, leaseID clientv3.LeaseID) error {
+func (c *Client) UpdateStats(etcdPath string, stats map[int]uint64) error {
 	bytes, err := json.Marshal(&Bandwidth{
 		In:  int(stats[0] / 1024),
 		Out: int((stats[1] + stats[2]) / 1024),
@@ -243,7 +261,7 @@ func (c *Client) UpdateStats(etcdPath string, stats map[int]uint64, leaseID clie
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), c.requestTimeout)
-	_, err = c.cli.Put(ctx, etcdPath, string(bytes), clientv3.WithLease(leaseID))
+	_, err = c.cli.Put(ctx, etcdPath, string(bytes), clientv3.WithLease(c.leaseID))
 	cancel()
 	if err != nil {
 		return err
@@ -567,5 +585,15 @@ func (c *Client) GetAllGroups(allFlag bool) (map[string]*FlatGroup, error) {
 
 // IsAlive checks if etcd client is still usable.
 func (c *Client) IsAlive() bool {
-	return c.cli != nil
+	return c.cli != nil && c.active == true
+}
+
+// Enable enables the etcd client.
+func (c *Client) Enable() {
+	c.active = true
+}
+
+// Disable disables the etcd client.
+func (c *Client) Disable() {
+	c.active = false
 }