func (k *Keeper) sendLoop(t *tomb.Tomb) error {
var timeout <-chan time.Time = nil
for {
select {
case rep := <-k.sendChan:
var buf []byte = k.alloc()
defer k.free(buf)
bytesWritten, err := EncodePacket(buf[4:], rep)
if err != nil {
return err
}
// write frame size
binary.BigEndian.PutUint32(buf[:4], uint32(bytesWritten))
// write buffer to connection
_, err = k.conn.Write(buf[:4+bytesWritten])
if err != nil {
return err
}
// log output
log.Debug("-> ", fmt.Sprintf("%x", buf[:4+bytesWritten]))
case <-t.Dying():
// create a timeout in order to send pending requests (close reply)
if timeout == nil {
timeout = time.After(100 * time.Millisecond)
}
case <-timeout:
close(k.sendChan)
return nil
}
}
}
func (s *shellHandler) startTerminal(parentTomb tomb.Tomb, sshConn *ssh.ServerConn, channel ssh.Channel) error {
defer channel.Close()
prompt := ">>> "
term := terminal.NewTerminal(channel, prompt)
// // Try to make the terminal raw
// oldState, err := terminal.MakeRaw(0)
// if err != nil {
// logger.Warn("Error making terminal raw: ", err.Error())
// }
// defer terminal.Restore(0, oldState)
// Get username
username, ok := sshConn.Permissions.Extensions["username"]
if !ok {
username = "******"
}
// Write ascii text
term.Write([]byte(fmt.Sprintf("\r\n Nice job, %s! You are connected!\r\n", username)))
defer term.Write([]byte(fmt.Sprintf("\r\nGoodbye, %s!\r\n", username)))
// Start REPL
for {
select {
case <-parentTomb.Dying():
return nil
default:
s.logger.Info("Reading line...")
input, err := term.ReadLine()
if err != nil {
fmt.Errorf("Readline() error")
return err
}
// Process line
line := strings.TrimSpace(input)
if len(line) > 0 {
// Log input and handle exit requests
if line == "exit" || line == "quit" {
s.logger.Info("Closing connection")
return nil
}
// Echo input
channel.Write(term.Escape.Green)
channel.Write([]byte(line + "\r\n"))
channel.Write(term.Escape.Reset)
}
}
}
return nil
}
// NewContext returns a Context that is canceled either when parent is canceled
// or when t is Killed.
func NewContext(parent context.Context, t *tomb.Tomb) context.Context {
ctx, cancel := context.WithCancel(parent)
go func() {
select {
case <-t.Dying():
cancel()
case <-ctx.Done():
}
}()
return ctx
}
func (k *Keeper) processorLoop(t *tomb.Tomb) error {
for {
select {
case f := <-k.processorChan:
if err := f(); err != nil {
return err
}
case <-t.Dying():
close(k.processorChan)
return nil
}
}
}
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 (k *Keeper) requestLoop(t *tomb.Tomb) error {
for {
select {
case buf := <-k.recvChan:
// current & total bytes read
bytesRead := 0
totalBytesRead := 0
// parse request header
reqHdr := &OpReqHeader{}
bytesRead, err := DecodePacket(buf, reqHdr)
totalBytesRead = totalBytesRead + bytesRead
if err != nil {
log.Error(fmt.Sprintf("unable to decode request header: %s", err.Error()))
return err
}
// create request
creator, found := creatorByOpCode[reqHdr.OpCode]
if !found {
log.Error(fmt.Sprintf("cannot create opcode: %d", reqHdr.OpCode))
return err
}
// parse request
req := creator()
bytesRead, err = DecodePacket(buf[bytesRead:], req)
totalBytesRead = totalBytesRead + bytesRead
if err != nil {
log.Error(fmt.Sprintf("unable to decode request: %s", err.Error()))
return err
}
// queue processor
k.processorChan <- func() error {
processOpReq(OpReq{Hdr: reqHdr, Req: req}, k.storeClient, k.sendChan)
if reqHdr.OpCode == opClose {
return errors.New("graceful connection close requested")
}
return nil
}
case <-t.Dying():
close(k.recvChan)
return nil
}
}
}
// 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 (s *shellHandler) Handle(parentTomb tomb.Tomb, sshConn *ssh.ServerConn, channel ssh.Channel, requests <-chan *ssh.Request) error {
defer channel.Close()
users, err := s.system.Users()
if err != nil {
return err
}
user, err := users.Get(sshConn.Permissions.Extensions["username"])
if err != nil {
return err
}
// 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.
for {
select {
case <-parentTomb.Dying():
t.Kill(nil)
return t.Wait()
case req := <-requests:
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
go s.startTerminal(t, channel, s.system, user)
default:
// fmt.Println("default req: ", req)
}
req.Reply(ok, nil)
}
}
return nil
}
func (k *Keeper) recvLoop(t *tomb.Tomb) error {
for {
_, err := k.read(k.temp[:4])
if err != nil {
return err
}
// parse frame size
len := binary.BigEndian.Uint32(k.temp[:4])
if len > bufferSize {
return errors.New(fmt.Sprintf("length should be at most %d bytes (received %d)", bufferSize, len))
}
// alloc buffer (freeing if full read from conn wasn't possible)
var buf = k.alloc()
defer k.free(buf)
// read frame
_, err = k.read(buf[:len])
if err != nil {
return err
}
// log input
log.Debug("<- ", fmt.Sprintf("%x%x", k.temp[:4], buf[:len]))
// send frame to channel
select {
case k.recvChan <- buf[:len]:
case <-t.Dying():
return nil
}
}
return nil
}
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 (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()
}
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()
}
func (s *shellHandler) startTerminal(parentTomb tomb.Tomb, channel ssh.Channel, system datamodel.System, user datamodel.User) {
defer channel.Close()
prompt := "kappa> "
term := terminal.NewTerminal(channel, prompt)
// // Try to make the terminal raw
// oldState, err := terminal.MakeRaw(0)
// if err != nil {
// logger.Warn("Error making terminal raw: ", err.Error())
// }
// defer terminal.Restore(0, oldState)
// Write ascii text
term.Write([]byte("\r\n"))
for _, line := range common.ASCII {
term.Write([]byte(line))
term.Write([]byte("\r\n"))
}
// Write login message
term.Write([]byte("\r\n\n"))
client.GetMessage(channel, common.DefaultColorCodes)
term.Write([]byte("\n"))
// Create query executor
executor := executor.NewExecutor(executor.NewSession("", user), common.NewTerminal(term, prompt), system)
// Start REPL
for {
select {
case <-parentTomb.Dying():
return
default:
input, err := term.ReadLine()
if err != nil {
fmt.Errorf("Readline() error")
break
}
// Process line
line := strings.TrimSpace(input)
if len(line) > 0 {
// Log input and handle exit requests
if line == "exit" || line == "quit" {
s.logger.Info("Closing connection")
break
} else if line == "quote me" {
term.Write([]byte("\r\n"))
client.GetMessage(channel, common.DefaultColorCodes)
term.Write([]byte("\r\n"))
continue
} else if strings.HasPrefix(line, "//") || strings.HasPrefix(line, "--") {
channel.Write(common.DefaultColorCodes.LightGrey)
channel.Write([]byte(line + "\r\n"))
channel.Write(common.DefaultColorCodes.Reset)
continue
}
// Parse statement
stmt, err := skl.ParseStatement(line)
// Return parse error in red
if err != nil {
s.logger.Warn("Bad Statement", "statement", line, "error", err)
channel.Write(common.DefaultColorCodes.LightRed)
channel.Write([]byte(err.Error()))
channel.Write([]byte("\r\n"))
channel.Write(common.DefaultColorCodes.Reset)
continue
}
// Execute statements
w := common.ResponseWriter{common.DefaultColorCodes, channel}
executor.Execute(&w, stmt)
}
}
}
}