merge 4.0.0-RC

3 files changed, 338 insertions, 0 deletions

diff --git a/internal/clients/connectors/connector.go b/internal/clients/connectors/connector.go
new file mode 100644
index 0000000..3ab6a08
--- /dev/null
+++ b/internal/clients/connectors/connector.go
@@ -0,0 +1,17 @@
+package connectors
+
+import (
+	"context"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+)
+
+// Connector interface.
+type Connector interface {
+	// Start the connection.
+	Start(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{})
+	// Server hostname.
+	Server() string
+	// Handler for the connection.
+	Handler() handlers.Handler
+}
diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go
new file mode 100644
index 0000000..aeb2a41
--- /dev/null
+++ b/internal/clients/connectors/serverconnection.go
@@ -0,0 +1,206 @@
+package connectors
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/ssh/client"
+
+	"golang.org/x/crypto/ssh"
+)
+
+// ServerConnection represents a connection to a single remote dtail server via
+// SSH protocol.
+type ServerConnection struct {
+	// The full server string as received from the server discovery (can be with port number)
+	server string
+	// Only the hostname or FQDN (without the port number)
+	hostname string
+	// Only the port number.
+	port            int
+	config          *ssh.ClientConfig
+	handler         handlers.Handler
+	commands        []string
+	hostKeyCallback client.HostKeyCallback
+	throttlingDone  bool
+}
+
+// NewServerConnection returns a new DTail SSH server connection.
+func NewServerConnection(server string, userName string,
+	authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback,
+	handler handlers.Handler, commands []string) *ServerConnection {
+
+	dlog.Client.Debug(server, "Creating new connection", server, handler, commands)
+	c := ServerConnection{
+		hostKeyCallback: hostKeyCallback,
+		server:          server,
+		handler:         handler,
+		commands:        commands,
+		config: &ssh.ClientConfig{
+			User:            userName,
+			Auth:            authMethods,
+			HostKeyCallback: hostKeyCallback.Wrap(),
+			Timeout:         time.Second * 2,
+		},
+	}
+
+	c.initServerPort()
+	return &c
+}
+
+// Server returns the server hostname connected to.
+func (c *ServerConnection) Server() string { return c.server }
+
+// Handler returns the handler used for the connection.
+func (c *ServerConnection) Handler() handlers.Handler { return c.handler }
+
+// Attempt to parse the server port address from the provided server FQDN.
+func (c *ServerConnection) initServerPort() {
+	parts := strings.Split(c.server, ":")
+	if len(parts) == 1 {
+		c.hostname = c.server
+		c.port = config.Common.SSHPort
+		return
+	}
+
+	dlog.Client.Debug("Parsing port from hostname", parts)
+	port, err := strconv.Atoi(parts[1])
+	if err != nil {
+		dlog.Client.FatalPanic("Unable to parse client port", c.server, parts, err)
+	}
+	c.hostname = parts[0]
+	c.port = port
+}
+
+// Start the connection to the server.
+func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc,
+	throttleCh, statsCh chan struct{}) {
+
+	// Throttle how many connections can be established concurrently (based on ch length)
+	dlog.Client.Debug(c.server, "Throttling connection", len(throttleCh), cap(throttleCh))
+
+	select {
+	case throttleCh <- struct{}{}:
+	case <-ctx.Done():
+		dlog.Client.Debug(c.server, "Not establishing connection as context is done",
+			len(throttleCh), cap(throttleCh))
+		return
+	}
+
+	dlog.Client.Debug(c.server, "Throttling says that the connection can be established",
+		len(throttleCh), cap(throttleCh))
+
+	go func() {
+		defer func() {
+			if !c.throttlingDone {
+				dlog.Client.Debug(c.server, "Unthrottling connection (1)",
+					len(throttleCh), cap(throttleCh))
+				c.throttlingDone = true
+				<-throttleCh
+			}
+			cancel()
+		}()
+
+		if err := c.dial(ctx, cancel, throttleCh, statsCh); err != nil {
+			dlog.Client.Warn(c.server, err)
+			if c.hostKeyCallback.Untrusted(c.server) {
+				dlog.Client.Debug(c.server, "Not trusting host")
+			}
+		}
+	}()
+
+	<-ctx.Done()
+}
+
+// Dail into a new SSH connection. Close connection in case of an error.
+func (c *ServerConnection) dial(ctx context.Context, cancel context.CancelFunc,
+	throttleCh, statsCh chan struct{}) error {
+
+	dlog.Client.Debug(c.server, "Incrementing connection stats")
+	statsCh <- struct{}{}
+	defer func() {
+		dlog.Client.Debug(c.server, "Decrementing connection stats")
+		<-statsCh
+	}()
+
+	address := fmt.Sprintf("%s:%d", c.hostname, c.port)
+	dlog.Client.Debug(c.server, "Dialing into the connection", address)
+
+	client, err := ssh.Dial("tcp", address, c.config)
+	if err != nil {
+		return err
+	}
+	defer client.Close()
+
+	return c.session(ctx, cancel, client, throttleCh)
+}
+
+// Create the SSH session. Close the session in case of an error.
+func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFunc,
+	client *ssh.Client, throttleCh chan struct{}) error {
+
+	dlog.Client.Debug(c.server, "Creating SSH session")
+	session, err := client.NewSession()
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+	return c.handle(ctx, cancel, session, throttleCh)
+}
+
+func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc,
+	session *ssh.Session, throttleCh chan struct{}) error {
+
+	dlog.Client.Debug(c.server, "Creating handler for SSH session")
+	stdinPipe, err := session.StdinPipe()
+	if err != nil {
+		return err
+	}
+	stdoutPipe, err := session.StdoutPipe()
+	if err != nil {
+		return err
+	}
+	if err := session.Shell(); err != nil {
+		return err
+	}
+
+	go func() {
+		io.Copy(stdinPipe, c.handler)
+		cancel()
+	}()
+	go func() {
+		io.Copy(c.handler, stdoutPipe)
+		cancel()
+	}()
+	go func() {
+		select {
+		case <-c.handler.Done():
+		case <-ctx.Done():
+		}
+		cancel()
+	}()
+
+	// Send all commands to client.
+	for _, command := range c.commands {
+		dlog.Client.Debug(command)
+		c.handler.SendMessage(command)
+	}
+
+	if !c.throttlingDone {
+		dlog.Client.Debug(c.server, "Unthrottling connection (2)",
+			len(throttleCh), cap(throttleCh))
+		c.throttlingDone = true
+		<-throttleCh
+	}
+
+	<-ctx.Done()
+	c.handler.Shutdown()
+	return nil
+}
diff --git a/internal/clients/connectors/serverless.go b/internal/clients/connectors/serverless.go
new file mode 100644
index 0000000..431247a
--- /dev/null
+++ b/internal/clients/connectors/serverless.go
@@ -0,0 +1,115 @@
+package connectors
+
+import (
+	"context"
+	"io"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	serverHandlers "github.com/mimecast/dtail/internal/server/handlers"
+	user "github.com/mimecast/dtail/internal/user/server"
+)
+
+// Serverless creates a server object directly without TCP.
+type Serverless struct {
+	handler  handlers.Handler
+	commands []string
+	userName string
+}
+
+// NewServerless starts a new serverless session.
+func NewServerless(userName string, handler handlers.Handler,
+	commands []string) *Serverless {
+
+	dlog.Client.Debug("Creating new serverless connector", handler, commands)
+	return &Serverless{
+		userName: userName,
+		handler:  handler,
+		commands: commands,
+	}
+}
+
+// Server returns serverless server indicator.
+func (s *Serverless) Server() string {
+	return "local(serverless)"
+}
+
+// Handler returns the handler used for the serverless connection.
+func (s *Serverless) Handler() handlers.Handler {
+	return s.handler
+}
+
+// Start the serverless connection.
+func (s *Serverless) Start(ctx context.Context, cancel context.CancelFunc,
+	throttleCh, statsCh chan struct{}) {
+
+	dlog.Client.Debug("Starting serverless connector")
+	go func() {
+		defer cancel()
+		if err := s.handle(ctx, cancel); err != nil {
+			dlog.Client.Warn(err)
+		}
+	}()
+	<-ctx.Done()
+}
+
+func (s *Serverless) handle(ctx context.Context, cancel context.CancelFunc) error {
+	dlog.Client.Debug("Creating server handler for a serverless session")
+
+	user, err := user.New(s.userName, s.Server())
+	if err != nil {
+		return err
+	}
+
+	var serverHandler serverHandlers.Handler
+	switch s.userName {
+	case config.HealthUser:
+		dlog.Client.Debug("Creating serverless health handler")
+		serverHandler = serverHandlers.NewHealthHandler(user)
+	default:
+		dlog.Client.Debug("Creating serverless server handler")
+		serverHandler = serverHandlers.NewServerHandler(
+			user,
+			make(chan struct{}, config.Server.MaxConcurrentCats),
+			make(chan struct{}, config.Server.MaxConcurrentTails),
+		)
+	}
+
+	terminate := func() {
+		dlog.Client.Debug("Terminating serverless connection")
+		serverHandler.Shutdown()
+		cancel()
+	}
+
+	go func() {
+		io.Copy(serverHandler, s.handler)
+		dlog.Client.Trace("io.Copy(serverHandler, s.handler) => done")
+		terminate()
+	}()
+	go func() {
+		io.Copy(s.handler, serverHandler)
+		dlog.Client.Trace("io.Copy(s.handler, serverHandler) => done")
+		terminate()
+	}()
+	go func() {
+		select {
+		case <-s.handler.Done():
+			dlog.Client.Trace("<-s.handler.Done()")
+		case <-ctx.Done():
+			dlog.Client.Trace("<-ctx.Done()")
+		}
+		terminate()
+	}()
+
+	// Send all commands to client.
+	for _, command := range s.commands {
+		dlog.Client.Debug("Sending command to serverless server", command)
+		s.handler.SendMessage(command)
+	}
+
+	<-ctx.Done()
+	dlog.Client.Trace("s.handler.Shutdown()")
+	s.handler.Shutdown()
+	return nil
+}