github.com/mutagen-io/mutagen@v0.18.0-rc1/cmd/mutagen-agent/forwarder.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"os/signal"
     7  
     8  	"github.com/spf13/cobra"
     9  
    10  	"github.com/mutagen-io/mutagen/cmd"
    11  
    12  	"github.com/mutagen-io/mutagen/pkg/agent"
    13  	"github.com/mutagen-io/mutagen/pkg/forwarding/endpoint/remote"
    14  	"github.com/mutagen-io/mutagen/pkg/logging"
    15  	"github.com/mutagen-io/mutagen/pkg/mutagen"
    16  )
    17  
    18  // forwarderMain is the entry point for the forwarder command.
    19  func forwarderMain(_ *cobra.Command, _ []string) error {
    20  	// Create a channel to track termination signals. We do this before creating
    21  	// and starting other infrastructure so that we can ensure things terminate
    22  	// smoothly, not mid-initialization.
    23  	signalTermination := make(chan os.Signal, 1)
    24  	signal.Notify(signalTermination, cmd.TerminationSignals...)
    25  
    26  	// Set up a logger on the standard error stream.
    27  	logLevel := logging.LevelInfo
    28  	if forwarderConfiguration.logLevel != "" {
    29  		if l, ok := logging.NameToLevel(forwarderConfiguration.logLevel); !ok {
    30  			return fmt.Errorf("invalid log level specified: %s", forwarderConfiguration.logLevel)
    31  		} else {
    32  			logLevel = l
    33  		}
    34  	}
    35  	logger := logging.NewLogger(logLevel, os.Stderr)
    36  
    37  	// Create a stream using standard input/output.
    38  	stream := newStdioStream()
    39  
    40  	// Perform an agent handshake.
    41  	if err := agent.ServerHandshake(stream); err != nil {
    42  		return fmt.Errorf("server handshake failed: %w", err)
    43  	}
    44  
    45  	// Perform a version handshake.
    46  	if err := mutagen.ServerVersionHandshake(stream); err != nil {
    47  		return fmt.Errorf("version handshake error: %w", err)
    48  	}
    49  
    50  	// Serve a forwarder on standard input/output and monitor for its
    51  	// termination.
    52  	forwardingTermination := make(chan error, 1)
    53  	go func() {
    54  		forwardingTermination <- remote.ServeEndpoint(logger, stream)
    55  	}()
    56  
    57  	// Wait for termination from a signal or the forwarder.
    58  	select {
    59  	case s := <-signalTermination:
    60  		return fmt.Errorf("terminated by signal: %s", s)
    61  	case err := <-forwardingTermination:
    62  		return fmt.Errorf("forwarding terminated: %w", err)
    63  	}
    64  }
    65  
    66  // forwarderCommand is the forwarder command.
    67  var forwarderCommand = &cobra.Command{
    68  	Use:          agent.CommandForwarder,
    69  	Short:        "Run the agent in forwarder mode",
    70  	Args:         cmd.DisallowArguments,
    71  	RunE:         forwarderMain,
    72  	SilenceUsage: true,
    73  }
    74  
    75  // forwarderConfiguration stores configuration for the forwarder command.
    76  var forwarderConfiguration struct {
    77  	// help indicates whether or not to show help information and exit.
    78  	help bool
    79  	// logLevel indicates the log level to use.
    80  	logLevel string
    81  }
    82  
    83  func init() {
    84  	// Grab a handle for the command line flags.
    85  	flags := forwarderCommand.Flags()
    86  
    87  	// Disable alphabetical sorting of flags in help output.
    88  	flags.SortFlags = false
    89  
    90  	// Manually add a help flag to override the default message. Cobra will
    91  	// still implement its logic automatically.
    92  	flags.BoolVarP(&forwarderConfiguration.help, "help", "h", false, "Show help information")
    93  
    94  	// Wire up logging flags.
    95  	flags.StringVar(&forwarderConfiguration.logLevel, agent.FlagLogLevel, "", "Set the log level")
    96  }