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) } } } }