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

     1  // Copyright 2023 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  //go:build !plan9 && !windows
     6  
     7  package main
     8  
     9  import (
    10  	"fmt"
    11  	"io"
    12  	"log"
    13  	"os/exec"
    14  
    15  	"github.com/creack/pty"
    16  	"github.com/gliderlabs/ssh"
    17  	"golang.org/x/build/internal/envutil"
    18  )
    19  
    20  func startSSHServerSwarming() {
    21  	buildletSSHServer = &ssh.Server{
    22  		Addr:    "localhost:" + sshPort(),
    23  		Handler: sshHandler,
    24  		PublicKeyHandler: func(ctx ssh.Context, key ssh.PublicKey) bool {
    25  			allowed, _, _, _, err := ssh.ParseAuthorizedKey(buldletAuthKeys)
    26  			if err != nil {
    27  				log.Printf("error parsing authorized key: %s", err)
    28  				return false
    29  			}
    30  			return ssh.KeysEqual(key, allowed)
    31  		},
    32  	}
    33  	go func() {
    34  		err := buildletSSHServer.ListenAndServe()
    35  		if err != nil {
    36  			log.Printf("buildlet SSH Server stopped: %s", err)
    37  		}
    38  	}()
    39  	teardownFuncs = append(teardownFuncs, func() {
    40  		buildletSSHServer.Close()
    41  		log.Println("shutting down SSH Server")
    42  	})
    43  }
    44  
    45  func sshHandler(s ssh.Session) {
    46  	ptyReq, winCh, isPty := s.Pty()
    47  	if !isPty {
    48  		fmt.Fprint(s, "scp is not supported\n")
    49  		return
    50  	}
    51  	var cmd *exec.Cmd
    52  	cmd = exec.Command(shell())
    53  
    54  	envutil.SetEnv(cmd, "TERM="+ptyReq.Term)
    55  	f, err := pty.Start(cmd)
    56  	if err != nil {
    57  		fmt.Fprintf(s, "unable to start shell %q: %s\n", shell(), err)
    58  		log.Printf("unable to start shell: %s", err)
    59  		return
    60  	}
    61  	defer f.Close()
    62  	go func() {
    63  		for win := range winCh {
    64  			pty.Setsize(f, &pty.Winsize{
    65  				Rows: uint16(win.Height),
    66  				Cols: uint16(win.Width),
    67  			})
    68  		}
    69  	}()
    70  	go func() {
    71  		io.Copy(f, s) // stdin
    72  	}()
    73  	io.Copy(s, f) // stdout
    74  	cmd.Process.Kill()
    75  	cmd.Wait()
    76  }