zrepl/sshbytestream/ssh.go
2017-04-29 20:09:09 +02:00

117 lines
2.4 KiB
Go

package sshbytestream
import (
"bytes"
"context"
"fmt"
"io"
"os"
"os/exec"
"sync"
)
type SSHTransport struct {
Host string
User string
Port uint16
TransportOpenCommand []string
Options []string
}
var SSHCommand string = "ssh"
func Incoming() (wc io.ReadWriteCloser, err error) {
// derivce ReadWriteCloser from stdin & stdout
return IncomingReadWriteCloser{}, nil
}
type IncomingReadWriteCloser struct{}
func (f IncomingReadWriteCloser) Read(p []byte) (n int, err error) {
return os.Stdin.Read(p)
}
func (f IncomingReadWriteCloser) Write(p []byte) (n int, err error) {
return os.Stdout.Write(p)
}
func (f IncomingReadWriteCloser) Close() (err error) {
os.Exit(0)
return nil
}
func Outgoing(remote SSHTransport) (conn io.ReadWriteCloser, err error) {
ctx, cancel := context.WithCancel(context.Background())
sshArgs := make([]string, 0, 2*len(remote.Options)+len(remote.TransportOpenCommand)+4)
sshArgs = append(sshArgs,
"-p", fmt.Sprintf("%d", remote.Port),
"-o", "BatchMode=yes",
)
for _, option := range remote.Options {
sshArgs = append(sshArgs, "-o", option)
}
sshArgs = append(sshArgs, fmt.Sprintf("%s@%s", remote.User, remote.Host))
sshArgs = append(sshArgs, remote.TransportOpenCommand...)
cmd := exec.CommandContext(ctx, SSHCommand, sshArgs...)
var in io.WriteCloser
var out io.ReadCloser
if in, err = cmd.StdinPipe(); err != nil {
return
}
if out, err = cmd.StdoutPipe(); err != nil {
return
}
f := ForkedSSHReadWriteCloser{
RemoteStdin: in,
RemoteStdout: out,
Cancel: cancel,
Command: cmd,
exitWaitGroup: &sync.WaitGroup{},
}
f.exitWaitGroup.Add(1)
go func() {
defer f.exitWaitGroup.Done()
stderr, _ := cmd.StderrPipe()
var b bytes.Buffer
if err := cmd.Run(); err != nil {
io.Copy(&b, stderr)
fmt.Println(b.String())
fmt.Printf("%v\n", cmd.ProcessState)
//panic(err)
}
}()
return f, nil
}
type ForkedSSHReadWriteCloser struct {
RemoteStdin io.Writer
RemoteStdout io.Reader
Command *exec.Cmd
Cancel context.CancelFunc
exitWaitGroup *sync.WaitGroup
}
func (f ForkedSSHReadWriteCloser) Read(p []byte) (n int, err error) {
return f.RemoteStdout.Read(p)
}
func (f ForkedSSHReadWriteCloser) Write(p []byte) (n int, err error) {
return f.RemoteStdin.Write(p)
}
func (f ForkedSSHReadWriteCloser) Close() (err error) {
f.Cancel()
f.exitWaitGroup.Wait()
return nil
}