chisel/main.go

package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"strconv"
	"strings"
	"time"

	chclient "github.com/jpillora/chisel/client"
	chserver "github.com/jpillora/chisel/server"
	chshare "github.com/jpillora/chisel/share"
	"github.com/jpillora/chisel/share/ccrypto"
	"github.com/jpillora/chisel/share/cos"
	"github.com/jpillora/chisel/share/settings"
	"github.com/jpillora/opts"
)

// chiselRepo is rendered in the "Read more:" section of every help screen.
const chiselRepo = "https://github.com/jpillora/chisel"

// versionTemplate keeps the Go runtime version alongside the build version in
// the help text. The --version flag still prints the build version alone.
var versionTemplate = "\n{{bold \"Version:\"}}\n{{.Pad}}{{accent \"" +
	chshare.BuildVersion + " (" + runtime.Version() + ")\"}}\n"

type rootOptions struct{}

func main() {
	serverConfig := newServerOptions()
	clientConfig := newClientOptions()
	cli := newCLI(serverConfig, clientConfig)

	parsed, err := cli.ParseArgsError(os.Args)
	if err != nil {
		// opts deliberately models help and version as successful parse exits.
		// Keep Chisel's historical stdout and trailing-newline behaviour.
		if isHelpExit(err, os.Args[1:]) {
			fmt.Print(err)
			return
		}
		if err.Error() == chshare.BuildVersion && hasRootVersionFlag(os.Args[1:]) {
			fmt.Println(err)
			return
		}
		fmt.Fprint(os.Stderr, parsed.Selected().Help())
		os.Exit(1)
	}
	if !parsed.IsRunnable() {
		fmt.Print(parsed.Help())
		return
	}
	if err := parsed.Run(); err != nil {
		log.Fatal(err)
	}
}

func newCLI(serverConfig *serverOptions, clientConfig *clientOptions) opts.Opts {
	return opts.New(&rootOptions{}).
		Name("chisel").
		Version(chshare.BuildVersion).
		Repo(chiselRepo).
		DocSet("version", versionTemplate).
		AddCommand(newServerCLI(serverConfig)).
		AddCommand(newClientCLI(clientConfig))
}

func newServerCLI(serverConfig *serverOptions) opts.Opts {
	return opts.New(serverConfig).
		Name("server").
		Summary("runs chisel in server mode").
		Repo(chiselRepo).
		DocSet("version", versionTemplate).
		DocAfter("flaggroups", "signals", signalsTemplate)
}

func newClientCLI(clientConfig *clientOptions) opts.Opts {
	return opts.New(clientConfig).
		Name("client").
		Summary("runs chisel in client mode").
		Repo(chiselRepo).
		DocSet("version", versionTemplate).
		DocAfter("summary", "remotes", remotesTemplate).
		DocAfter("flaggroups", "signals", signalsTemplate)
}

func hasHelpFlag(args []string) bool {
	for _, arg := range args {
		if arg == "--help" || arg == "-help" || arg == "--h" || arg == "-h" {
			return true
		}
	}
	return false
}

func isHelpExit(err error, args []string) bool {
	message := err.Error()
	return hasHelpFlag(args) && strings.Contains(message, "Usage:") && !strings.Contains(message, "Error:")
}

func hasRootVersionFlag(args []string) bool {
	for _, arg := range args {
		if arg == "--" || arg == "-" || !strings.HasPrefix(arg, "-") {
			return false
		}
		if arg == "--version" || arg == "-version" || arg == "--v" || arg == "-v" {
			return true
		}
	}
	return false
}

func optsTagValue(tag, key string) string {
	for _, part := range strings.Split(tag, ",") {
		name, value, ok := strings.Cut(strings.TrimSpace(part), "=")
		if ok && name == key {
			return value
		}
	}
	return ""
}

// signalsTemplate documents the signals both commands handle. The layout is
// deliberate, so it is rendered verbatim rather than reflowed.
const signalsTemplate = "\n{{bold \"Signals:\"}}\n" +
	"{{.Pad}}The chisel process is listening for:\n" +
	"{{.Pad}}  a SIGINT or SIGTERM to begin a graceful shutdown\n" +
	"{{.Pad}}    (a second signal forces an immediate exit),\n" +
	"{{.Pad}}  a SIGUSR2 to print process stats, and\n" +
	"{{.Pad}}  a SIGHUP to short-circuit the client reconnect timer\n"

