[ Web Proxy ]
URL:
Viewing: https://raw.githubusercontent.com/pingf/sshcode/reuse-ssh-connection/sshcode.go [Back]  [Original]

package main

import (
	"context"
	"fmt"
	"math/rand"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/pkg/browser"
	"go.coder.com/flog"
	"golang.org/x/xerrors"
)

const codeServerPath = "~/.cache/sshcode/sshcode-server"
const sshDirectory = "~/.ssh"
const sshDirectoryUnsafeModeMask = 0022
const sshControlPath = sshDirectory + "/control-%h-%p-%r"

type options struct {
	skipSync          bool
	syncBack          bool
	noOpen            bool
	noReuseConnection bool
	bindAddr          string
	remotePort        string
	sshFlags          string
}

func sshCode(host, dir string, o options) error {
	host, extraSSHFlags, err := parseHost(host)
	if err != nil {
		return xerrors.Errorf("failed to parse host IP: %w", err)
	}
	if extraSSHFlags != "" {
		o.sshFlags = strings.Join([]string{extraSSHFlags, o.sshFlags}, " ")
	}

	o.bindAddr, err = parseBindAddr(o.bindAddr)
	if err != nil {
		return xerrors.Errorf("failed to parse bind address: %w", err)
	}

	if o.remotePort == "" {
		o.remotePort, err = randomPort()
	}
	if err != nil {
		return xerrors.Errorf("failed to find available remote port: %w", err)
	}

	// Check the SSH directory's permissions and warn the user if it is not safe.
	sshDirectoryMode, err := os.Lstat(expandPath(sshDirectory))
	if err != nil {
		if !o.noReuseConnection {
			flog.Info("failed to stat %v directory, disabling connection reuse feature: %v", sshDirectory, err)
			o.noReuseConnection = true
		}
	} else {
		if !sshDirectoryMode.IsDir() {
			if !o.noReuseConnection {
				flog.Info("%v is not a directory, disabling connection reuse feature", sshDirectory)
				o.noReuseConnection = true
			} else {
				flog.Info("warning: %v is not a directory", sshDirectory)
			}
		}
		if sshDirectoryMode.Mode().Perm()&sshDirectoryUnsafeModeMask != 0 {
			flog.Info("warning: the %v directory has unsafe permissions, they should only be writable by "+
				"the owner (and files inside should be set to 0600)", sshDirectory)
		}
	}

	// Start SSH master connection socket. This prevents multiple password prompts from appearing as authentication
	// only happens on the initial connection.
	if !o.noReuseConnection {
		newSSHFlags := fmt.Sprintf(`%v -o "ControlPath=%v"`, o.sshFlags, sshControlPath)

		// -MN means "start a master socket and don't open a session, just connect".
		sshCmdStr := fmt.Sprintf(`exec ssh %v -MN %v`, newSSHFlags, host)
		sshMasterCmd := exec.Command("sh", "-c", sshCmdStr)
		sshMasterCmd.Stdin = os.Stdin
		sshMasterCmd.Stdout = os.Stdout
		sshMasterCmd.Stderr = os.Stderr
		stopSSHMaster := func() {
			if sshMasterCmd.Process != nil {
				err := sshMasterCmd.Process.Signal(syscall.Signal(0))
				if err != nil {
					return
				}
				err = sshMasterCmd.Process.Signal(syscall.SIGTERM)
				if err != nil {
					flog.Error("failed to send SIGTERM to SSH master process: %v", err)
				}
			}
		}
		defer stopSSHMaster()

		err = sshMasterCmd.Start()
		go sshMasterCmd.Wait()
		if err != nil {
			flog.Error("failed to start SSH master connection, disabling connection reuse feature: %v", err)
			o.noReuseConnection = true
			stopSSHMaster()
		} else {
			err = checkSSHMaster(sshMasterCmd, newSSHFlags, host)
			if err != nil {
				flog.Error("SSH master failed to be ready in time, disabling connection reuse feature: %v", err)
				o.noReuseConnection = true
				stopSSHMaster()
			} else {
				sshMasterCmd.Stdin = nil
				o.sshFlags = newSSHFlags
			}
		}
	}

	flog.Info("ensuring code-server is updated...")
	dlScript := downloadScript(codeServerPath)

	// Downloads the latest code-server and allows it to be executed.
	sshCmdStr := fmt.Sprintf("ssh %v %v '/usr/bin/env bash'", o.sshFlags, host)

	sshCmd := exec.Command("sh", "-c", sshCmdStr)
	sshCmd.Stdout = os.Stdout
	sshCmd.Stderr = os.Stderr
	sshCmd.Stdin = strings.NewReader(dlScript)
	err = sshCmd.Run()
	if err != nil {
		return xerrors.Errorf("failed to update code-server: \n---ssh cmd---\n%s\n---download script---\n%s: %w",
			sshCmdStr,
			dlScript,
			err,
		)
	}

	if !o.skipSync {
		start := time.Now()
		flog.Info("syncing settings")
		err = syncUserSettings(o.sshFlags, host, false)
		if err != nil {
			return xerrors.Errorf("failed to sync settings: %w", err)
		}

		flog.Info("synced settings in %s", time.Since(start))

		flog.Info("syncing extensions")
		err = syncExtensions(o.sshFlags, host, false)
		if err != nil {
			return xerrors.Errorf("failed to sync extensions: %w", err)
		}
		flog.Info("synced extensions in %s", time.Since(start))
	}

	flog.Info("starting code-server...")

	flog.Info("Tunneling remote port %v to %v", o.remotePort, o.bindAddr)

	sshCmdStr =
		fmt.Sprintf("ssh -tt -q -L %v:localhost:%v %v %v 'cd %v; %v --host 127.0.0.1 --allow-http --no-auth --port=%v'",
			o.bindAddr, o.remotePort, o.sshFlags, host, dir, codeServerPath, o.remotePort,
		)

	// Starts code-server and forwards the remote port.
	sshCmd = exec.Command("sh", "-c", sshCmdStr)
	sshCmd.Stdin = os.Stdin
	sshCmd.Stdout = os.Stdout
	sshCmd.Stderr = os.Stderr
	err = sshCmd.Start()
	if err != nil {
		return xerrors.Errorf("failed to start code-server: %w", err)
	}

	url := fmt.Sprintf("http://%s", o.bindAddr)
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	client := http.Client{
		Timeout: time.Second * 3,
	}
	for {
		if ctx.Err() != nil {
			return xerrors.Errorf("code-server didn't start in time: %w", ctx.Err())
		}
		// Waits for code-server to be available before opening the browser.
		resp, err := client.Get(url)
		if err != nil {
			continue
		}
		resp.Body.Close()
		break
	}

	ctx, cancel = context.WithCancel(context.Background())

	if !o.noOpen {
		openBrowser(url)
	}

	go func() {
		defer cancel()
		sshCmd.Wait()
	}()

	c := make(chan os.Signal)
	signal.Notify(c, os.Interrupt)

	select {
	case 

Web Proxy Viewer  |  New URL  |  Original Page