SSH Keep Alive

resolves #8

connector/connection.go

@@ -2,83 +2,34 @@ package connector
 
 import (
 	"bytes"
-	"io/ioutil"
-	"strings"
-
+	"github.com/pkg/errors"
+	"net"
 	"sync"
 
-	"time"
-
 	"golang.org/x/crypto/ssh"
 )
 
-const timeoutInSeconds = 5
-
-var (
-	cachedConfig *ssh.ClientConfig
-	lock         = &sync.Mutex{}
-)
-
-func config(user, keyFile string) (*ssh.ClientConfig, error) {
-	lock.Lock()
-	defer lock.Unlock()
-
-	if cachedConfig != nil {
-		return cachedConfig, nil
-	}
-
-	pk, err := loadPublicKeyFile(keyFile)
-	if err != nil {
-		return nil, err
-	}
-
-	cachedConfig = &ssh.ClientConfig{
-		User:            user,
-		Auth:            []ssh.AuthMethod{pk},
-		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
-		Timeout:         timeoutInSeconds * time.Second,
-	}
-
-	return cachedConfig, nil
-}
-
-// NewSSSHConnection connects to device
-func NewSSSHConnection(host, user, keyFile string) (*SSHConnection, error) {
-	if !strings.Contains(host, ":") {
-		host = host + ":22"
-	}
-
-	c := &SSHConnection{Host: host}
-	err := c.Connect(user, keyFile)
-	if err != nil {
-		return nil, err
-	}
-
-	return c, nil
-}
-
 // SSHConnection encapsulates the connection to the device
 type SSHConnection struct {
-	conn *ssh.Client
-	Host string
-}
-
-// Connect connects to the device
-func (c *SSHConnection) Connect(user, keyFile string) error {
-	config, err := config(user, keyFile)
-	if err != nil {
-		return err
-	}
-
-	c.conn, err = ssh.Dial("tcp", c.Host, config)
-	return err
+	host   string
+	client *ssh.Client
+	conn   net.Conn
+	mu     sync.Mutex
+	done   chan struct{}
 }
 
 // RunCommand runs a command against the device
 func (c *SSHConnection) RunCommand(cmd string) ([]byte, error) {
-	session, err := c.conn.NewSession()
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.client == nil {
+		return nil, errors.New("not conneted")
+	}
+
+	session, err := c.client.NewSession()
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "could not open session")
 	}
 	defer session.Close()
 
@@ -87,32 +38,40 @@ func (c *SSHConnection) RunCommand(cmd string) ([]byte, error) {
 
 	err = session.Run(cmd)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "could not run command")
 	}
 
 	return b.Bytes(), nil
 }
 
-// Close closes connection
-func (c *SSHConnection) Close() {
-	if c.conn == nil {
-		return
-	}
+func (c *SSHConnection) isConnected() bool {
+	return c.conn != nil
+}
+
+func (c *SSHConnection) terminate() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 
 	c.conn.Close()
+
+	c.client = nil
 	c.conn = nil
 }
 
