Simple tcp to socket multiplexing reverse proxy (#438)

Signed-off-by: David Chung <[email protected]>

Author: David Chung

Makefile

@@ -126,6 +126,14 @@ endif
 	$(call build_binary,infrakit-instance-maas,github.com/docker/infrakit/examples/instance/maas)
 	$(call build_binary,infrakit-event-time,github.com/docker/infrakit/examples/event/time)
 
+cli: build-cli
+build-cli:
+	@echo "+ $@"
+ifneq (,$(findstring .m,$(VERSION)))
+	@echo "\nWARNING - repository contains uncommitted changes, tagging binaries as dirty\n"
+endif
+	$(call build_binary,infrakit,github.com/docker/infrakit/cmd/cli)
+
 install:
 	@echo "+ $@"
 	@go install ${GO_LDFLAGS} $(PKGS)

cmd/cli/main.go

@@ -2,38 +2,73 @@ package main
 
 import (
 	"errors"
+	"flag"
+	"net/url"
 	"os"
 
-	log "github.com/Sirupsen/logrus"
 	"github.com/docker/infrakit/pkg/cli"
 	"github.com/docker/infrakit/pkg/discovery"
+	"github.com/docker/infrakit/pkg/discovery/remote"
+	logutil "github.com/docker/infrakit/pkg/log"
 	"github.com/spf13/cobra"
 )
 
 // A generic client for infrakit
-
 func main() {
 
+	log := logutil.New("module", "main")
+
 	cmd := &cobra.Command{
 		Use:   os.Args[0],
 		Short: "infrakit cli",
 	}
-	logLevel := cmd.PersistentFlags().Int("log", cli.DefaultLogLevel, "Logging level. 0 is least verbose. Max is 5")
+
+	// Log setup
+	logOptions := &logutil.ProdDefaults
+	ulist := []*url.URL{}
+	remotes := []string{}
+
+	cmd.PersistentFlags().AddFlagSet(cli.Flags(logOptions))
+	cmd.PersistentFlags().AddGoFlagSet(flag.CommandLine)
+
+	cmd.PersistentFlags().StringSliceVarP(&remotes, "host", "H", remotes, "host list. Default is local sockets")
+
+	// parse the list of hosts
 	cmd.PersistentPreRunE = func(c *cobra.Command, args []string) error {
-		cli.SetLogLevel(*logLevel)
+		logutil.Configure(logOptions)
+
+		if len(remotes) > 0 {
+			for _, h := range remotes {
+				u, err := url.Parse(h)
+				if err != nil {
+					return err
+				}
+				ulist = append(ulist, u)
+			}
+		}
 		return nil
 	}
 
-	// Don't print usage text for any error returned from a RunE function.  Only print it when explicitly requested.
+	// Don't print usage text for any error returned from a RunE function.
+	// Only print it when explicitly requested.
 	cmd.SilenceUsage = true
 
-	// Don't automatically print errors returned from a RunE function.  They are returned from cmd.Execute() below
-	// and we print it ourselves.
+	// Don't automatically print errors returned from a RunE function.
+	// They are returned from cmd.Execute() below and we print it ourselves.
 	cmd.SilenceErrors = true
 	f := func() discovery.Plugins {
-		d, err := discovery.NewPluginDiscovery()
+		if len(ulist) == 0 {
+			d, err := discovery.NewPluginDiscovery()
+			if err != nil {
+				log.Crit("Failed to initialize plugin discovery", "err", err)
+				os.Exit(1)
+			}
+			return d
+		}
+
+		d, err := remote.NewPluginDiscovery(ulist)
 		if err != nil {
-			log.Fatalf("Failed to initialize plugin discovery: %s", err)
+			log.Crit("Failed to initialize remote plugin discovery", "err", err)
 			os.Exit(1)
 		}
 		return d
@@ -46,19 +81,20 @@ func main() {
 	cmd.AddCommand(metadataCommand(f))
 	cmd.AddCommand(eventCommand(f))
 	cmd.AddCommand(pluginCommand(f))
+	cmd.AddCommand(utilCommand(f))
 
 	cmd.AddCommand(instancePluginCommand(f), groupPluginCommand(f), flavorPluginCommand(f), resourcePluginCommand(f))
 
 	err := cmd.Execute()
 	if err != nil {
-		log.Fatal(err)
+		log.Crit("error executing", "err", err)
 		os.Exit(1)
 	}
 }
 
 func assertNotNil(message string, f interface{}) {
 	if f == nil {
-		log.Error(errors.New(message))
+		logutil.New("main", "assert").Error("assert failed", "err", errors.New(message))
 		os.Exit(1)
 	}
 }

cmd/cli/util.go

@@ -0,0 +1,74 @@
+package main
+
+import (
+	"time"
+
+	"github.com/docker/infrakit/pkg/discovery"
+	"github.com/docker/infrakit/pkg/log"
+	"github.com/docker/infrakit/pkg/rpc/mux"
+	"github.com/spf13/cobra"
+)
+
+var logger = log.New("module", "cli/util")
+
+func muxCommand(plugins func() discovery.Plugins) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "mux",
+		Short: "API mux service",
+	}
+
+	listen := cmd.Flags().StringP("listen", "l", ":8080", "Listening port")
+	autoStop := cmd.Flags().BoolP("auto-stop", "a", false, "True to stop when no plugins are running")
+	interval := cmd.Flags().DurationP("scan", "s", 1*time.Minute, "Scan interval to check for plugins")
+
+	cmd.RunE = func(c *cobra.Command, args []string) error {
+		logger.Info("Starting mux server", "listen", *listen)
+		server, err := mux.NewServer(*listen, plugins)
+		if err != nil {
+			return err
+		}
+		defer server.Stop()
+
+		block := make(chan struct{})
+		// If the sockets are gone, then we can safely exit.
+		go func() {
+			checkNow := time.Tick(*interval)
+			for {
+				select {
+				case <-server.Wait():
+					logger.Info("server stopped")
+					close(block)
+					return
+
+				case <-checkNow:
+					if m, err := plugins().List(); err == nil {
+						if len(m) == 0 && *autoStop {
+							logger.Info("scan found no plugins.")
+							close(block)
+							return
+						}
+					}
+
+				}
+			}
+		}()
+
+		<-block
+
+		return nil
+	}
+
+	return cmd
+}
+
+func utilCommand(plugins func() discovery.Plugins) *cobra.Command {
+
+	util := &cobra.Command{
+		Use:   "util",
+		Short: "Utilties",
+	}
+
+	util.AddCommand(muxCommand(plugins))
+
+	return util
+}

pkg/discovery/dir.go

@@ -4,23 +4,51 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"os/user"
+	"path"
 	"path/filepath"
 	"sync"
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/docker/infrakit/pkg/plugin"
 )
 
