session.go

package client

import (
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/docker/docker/pkg/term"
	"github.com/gravitational/teleport"
	"github.com/gravitational/teleport/lib/defaults"
	"github.com/gravitational/teleport/lib/session"
	"github.com/gravitational/teleport/lib/sshutils"
	"github.com/gravitational/teleport/lib/utils"
	"github.com/gravitational/trace"
	"golang.org/x/crypto/ssh"

	log "github.com/Sirupsen/logrus"
)

type NodeSession struct {
	// id is the Teleport session ID
	id session.ID

	// env is the environment variables that need to be created
	// on the server for this session
	env map[string]string

	// nodeClient is the parent of this session: the client connected to an
	// SSH node
	nodeClient *NodeClient

	// Standard input/outputs for this session
	stdin  io.Reader
	stdout io.Writer
	stderr io.Writer

	// closer is used to simultaneously close all goroutines created by
	// this session. It's also used to wait for everyone to close
	closer *utils.CloseBroadcaster

	ExitMsg string
}

// newSession creates a new Teleport session with the given remote node
// if 'joinSessin' is given, the session will join the existing session
// of another user
func newSession(client *NodeClient,
	joinSession *session.Session,
	env map[string]string,
	stdin io.Reader,
	stdout io.Writer,
	stderr io.Writer) (*NodeSession, error) {

	if stdin == nil {
		stdin = os.Stdin
	}
	if stdout == nil {
		stdout = os.Stdout
	}
	if stderr == nil {
		stderr = os.Stderr
	}

	var err error
	ns := &NodeSession{
		env:        env,
		nodeClient: client,
		stdin:      stdin,
		stdout:     stdout,
		stderr:     stderr,
		closer:     utils.NewCloseBroadcaster(),
	}
	// if we're joining an existing session, we need to assume that session's
	// existing/current terminal size:
	if joinSession != nil {
		ns.id = joinSession.ID
		tsize := joinSession.TerminalParams.Winsize()
		if ns.isTerminalAttached() {
			err = term.SetWinsize(0, tsize)
			if err != nil {
				log.Error(err)
			}
			os.Stdout.Write([]byte(fmt.Sprintf("\x1b[8;%d;%dt", tsize.Height, tsize.Width)))
		}
		// new session!
	} else {
		ns.id = session.NewID()
	}
	if ns.env == nil {
		ns.env = make(map[string]string)
	}
	ns.env[sshutils.SessionEnvVar] = string(ns.id)
	return ns, nil
}

func (ns *NodeSession) regularSession(callback func(s *ssh.Session) error) error {
	session, err := ns.createServerSession()
	if err != nil {
		return trace.Wrap(err)
	}
	session.Stdout = ns.stdout
	session.Stderr = ns.stderr
	session.Stdin = ns.stdin
	return trace.Wrap(callback(session))
}

type interactiveCallback func(serverSession *ssh.Session, shell io.ReadWriteCloser) error

func (ns *NodeSession) createServerSession() (*ssh.Session, error) {
	sess, err := ns.nodeClient.Client.NewSession()
	if err != nil {
		return nil, trace.Wrap(err)
	}
	// pass language info into the remote session.
	evarsToPass := []string{"LANG", "LANGUAGE"}
	for _, evar := range evarsToPass {
		if value := os.Getenv(evar); value != "" {
			err = sess.Setenv(evar, value)
			if err != nil {
				log.Warn(err)
			}
		}
	}
	// pass environment variables set by client
	for key, val := range ns.env {
		err = sess.Setenv(key, val)
		if err != nil {
			log.Warn(err)
		}
	}
	return sess, nil
}

