micro/gateway/cli/client.go

package cli

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"runtime"
	"strings"
	"sync"
	"time"

	"git.nspix.com/golang/micro/helper/net/ip"
	bytepkg "git.nspix.com/golang/micro/helper/pool/byte"
	"github.com/peterh/liner"
)

type Client struct {
	seq           uint16
	appName       string
	conn          net.Conn
	context       context.Context
	readyChan     chan struct{}
	completerChan chan *Frame
	responseChan  chan *Frame
	mutex         sync.Mutex
}

func (client *Client) writePack(packet *Frame) (err error) {
	if packet.Timestamp == 0 {
		packet.Timestamp = time.Now().Unix()
	}
	client.mutex.Lock()
	client.seq++
	packet.Seq = client.seq
	client.mutex.Unlock()
	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 = bytepkg.Get(MaxReadBufferLength)
	defer func() {
		bytepkg.Put(buf)
	}()
	for {
		packet := &Frame{}
		if packet, err = nextFrame(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, os: %s, ip: %s", vs["version"], runtime.GOOS, ip.InternalIP()), EOL)
				client.stdout(fmt.Sprintf("Last login: %s from %s@%s", vs["login_at"], vs["client_addr"], vs["cid"]), EOL)
				client.stdout("Type 'help;' for help. Type 'cls' to clear the current input statement.", EOL)
				select {
				case client.readyChan <- struct{}{}:
				default:
				}
			}
		case PacketTypeData:
			select {
			case client.responseChan <- packet:
			case <-time.After(time.Second):
			}
		}
	}
	select {
	case c <- err:
	case <-client.context.Done():
	}
}

func (client *Client) waitResponse(ctx context.Context, seq uint16) (err error) {
	childCtx, cancelFunc := context.WithTimeout(ctx, time.Second*30)
	defer func() {
		cancelFunc()
	}()
	for {
		select {
		case <-childCtx.Done():
			err = client.context.Err()
			return
		case res, ok := <-client.responseChan:
			if !ok {
				break
			}
			if res.Seq == seq {
				if res.Error != "" {
					client.stdout(res.Error, EOL)
				} else {
					client.stdout(string(bytes.Trim(res.Data, "\r\n")))
				}
			}
			if res.Flag == FlagComplete {
				client.stdout(EOL)
				return
			}
		}
	}
}

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
		}
		reqFrame := newFrame(PacketTypeData, FlagComplete, []byte(line))
		if err = client.writePack(reqFrame); err != nil {
			break
		}
		state.AppendHistory(line)
		client.waitResponse(client.context, reqFrame.Seq)
	}
}

func (client *Client) CompleterHandleFunc(str string) (ss []string) {
	var err error
	if err = client.writePack(newFrame(PacketTypeCompleter, FlagComplete, []byte(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(newFrame(PacketTypeEcho, FlagComplete, nil)); 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 = nextFrame(conn); err != nil {
		return
	}
	if err = cli.writePack(newFrame(PacketTypeData, FlagComplete, []byte(cmd))); err != nil {
		return
	}
	if frame, err = nextFrame(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
}