« Back to Index

Go: local SSH tunnel using google cloud

View original Gist on GitHub

Tags: #go #ssh

proxy.go

// Package main is the entrypoint to the proxy CLI program.
package main

import (
	"bytes"
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"syscall"

	"golang.org/x/sys/unix"
)

// Open an SSH tunnel on a local port to the given proxy address.
func main() {
	if err := run(); err != nil {
		log.Fatalf("failed to run: %s", err)
	}
}

func run() error {
	proxyAddr := flag.String("connect", "", "open SSH tunnel to `IP address` (omit port)")
	login := flag.Bool("login", false, "log in instead of opening tunnel")
	port := flag.String("port", "1080", "local `port` number of tunnel")
	flag.Parse()

	gcpURL, err := getGCPURL(*proxyAddr)
	if gcpURL == "" || err != nil {
		if err != nil {
			return fmt.Errorf("failed to get GCP URL: %w", err)
		}
		msg := "no GCP proxy with IP address %v found; make sure you’ve run 'gcloud components install beta' first; check gcloud auth list and the IP address"
		return fmt.Errorf(msg, *proxyAddr)
	}

	if *login {
		loginReplaceProc(gcpURL)
		return nil
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cmd, err := connect(ctx, gcpURL, *port, os.Stderr)
	if err != nil {
		return nil
	}
	if cmd == nil {
		return fmt.Errorf("no proxy with address: %#v", *proxyAddr)
	}

	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
		<-c
		cancel()
	}()

	fmt.Fprintf(os.Stdout, "forwarding localhost:%s to proxy: %s\n", *port, *proxyAddr)
	err = cmd.Wait()
	if err != nil {
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			err = fmt.Errorf("exit error: %s: %w", exitErr.Stderr, err)
		}
		return err
	}

	return nil
}

func connect(ctx context.Context, url, port string, out io.Writer) (*exec.Cmd, error) {
	cmd := exec.CommandContext(ctx, "gcloud", "compute", "ssh", "--ssh-flag=-nNTC -D localhost:"+port, "--project=<REDACTED>", url) // #nosec: G204
	if out != nil {
		cmd.Stdout = out
		cmd.Stderr = out
	}
	err := cmd.Start()
	if err != nil {
		return nil, fmt.Errorf("failed to start gcloud command: %w", err)
	}
	return cmd, nil
}

func loginReplaceProc(url string) {
	cmd := exec.Command("gcloud", "beta", "compute", "ssh", "--project=<REDACTED>", url)
	err := unix.Exec(cmd.Path, cmd.Args, os.Environ())
	fmt.Fprintln(os.Stderr, err)
}

func getGCPURL(addr string) (string, error) {
	cmd := exec.Command("gcloud", "beta", "compute", "addresses", "list", "--format=value(users[0].flatten())", "--project=<REDACTED>", "--filter=labels.production=ssh-proxy AND address="+addr) // #nosec: G204
	out, err := cmd.Output()
	out = bytes.TrimSpace(out)
	return string(out), err
}