golang.org/x/build@v0.0.0-20240506185731-218518f32b70/cmd/buildlet/testssh/testssh.go (about)

     1  // Copyright 2019 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // The testssh binary exists to verify that a buildlet container's
     6  // ssh works, without running the whole coordinator binary in the
     7  // staging environment.
     8  package main
     9  
    10  import (
    11  	"bytes"
    12  	"flag"
    13  	"fmt"
    14  	"io"
    15  	"log"
    16  	"net"
    17  	"os"
    18  	"os/exec"
    19  	"path/filepath"
    20  	"strings"
    21  	"time"
    22  
    23  	"golang.org/x/build/buildenv"
    24  	"golang.org/x/build/buildlet"
    25  )
    26  
    27  var (
    28  	container  = flag.String("container", "", "if non-empty, the ID of a running docker container")
    29  	startImage = flag.String("start-image", "", "if non-empty, the Docker image to start a buildlet of locally, and use its container ID for the -container value")
    30  	user       = flag.String("user", "root", "SSH user")
    31  )
    32  
    33  func main() {
    34  	flag.Parse()
    35  	ipPort := getIPPort()
    36  	defer cleanContainer()
    37  
    38  	bc := buildlet.NewClient(ipPort, buildlet.NoKeyPair)
    39  	for {
    40  		c, err := net.Dial("tcp", ipPort)
    41  		if err == nil {
    42  			c.Close()
    43  			break
    44  		}
    45  		log.Printf("waiting for %v to come up...", ipPort)
    46  		time.Sleep(time.Second)
    47  	}
    48  
    49  	pubKey, privPath := genKey()
    50  
    51  	log.Printf("hitting buildlet's /connect-ssh ...")
    52  	buildletConn, err := bc.ConnectSSH(*user, pubKey)
    53  	if err != nil {
    54  		var out []byte
    55  		if *container != "" {
    56  			var err error
    57  			out, err = exec.Command("docker", "logs", *container).CombinedOutput()
    58  			if err != nil {
    59  				log.Printf("failed to fetch docker logs: %v", err)
    60  			}
    61  		}
    62  		cleanContainer()
    63  		log.Printf("image logs: %s", out)
    64  		log.Fatalf("ConnectSSH: %v (logs above)", err)
    65  	}
    66  	defer buildletConn.Close()
    67  	log.Printf("ConnectSSH succeeded; testing connection...")
    68  
    69  	ln, err := net.Listen("tcp", "localhost:0")
    70  	if err != nil {
    71  		log.Fatal(err)
    72  	}
    73  	go func() {
    74  		c, err := ln.Accept()
    75  		if err != nil {
    76  			log.Fatal(err)
    77  		}
    78  		go io.Copy(buildletConn, c)
    79  		go io.Copy(c, buildletConn)
    80  	}()
    81  	ip, port, err := net.SplitHostPort(ln.Addr().String())
    82  	if err != nil {
    83  		log.Fatal(err)
    84  	}
    85  
    86  	cmd := exec.Command("ssh",
    87  		"-v",
    88  		"-i", privPath,
    89  		"-o", "StrictHostKeyChecking=no",
    90  		"-o", "UserKnownHostsFile=/dev/null",
    91  		"-o", "LogLevel=ERROR",
    92  		"-p", port,
    93  		*user+"@"+ip,
    94  		"echo", "SSH works")
    95  	stdout := new(bytes.Buffer)
    96  	stderr := new(bytes.Buffer)
    97  	cmd.Stderr = stderr
    98  	cmd.Stdout = stdout
    99  	if err := cmd.Run(); err != nil {
   100  		cleanContainer()
   101  		log.Fatalf("ssh client: %v, %s", err, stderr)
   102  	}
   103  	fmt.Print(stdout.String())
   104  }
   105  
   106  func cleanContainer() {
   107  	if *startImage == "" {
   108  		return
   109  	}
   110  	out, err := exec.Command("docker", "rm", "-f", *container).CombinedOutput()
   111  	if err != nil {
   112  		log.Printf("docker rm: %v, %s", err, out)
   113  	}
   114  }
   115  
   116  func genKey() (pubKey, privateKeyPath string) {
   117  	cache, err := os.UserCacheDir()
   118  	if err != nil {
   119  		log.Fatal(err)
   120  	}
   121  	cache = filepath.Join(cache, "testssh")
   122  	os.MkdirAll(cache, 0755)
   123  	privateKeyPath = filepath.Join(cache, "testkey")
   124  	pubKeyPath := filepath.Join(cache, "testkey.pub")
   125  	if _, err := os.Stat(pubKeyPath); err != nil {
   126  		out, err := exec.Command("ssh-keygen", "-t", "ed25519", "-f", privateKeyPath, "-N", "").CombinedOutput()
   127  		if err != nil {
   128  			log.Fatalf("ssh-keygen: %v, %s", err, out)
   129  		}
   130  	}
   131  	slurp, err := os.ReadFile(pubKeyPath)
   132  	if err != nil {
   133  		log.Fatal(err)
   134  	}
   135  	return strings.TrimSpace(string(slurp)), privateKeyPath
   136  }
   137  
   138  func getIPPort() string {
   139  	if *startImage != "" {
   140  		buildlet := "buildlet.linux-amd64"
   141  		log.Printf("creating container with image %s ...", *startImage)
   142  		out, err := exec.Command("docker", "run", "-d",
   143  			"--stop-timeout=300",
   144  			"-e", "META_BUILDLET_BINARY_URL=https://storage.googleapis.com/"+buildenv.Production.BuildletBucket+"/"+buildlet,
   145  			*startImage).CombinedOutput()
   146  		if err != nil {
   147  			log.Fatalf("docker run: %v, %s", err, out)
   148  		}
   149  		*container = strings.TrimSpace(string(out))
   150  		log.Printf("created container %s ...", *container)
   151  	}
   152  	if *container != "" {
   153  		out, err := exec.Command("bash", "-c", "docker inspect "+*container+" | jq -r '.[0].NetworkSettings.IPAddress'").CombinedOutput()
   154  		if err != nil {
   155  			log.Fatalf("%v: %s", err, out)
   156  		}
   157  		return strings.TrimSpace(string(out)) + ":80"
   158  	}
   159  	log.Fatalf("no address specified")
   160  	return ""
   161  }