func generatePidFile() error {
	pid := []byte(strconv.Itoa(os.Getpid()))
	return os.WriteFile("chisel.pid", pid, 0644)
}

type sharedStringFlag struct {
	target *string
}

func (f sharedStringFlag) String() string {
	if f.target == nil {
		return ""
	}
	return *f.target
}

func (f *sharedStringFlag) Set(value string) error {
	*f.target = value
	return nil
}

// Field order determines the order flags appear in the help text.
type serverOptions struct {
	Host   string `opts:"name=host,short=-" help:"Defines the HTTP listening host – the network interface (defaults the environment variable HOST and falls back to 0.0.0.0)."`
	Port   string `opts:"name=port,short=p" help:"Defines the HTTP listening port (defaults to the environment variable PORT and falls back to port 8080)."`
	KeyGen string `opts:"name=keygen,short=-" help:"A path to write a newly generated PEM-encoded SSH private key file. If users depend on your --key fingerprint, you may also include your --key to output your existing key. Use - (dash) to output the generated key to stdout."`
	chserver.Config
	ProxyAlias sharedStringFlag `opts:"name=proxy,mode=flag,short=-" help:"Alias for --backend."`
	PID        bool             `opts:"name=pid,short=-" help:"Generate pid file in current working directory"`
	Verbose    bool             `opts:"name=verbose,short=v" help:"Enable verbose logging"`
}

func newServerOptions() *serverOptions {
	o := &serverOptions{}
	o.Config.KeepAlive = 25 * time.Second
	o.ProxyAlias.target = &o.Config.Proxy
	return o
}

func (o *serverOptions) Run() error {
	config := &o.Config

	if o.KeyGen != "" {
		return ccrypto.GenerateKeyFile(o.KeyGen, config.KeySeed)
	}

	if config.KeySeed != "" {
		log.Print("Option `--key` is deprecated and will be removed in a future version of chisel.")
		log.Print("Please use `chisel server --keygen /file/path`, followed by `chisel server --keyfile /file/path` to specify the SSH private key")
	}

	if o.Host == "" {
		o.Host = os.Getenv("HOST")
	}
	if o.Host == "" {
		o.Host = "0.0.0.0"
	}
	if o.Port == "" {
		o.Port = os.Getenv("PORT")
	}
	if o.Port == "" {
		o.Port = "8080"
	}
	if config.KeyFile == "" {
		config.KeyFile = settings.Env("KEY_FILE")
	}
	if config.KeySeed == "" {
		config.KeySeed = settings.Env("KEY")
	}
	if config.Auth == "" {
		config.Auth = os.Getenv("AUTH")
	}
	s, err := chserver.NewServer(config)
	if err != nil {
		return err
	}
	s.Debug = o.Verbose
	if o.PID {
		if err := generatePidFile(); err != nil {
			return err
		}
	}
	go cos.GoStats()
	ctx := cos.InterruptContext()
	if err := s.StartContext(ctx, o.Host, o.Port); err != nil {
		return err
	}
	return s.Wait()
}

func setHeader(headers http.Header, arg string) error {
	index := strings.Index(arg, ":")
	if index < 0 {
		return fmt.Errorf(`Invalid header (%s). Should be in the format "HeaderName: HeaderContent"`, arg)
	}
	key := arg[0:index]
	value := arg[index+1:]
	headers.Set(key, strings.TrimSpace(value))
	return nil
}

