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  }