micro/gateway/cli/client.go

package cli

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	byte2 "git.nspix.com/golang/micro/helper/pool/byte"
	"git.nspix.com/golang/micro/helper/unsafestr"
	"github.com/peterh/liner"
	"net"
	"os"
	"strings"
	"time"
)

type Client struct {
	appName       string
	conn          net.Conn
	context       context.Context
	readyChan     chan struct{}
	completerChan chan *Frame
	responseChan  chan *Frame
}

func (client *Client) writePack(packet *Frame) (err error) {
	if packet.Timestamp == 0 {
		packet.Timestamp = time.Now().Unix()
	}
	err = writeFrame(client.conn, packet)
	return
}

func (client *Client) stdout(args ...interface{}) {
	fmt.Fprint(os.Stdout, args...)
}

func (client *Client) rdyLoop(c chan error) {
	var (
		err error
		buf []byte
	)
	buf = byte2.Get(MaxReadBufferLength)
	defer func() {
		byte2.Put(buf)
	}()
	for {
		packet := &Frame{}
		if packet, err = readFrame(client.conn); err != nil {
			c <- err
			break
		}
		switch packet.Type {
		case PacketTypeCompleter:
			client.completerChan <- packet
		case PacketTypeEcho:
			vs := make(map[string]string)
			if err = json.Unmarshal(packet.Data, &vs); err == nil {
				time.Sleep(time.Millisecond * 200)
				client.stdout("\u001B[2K")
				client.stdout("\u001B[0G")
				client.appName = vs["name"]
				client.stdout(fmt.Sprintf("Welcome to the %s monitor.", vs["name"]), EOL)
				client.stdout(fmt.Sprintf("Server Version: %s, connection id is %s", vs["version"], vs["cid"]), EOL)
				client.stdout(fmt.Sprintf("Last login: %s from %s", vs["login_at"], vs["client_addr"]), EOL)
				client.stdout("Type 'help;' for help. Type 'cls' to clear the current input statement.", EOL)
				select {
				case client.readyChan <- struct{}{}:
				default:
				}
			}
		case PacketTypeData:
			client.responseChan <- packet
		}
	}
	select {
	case c <- err:
	case <-client.context.Done():
	}
}

func (client *Client) interactive(c chan error) {
	var (
		err   error
		line  string
		state *liner.State
	)
	state = liner.NewLiner()
	state.SetCompleter(client.CompleterHandleFunc)
	select {
	case <-client.readyChan:
	case <-client.context.Done():
		return
	}
	defer func() {
		_ = state.Close()
		c <- err
	}()
	for {
		if line, err = state.Prompt(client.appName + "> "); err != nil {
			break
		}
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		if strings.ToLower(line) == "exit" || strings.ToLower(line) == "quit" {
			client.stdout("Bye Bye", EOL)
			return
		}
		if strings.ToLower(line) == "clear" || strings.ToLower(line) == "cls" {
			client.stdout("\033[2J")
			continue
		}
		if err = client.writePack(&Frame{Type: PacketTypeData, Data: []byte(line)}); err != nil {
			break
		}
		state.AppendHistory(line)
		select {
		case <-client.context.Done():
			err = client.context.Err()
			return
		case <-time.After(time.Second * 30):
			client.stdout("request timeout", EOL)
		case res, ok := <-client.responseChan:
			if !ok {
				break
			}
			if res.Error != "" {
				client.stdout(res.Error, EOL)
			} else {
				client.stdout(string(bytes.Trim(res.Data, "\r\n")), EOL)
			}
		}
	}
}

func (client *Client) CompleterHandleFunc(str string) (ss []string) {
	var err error
	if err = client.writePack(&Frame{
		Type: PacketTypeCompleter,
		Data: unsafestr.StringToBytes(str),
	}); err != nil {
		return
	}
	select {
	case <-time.After(time.Second * 5):
		return nil
	case <-client.context.Done():
		return nil
	case resp, ok := <-client.completerChan:
		if ok {
			ss = make([]string, 0)
			err = json.Unmarshal(resp.Data, &ss)
		}
	}
	return
}

func (client *Client) Run() (err error) {
	var (
		errChan chan error
	)
	errChan = make(chan error)
	defer func() {
		_ = client.conn.Close()
	}()
	if err = client.writePack(&Frame{
		Type: PacketTypeEcho,
	}); err != nil {
		return
	}
	client.stdout("connecting")
	go client.rdyLoop(errChan)
	go client.interactive(errChan)
	err = <-errChan
	return
}

func ExecuteContext(ctx context.Context, conn net.Conn, cmd string) (s string, err error) {
	if ctx == nil {
		ctx = context.Background()
	}
	cli := &Client{
		conn:          conn,
		context:       ctx,
		readyChan:     make(chan struct{}),
		completerChan: make(chan *Frame),
		responseChan:  make(chan *Frame, 1),
	}
	if err = cli.writePack(&Frame{
		Type: PacketTypeEcho,
	}); err != nil {
		return
	}
	frame := &Frame{}
	if frame, err = readFrame(conn); err != nil {
		return
	}
	if err = cli.writePack(&Frame{Type: PacketTypeData, Data: []byte(cmd)}); err != nil {
		return
	}
	if frame, err = readFrame(conn); err != nil {
		return
	}
	if frame.Error != "" {
		err = errors.New(frame.Error)
	} else {
		s = string(frame.Data)
	}
	return
}

func OpenInteractive(ctx context.Context, conn net.Conn) (err error) {
	if ctx == nil {
		ctx = context.Background()
	}
	cli := &Client{
		conn:          conn,
		context:       ctx,
		readyChan:     make(chan struct{}),
		completerChan: make(chan *Frame),
		responseChan:  make(chan *Frame, 1),
	}
	err = cli.Run()
	return
}