github.com/nya3jp/tast@v0.0.0-20230601000426-85c8e4d83a9b/src/go.chromium.org/tast/core/internal/fakesshserver/fakesshserver.go (about)

     1  // Copyright 2021 The ChromiumOS Authors
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  // Package fakesshserver implements a fake SSH server.
     6  package fakesshserver
     7  
     8  import (
     9  	"crypto/rsa"
    10  	"io"
    11  	"net"
    12  	"os/exec"
    13  	"strings"
    14  
    15  	"go.chromium.org/tast/core/internal/sshtest"
    16  )
    17  
    18  // Process implements a simulated process started by a fake SSH server.
    19  type Process func(stdin io.Reader, stdout, stderr io.Writer) int
    20  
    21  // Handler receives a command requested by an SSH client and decides whether to
    22  // handle the request.
    23  // If it returns true, a reply is sent to the client indicating that the command
    24  // is accepted, and returned Process is called with stdin/stdout/stderr.
    25  // If it returns false, an unsuccessful reply is sent to the client.
    26  type Handler func(cmd string) (Process, bool)
    27  
    28  // ExactMatchHandler constructs a Handler that replies to a command request by
    29  // proc if it exactly matches with cmd.
    30  func ExactMatchHandler(cmd string, proc Process) Handler {
    31  	return func(c string) (Process, bool) {
    32  		if c != cmd {
    33  			return nil, false
    34  		}
    35  		return proc, true
    36  	}
    37  }
    38  
    39  // ShellHandler constructs a Handler that replies to a command request by
    40  // running it as is with "sh -c" if its prefix matches with the given prefix.
    41  func ShellHandler(prefix string) Handler {
    42  	return func(c string) (Process, bool) {
    43  		if !strings.HasPrefix(c, prefix) {
    44  			return nil, false
    45  		}
    46  		return func(stdin io.Reader, stdout, stderr io.Writer) int {
    47  			cmd := exec.Command("sh", "-c", c)
    48  			cmd.Stdin = stdin
    49  			cmd.Stdout = stdout
    50  			cmd.Stderr = stderr
    51  			err := cmd.Run()
    52  			if err != nil {
    53  				if xerr, ok := err.(*exec.ExitError); ok {
    54  					return xerr.ExitCode()
    55  				}
    56  				return 255
    57  			}
    58  			return 0
    59  		}, true
    60  	}
    61  }
    62  
    63  // Server maintains resources related to a fake SSH server.
    64  type Server struct {
    65  	server *sshtest.SSHServer
    66  }
    67  
    68  // Start starts a new fake SSH server.
    69  func Start(userKey *rsa.PublicKey, hostKey *rsa.PrivateKey, handlers []Handler) (*Server, error) {
    70  	server, err := sshtest.NewSSHServer(userKey, hostKey, func(req *sshtest.ExecReq) {
    71  		for _, handler := range handlers {
    72  			cmd, ok := handler(req.Cmd)
    73  			if !ok {
    74  				continue
    75  			}
    76  			req.Start(true)
    77  			status := cmd(req, req, req.Stderr())
    78  			req.End(status)
    79  			return
    80  		}
    81  		req.Start(false)
    82  	})
    83  	if err != nil {
    84  		return nil, err
    85  	}
    86  	return &Server{server: server}, nil
    87  }
    88  
    89  // Stop stops the fake SSH server.
    90  func (s *Server) Stop() {
    91  	s.server.Close()
    92  }
    93  
    94  // Addr returns the address the server listens to.
    95  func (s *Server) Addr() net.Addr { return s.server.Addr() }