// remotesTemplate documents the <remote> syntax. Its bullet lists, examples and
// ProxyCommand snippet are laid out by hand, so it is rendered verbatim.
const remotesTemplate = `
<server> is the URL to the chisel server.

<remote>s are remote connections tunneled through the server, each of
which come in the form:

  <local-host>:<local-port>:<remote-host>:<remote-port>/<protocol>

  ■ local-host defaults to 0.0.0.0 (all interfaces).
  ■ local-port defaults to remote-port.
  ■ remote-port is required*.
  ■ remote-host defaults to 127.0.0.1 (server localhost).
  ■ protocol defaults to tcp.

which shares <remote-host>:<remote-port> from the server to the client
as <local-host>:<local-port>, or:

  R:<local-interface>:<local-port>:<remote-host>:<remote-port>/<protocol>

which does reverse port forwarding, sharing <remote-host>:<remote-port>
from the client to the server's <local-interface>:<local-port>.

  example remotes

    3000
    example.com:3000
    3000:google.com:80
    192.168.0.5:3000:google.com:80
    socks
    5000:socks
    udpsocks
    5000:udpsocks
    R:2222:localhost:22
    R:socks
    R:5000:socks
    R:5000:udpsocks
    stdio:example.com:22
    1.1.1.1:53/udp

  When the chisel server has --socks5 enabled, remotes can
  specify "socks" or "udpsocks" in place of remote-host and
  remote-port. Both default to the local host and port
  127.0.0.1:1080. Connections to this remote will terminate
  at the server's internal SOCKS5 proxy. When the server also
  has --authfile set, SOCKS5 access requires an entry matching
  the token "socks" in the user's address list.

  When the chisel server has --reverse enabled, remotes can
  be prefixed with R to denote that they are reversed. That
  is, the server will listen and accept connections, and they
  will be proxied through the client which specified the remote.
  Reverse remotes specifying "R:socks" will listen on the server's
  default socks port (1080) and terminate the connection at the
  client's internal SOCKS5 proxy.

  The "socks" remote supports TCP CONNECT only. Use "udpsocks"
  for TCP CONNECT and UDP ASSOCIATE, including "R:udpsocks".
  It binds TCP and UDP on the same configured interface and port;
  the UDP relay never uses a randomly allocated listening port.
  UDP requires support on both Chisel peers. Destination names are
  resolved at the exit node. If UDP cannot bind, TCP remains available
  and UDP ASSOCIATE fails; no alternate port is used. Clients using
  a zero UDP source port allow one pending association per source IP
  until the first valid datagram pins it. BIND, UDP fragmentation,
  and udpsocks over stdio are not supported.

  When stdio is used as local-host, the tunnel will connect standard
  input/output of this program with the remote. This is useful when
  combined with ssh ProxyCommand. You can use
    ssh -o ProxyCommand='chisel client chiselserver stdio:%h:%p' \
        [email protected]
  to connect to an SSH server through the tunnel.
`

// Field order determines the order flags appear in the help text.
type clientOptions struct {
	chclient.Config
	Header   []string `opts:"name=header,short=-" help:"Set a custom header in the form \"HeaderName: HeaderContent\". Can be used multiple times. (e.g --header \"Foo: Bar\" --header \"Hello: World\")"`
	Hostname string   `opts:"name=hostname,short=-" help:"Optionally set the 'Host' header (defaults to the host found in the server url)."`
	SNI      string   `opts:"name=sni,short=-" help:"Override the ServerName when using TLS (defaults to the hostname)."`
	PID      bool     `opts:"name=pid,short=-" help:"Generate pid file in current working directory"`
}

func newClientOptions() *clientOptions {
	o := &clientOptions{}
	o.Config.KeepAlive = 25 * time.Second
	o.Config.MaxRetryCount = -1
	o.Config.Headers = http.Header{}
	return o
}

func (o *clientOptions) Run() error {
	config := &o.Config
	for _, header := range o.Header {
		if err := setHeader(config.Headers, header); err != nil {
			return err
		}
	}
	//default auth
	if config.Auth == "" {
		config.Auth = os.Getenv("AUTH")
	}
	//move hostname onto headers
	if o.Hostname != "" {
		config.Headers.Set("Host", o.Hostname)
		config.TLS.ServerName = o.Hostname
	}

	if o.SNI != "" {
		config.TLS.ServerName = o.SNI
	}

	//ready
	c, err := chclient.NewClient(config)
	if err != nil {
		return err
	}
	c.Debug = config.Verbose
	if o.PID {
		if err := generatePidFile(); err != nil {
			return err
		}
	}
	go cos.GoStats()
	ctx := cos.InterruptContext()
	if err := c.Start(ctx); err != nil {
		return err
	}
	return c.Wait()
}