sshtunnel.go

// Copyright 2021 Jürgen Enge, info-age GmbH, Basel
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package sshtunnel

/*
   https://gist.github.com/svett/5d695dcc4cc6ad5dd275

*/

import (
	"fmt"
	"github.com/goph/emperror"
	"github.com/op/go-logging"
	"golang.org/x/crypto/ssh"
	"io"
	"io/ioutil"
	"net"
	"sync"
)

type Endpoint struct {
	Host string
	Port int
}

type SourceDestination struct {
	Local  *Endpoint
	Remote *Endpoint
}

func (endpoint *Endpoint) String() string {
	return fmt.Sprintf("%s:%d", endpoint.Host, endpoint.Port)
}

type SSHtunnel struct {
	server   *Endpoint
	tunnels  map[string]*SourceDestination
	quit     map[string]chan interface{}
	listener map[string]net.Listener
	config   *ssh.ClientConfig
	client   *ssh.Client
	log      *logging.Logger
	wg       sync.WaitGroup
}

func NewSSHTunnel(user, privateKey string, serverEndpoint *Endpoint, tunnels map[string]*SourceDestination, log *logging.Logger) (*SSHtunnel, error) {
	key, err := ioutil.ReadFile(privateKey)
	if err != nil {
		return nil, emperror.Wrapf(err, "Unable to read private key %s", privateKey)
	}
	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil, emperror.Wrapf(err, "Unable to parse private key")
	}

	sshConfig := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	tunnel := &SSHtunnel{
		config:   sshConfig,
		server:   serverEndpoint,
		tunnels:  tunnels,
		quit:     make(map[string]chan interface{}),
		listener: make(map[string]net.Listener),
		log:      log,
	}

	return tunnel, nil
}

func (tunnel *SSHtunnel) String() string {
	str := fmt.Sprintf("%v@%v:%v",
		tunnel.config.User,
		tunnel.server.Host, tunnel.server.Port,
	)
	for _, srcdests := range tunnel.tunnels {
		str += fmt.Sprintf(" - (%v:%v -> %v:%v)",
			srcdests.Local.Host, srcdests.Local.Port,
			srcdests.Remote.Host, srcdests.Remote.Port,
		)
	}
	return str
}

func (tunnel *SSHtunnel) Close() {
	tunnel.client.Close()
	for key, listener := range tunnel.listener {
		if q, ok := tunnel.quit[key]; ok {
			close(q)
		}
		listener.Close()
	}
	tunnel.wg.Wait()
	tunnel.log.Debugf("tunnel %s closed", tunnel.String())
}

func (tunnel *SSHtunnel) Start() error {
	var err error
	tunnel.log.Debugf("dialing ssh: %v", tunnel.String())
	tunnel.client, err = ssh.Dial("tcp", tunnel.server.String(), tunnel.config)
	if err != nil {
		return emperror.Wrapf(err, "server dial error to %v", tunnel.server.String())
	}

	for key, t := range tunnel.tunnels {
		tunnel.log.Debugf("starting listener %s -> %s", t.Local.String(), t.Remote.String())
		tunnel.listener[key], err = net.Listen("tcp", t.Local.String())
		if err != nil {
			return emperror.Wrapf(err, "cannot start listener on %v", t.Local.String())
		}
		tunnel.quit[key] = make(chan interface{})

		go func(k string) {
			//defer tunnel.wg.Done()
			conn, err := tunnel.listener[k].Accept()
			if err != nil {
				select {
				case <-tunnel.quit[k]:
					return
				default:
					tunnel.log.Errorf("error accepting connection on %v", tunnel.tunnels[k].Local.String())
				}
			} else {
				tunnel.wg.Add(1)
				go func() {
					tunnel.forward(conn, tunnel.tunnels[k].Remote)
					tunnel.wg.Done()
				}()
			}
		}(key)
	}
	return nil
}

func (tunnel *SSHtunnel) forward(localConn net.Conn, endpoint *Endpoint) {
	var err error

	remoteConn, err := tunnel.client.Dial("tcp", endpoint.String())
	if err != nil {
		tunnel.log.Errorf("Remote dial error %v: %v", endpoint.String(), err)
		return
	}

	copyConn := func(writer, reader net.Conn) {
		defer writer.Close()
		defer reader.Close()

		_, err := io.Copy(writer, reader)
		if err != nil {
			tunnel.log.Errorf("io.Copy error %v: %v", endpoint.String(), err)
		}
	}
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		copyConn(localConn, remoteConn)
		wg.Done()
	}()
	go func() {
		copyConn(remoteConn, localConn)
		wg.Done()
	}()
	wg.Wait()
}