func (s *SSHServer) handleTCPConnection(parentTomb tomb.Tomb, conn net.Conn) error {
// Convert to SSH connection
sshConn, channels, requests, err := ssh.NewServerConn(conn, s.config.sshConfig)
if err != nil {
s.config.Logger.Warn("SSH handshake failed: %s", conn.RemoteAddr())
return err
}
// Close connection on exit
s.config.Logger.Debug("Handshake successful")
defer sshConn.Conn.Close()
// Discard requests
go ssh.DiscardRequests(requests)
// Create new tomb stone
var t tomb.Tomb
for {
select {
case ch := <-channels:
chType := ch.ChannelType()
// Determine if channel is acceptable (has a registered handler)
handler, ok := s.config.Handler(chType)
if !ok {
s.config.Logger.Info("UnknownChannelType", "type", chType)
ch.Reject(ssh.UnknownChannelType, chType)
break
}
// Accept channel
channel, requests, err := ch.Accept()
if err != nil {
s.config.Logger.Warn("Error creating channel")
continue
}
t.Go(func() error {
return handler.Handle(t, sshConn, channel, requests)
})
case <-parentTomb.Dying():
t.Kill(nil)
if err := t.Wait(); err != nil {
s.config.Logger.Warn("ssh handler error: %s", err)
}
sshConn.Close()
return sshConn.Wait()
}
}
return nil
}
func (s *shellHandler) Handle(parentTomb tomb.Tomb, sshConn *ssh.ServerConn, channel ssh.Channel, requests <-chan *ssh.Request) error {
defer channel.Close()
s.logger.Info("WooHoo!!! Inside Handler!")
// Create tomb for terminal goroutines
var t tomb.Tomb
// Sessions have out-of-band requests such as "shell",
// "pty-req" and "env". Here we handle only the
// "shell" request.
t.Go(func() error {
OUTER:
for {
select {
case <-parentTomb.Dying():
t.Kill(nil)
break OUTER
case req := <-requests:
if req == nil {
break OUTER
}
ok := false
switch req.Type {
case "shell":
ok = true
if len(req.Payload) > 0 {
// fmt.Println(string(req.Payload))
// We don't accept any
// commands, only the
// default shell.
ok = false
}
case "pty-req":
// Responding 'ok' here will let the client
// know we have a pty ready for input
ok = true
t.Go(func() error {
return s.startTerminal(t, sshConn, channel)
})
}
req.Reply(ok, nil)
}
}
return nil
})
return t.Wait()
}
// listen accepts new connections and handles the conversion from TCP to SSH connections.
func (s *SSHServer) listen() error {
defer s.listener.Close()
// Create tomb for connection goroutines
var t tomb.Tomb
t.Go(func() error {
OUTER:
for {
// Accepts will only block for 1s
s.listener.SetDeadline(time.Now().Add(s.config.Deadline))
select {
// Stop server on channel receive
case <-s.t.Dying():
t.Kill(nil)
break OUTER
default:
// Accept new connection
tcpConn, err := s.listener.Accept()
if err != nil {
if neterr, ok := err.(net.Error); ok && neterr.Timeout() {
s.config.Logger.Trace("Connection timeout...")
} else {
s.config.Logger.Warn("Connection failed", "error", err)
}
continue
}
// Handle connection
s.config.Logger.Info("Successful TCP connection:", tcpConn.RemoteAddr().String())
t.Go(func() error {
// Return the error for the
err := s.handleTCPConnection(t, tcpConn)
if err != io.EOF {
return err
}
return nil
})
}
}
return nil
})
return t.Wait()
}
func (e *EchoHandler) Handle(t tomb.Tomb, conn *ssh.ServerConn, channel ssh.Channel, requests <-chan *ssh.Request) error {
defer channel.Close()
e.logger.Info("echo handle called!")
// Create tomb for terminal goroutines
var tmb tomb.Tomb
type msg struct {
line []byte
isPrefix bool
err error
}
in := make(chan msg)
defer close(in)
reader := bufio.NewReader(channel)
tmb.Go(func() error {
tmb.Go(func() error {
for {
line, pre, err := reader.ReadLine()
if err != nil {
tmb.Kill(nil)
return nil
}
select {
case in <- msg{line, pre, err}:
case <-t.Dying():
tmb.Kill(nil)
return nil
case <-tmb.Dying():
return nil
}
}
})
tmb.Go(func() error {
for {
e.logger.Info("time: ", time.Now())
select {
case <-tmb.Dying():
return nil
case <-t.Dying():
tmb.Kill(nil)
return nil
case m := <-in:
if m.err != nil {
tmb.Kill(m.err)
return m.err
}
// Send echo
channel.Write(m.line)
}
}
})
return nil
})
return tmb.Wait()
}
// handleTCPConnection converts the TCP connection into an ssh.ServerConn and processes all incoming channel requests
// and dispatches them to defined handlers if they exist.
func (s *SSHServer) handleTCPConnection(parentTomb tomb.Tomb, conn net.Conn) error {
// Convert to SSH connection
sshConn, channels, requests, err := ssh.NewServerConn(conn, s.config.sshConfig)
if err != nil {
s.config.Logger.Warn("SSH handshake failed:", "addr", conn.RemoteAddr().String(), "error", err)
conn.Close()
return err
}
// Close connection on exit
s.config.Logger.Debug("Handshake successful")
defer sshConn.Close()
defer sshConn.Wait()
// Discard requests
go ssh.DiscardRequests(requests)
// Create new tomb stone
var t tomb.Tomb
t.Go(func() error {
OUTER:
for {
select {
case ch := <-channels:
// Check if chan was closed
if ch == nil {
t.Kill(nil)
break OUTER
}
// Get channel type
chType := ch.ChannelType()
// Determine if channel is acceptable (has a registered handler)
handler, ok := s.config.Handler(chType)
if !ok {
s.config.Logger.Info("UnknownChannelType", "type", chType)
ch.Reject(ssh.UnknownChannelType, chType)
t.Kill(nil)
break OUTER
}
// Accept channel
channel, requests, err := ch.Accept()
if err != nil {
s.config.Logger.Warn("Error creating channel")
continue
}
t.Go(func() error {
err := handler.Handle(t, sshConn, channel, requests)
if err != nil {
s.config.Logger.Warn("Handler raised an error", err)
}
s.config.Logger.Warn("Exiting channel", chType)
t.Kill(err)
return err
})
case <-parentTomb.Dying():
t.Kill(nil)
break OUTER
case <-t.Dying():
break OUTER
}
}
return nil
})
// Wait for all goroutines to finish
return t.Wait()
}
func (e *EchoHandler) Handle(parentTomb tomb.Tomb, sshConn *ssh.ServerConn, channel ssh.Channel, requests <-chan *ssh.Request) error {
defer channel.Close()
// Create tomb for terminal goroutines
var t tomb.Tomb
type msg struct {
length uint32
data []byte
}
in := make(chan msg)
defer close(in)
// Sessions have out-of-band requests such as "shell",
// "pty-req" and "env". Here we handle only the
// "shell" request.
t.Go(func() error {
var buffer bytes.Buffer
// Read channel
t.Go(func() error {
length := make([]byte, 4)
for {
n, err := channel.Read(length)
if err != nil {
return err
} else if n != 4 {
return errors.New("Invalid message length")
}
// Decode length
l, err := xbinary.LittleEndian.Uint32(length, 0)
if err != nil {
return err
}
// Read data
n64, err := buffer.ReadFrom(io.LimitReader(channel, int64(l)))
if err != nil {
return err
} else if n64 != int64(l) {
return errors.New("error: reading message")
}
select {
case <-parentTomb.Dying():
return nil
case <-t.Dying():
return nil
case in <- msg{l, buffer.Bytes()}:
}
}
})
length := make([]byte, 4)
OUTER:
for {
select {
case <-parentTomb.Dying():
t.Kill(nil)
break OUTER
case m := <-in:
if m.length == 0 {
return nil
}
// Encode length
_, err := xbinary.LittleEndian.PutUint32(length, 0, m.length)
if err != nil {
t.Kill(err)
return nil
}
// Write echo response
channel.Write(length)
channel.Write(m.data)
}
}
return nil
})
return t.Wait()
}