/////////////////////////////////////////////////////////////////////////////////////
// TestAttachMultiple sets up the config for attach testing - tests launching and
// attaching to multiple processes simultaneously
//
func TestAttachMultiple(t *testing.T) {
_, mocker := testSetup(t)
defer testTeardown(t, mocker)
testServer, _ := server.(*testAttachServer)
cfg := executor.ExecutorConfig{
Common: executor.Common{
ID: "tee1",
Name: "tether_test_executor",
},
Sessions: map[string]*executor.SessionConfig{
"tee1": &executor.SessionConfig{
Common: executor.Common{
ID: "tee1",
Name: "tether_test_session",
},
Tty: false,
Attach: true,
Cmd: executor.Cmd{
Path: "/usr/bin/tee",
// grep, matching everything, reading from stdin
Args: []string{"/usr/bin/tee", pathPrefix + "/tee1.out"},
Env: []string{},
Dir: "/",
},
},
"tee2": &executor.SessionConfig{
Common: executor.Common{
ID: "tee2",
Name: "tether_test_session2",
},
Tty: false,
Attach: true,
Cmd: executor.Cmd{
Path: "/usr/bin/tee",
// grep, matching everything, reading from stdin
Args: []string{"/usr/bin/tee", pathPrefix + "/tee2.out"},
Env: []string{},
Dir: "/",
},
},
"tee3": &executor.SessionConfig{
Common: executor.Common{
ID: "tee3",
Name: "tether_test_session2",
},
Tty: false,
Attach: false,
Cmd: executor.Cmd{
Path: "/usr/bin/tee",
// grep, matching everything, reading from stdin
Args: []string{"/usr/bin/tee", pathPrefix + "/tee3.out"},
Env: []string{},
Dir: "/",
},
},
},
Key: genKey(),
Diagnostics: executor.Diagnostics{
DebugLevel: 2,
},
}
_, _, conn := StartAttachTether(t, &cfg, mocker)
defer conn.Close()
// wait for updates to occur
<-mocker.Started
<-testServer.updated
if !testServer.enabled {
t.Errorf("attach server was not enabled")
}
cconfig := &ssh.ClientConfig{
User: "******",
HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
return nil
},
}
// create the SSH client
sConn, chans, reqs, err := ssh.NewClientConn(conn, "notappliable", cconfig)
assert.NoError(t, err)
defer sConn.Close()
client := ssh.NewClient(sConn, chans, reqs)
ids, err := attach.SSHls(client)
assert.NoError(t, err)
// there's no ordering guarantee in the returned ids
if len(ids) != len(cfg.Sessions) {
t.Errorf("ID list - expected %d, got %d", len(cfg.Sessions), len(ids))
}
// check the ids we got correspond to those in the config
for _, id := range ids {
if _, ok := cfg.Sessions[id]; !ok {
t.Errorf("Expected sessions to have an entry for %s", id)
}
}
sessionA, err := attach.SSHAttach(client, "tee1")
assert.NoError(t, err)
sessionB, err := attach.SSHAttach(client, "tee2")
assert.NoError(t, err)
stdoutA := sessionA.Stdout()
stdoutB := sessionB.Stdout()
// FIXME: this is line buffered - how do I disable that so we don't have odd hangs to diagnose
// when the trailing \n is missed
testBytesA := []byte("hello world!\n")
testBytesB := []byte("goodbye world!\n")
// read from session into buffer
bufA := &bytes.Buffer{}
bufB := &bytes.Buffer{}
var wg sync.WaitGroup
// wg.Add cannot go inside the go routines as the Add may not have happened by the time we call Wait
wg.Add(2)
go func() {
io.CopyN(bufA, stdoutA, int64(len(testBytesA)))
wg.Done()
}()
go func() {
io.CopyN(bufB, stdoutB, int64(len(testBytesB)))
wg.Done()
}()
// write something to echo
log.Debug("sending test data")
sessionA.Stdin().Write(testBytesA)
sessionB.Stdin().Write(testBytesB)
log.Debug("sent test data")
// wait for the close to propagate
wg.Wait()
sessionA.CloseStdin()
sessionB.CloseStdin()
<-mocker.Cleaned
assert.Equal(t, bufA.Bytes(), testBytesA)
assert.Equal(t, bufB.Bytes(), testBytesB)
}
func TestReattach(t *testing.T) {
_, mocker := testSetup(t)
defer testTeardown(t, mocker)
testServer, _ := server.(*testAttachServer)
cfg := executor.ExecutorConfig{
Common: executor.Common{
ID: "attach",
Name: "tether_test_executor",
},
Sessions: map[string]*executor.SessionConfig{
"attach": &executor.SessionConfig{
Common: executor.Common{
ID: "attach",
Name: "tether_test_session",
},
Tty: false,
Attach: true,
RunBlock: true,
Cmd: executor.Cmd{
Path: "/usr/bin/tee",
// grep, matching everything, reading from stdin
Args: []string{"/usr/bin/tee", pathPrefix + "/tee.out"},
Env: []string{},
Dir: "/",
},
},
},
Key: genKey(),
}
_, _, conn := StartAttachTether(t, &cfg, mocker)
defer conn.Close()
// wait for updates to occur
<-testServer.updated
if !testServer.enabled {
t.Errorf("attach server was not enabled")
}
containerConfig := &ssh.ClientConfig{
User: "******",
HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
return nil
},
}
// create the SSH client from the mocked connection
sshConn, chans, reqs, err := ssh.NewClientConn(conn, "notappliable", containerConfig)
assert.NoError(t, err)
defer sshConn.Close()
var sshSession attach.SessionInteraction
done := make(chan bool)
buf := &bytes.Buffer{}
testBytes := []byte("\x1b[32mhello world\x1b[39m!\n")
attachFunc := func() {
attachClient := ssh.NewClient(sshConn, chans, reqs)
if attachClient == nil {
t.Errorf("Failed to get ssh.NewClient")
}
_, err = attach.SSHls(attachClient)
assert.NoError(t, err)
sshSession, err = attach.SSHAttach(attachClient, cfg.ID)
assert.NoError(t, err)
stdout := sshSession.Stdout()
// read from session into buffer
go func() {
io.CopyN(buf, stdout, int64(len(testBytes)))
done <- true
}()
// write something to echo
log.Debug("sending test data")
sshSession.Stdin().Write(testBytes)
log.Debug("sent test data")
}
limit := 100
for i := 0; i <= limit; i++ {
if i > 0 {
// truncate the buffer for the retach
buf.Reset()
testBytes = []byte(fmt.Sprintf("\x1b[32mhello world - again %dth time \x1b[39m!\n", i))
}
// attach
attachFunc()
// wait for the close to propagate
<-done
// send close-stdin if this is the last iteration
if i == limit {
// exit
sshSession.CloseStdin()
} else {
// detach
sshSession.Stdin().Close()
}
assert.Equal(t, buf.Bytes(), testBytes)
}
}
func attachCase(t *testing.T, runblock bool) {
_, mocker := testSetup(t)
defer testTeardown(t, mocker)
testServer, _ := server.(*testAttachServer)
cfg := executor.ExecutorConfig{
Common: executor.Common{
ID: "attach",
Name: "tether_test_executor",
},
Sessions: map[string]*executor.SessionConfig{
"attach": {
Common: executor.Common{
ID: "attach",
Name: "tether_test_session",
},
Tty: false,
Attach: true,
RunBlock: runblock,
Cmd: executor.Cmd{
Path: "/usr/bin/tee",
// grep, matching everything, reading from stdin
Args: []string{"/usr/bin/tee", pathPrefix + "/tee.out"},
Env: []string{},
Dir: "/",
},
},
},
Key: genKey(),
}
_, _, conn := StartAttachTether(t, &cfg, mocker)
defer conn.Close()
// wait for updates to occur
<-testServer.updated
if !testServer.enabled {
t.Errorf("attach server was not enabled")
}
containerConfig := &ssh.ClientConfig{
User: "******",
HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
return nil
},
}
// create the SSH client from the mocked connection
sshConn, chans, reqs, err := ssh.NewClientConn(conn, "notappliable", containerConfig)
assert.NoError(t, err)
defer sshConn.Close()
attachClient := ssh.NewClient(sshConn, chans, reqs)
if runblock {
_, err = attach.SSHls(attachClient)
assert.NoError(t, err)
}
sshSession, err := attach.SSHAttach(attachClient, cfg.ID)
assert.NoError(t, err)
stdout := sshSession.Stdout()
// FIXME: the pipe pair are line buffered - how do I disable that so we
// don't have odd hangs to diagnose when the trailing \n is missed
testBytes := []byte("\x1b[32mhello world\x1b[39m!\n")
// read from session into buffer
buf := &bytes.Buffer{}
done := make(chan bool)
go func() { io.CopyN(buf, stdout, int64(len(testBytes))); done <- true }()
// write something to echo
log.Debug("sending test data")
sshSession.Stdin().Write(testBytes)
log.Debug("sent test data")
// wait for the close to propagate
<-done
sshSession.CloseStdin()
}