github.com/jlmeeker/kismatic@v1.10.1-0.20180612190640-57f9005a1f1a/integration-tests/ssh.go (about) 1 package integration_tests 2 3 import ( 4 "errors" 5 "fmt" 6 "os" 7 "os/exec" 8 "time" 9 ) 10 11 func runViaSSH(cmds []string, hosts []NodeDeets, sshKey string, period time.Duration) error { 12 timeout := time.After(period) 13 bail := make(chan struct{}) 14 cmdSuccess := make(chan bool) 15 // Create a goroutine per host. Each goroutine runs the commands serially on the host 16 // until one of these is true: 17 // a) all commands were executed successfully, 18 // b) an error occurred when running a command, 19 // c) the goroutine got a signal to bail 20 for _, host := range hosts { 21 go func(node NodeDeets) { 22 for _, cmd := range cmds { 23 res, err := executeCmd(cmd, node.PublicIP, node.SSHUser, sshKey) 24 fmt.Println(res) 25 select { 26 case cmdSuccess <- err == nil: 27 if err != nil { 28 return 29 } 30 case <-bail: 31 return 32 } 33 } 34 }(host) 35 } 36 37 // The bail channel is closed if we encounter an error, or if the timeout is reached. 38 // This will signal all goroutines to return. 39 defer close(bail) 40 41 // At most, we will get a total of hosts * cmds status messages in the channel. 42 for i := 0; i < len(hosts)*len(cmds); i++ { 43 select { 44 case ok := <-cmdSuccess: 45 if !ok { 46 return fmt.Errorf("error running command on node") 47 } 48 case <-timeout: 49 return fmt.Errorf("timed out running commands on nodes") 50 } 51 } 52 return nil 53 } 54 55 func executeCmd(cmd, hostname, user, sshKey string) (string, error) { 56 sshCmd := exec.Command("ssh", "-o", "StrictHostKeyChecking no", "-t", "-t", "-i", sshKey, user+"@"+hostname, cmd) 57 sshCmd.Stdin = os.Stdin 58 sshOut, sshErr := sshCmd.CombinedOutput() 59 return hostname + ": " + string(sshOut), sshErr 60 } 61 62 func copyFileToRemote(file string, destFile string, node NodeDeets, sshKey string, period time.Duration) error { 63 timeout := time.After(period) 64 success := make(chan bool) 65 go func() { 66 out, err := scpFile(file, destFile, node.SSHUser, node.PublicIP, sshKey) 67 fmt.Println(out) 68 success <- err == nil 69 }() 70 select { 71 case ok := <-success: 72 if !ok { 73 return errors.New("failed to copy file to node") 74 } 75 case <-timeout: 76 return errors.New("timed out copying file to node") 77 } 78 return nil 79 } 80 81 func scpFile(filePath string, destFilePath string, user, hostname, sshKey string) (string, error) { 82 ver := exec.Command("scp", "-o", "StrictHostKeyChecking no", "-i", sshKey, filePath, user+"@"+hostname+":"+destFilePath) 83 out, err := ver.CombinedOutput() 84 return string(out), err 85 } 86 87 // WaitUntilSSHOpen waits up to the given timeout for a successful SSH connection to 88 // the given node. If the connection is open, returns true. If the timeout is reached, returns false. 89 func WaitUntilSSHOpen(publicIP, sshUser, sshKey string, timeout time.Duration) bool { 90 tout := time.After(timeout) 91 tick := time.Tick(3 * time.Second) 92 for { 93 select { 94 case <-tout: 95 return false 96 case <-tick: 97 cmd := exec.Command("ssh") 98 cmd.Args = append(cmd.Args, "-i", sshKey) 99 cmd.Args = append(cmd.Args, "-o", "ConnectTimeout=5") 100 cmd.Args = append(cmd.Args, "-o", "BatchMode=yes") 101 cmd.Args = append(cmd.Args, "-o", "StrictHostKeyChecking=no") 102 cmd.Args = append(cmd.Args, fmt.Sprintf("%s@%s", sshUser, publicIP), "exit") // just call exit if we are able to connect 103 if err := cmd.Run(); err == nil { 104 // command succeeded 105 fmt.Println() 106 return true 107 } 108 fmt.Printf("?") 109 } 110 } 111 }