laptop-booter/ssh.go at master · milanaleksic/laptop-booter · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
package laptop_booter

// based on http://blog.ralch.com/tutorial/golang-ssh-tunneling/

import (
	"fmt"
	"io"
	"net"
	"strconv"
	"time"

	"log"

	"golang.org/x/crypto/ssh"
)

type Endpoint struct {
	// Server host address
	Host string
	// Server port
	Port int
}

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

func (endpoint *Endpoint) IsSet() bool {
	return endpoint.Host != ""
}

type SSHTunnel struct {
	Local    *Endpoint
	Mediator *Endpoint
	Remote   *Endpoint

	Config *ssh.ClientConfig
}

func (tunnel *SSHTunnel) BlockingListen() error {
	listener, err := net.Listen("tcp", tunnel.Local.String())
	if err != nil {
		return err
	}
	defer listener.Close()

	for {
		conn, err := listener.Accept()
		if err != nil {
			return err
		}
		go tunnel.forward(conn)
	}
}

func (tunnel *SSHTunnel) forward(localConn net.Conn) {
	var remoteConn net.Conn
	var err error
	if tunnel.Mediator.IsSet() {
		serverConn, err := ssh.Dial("tcp", tunnel.Mediator.String(), tunnel.Config)
		if err != nil {
			fmt.Printf("Server dial error to %s:%d, %s\n", tunnel.Remote.Host, tunnel.Remote.Port, err)
			return
		}

		remoteConn, err = serverConn.Dial("tcp", tunnel.Remote.String())
	} else {
		remoteConn, err = net.Dial("tcp", tunnel.Remote.String())
	}

	if err != nil {
		fmt.Printf("Remote dial error to %s:%d, %+v\n", tunnel.Remote.Host, tunnel.Remote.Port, err)
		return
	}

	copyConn := func(writer, reader net.Conn) {
		_, err := io.Copy(writer, reader)
		if err != nil {
			fmt.Printf("io.Copy error to %s:%d, %s", tunnel.Remote.Host, tunnel.Remote.Port, err)
		}
	}

	go copyConn(localConn, remoteConn)
	go copyConn(remoteConn, localConn)
}

func (tunnel *SSHTunnel) Activate() {
	log.Printf("Activating local port %v for tunnel to %v (user: %s)", tunnel.Local.Port, tunnel.Remote, tunnel.Config.User)
	go tunnel.BlockingListen()
	tunnel.waitForLocalHostOpen()
}

func (tunnel *SSHTunnel) waitForLocalHostOpen() {
	for {
		conn, _ := net.DialTimeout("tcp", net.JoinHostPort("localhost", strconv.Itoa(tunnel.Local.Port)), 100*time.Millisecond)
		if conn != nil {
			conn.Close()
			break
		}
	}
}