light-apiengine/engine/auth/auth_client.go

package auth

import (
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"git.qianqiusoft.com/qianqiusoft/light-apiengine/config"
	"git.qianqiusoft.com/qianqiusoft/light-apiengine/entitys"
	sysutils "git.qianqiusoft.com/qianqiusoft/light-apiengine/utils"
	"net"
	"strconv"
	"time"
)

const (
	__KEY = "Light#dauth-@*I2"

	CMD_NEW = "new"
	CMD_REMOVE = "remove"
)

var TCPClient *TcpClient

type authPackage struct {
	Cmd string
	Content []byte
}

func (ap *authPackage)toBytes()[]byte{
	buf := bytes.NewBuffer([]byte{})

	b := []byte(ap.Cmd)
	buf.Write(uint64ToBytes(len(b)))
	buf.Write(b)

	buf.Write(uint64ToBytes(len(ap.Content)))
	buf.Write(ap.Content)

	return buf.Bytes()
}

type TcpClient struct {
	conn    net.Conn
	pchan   chan *authPackage
	done    chan bool
	exited bool
	verified   bool
}

// 创建client
func NewTcpClient()*TcpClient{
	c := &TcpClient{
		pchan: make(chan *authPackage, 100),
		done:  make(chan bool),
		exited: false,
		verified: false,
	}
	return c
}

// 启动
func(c *TcpClient)Start() {
	go func() {
		defer func() {
			if p := recover(); p != nil {
				fmt.Println("ecover", p)
			}
		}()

		var err error = nil
		address := config.AppConfig.GetKey("auth_server")
		fmt.Println("auth client start, dial address is", address)
		c.conn, err = net.Dial("tcp", address)
		if err != nil {
			fmt.Println("Error dialing", err.Error())
			c.restart()
			return
		}

		fmt.Println("发送验证")
		sendVerify(c.conn) // 发送验证，不需要读取返回值，如果验证错误立刻关掉

		fmt.Println("读取验证结果")
		vresp, err := readString(c.conn)
		if err != nil {
			fmt.Println("Error dialing", err.Error())
			c.restart()
			return
		}
		if vresp != "ok"{
			// 验证失败
			fmt.Println("verify is not ok", vresp)
			return
		}
		c.verified = true

		// send
		go func() {
			for {
				select {
				case data := <-c.pchan:
					_, err := c.conn.Write(data.toBytes())
					if err != nil {
						fmt.Println("写入内容错误", err.Error())
						return
					}
				case <-c.done:
					return
				}
			}
		}()

		// receive
		for {
			cmd, err := readString(c.conn) // 读取命令
			if err != nil {
				c.done <- true
				fmt.Println("读取命令错误", err.Error())
				break
			}
			bytess, err := readBytes(c.conn)
			if err != nil {
				c.done <- true
				fmt.Println("读取token内容错误", err.Error())
				break
			}
			if cmd == "remove" {
				// 移除，此时bytess为tokenstring
				sysutils.GetGlobalTokenStore().Remove(string(bytess))
			} else if cmd == "new" {
				// 新建
				token, _ := bytesToToken(bytess)
				sysutils.GetGlobalTokenStore().Set(token.AccessToken, token)
			}
		}

		c.restart()
	}()
}

// 停止
func(c *TcpClient)Stop(){
	c.exited = true
	c.conn.Close()
}

// 检测
func(c *TcpClient)restart(){
	if c.exited{
		// 已退出则不重启
		return
	}

	go func(){
		time.Sleep(3 * time.Second)
		c.verified = false
		c.Start()
	}()
}

// 发送bytes
func(c *TcpClient)Send(cmd string, bytess []byte){
	if !c.verified{
		fmt.Println("未认证")
		return
	}

	c.pchan <- &authPackage{
		Cmd: cmd,
		Content: bytess,
	}
}

// 发送token
func(c *TcpClient)SendToken(token *entitys.Token){
	c.Send(CMD_NEW, tokenToBytes(token))
}

// 读取字符串
func readString(conn net.Conn)(string, error){
	// 读长度
	size, err := readUInt32(conn)
	if err != nil{
		fmt.Println("读取长度失败，", err.Error())
		return "", err
	}

	// 读字符串
	b := make([]byte, size)
	n, err := conn.Read(b)
	if n != int(size){
		return "", errors.New("读取长度不是" + strconv.Itoa(int(size)))
	}
	return string(b), nil
}

// 写入字符串
func writeString(conn net.Conn, str string)error{
	if str == ""{
		return errors.New("字符串为空")
	}else{
		fmt.Println("---->1", str)
	}

	bytess := []byte(str)
	size := len(bytess)

	fmt.Println("---->2", str)
	// 发送长度
	err := writeUInt32(conn, uint32(size))
	if err != nil{
		fmt.Println("发送长度失败，", err.Error())
		return err
	}

	fmt.Println("---->3", str)
	// 发送内容
	n, err := conn.Write(bytess)
	if err != nil{
		fmt.Println("发送内容失败，", err.Error())
		return err
	}
	if n != size{
		return errors.New("发送长度不是" + strconv.Itoa(int(size)))
	}
	fmt.Println("---->4", str)

	return nil
}