-type errNotUnixSocket string
+// Dir returns the directory to use for plugin discovery, which may be customized by the environment.
+func Dir() string {
+	if pluginDir := os.Getenv(PluginDirEnvVar); pluginDir != "" {
+		return pluginDir
+	}
 
-func (e errNotUnixSocket) Error() string {
-	return string(e)
+	home := os.Getenv("HOME")
+	if usr, err := user.Current(); err == nil {
+		home = usr.HomeDir
+	}
+	return path.Join(home, ".infrakit/plugins")
 }
 
-// IsErrNotUnixSocket returns true if the error is due to the file not being a valid unix socket.
-func IsErrNotUnixSocket(e error) bool {
-	_, is := e.(errNotUnixSocket)
-	return is
+// NewPluginDiscovery creates a plugin discovery based on the environment configuration.
+func NewPluginDiscovery() (Plugins, error) {
+	return NewPluginDiscoveryWithDirectory(Dir())
+}
+
+// NewPluginDiscoveryWithDirectory creates a plugin discovery based on the directory given.
+func NewPluginDiscoveryWithDirectory(pluginDir string) (Plugins, error) {
+	stat, err := os.Stat(pluginDir)
+	if err == nil {
+		if !stat.IsDir() {
+			return nil, fmt.Errorf("Plugin dir %s is a file", pluginDir)
+		}
+	} else {
+		if os.IsNotExist(err) {
+			if err := os.MkdirAll(pluginDir, 0700); err != nil {
+				return nil, fmt.Errorf("Failed to create plugin dir %s: %s", pluginDir, err)
+			}
+		} else {
+			return nil, fmt.Errorf("Failed to access plugin dir %s: %s", pluginDir, err)
+		}
+	}
+
+	return newDirPluginDiscovery(pluginDir)
 }
 
 type dirPluginDiscovery struct {
@@ -38,7 +66,7 @@ func (r *dirPluginDiscovery) Find(name plugin.Name) (*plugin.Endpoint, error) {
 
 	p, exists := plugins[lookup]
 	if !exists {
-		return nil, fmt.Errorf("Plugin not found: %s (looked up using %s)", name, lookup)
+		return nil, ErrNotFound(string(name))
 	}
 
 	return p, nil
@@ -54,16 +82,16 @@ func newDirPluginDiscovery(dir string) (*dirPluginDiscovery, error) {
 }
 
 func (r *dirPluginDiscovery) dirLookup(entry os.FileInfo) (*plugin.Endpoint, error) {
+	socketPath := filepath.Join(r.dir, entry.Name())
 	if entry.Mode()&os.ModeSocket != 0 {
-		socketPath := filepath.Join(r.dir, entry.Name())
 		return &plugin.Endpoint{
 			Protocol: "unix",
 			Address:  socketPath,
 			Name:     entry.Name(),
 		}, nil
 	}
 
-	return nil, errNotUnixSocket(fmt.Sprintf("File is not a socket: %s", entry))
+	return nil, ErrNotUnixSocket(socketPath)
 }
 
 // List returns a list of plugins known, keyed by the name

pkg/discovery/dir_test.go

@@ -14,7 +14,7 @@ import (
 )
 
 func TestErrNotUnixSocket(t *testing.T) {
-	err := errNotUnixSocket("no socket!")
+	err := ErrNotUnixSocket("no socket!")
 	require.Error(t, err)
 	require.True(t, IsErrNotUnixSocket(err))
 }

pkg/discovery/discovery.go

@@ -2,9 +2,6 @@ package discovery
 
 import (
 	"fmt"
-	"os"
-	"os/user"
-	"path"
 
 	"github.com/docker/infrakit/pkg/plugin"
 )
@@ -21,40 +18,28 @@ const (
 	PluginDirEnvVar = "INFRAKIT_PLUGINS_DIR"
 )
 
-// Dir returns the directory to use for plugin discovery, which may be customized by the environment.
-func Dir() string {
-	if pluginDir := os.Getenv(PluginDirEnvVar); pluginDir != "" {
-		return pluginDir
-	}
-
-	home := os.Getenv("HOME")
-	if usr, err := user.Current(); err == nil {
-		home = usr.HomeDir
-	}
-	return path.Join(home, ".infrakit/plugins")
+// ErrNotUnixSocket is the error raised when the file is not a unix socket
+type ErrNotUnixSocket string
+
+func (e ErrNotUnixSocket) Error() string {
+	return fmt.Sprintf("not a unix socket:%s", string(e))
+}
+
+// IsErrNotUnixSocket returns true if the error is due to the file not being a valid unix socket.
+func IsErrNotUnixSocket(e error) bool {
+	_, is := e.(ErrNotUnixSocket)
+	return is
 }
 
-// NewPluginDiscovery creates a plugin discovery based on the environment configuration.
-func NewPluginDiscovery() (Plugins, error) {
-	return NewPluginDiscoveryWithDirectory(Dir())
+// ErrNotFound is the error raised when the plugin is not found
+type ErrNotFound string
+
+func (e ErrNotFound) Error() string {
+	return fmt.Sprintf("plugin not found:%s", string(e))
 }
 
-// NewPluginDiscoveryWithDirectory creates a plugin discovery based on the directory given.
-func NewPluginDiscoveryWithDirectory(pluginDir string) (Plugins, error) {
-	stat, err := os.Stat(pluginDir)
-	if err == nil {
-		if !stat.IsDir() {
-			return nil, fmt.Errorf("Plugin dir %s is a file", pluginDir)
-		}
-	} else {
-		if os.IsNotExist(err) {
-			if err := os.MkdirAll(pluginDir, 0700); err != nil {
-				return nil, fmt.Errorf("Failed to create plugin dir %s: %s", pluginDir, err)
-			}
-		} else {
-			return nil, fmt.Errorf("Failed to access plugin dir %s: %s", pluginDir, err)
-		}
-	}
-
-	return newDirPluginDiscovery(pluginDir)
+// IsErrNotFound returns true if the error is due to a plugin not found.
+func IsErrNotFound(e error) bool {
+	_, is := e.(ErrNotFound)
+	return is
 }