Merge pull request gravitational#509 from gravitational/ev/leak

Ev/leak

integration/helpers.go

@@ -194,7 +194,7 @@ func (i *TeleInstance) GetSiteAPI(siteName string) auth.ClientI {
 func (i *TeleInstance) Create(trustedSecrets []*InstanceSecrets, enableSSH bool, console io.Writer) error {
 	dataDir, err := ioutil.TempDir("", "cluster-"+i.Secrets.SiteName)
 	if err != nil {
-		return err
+		return trace.Wrap(err)
 	}
 	tconf := service.MakeDefaultConfig()
 	tconf.SeedConfig = true
@@ -229,7 +229,7 @@ func (i *TeleInstance) Create(trustedSecrets []*InstanceSecrets, enableSSH bool,
 	i.Config = tconf
 	i.Process, err = service.NewTeleport(tconf)
 	if err != nil {
-		return err
+		return trace.Wrap(err)
 	}
 	// create users:
 	auth := i.Process.GetAuthServer()
@@ -239,7 +239,7 @@ func (i *TeleInstance) Create(trustedSecrets []*InstanceSecrets, enableSSH bool,
 			AllowedLogins: user.AllowedLogins,
 		})
 		if err != nil {
-			return err
+			return trace.Wrap(err)
 		}
 		priv, pub, _ := tconf.Keygen.GenerateKeyPair("")
 		//priv, pub := makeKey()
@@ -262,7 +262,7 @@ func (i *TeleInstance) Create(trustedSecrets []*InstanceSecrets, enableSSH bool,
 func (i *TeleInstance) Reset() (err error) {
 	i.Process, err = service.NewTeleport(i.Config)
 	if err != nil {
-		return err
+		return trace.Wrap(err)
 	}
 	return nil
 }

lib/client/api.go

@@ -187,8 +187,22 @@ func (c *Config) ProxySpecified() bool {
 type TeleportClient struct {
 	Config
 	localAgent *LocalKeyAgent
+
+	// OnShellCreated gets called when the shell is created. It's
+	// safe to keep it nil
+	OnShellCreated ShellCreatedCallback
+
+	// ExitMsg (if set) will be printed at the end of the SSH session
+	ExitMsg string
 }
 
+// ShellCreatedCallback can be supplied for every teleport client. It will
+// be called right after the remote shell is created, but the session
+// hasn't begun yet.
+//
+// It allows clients to cancel SSH action
+type ShellCreatedCallback func(shell io.ReadWriteCloser) (exit bool, err error)
+
 func (tc *TeleportClient) authMethods() []ssh.AuthMethod {
 	return tc.Config.AuthMethods
 }
@@ -669,6 +683,10 @@ func (tc *TeleportClient) runCommand(siteName string, nodeAddresses []string, pr
 // sessionID : when empty, creates a new shell. otherwise it tries to join the existing session.
 // stdin  : standard input to use. if nil, uses os.Stdin
 func (tc *TeleportClient) runShell(nodeClient *NodeClient, sessionID session.ID, stdin io.Reader) error {
+	var (
+		state *term.State
+		err   error
+	)
 	defer nodeClient.Close()
 	address := tc.NodeHostPort()
 
@@ -681,74 +699,95 @@ func (tc *TeleportClient) runShell(nodeClient *NodeClient, sessionID session.ID,
 	if tc.Stderr == nil {
 		tc.Stderr = os.Stderr
 	}
+	attachedTerm := (stdin == os.Stdin && term.IsTerminal(0))
+
+	winSize := &term.Winsize{Width: 80, Height: 25}
+	if attachedTerm {
+		winSize, err = term.GetWinsize(0)
+		if err != nil {
+			log.Error(err)
+		}
+	}
+
+	shell, err := nodeClient.Shell(
+		int(winSize.Width),
+		int(winSize.Height),
+		sessionID,
+		tc.Config.Env,
+		attachedTerm)
+
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	defer shell.Close()
+
+	// user-supplied callback
+	if tc.OnShellCreated != nil {
+		exit, err := tc.OnShellCreated(shell)
+		if exit {
+			return trace.Wrap(err)
+		}
+	}
+
 	// terminal must be in raw mode
-	var (
-		state   *term.State
-		err     error
-		exitMsg string
-	)
-	if stdin == os.Stdin && term.IsTerminal(0) {
+	if attachedTerm {
 		state, err = term.SetRawTerminal(0)
 		if err != nil {
 			return trace.Wrap(err)
 		}
+		log.Infof("connecting to remote shell using stdin")
+	} else {
+		log.Infof("connecting to remote shell NOT using stdin")
 	}
 	defer func() {
 		if state != nil {
 			term.RestoreTerminal(0, state)
 		}
-		if exitMsg != "" {
-			fmt.Println(exitMsg)
+		if tc.ExitMsg != "" {
+			fmt.Println(tc.ExitMsg)
 		}
 	}()
 
 	broadcastClose := utils.NewCloseBroadcaster()
 
-	// Catch term signals
-	exitSignals := make(chan os.Signal, 1)
-	signal.Notify(exitSignals, syscall.SIGTERM)
-	go func() {
-		defer broadcastClose.Close()
-		<-exitSignals
-		exitMsg = fmt.Sprintf("Connection to %s closed\n", address)
-	}()
-
-	winSize, err := term.GetWinsize(0)
-	if err != nil {
-		log.Error(err)
-		winSize = &term.Winsize{Width: 80, Height: 25}
-	}
-
-	shell, err := nodeClient.Shell(int(winSize.Width), int(winSize.Height), sessionID, tc.Config.Env)
-	if err != nil {
-		return trace.Wrap(err)
-	}
+	// Catch term signals, but only if we're attached to a real terminal
+	if attachedTerm {
+		exitSignals := make(chan os.Signal, 1)
+		signal.Notify(exitSignals, syscall.SIGTERM)
+		go func() {
+			defer broadcastClose.Close()
+			<-exitSignals
+			if tc.ExitMsg == "" {
+				tc.ExitMsg = fmt.Sprintf("Connection to %s closed\n", address)
+			}
+		}()
 
-	// Catch Ctrl-C signal
-	ctrlCSignal := make(chan os.Signal, 1)
-	signal.Notify(ctrlCSignal, syscall.SIGINT)
-	go func() {
-		for {
-			<-ctrlCSignal
-			_, err := shell.Write([]byte{3})
-			if err != nil {
-				log.Errorf(err.Error())
+		// Catch Ctrl-C signal
+		ctrlCSignal := make(chan os.Signal, 1)
+		signal.Notify(ctrlCSignal, syscall.SIGINT)
+		go func() {
+			for {
+				<-ctrlCSignal
+				_, err := shell.Write([]byte{3})
+				if err != nil {
+					log.Errorf(err.Error())
+				}
 			}
-		}
-	}()
+		}()
 
-	// Catch Ctrl-Z signal
-	ctrlZSignal := make(chan os.Signal, 1)
-	signal.Notify(ctrlZSignal, syscall.SIGTSTP)
-	go func() {
-		for {
-			<-ctrlZSignal
-			_, err := shell.Write([]byte{26})
-			if err != nil {
-				log.Errorf(err.Error())
+		// Catch Ctrl-Z signal
+		ctrlZSignal := make(chan os.Signal, 1)
+		signal.Notify(ctrlZSignal, syscall.SIGTSTP)
+		go func() {
+			for {
+				<-ctrlZSignal
+				_, err := shell.Write([]byte{26})
+				if err != nil {
+					log.Errorf(err.Error())
+				}
 			}
-		}
-	}()
+		}()
+	}
 
 	// copy from the remote shell to the local
 	go func() {
@@ -757,7 +796,9 @@ func (tc *TeleportClient) runShell(nodeClient *NodeClient, sessionID session.ID,
 		if err != nil {
 			log.Errorf(err.Error())
 		}
-		exitMsg = fmt.Sprintf("Connection to %s closed from the remote side", address)
+		if tc.ExitMsg == "" {
+			tc.ExitMsg = fmt.Sprintf("Connection to %s closed from the remote side", address)
+		}
 	}()
 
 	// copy from the local shell to the remote
@@ -773,7 +814,7 @@ func (tc *TeleportClient) runShell(nodeClient *NodeClient, sessionID session.ID,
 			if n > 0 {
 				_, err = shell.Write(buf[:n])
 				if err != nil {
-					exitMsg = err.Error()
+					tc.ExitMsg = err.Error()
 					return
 				}
 			}

lib/client/client.go

@@ -262,8 +262,21 @@ func (proxy *ProxyClient) Close() error {
 	return proxy.Client.Close()
 }
 
-// Shell returns remote shell as io.ReadWriterCloser object
-func (client *NodeClient) Shell(width, height int, sessionID session.ID, env map[string]string) (io.ReadWriteCloser, error) {
+// Shell returns a configured remote shell (for a window of a requested size)
+// as io.ReadWriterCloser object
+//
+// Parameters:
+//	- width & height : size of the window
+//  - session id     : ID of a session (if joining existing) or empty to create
+//                     a new session
+//  - env            : list of environment variables to set for a new session
+//  - attachedTerm   : boolean indicating if this client is attached to a real terminal
+func (client *NodeClient) Shell(
+	width, height int,
+	sessionID session.ID,
+	env map[string]string,
+	attachedTerm bool) (io.ReadWriteCloser, error) {
+
 	if sessionID == "" {
 		// initiate a new session if not passed
 		sessionID = session.NewID()
@@ -333,7 +346,7 @@ func (client *NodeClient) Shell(width, height int, sessionID session.ID, env map
 	// this goroutine sleeps until a terminal size changes (it receives an OS signal)
 	sigC := make(chan os.Signal, 1)
 	signal.Notify(sigC, syscall.SIGWINCH)
-	go func() {
+	broadcastTerminalSize := func() {
 		for {
 			select {
 			case sig := <-sigC:
@@ -343,8 +356,8 @@ func (client *NodeClient) Shell(width, height int, sessionID session.ID, env map
 				// get the size:
 				winSize, err := term.GetWinsize(0)
 				if err != nil {
-					log.Infof("error getting size: %s", err)
-					continue
+					log.Errorf("Error getting size: %s", err)
+					break
 				}
 				// send the new window size to the server
 				_, err = clientSession.SendRequest(
@@ -357,22 +370,22 @@ func (client *NodeClient) Shell(width, height int, sessionID session.ID, env map
 					log.Infof("failed to send window change reqest: %v", err)
 				}
 			case <-broadcastClose.C:
-				log.Infof("detected close")
 				return
 			}
 		}
-	}()
+	}
 
-	tick := time.NewTicker(defaults.SessionRefreshPeriod)
 	// detect changes of the session's terminal
-	go func() error {
+	updateTerminalSize := func() {
+		tick := time.NewTicker(defaults.SessionRefreshPeriod)
 		defer tick.Stop()
 		var prevSess *session.Session
 		for {
 			select {
 			case <-tick.C:
 				sess, err := siteClient.GetSession(sessionID)
 				if err != nil {
+					log.Error(err)
 					continue
 				}
 				// no previous session
@@ -386,7 +399,6 @@ func (client *NodeClient) Shell(width, height int, sessionID session.ID, env map
 				}
 
 				newSize := sess.TerminalParams.Winsize()
-
 				currentSize, err := term.GetWinsize(0)
 				if err != nil {
 					log.Error(err)
@@ -401,10 +413,15 @@ func (client *NodeClient) Shell(width, height int, sessionID session.ID, env map
 				}
 				prevSess = sess
 			case <-broadcastClose.C:
-				return nil
+				return
 			}
 		}
-	}()
+	}
+
+	if attachedTerm {
+		go broadcastTerminalSize()
+		go updateTerminalSize()
+	}
 
 	go func() {
 		io.Copy(os.Stderr, stderr)

lib/service/service.go

@@ -656,6 +656,7 @@ func (process *TeleportProcess) initProxyEndpoint(conn *Connector) error {
 			utils.Consolef(cfg.Console, "[PROXY] starting the web server: %v", err)
 			return trace.Wrap(err)
 		}
+		defer webHandler.Close()
 
 		proxyLimiter.WrapHandle(webHandler)
 		process.BroadcastEvent(Event{Name: ProxyWebServerEvent, Payload: webHandler})

lib/web/sessions.go

@@ -164,6 +164,7 @@ type sessionCache struct {
 
 // Close closes all allocated resources and stops goroutines
 func (s *sessionCache) Close() error {
+	log.Infof("[WEB] closing session cache")
 	return s.closer.Close()
 }
 

lib/web/web.go

@@ -1249,17 +1249,6 @@ type Server struct {
 	http.Server
 }
 
-func New(addr utils.NetAddr, cfg Config) (*Server, error) {
-	h, err := NewHandler(cfg)
-	if err != nil {
-		return nil, err
-	}
-	srv := &Server{}
-	srv.Server.Addr = addr.Addr
-	srv.Server.Handler = h
-	return srv, nil
-}
-
 // CreateSignupLink generates and returns a URL which is given to a new
 // user to complete registration with Teleport via Web UI
 func CreateSignupLink(hostPort string, token string) string {