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

     1  // Copyright 2017 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 sshtest
     6  
     7  import (
     8  	"crypto/rsa"
     9  	"errors"
    10  	"os"
    11  )
    12  
    13  const (
    14  	defaultKeyBits = 1024
    15  )
    16  
    17  // MustGenerateKeys can be called from a test file's init function to generate
    18  // 1024-bit user and host keys. Panics on error.
    19  func MustGenerateKeys() (userKey, hostKey *rsa.PrivateKey) {
    20  	var err error
    21  	if userKey, hostKey, err = GenerateKeys(defaultKeyBits); err != nil {
    22  		panic(err)
    23  	}
    24  	return userKey, hostKey
    25  }
    26  
    27  // TestData contains common data that can be used by tests that interact with an SSHServer.
    28  type TestData struct { // NOLINT
    29  	Srvs        []*SSHServer
    30  	UserKeyFile string
    31  }
    32  
    33  // NewTestData initializes and returns a TestData struct. Panics on error.
    34  func NewTestData(handlers ...ExecHandler) *TestData {
    35  	if len(handlers) == 0 {
    36  		panic(errors.New("no handler is specfied"))
    37  	}
    38  	userKey, hostKey := MustGenerateKeys()
    39  	var servers []*SSHServer
    40  	for _, handler := range handlers {
    41  		srv, err := NewSSHServer(&userKey.PublicKey, hostKey, handler)
    42  		if err != nil {
    43  			for _, srv := range servers {
    44  				srv.Close()
    45  			}
    46  			panic(err)
    47  		}
    48  		servers = append(servers, srv)
    49  	}
    50  	userKeyFile, err := WriteKey(userKey)
    51  	if err != nil {
    52  		for _, srv := range servers {
    53  			srv.Close()
    54  		}
    55  		panic(err)
    56  	}
    57  	return &TestData{
    58  		Srvs:        servers,
    59  		UserKeyFile: userKeyFile,
    60  	}
    61  }
    62  
    63  // Close stops the SSHServer and deletes the user key file.
    64  func (d *TestData) Close() {
    65  	for _, s := range d.Srvs {
    66  		s.Close()
    67  	}
    68  	os.Remove(d.UserKeyFile)
    69  }