handy-sshd/cmd/test_util_test.go
2023-08-10 22:53:02 +09:00

103 lines
2.6 KiB
Go

package cmd
import (
"github.com/stretchr/testify/assert"
"golang.org/x/crypto/ssh"
"io"
"net"
"os/exec"
"strconv"
"testing"
)
func getAvailableTcpPort() int {
ln, err := net.Listen("tcp", ":0")
if err != nil {
panic(err)
}
defer ln.Close()
return ln.Addr().(*net.TCPAddr).Port
}
func waitTCPServer(port int) {
for {
conn, err := net.Dial("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(port)))
if err == nil {
conn.Close()
break
}
}
}
func assertExec(t *testing.T, client *ssh.Client) {
session, err := client.NewSession()
assert.NoError(t, err)
defer session.Close()
whoamiBytes, err := session.Output("whoami")
assert.NoError(t, err)
expectedWhoamiBytes, err := exec.Command("whoami").Output()
assert.NoError(t, err)
assert.Equal(t, string(whoamiBytes), string(expectedWhoamiBytes))
}
func assertLocalPortForwarding(t *testing.T, client *ssh.Client) {
var remoteTcpPort int
acceptedConnChan := make(chan net.Conn)
{
ln, err := net.Listen("tcp", ":0")
assert.NoError(t, err)
remoteTcpPort = ln.Addr().(*net.TCPAddr).Port
go func() {
conn, err := ln.Accept()
assert.NoError(t, err)
acceptedConnChan <- conn
}()
}
raddr := &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: remoteTcpPort}
conn, err := client.DialTCP("tcp", nil, raddr)
assert.NoError(t, err)
defer conn.Close()
acceptedConn := <-acceptedConnChan
defer acceptedConn.Close()
{
localToRemote := [3]byte{1, 2, 3}
_, err = conn.Write(localToRemote[:])
assert.NoError(t, err)
var buf [len(localToRemote)]byte
_, err = io.ReadFull(acceptedConn, buf[:])
assert.NoError(t, err)
assert.Equal(t, buf, localToRemote)
}
{
remoteToLocal := [4]byte{10, 20, 30, 40}
_, err = acceptedConn.Write(remoteToLocal[:])
assert.NoError(t, err)
var buf [len(remoteToLocal)]byte
_, err = io.ReadFull(conn, buf[:])
assert.NoError(t, err)
assert.Equal(t, buf, remoteToLocal)
}
}
func assertRemotePortForwardingTODO(t *testing.T, client *ssh.Client) {
remotePort := getAvailableTcpPort()
acceptedConnChan := make(chan net.Conn)
var _ = acceptedConnChan
ln, err := client.Listen("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(remotePort)))
var _ = ln
assert.NoError(t, err)
go func() {
//conn, err := ln.Accept()
//assert.NoError(t, err)
//acceptedConnChan <- conn
}()
conn, err := net.Dial("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(remotePort)))
assert.NoError(t, err)
defer conn.Close()
// FIXME: implement but the following suspends
//acceptedConn := <-acceptedConnChan
//defer acceptedConn.Close()
// TODO: conn <--> acceptedConn communication
}