// interactiveSession creates an interactive session on the remote node, executes
// the given callback on it, and waits for the session to end
func (ns *NodeSession) interactiveSession(callback interactiveCallback) error {
	// determine what kind of a terminal we need
	termType := os.Getenv("TERM")
	if termType == "" {
		termType = teleport.SafeTerminalType
	}
	ns.env["TERM"] = termType
	// create the server-side session:
	sess, err := ns.createServerSession()
	if err != nil {
		return trace.Wrap(err)
	}
	// allocate terminal on the server:
	remoteTerm, err := ns.allocateTerminal(termType, sess)
	if err != nil {
		return trace.Wrap(err)
	}
	defer remoteTerm.Close()

	// call the passed callback and give them the established
	// ssh session:
	if err := callback(sess, remoteTerm); err != nil {
		return trace.Wrap(err)
	}

	// Catch term signals, but only if we're attached to a real terminal
	if ns.isTerminalAttached() {
		ns.watchSignals(remoteTerm)
	}

	// start piping input into the remote shell and pipe the output from
	// the remote shell into stdout:
	ns.pipeInOut(remoteTerm)

	// switch the terminal to raw mode (and switch back on exit!)
	if ns.isTerminalAttached() {
		ts, err := term.SetRawTerminal(0)
		if err != nil {
			log.Warn(err)
		} else {
			defer term.RestoreTerminal(0, ts)
		}
	}
	// wait for the session to end
	<-ns.closer.C
	return nil
}

// allocateTerminal creates (allocates) a server-side terminal for this session.
func (ns *NodeSession) allocateTerminal(termType string, s *ssh.Session) (io.ReadWriteCloser, error) {
	var err error
	// read the size of the terminal window:
	tsize := &term.Winsize{
		Width:  teleport.DefaultTerminalWidth,
		Height: teleport.DefaultTerminalHeight,
	}
	if ns.isTerminalAttached() {
		tsize, err = term.GetWinsize(0)
		if err != nil {
			log.Error(err)
		}
	}
	// ... and request a server-side terminal of the same size:
	err = s.RequestPty(termType,
		int(tsize.Height),
		int(tsize.Width),
		ssh.TerminalModes{})
	if err != nil {
		return nil, trace.Wrap(err)
	}
	writer, err := s.StdinPipe()
	if err != nil {
		return nil, trace.Wrap(err)
	}
	reader, err := s.StdoutPipe()
	if err != nil {
		return nil, trace.Wrap(err)
	}
	stderr, err := s.StderrPipe()
	if err != nil {
		return nil, trace.Wrap(err)
	}
	if ns.isTerminalAttached() {
		go ns.updateTerminalSize(s)
	}
	go func() {
		io.Copy(os.Stderr, stderr)
	}()
	return utils.NewPipeNetConn(
		reader,
		writer,
		utils.MultiCloser(writer, s, ns.closer),
		&net.IPAddr{},
		&net.IPAddr{},
	), nil
}

func (ns *NodeSession) updateTerminalSize(s *ssh.Session) {
	// sibscribe for "terminal resized" signal:
	sigC := make(chan os.Signal, 1)
	signal.Notify(sigC, syscall.SIGWINCH)
	currentSize, _ := term.GetWinsize(0)

	// start the timer which asks for server-side window size changes:
	siteClient, err := ns.nodeClient.Proxy.ConnectToSite()
	if err != nil {
		log.Error(err)
	}
	tick := time.NewTicker(defaults.SessionRefreshPeriod)
	defer tick.Stop()

	var prevSess *session.Session
	for {
		select {
		// our own terminal window got resized:
		case sig := <-sigC:
			if sig == nil {
				return
			}
			// get the size:
			winSize, err := term.GetWinsize(0)
			if err != nil {
				log.Warnf("[CLIENT] Error getting size: %s", err)
				break
			}
			// it's the result of our own size change (see below)
			if winSize.Height == currentSize.Height && winSize.Width == currentSize.Width {
				continue
			}
			// send the new window size to the server
			_, err = s.SendRequest(
				sshutils.WindowChangeReq, false,
				ssh.Marshal(sshutils.WinChangeReqParams{
					W: uint32(winSize.Width),
					H: uint32(winSize.Height),
				}))
			if err != nil {
				log.Warnf("[CLIENT] failed to send window change reqest: %v", err)
			}
		case <-tick.C:
			sess, err := siteClient.GetSession(ns.id)
			if err != nil {
				log.Error(err)
				continue
			}
			// no previous session
			if prevSess == nil || sess == nil {
				prevSess = sess
				continue
			}
			log.Infof("[CLIENT] updating the session %v with %d parties", sess.ID, len(sess.Parties))
			// nothing changed
			if prevSess.TerminalParams.W == sess.TerminalParams.W && prevSess.TerminalParams.H == sess.TerminalParams.H {
				continue
			}

			newSize := sess.TerminalParams.Winsize()
			currentSize, err = term.GetWinsize(0)
			if err != nil {
				log.Error(err)
			}
			if currentSize.Width != newSize.Width || currentSize.Height != newSize.Height {
				// ok, something have changed, let's resize to the new parameters
				err = term.SetWinsize(0, newSize)
				if err != nil {
					log.Error(err)
				}
				os.Stdout.Write([]byte(fmt.Sprintf("\x1b[8;%d;%dt", newSize.Height, newSize.Width)))
			}
			prevSess = sess
		case <-ns.closer.C:
			return
		}
	}
}

