Golang v2.Tomb類代碼示例

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 = "user"
	}

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