-func loadPublicKeyFile(file string) (ssh.AuthMethod, error) {
-	b, err := ioutil.ReadFile(file)
-	if err != nil {
-		return nil, err
-	}
+func (c *SSHConnection) close() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
 
-	key, err := ssh.ParsePrivateKey(b)
-	if err != nil {
-		return nil, err
+	if c.client != nil {
+		c.client.Close()
 	}
 
-	return ssh.PublicKeys(key), nil
+	c.done <- struct{}{}
+	c.conn = nil
+	c.client = nil
+}
+
+// Host returns the hostname connected to
+func (c *SSHConnection) Host() string {
+	return c.host
 }

connector/connection_manager.go

@@ -0,0 +1,187 @@
+package connector
+
+import (
+	"github.com/prometheus/common/log"
+	"io"
+	"io/ioutil"
+	"net"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/pkg/errors"
+	"golang.org/x/crypto/ssh"
+)
+
+const timeoutInSeconds = 5
+
+// Option defines options for the manager which are applied on creation
+type Option func(*SSHConnectionManager)
+
+// WithReconnectInterval sets the reconnect interval (default 10 seconds)
+func WithReconnectInterval(d time.Duration) Option {
+	return func(m *SSHConnectionManager) {
+		m.reconnectInterval = d
+	}
+}
+
+// WithKeepAliveInterval sets the keep alive interval (default 10 seconds)
+func WithKeepAliveInterval(d time.Duration) Option {
+	return func(m *SSHConnectionManager) {
+		m.keepAliveInterval = d
+	}
+}
+
+// WithKeepAliveTimeout sets the timeout after an ssh connection to be determined dead (default 15 seconds)
+func WithKeepAliveTimeout(d time.Duration) Option {
+	return func(m *SSHConnectionManager) {
+		m.keepAliveTimeout = d
+	}
+}
+
+// SSHConnectionManager manages SSH connections to different devices
+type SSHConnectionManager struct {
+	config            *ssh.ClientConfig
+	connections       map[string]*SSHConnection
+	reconnectInterval time.Duration
+	keepAliveInterval time.Duration
+	keepAliveTimeout  time.Duration
+	mu                sync.Mutex
+}
+
+// NewConnectionManager creates a new connection manager
+func NewConnectionManager(user string, key io.Reader, opts ...Option) (*SSHConnectionManager, error) {
+	pk, err := loadPublicKeyFile(key)
+	if err != nil {
+		return nil, err
+	}
+
+	cfg := &ssh.ClientConfig{
+		User:            user,
+		Auth:            []ssh.AuthMethod{pk},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		Timeout:         timeoutInSeconds * time.Second,
+	}
+
+	m := &SSHConnectionManager{
+		config:            cfg,
+		connections:       make(map[string]*SSHConnection),
+		reconnectInterval: 30 * time.Second,
+		keepAliveInterval: 10 * time.Second,
+		keepAliveTimeout:  15 * time.Second,
+	}
+
+	for _, opt := range opts {
+		opt(m)
+	}
+
+	return m, nil
+}
+
+// Connect connects to a device or returns an long living connection
+func (m *SSHConnectionManager) Connect(host string) (*SSHConnection, error) {
+	if !strings.Contains(host, ":") {
+		host = host + ":22"
+	}
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if connection, found := m.connections[host]; found {
+		if !connection.isConnected() {
+			return nil, errors.New("not connected")
+		}
+
+		return connection, nil
+	}
+
+	return m.connect(host)
+}
+
+func (m *SSHConnectionManager) connect(host string) (*SSHConnection, error) {
+	client, conn, err := m.connectToServer(host)
+	if err != nil {
+		return nil, err
+	}
+
+	c := &SSHConnection{
+		conn:   conn,
+		client: client,
+		host:   host,
+		done:   make(chan struct{}),
+	}
+	go m.keepAlive(c)
+
+	m.connections[host] = c
+
+	return c, nil
+}
+
+func (m *SSHConnectionManager) connectToServer(host string) (*ssh.Client, net.Conn, error) {
+	conn, err := net.DialTimeout("tcp", host, m.config.Timeout)
+	if err != nil {
+		return nil, nil, errors.Wrap(err, "could not open tcp connection")
+	}
+
+	c, chans, reqs, err := ssh.NewClientConn(conn, host, m.config)
+	if err != nil {
+		return nil, nil, errors.Wrap(err, "could not connect to device")
+	}
+
+	return ssh.NewClient(c, chans, reqs), conn, nil
+}
+
+func (m *SSHConnectionManager) keepAlive(connection *SSHConnection) {
+	for {
+		select {
+		case <-time.After(m.keepAliveInterval):
+			log.Debugf("Sending keepalive for ")
+			connection.conn.SetDeadline(time.Now().Add(m.keepAliveTimeout))
+			_, _, err := connection.client.SendRequest("keepalive@golang.org", true, nil)
+			if err != nil {
+				log.Infof("Lost connection to %s (%v). Trying to reconnect...", connection.Host(), err)
+				connection.terminate()
+				m.reconnect(connection)
+			}
+		case <-connection.done:
+			return
+		}
+	}
+}
+
+func (m *SSHConnectionManager) reconnect(connection *SSHConnection) {
+	for {
+		client, conn, err := m.connectToServer(connection.Host())
+		if err == nil {
+			connection.client = client
+			connection.conn = conn
+			return
+		}
+
+		log.Infof("Reconnect to %s failed: %v", connection.Host(), err)
+		time.Sleep(m.reconnectInterval)
+	}
+}
+
+// Close closes all TCP connections and stop keep alives
+func (m *SSHConnectionManager) Close() error {
+	for _, c := range m.connections {
+		c.close()
+	}
+
+	return nil
+}
+
+func loadPublicKeyFile(r io.Reader) (ssh.AuthMethod, error) {
+	b, err := ioutil.ReadAll(r)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not read from reader")
+	}
+
+	key, err := ssh.ParsePrivateKey(b)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not parse private key")
+	}
+
+	return ssh.PublicKeys(key), nil
+}

go.mod

