mirror of
https://github.com/zrepl/zrepl.git
synced 2024-11-26 02:14:44 +01:00
69f8e7cfc3
Move from rpc to separate util package.
110 lines
2.3 KiB
Go
110 lines
2.3 KiB
Go
package sshbytestream
|
|
|
|
import (
|
|
"io"
|
|
"os"
|
|
"os/exec"
|
|
"github.com/zrepl/zrepl/model"
|
|
"context"
|
|
"fmt"
|
|
"bytes"
|
|
"sync"
|
|
)
|
|
|
|
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(name string, remote model.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
|
|
} |