i2-bill-erp/internal/utils/token.go

// ERP 代码, light的逻辑， 废弃！！！！！
// ERP 代码, light的逻辑， 废弃！！！！！

package utils

import (
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.i2edu.net/i2/go-zero/core/logx"
	"git.i2edu.net/i2/i2-bill-erp/internal/config"
)

type Token struct {
	Lock         *sync.RWMutex
	Result       int                    `json:"-"`
	UserId       string                 `json:"user_id"`
	AccessToken  string                 `json:"access_token"`
	RefreshToken string                 `json:"refresh_token"`
	LoginID      string                 `json:"login_id"`
	TimeStamp    uint64                 `json:"time_stamp"`
	ServerIp     string                 `json:"server_ip"`
	Domain       string                 `json:"domain"`
	Password     string                 `json:"-"`
	Map          map[string]interface{} `json:"map"`
}

type TokenStore struct {
	storeServiceUrl string

	name   string
	lock   *sync.RWMutex
	tokens map[string]*Token
}

type IAuth interface {
	Init()
}

var globalTokenStore *TokenStore = nil

func Init(c *config.Config) {
	globalTokenStore = &TokenStore{name: "sso", lock: new(sync.RWMutex), tokens: make(map[string]*Token)}
	if strings.HasPrefix(c.Erp.AuthServer, "http") {
		globalTokenStore.storeServiceUrl = c.Erp.AuthServer
	} else {
		iauthMap = make(map[string]IAuth)

		go globalTokenStore.startTokenCheckProcess()

		lightAuth := &LightAuth{}
		RegisterAuth("qianqiusoft.com", lightAuth)

		erpClient := NewTcpClient(c)
		erpClient.Start()
	}
}

type LightAuth struct {
}

func (la *LightAuth) Init() {

}

var iauthMap map[string]IAuth

/**
 * @brief: 注册自定登录对象, 会调用Init进行初始化
 * @param1 t: key
 * @param2 a: 认证对象
 * @return none
 */
func RegisterAuth(t string, a IAuth) {
	if t == "" || a == nil {
		return
	}

	var ok bool
	if _, ok = iauthMap[t]; ok {
		return
	} else {
		a.Init()
		iauthMap[t] = a
	}
}

func GetGlobalTokenStore() *TokenStore {
	return globalTokenStore
}

func (t *TokenStore) Get(key string) string {
	if t.storeServiceUrl == "" {
		t.lock.RLock()
		defer t.lock.RUnlock()
		if val, ok := t.tokens[key]; ok {
			//log.Println(key, "获取Token:", val.AccessToken, val.RefreshToken, val.LoginID)
			return val.UserId
		}
	} else {
		resp, err := http.Get(t.storeServiceUrl + key)
		if err != nil {
			return ""
		}
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return ""
		}
		token := &Token{}
		err = json.Unmarshal(b, &token)
		if err != nil {
			return ""
		}
		return token.UserId
	}

	return ""
}

func (t *TokenStore) Set(key string, v *Token) {
	t.lock.Lock()
	defer t.lock.Unlock()
	if val, ok := t.tokens[key]; !ok {
		t.tokens[key] = v
	} else if val != v {
		t.tokens[key] = v
	}
}

func (t *TokenStore) Remove(key string) {
	t.lock.Lock()
	defer t.lock.Unlock()
	delete(t.tokens, key)
	log.Println(key, "删除Key")
}

func (t *TokenStore) Refresh(key string) {
	t.lock.Lock()
	defer t.lock.Unlock()

	val, ok := t.tokens[key]
	if ok {
		val.TimeStamp = uint64(time.Now().UnixNano())
	}
}

func (t *TokenStore) startTokenCheckProcess() {
}

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

	CMD_NEW           = "new"
	CMD_REMOVE        = "remove"
	CMD_PINGPONG      = "pg"
	CMD_PINGPONG_RESP = "pg_resp"
)

var TCPClient *TcpClient

type authPackage struct {
	Cmd      string
	TokenStr string
	Content  []byte
}

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

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

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

	return buf.Bytes()
}

type TcpClient struct {
	conn     net.Conn // 连接
	c        *config.Config
	pchan    chan *authPackage // 包chan
	done     chan bool         // 是否done
	exited   bool              // 退出
	verified bool              // 验证是否
}

