github.com/looshlee/beatles@v0.0.0-20220727174639-742810ab631c/test/helpers/node.go (about) 1 // Copyright 2017 Authors of Cilium 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package helpers 16 17 import ( 18 "context" 19 "fmt" 20 "io" 21 "os" 22 "path/filepath" 23 "strings" 24 "sync" 25 "time" 26 27 "github.com/cilium/cilium/test/config" 28 ginkgoext "github.com/cilium/cilium/test/ginkgo-ext" 29 30 "github.com/sirupsen/logrus" 31 "golang.org/x/crypto/ssh" 32 ) 33 34 var ( 35 //SSHMetaLogs is a buffer where all commands sent over ssh are saved. 36 SSHMetaLogs = ginkgoext.NewWriter(new(Buffer)) 37 ) 38 39 // SSHMeta contains metadata to SSH into a remote location to run tests 40 type SSHMeta struct { 41 sshClient *SSHClient 42 env []string 43 rawConfig []byte 44 nodeName string 45 logger *logrus.Entry 46 } 47 48 // CreateSSHMeta returns an SSHMeta with the specified host, port, and user, as 49 // well as an according SSHClient. 50 func CreateSSHMeta(host string, port int, user string) *SSHMeta { 51 return &SSHMeta{ 52 sshClient: GetSSHClient(host, port, user), 53 } 54 } 55 56 func (s *SSHMeta) String() string { 57 return fmt.Sprintf("environment: %s, SSHClient: %s", s.env, s.sshClient.String()) 58 59 } 60 61 // CloseSSHClient closes all of the connections made by the SSH Client for this 62 // SSHMeta. 63 func (s *SSHMeta) CloseSSHClient() { 64 if s.sshClient == nil || s.sshClient.client == nil { 65 log.Error("SSH client is nil; cannot close") 66 } 67 if err := s.sshClient.client.Close(); err != nil { 68 log.WithError(err).Error("error closing SSH client") 69 } 70 } 71 72 // GetVagrantSSHMeta returns a SSHMeta initialized based on the provided 73 // SSH-config target. 74 func GetVagrantSSHMeta(vmName string) *SSHMeta { 75 config, err := GetVagrantSSHMetadata(vmName) 76 if err != nil { 77 return nil 78 } 79 80 log.Debugf("generated SSHConfig for node %s", vmName) 81 nodes, err := ImportSSHconfig(config) 82 if err != nil { 83 log.WithError(err).Error("Error importing ssh config") 84 return nil 85 } 86 var node *SSHConfig 87 log.Debugf("done importing ssh config") 88 for name := range nodes { 89 if strings.HasPrefix(name, vmName) { 90 node = nodes[name] 91 break 92 } 93 } 94 if node == nil { 95 log.Errorf("Node %s not found in ssh config", vmName) 96 return nil 97 } 98 sshMeta := &SSHMeta{ 99 sshClient: node.GetSSHClient(), 100 rawConfig: config, 101 nodeName: vmName, 102 } 103 104 sshMeta.setBasePath() 105 return sshMeta 106 } 107 108 // setBasePath if the SSHConfig is defined we set the BasePath to the GOPATH, 109 // from golang 1.8 GOPATH is by default $HOME/go so we also check that. 110 func (s *SSHMeta) setBasePath() { 111 if config.CiliumTestConfig.SSHConfig == "" { 112 return 113 } 114 115 gopath := s.Exec("echo $GOPATH").SingleOut() 116 if gopath != "" { 117 BasePath = filepath.Join(gopath, CiliumPath) 118 return 119 } 120 121 home := s.Exec("echo $HOME").SingleOut() 122 if home == "" { 123 return 124 } 125 126 BasePath = filepath.Join(home, "go", CiliumPath) 127 return 128 } 129 130 // ExecuteContext executes the given `cmd` and writes the cmd's stdout and 131 // stderr into the given io.Writers. 132 // Returns an error if context Deadline() is reached or if there was an error 133 // executing the command. 134 func (s *SSHMeta) ExecuteContext(ctx context.Context, cmd string, stdout io.Writer, stderr io.Writer) error { 135 if stdout == nil { 136 stdout = os.Stdout 137 } 138 139 if stderr == nil { 140 stderr = os.Stderr 141 } 142 fmt.Fprintln(SSHMetaLogs, cmd) 143 command := &SSHCommand{ 144 Path: cmd, 145 Stdin: os.Stdin, 146 Stdout: stdout, 147 Stderr: stderr, 148 } 149 return s.sshClient.RunCommandContext(ctx, command) 150 } 151 152 // ExecWithSudo returns the result of executing the provided cmd via SSH using 153 // sudo. 154 func (s *SSHMeta) ExecWithSudo(cmd string, options ...ExecOptions) *CmdRes { 155 command := fmt.Sprintf("sudo %s", cmd) 156 return s.Exec(command, options...) 157 } 158 159 // ExecOptions options to execute Exec and ExecWithContext 160 type ExecOptions struct { 161 SkipLog bool 162 } 163 164 // Exec returns the results of executing the provided cmd via SSH. 165 func (s *SSHMeta) Exec(cmd string, options ...ExecOptions) *CmdRes { 166 // Bound all command executions to be at most the timeout used by the CI 167 // so that commands do not block forever. 168 ctx, cancel := context.WithTimeout(context.Background(), HelperTimeout) 169 defer cancel() 170 return s.ExecContext(ctx, cmd, options...) 171 } 172 173 // ExecShort runs command with the provided options. It will take up to 174 // ShortCommandTimeout seconds to run the command before it times out. 175 func (s *SSHMeta) ExecShort(cmd string, options ...ExecOptions) *CmdRes { 176 ctx, cancel := context.WithTimeout(context.Background(), ShortCommandTimeout) 177 defer cancel() 178 return s.ExecContext(ctx, cmd, options...) 179 } 180 181 // ExecMiddle runs command with the provided options. It will take up to 182 // MidCommandTimeout seconds to run the command before it times out. 183 func (s *SSHMeta) ExecMiddle(cmd string, options ...ExecOptions) *CmdRes { 184 ctx, cancel := context.WithTimeout(context.Background(), MidCommandTimeout) 185 defer cancel() 186 return s.ExecContext(ctx, cmd, options...) 187 } 188 189 // ExecContextShort is a wrapper around ExecContext which creates a child 190 // context with a timeout of ShortCommandTimeout. 191 func (s *SSHMeta) ExecContextShort(ctx context.Context, cmd string, options ...ExecOptions) *CmdRes { 192 shortCtx, cancel := context.WithTimeout(ctx, ShortCommandTimeout) 193 defer cancel() 194 return s.ExecContext(shortCtx, cmd, options...) 195 } 196 197 // ExecContext returns the results of executing the provided cmd via SSH. 198 func (s *SSHMeta) ExecContext(ctx context.Context, cmd string, options ...ExecOptions) *CmdRes { 199 var ops ExecOptions 200 if len(options) > 0 { 201 ops = options[0] 202 } 203 204 log.Debugf("running command: %s", cmd) 205 stdout := new(Buffer) 206 stderr := new(Buffer) 207 start := time.Now() 208 err := s.ExecuteContext(ctx, cmd, stdout, stderr) 209 210 res := CmdRes{ 211 cmd: cmd, 212 stdout: stdout, 213 stderr: stderr, 214 success: true, // this may be toggled when err != nil below 215 duration: time.Since(start), 216 } 217 218 if err != nil { 219 res.success = false 220 // Set error code to 1 in case that it's another error to see that the 221 // command failed. If the default value (0) indicates that command 222 // works but it was not executed at all. 223 res.exitcode = 1 224 exiterr, isExitError := err.(*ssh.ExitError) 225 if isExitError { 226 // Set res's exitcode if the error is an ExitError 227 res.exitcode = exiterr.Waitmsg.ExitStatus() 228 } else { 229 // Log other error types. They are likely from SSH or the network 230 log.WithError(err).Errorf("Error executing command '%s'", cmd) 231 res.err = err 232 } 233 } 234 235 res.SendToLog(ops.SkipLog) 236 return &res 237 } 238 239 // GetCopy returns a copy of SSHMeta, useful for parallel requests 240 func (s *SSHMeta) GetCopy() *SSHMeta { 241 nodes, err := ImportSSHconfig(s.rawConfig) 242 if err != nil { 243 log.WithError(err).Error("while importing ssh config for meta copy") 244 return nil 245 } 246 247 config := nodes[s.nodeName] 248 if config == nil { 249 log.Errorf("no node %s in imported config", s.nodeName) 250 return nil 251 } 252 253 copy := &SSHMeta{ 254 sshClient: config.GetSSHClient(), 255 rawConfig: s.rawConfig, 256 nodeName: s.nodeName, 257 } 258 259 return copy 260 } 261 262 // ExecInBackground returns the results of running cmd via SSH in the specified 263 // context. The command will be executed in the background until context.Context 264 // is canceled or the command has finish its execution. 265 func (s *SSHMeta) ExecInBackground(ctx context.Context, cmd string, options ...ExecOptions) *CmdRes { 266 if ctx == nil { 267 panic("no context provided") 268 } 269 270 var ops ExecOptions 271 if len(options) > 0 { 272 ops = options[0] 273 } 274 275 fmt.Fprintln(SSHMetaLogs, cmd) 276 stdout := new(Buffer) 277 stderr := new(Buffer) 278 279 command := &SSHCommand{ 280 Path: cmd, 281 Stdin: os.Stdin, 282 Stdout: stdout, 283 Stderr: stderr, 284 } 285 var wg sync.WaitGroup 286 res := &CmdRes{ 287 cmd: cmd, 288 stdout: stdout, 289 stderr: stderr, 290 success: false, 291 wg: &wg, 292 } 293 294 res.wg.Add(1) 295 go func(res *CmdRes) { 296 defer res.wg.Done() 297 start := time.Now() 298 err := s.sshClient.RunCommandInBackground(ctx, command) 299 if err != nil { 300 exiterr, isExitError := err.(*ssh.ExitError) 301 if isExitError { 302 res.exitcode = exiterr.Waitmsg.ExitStatus() 303 // Set success as true if SIGINT signal was sent to command 304 if res.exitcode == 130 { 305 res.success = true 306 } 307 } 308 } else { 309 res.success = true 310 res.exitcode = 0 311 } 312 res.duration = time.Since(start) 313 res.SendToLog(ops.SkipLog) 314 }(res) 315 316 return res 317 }