// 读取bytes
func readBytes(conn net.Conn)([]byte, error){
	// 读长度
	size, err := readUInt32(conn)
	if err != nil{
		fmt.Println("读取长度失败，", err.Error())
		return nil, err
	}

	// 读字符串
	b := make([]byte, size)
	n, err := conn.Read(b)
	if n != int(size){
		return nil, errors.New("读取长度不是" + strconv.Itoa(int(size)))
	}
	return b, nil
}

// 读取uint64
func readUInt32(conn net.Conn)(uint32, error){
	b := make([]byte, 4)
	n, err := conn.Read(b)
	if err != nil{
		fmt.Println("读取长度失败，", err.Error())
		return 0, err
	}
	if n != 4{
		return 0, errors.New("读取长度不是4")
	}
	size := binary.BigEndian.Uint32(b)

	return size, nil
}

// 写入长度
func writeUInt32(conn net.Conn, v uint32)error{
	// 发送长度
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, v)
	n, err := conn.Write(b)
	if err != nil{
		fmt.Println("发送长度失败，", err.Error())
		return err
	}
	if n != 4{
		return errors.New("发送长度不是4")
	}

	return nil
}


// 写入长度
func writeUInt64(conn net.Conn, v uint64)error{
	// 发送长度
	b := make([]byte, 8)
	binary.BigEndian.PutUint64(b, v)
	n, err := conn.Write(b)
	if err != nil{
		fmt.Println("发送长度失败，", err.Error())
		return err
	}
	if n != 4{
		return errors.New("发送长度不是4")
	}

	return nil
}


// 读取uint64
func readStringByBytes(bytess []byte)(string, int, error) {
	size := binary.BigEndian.Uint32(bytess)
	return string(bytess[4 : 4+size]), int(size), nil
}

// 转token
func bytesToToken(content []byte)(*entitys.Token, error){
	token := &entitys.Token{}
	var index int = 0
	var size int
	var err error = nil
	token.UserId, size, err = readStringByBytes(content)
	if err != nil{
		fmt.Println("读取userid错误")
		return  nil, err
	}
	index += 4 + size

	token.AccessToken, size, err = readStringByBytes(content[index:])
	if err != nil{
		fmt.Println("读取AccessToken错误")
		return  nil, err
	}
	index += 4 + size

	token.RefreshToken, size, err = readStringByBytes(content[index:])
	if err != nil{
		fmt.Println("读取RefreshToken错误")
		return  nil, err
	}
	index += 4 + size

	token.LoginID, size, err = readStringByBytes(content[index:])
	if err != nil{
		fmt.Println("读取LoginID错误")
		return  nil, err
	}
	index += 4 + size

	token.TimeStamp = binary.BigEndian.Uint64(content[index:])
	index += 4

	token.ServerIp, size, err = readStringByBytes(content[index:])
	if err != nil{
		fmt.Println("读取ServerIp错误")
		return  nil, err
	}
	index += 4 + size

	token.Domain, size, err = readStringByBytes(content[index:])
	if err != nil{
		fmt.Println("读取Domain错误")
		return  nil, err
	}
	index += 4 + size

	return  token, nil
}

// int转bytes
func uint64ToBytes(v int)[]byte{
	b := make([]byte, 4)
	binary.BigEndian.PutUint32(b, uint32(v))

	return b
}

// 转bytes
func tokenToBytes(token *entitys.Token)[]byte{
	buf := bytes.NewBuffer([]byte{})

	t := []byte(token.UserId)
	buf.Write(uint64ToBytes(len(t)))
	buf.Write(t)

	t = []byte(token.AccessToken)
	buf.Write(uint64ToBytes(len(t)))
	buf.Write(t)

	t = []byte(token.RefreshToken)
	buf.Write(uint64ToBytes(len(t)))
	buf.Write(t)

	t = []byte(token.LoginID)
	buf.Write(uint64ToBytes(len(t)))
	buf.Write(t)

	buf.Write(uint64ToBytes(int(token.TimeStamp)))

	t = []byte(token.ServerIp)
	buf.Write(uint64ToBytes(len(t)))
	buf.Write(t)

	t = []byte(token.Domain)
	buf.Write(uint64ToBytes(len(t)))
	buf.Write(t)

	return buf.Bytes()
}

// 发送验证
func sendVerify(conn net.Conn){
	timestamp := time.Now().UnixNano()
	timestampStr := strconv.Itoa(int(timestamp))
	seed := timestampStr + __KEY
	hashVal := hash(seed)

	writeUInt64(conn, uint64(timestamp))
	writeString(conn, hashVal)
}

// md5 哈希
func hash(str string)string{
	h := md5.New()
	h.Write([]byte(str))
	return hex.EncodeToString(h.Sum(nil))
}