// 创建client
func NewTcpClient(cfg *config.Config) *TcpClient {
	c := &TcpClient{
		c:        cfg,
		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)
			}
			if c.conn != nil {
				c.conn.Close()
			}
			c.restart()
		}()

		var err error = nil
		address := c.c.Erp.AuthServer
		// fmt.Println("auth client start, dial address is", address)
		c.conn, err = net.Dial("tcp", address)
		if err != nil {
			fmt.Println("Error dialing erp client", err.Error())
			return
		}
		fmt.Println("发送验证")
		sendVerify(c.conn) // 发送验证，不需要读取返回值，如果验证错误立刻关掉

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

		// send
		go func() {
			for {
				select {
				case data := <-c.pchan:
					logx.Info("写入数据")
					c.conn.SetWriteDeadline(time.Now().Add(time.Second * 2))
					_, err := c.conn.Write(data.toBytes())
					if err != nil {
						fmt.Println("写入内容错误", err.Error())
						return
					}
				case <-c.done:
					logx.Info("发送数据done退出")
					return
				}
			}
		}()

		// receive
		for {
			cmd, err := readString(c.conn) // 读取命令
			if err != nil {
				c.done <- true
				fmt.Println("读取命令错误", err.Error())
				break
			}
			if cmd == CMD_NEW {
				err = c.newHandler()
			} else if cmd == CMD_REMOVE {
				err = c.removeHandler()
			} else if cmd == CMD_PINGPONG_RESP {

			} else {
				fmt.Println("未知cmd", cmd)
				continue
			}
			if err != nil {
				c.done <- true
				fmt.Println("处理错误", err.Error())
				break
			}
		}
	}()
}

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

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

	go func() {
		c.verified = false
		c.done = make(chan bool)
		c.pchan = make(chan *authPackage, 100)
		c.exited = false
		time.Sleep(3 * time.Second)
		c.Start()
	}()
}

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

	logx.Info("发送指令1", cmd)
	c.pchan <- &authPackage{
		Cmd:     cmd,
		Content: bytess,
	}
	logx.Info("发送指令2", cmd)
}

// 发送token
func (c *TcpClient) SendToken(token *Token) {
	// logx.Info("发送新建token")
	bytess := tokenToBytes(token)
	c.Send(CMD_NEW, bytess)
}

// 处理创建
func (c *TcpClient) newHandler() error {
	// fmt.Println("处理新建")
	bytess, err := readBytes(c.conn)
	if err != nil {
		fmt.Println("读取token内容错误", err.Error())
		return err
	}
	// 新建
	token, err := bytesToToken(bytess)
	if err != nil {
		logx.Info("bytesToToken 错误", err.Error())
		return err
	}
	GetGlobalTokenStore().Set(token.AccessToken, token)
	return nil
}

// 处理删除
func (c *TcpClient) removeHandler() error {
	fmt.Println("处理删除")
	bytess, err := readBytes(c.conn)
	if err != nil {
		fmt.Println("读取token内容错误", err.Error())
		return err
	}
	// 移除，此时bytess为tokenstring
	GetGlobalTokenStore().Remove(string(bytess))
	return nil
}

// 读取字符串
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("字符串为空")
	}

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

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

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

	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
}

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

	return b
}

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

	return b
}

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

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

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

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

	// fmt.Println("读取timestamp")
	token.TimeStamp = binary.BigEndian.Uint64(content[index:])
	index += 8

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

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

	return token, nil
}

// 转bytes，包括token开头
func tokenToBytes(token *Token) []byte {
	buf := bytes.NewBuffer([]byte{})

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

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

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

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

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

	fmt.Println(token.ServerIp)
	t = []byte(token.ServerIp)
	buf.Write(uint32ToBytes(len(t)))
	buf.Write(t)

	fmt.Println(token.Domain)
	t = []byte(token.Domain)
	buf.Write(uint32ToBytes(len(t)))
	buf.Write(t)

	bytess := buf.Bytes()
	buf = bytes.NewBuffer([]byte{}) // 这里用reset是错误的

	tokenstrbytes := []byte(token.AccessToken)
	buf.Write(uint32ToBytes(len(tokenstrbytes)))
	buf.Write(tokenstrbytes)
	buf.Write(uint32ToBytes(len(bytess)))
	buf.Write(bytess)

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