@@ -0,0 +1,18 @@
+module github.com/czerwonk/junos_exporter
+
+require (
+	github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc
+	github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf
+	github.com/beorn7/perks v0.0.0-20160804104726-4c0e84591b9a
+	github.com/golang/protobuf v0.0.0-20170920220647-130e6b02ab05
+	github.com/matttproud/golang_protobuf_extensions v1.0.0
+	github.com/prometheus/client_golang v0.8.0
+	github.com/prometheus/client_model v0.0.0-20170216185247-6f3806018612
+	github.com/prometheus/common v0.0.0-20171006141418-1bab55dd05db
+	github.com/prometheus/procfs v0.0.0-20170703101242-e645f4e5aaa8
+	github.com/sirupsen/logrus v1.0.3
+	golang.org/x/crypto v0.0.0-20170930174604-9419663f5a44
+	golang.org/x/sys v0.0.0-20171006175012-ebfc5b463182
+	gopkg.in/alecthomas/kingpin.v2 v2.2.5
+	gopkg.in/yaml.v2 v2.2.1
+)

go.sum

@@ -0,0 +1,29 @@
+github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc h1:cAKDfWh5VpdgMhJosfJnn5/FoN2SRZ4p7fJNX58YPaU=
+github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
+github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf h1:qet1QNfXsQxTZqLG4oE62mJzwPIB8+Tee4RNCL9ulrY=
+github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
+github.com/beorn7/perks v0.0.0-20160804104726-4c0e84591b9a h1:BtpsbiV638WQZwhA98cEZw2BsbnQJrbd0BI7tsy0W1c=
+github.com/beorn7/perks v0.0.0-20160804104726-4c0e84591b9a/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
+github.com/golang/protobuf v0.0.0-20170920220647-130e6b02ab05 h1:Kesru7U6Mhpf/x7rthxAKnr586VFmoE2NdEvkOKvfjg=
+github.com/golang/protobuf v0.0.0-20170920220647-130e6b02ab05/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
+github.com/matttproud/golang_protobuf_extensions v1.0.0 h1:YNOwxxSJzSUARoD9KRZLzM9Y858MNGCOACTvCW9TSAc=
+github.com/matttproud/golang_protobuf_extensions v1.0.0/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
+github.com/prometheus/client_golang v0.8.0 h1:1921Yw9Gc3iSc4VQh3PIoOqgPCZS7G/4xQNVUp8Mda8=
+github.com/prometheus/client_golang v0.8.0/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
+github.com/prometheus/client_model v0.0.0-20170216185247-6f3806018612 h1:13pIdM2tpaDi4OVe24fgoIS7ZTqMt0QI+bwQsX5hq+g=
+github.com/prometheus/client_model v0.0.0-20170216185247-6f3806018612/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
+github.com/prometheus/common v0.0.0-20171006141418-1bab55dd05db h1:PmL7nSW2mvuotGlJKuvUcSI/eE86zwYUcIAGoB6eHBk=
+github.com/prometheus/common v0.0.0-20171006141418-1bab55dd05db/go.mod h1:daVV7qP5qjZbuso7PdcryaAu0sAZbrN9i7WWcTMWvro=
+github.com/prometheus/procfs v0.0.0-20170703101242-e645f4e5aaa8 h1:uZfczEBIA1FZfOQo4/JWgGnMNd/4HVsM9A+B30wtlkA=
+github.com/prometheus/procfs v0.0.0-20170703101242-e645f4e5aaa8/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
+github.com/sirupsen/logrus v1.0.3 h1:B5C/igNWoiULof20pKfY4VntcIPqKuwEmoLZrabbUrc=
+github.com/sirupsen/logrus v1.0.3/go.mod h1:pMByvHTf9Beacp5x1UXfOR9xyW/9antXMhjMPG0dEzc=
+golang.org/x/crypto v0.0.0-20170930174604-9419663f5a44 h1:9lP3x0pW80sDI6t1UMSLA4to18W7R7imwAI/sWS9S8Q=
+golang.org/x/crypto v0.0.0-20170930174604-9419663f5a44/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4=
+golang.org/x/sys v0.0.0-20171006175012-ebfc5b463182 h1:7cKexPAAZFbkQtOZ/08DxRPYYxWzMBesz2/gC7esAtI=
+golang.org/x/sys v0.0.0-20171006175012-ebfc5b463182/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+gopkg.in/alecthomas/kingpin.v2 v2.2.5 h1:qskSCq465uEvC3oGocwvZNsO3RF3SpLVLumOAhL0bXo=
+gopkg.in/alecthomas/kingpin.v2 v2.2.5/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.1 h1:mUhvW9EsL+naU5Q3cakzfE91YhliOondGd6ZrsDBHQE=
+gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=