// isTerminalAttached returns true when this session is be controlled by
// a real terminal.
// It will return False for sessions initiated by the Web client or
// for non-interactive sessions (commands)
func (ns *NodeSession) isTerminalAttached() bool {
	return ns.stdin == os.Stdin && term.IsTerminal(os.Stdin.Fd())
}

// runShell executes user's shell on the remote node under an interactive session
func (ns *NodeSession) runShell(callback ShellCreatedCallback) error {
	return ns.interactiveSession(func(s *ssh.Session, shell io.ReadWriteCloser) error {
		// start the shell on the server:
		if err := s.Shell(); err != nil {
			return trace.Wrap(err)
		}
		// call the client-supplied callback
		if callback != nil {
			exit, err := callback(shell)
			if exit {
				return trace.Wrap(err)
			}
		}
		return nil
	})
}

// runCommand executes a given command either in interactive (with terminal attached)
// or non-intractive mode
func (ns *NodeSession) runCommand(cmd []string, callback ShellCreatedCallback, interactive bool) error {
	// stdin is not a terminal? refuse to allocate terminal on the server and go back
	// to "non-interactive":
	if interactive && ns.stdin == os.Stdin && !term.IsTerminal(os.Stdin.Fd()) {
		interactive = false
		fmt.Fprintf(os.Stderr, "TTY will not be allocated on the server because stdin is not a terminal\n")
	}

	// interactive session:
	if interactive {
		return ns.interactiveSession(func(s *ssh.Session, shell io.ReadWriteCloser) error {
			err := s.Start(strings.Join(cmd, " "))
			if err != nil {
				return trace.Wrap(err)
			}
			if callback != nil {
				exit, err := callback(shell)
				if exit {
					return trace.Wrap(err)
				}
			}
			return nil
		})
	}
	// non-interactive session:
	return ns.regularSession(func(s *ssh.Session) error {
		return s.Run(strings.Join(cmd, " "))
	})
}

// watchSignals register UNIX signal handlers and properly terminates a remote shell session
// must be called as a goroutine right after a remote shell is created
func (ns *NodeSession) watchSignals(shell io.Writer) {
	exitSignals := make(chan os.Signal, 1)
	// catch SIGTERM
	signal.Notify(exitSignals, syscall.SIGTERM)
	go func() {
		defer ns.closer.Close()
		<-exitSignals
	}()
	// 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())
			}
		}
	}()
}

// pipeInOut launches two goroutines: one to pipe the local input into the remote shell,
// and another to pipe the output of the remote shell into the local output
func (ns *NodeSession) pipeInOut(shell io.ReadWriteCloser) {
	// copy from the remote shell to the local output
	go func() {
		defer ns.closer.Close()
		_, err := io.Copy(ns.stdout, shell)
		if err != nil {
			log.Errorf(err.Error())
		}
	}()
	// copy from the local input to the remote shell:
	go func() {
		defer ns.closer.Close()
		buf := make([]byte, 128)
		for {
			n, err := ns.stdin.Read(buf)
			if err != nil {
				fmt.Println(trace.Wrap(err))
				return
			}
			if n > 0 {
				_, err = shell.Write(buf[:n])
				if err != nil {
					ns.ExitMsg = err.Error()
					return
				}
			}
		}
	}()
}