nps/utils/conn.go

package utils

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/golang/snappy"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const cryptKey = "1234567812345678"

type CryptConn struct {
	conn  net.Conn
	crypt bool
	rate  *Rate
}

func NewCryptConn(conn net.Conn, crypt bool, rate *Rate) *CryptConn {
	c := new(CryptConn)
	c.conn = conn
	c.crypt = crypt
	c.rate = rate
	return c
}

//加密写
func (s *CryptConn) Write(b []byte) (n int, err error) {
	n = len(b)
	if s.crypt {
		if b, err = AesEncrypt(b, []byte(cryptKey)); err != nil {
			return
		}
	}
	if b, err = GetLenBytes(b); err != nil {
		return
	}
	_, err = s.conn.Write(b)
	if s.rate != nil {
		s.rate.Get(int64(n))
	}
	return
}

//解密读
func (s *CryptConn) Read(b []byte) (n int, err error) {
	defer func() {
		if err == nil && n == len(IO_EOF) && string(b[:n]) == IO_EOF {
			err = io.EOF
			n = 0
		}
	}()
	var lens int
	var buf []byte
	var rb []byte
	c := NewConn(s.conn)
	if lens, err = c.GetLen(); err != nil {
		return
	}
	if buf, err = c.ReadLen(lens); err != nil {
		return
	}
	if s.crypt {
		if rb, err = AesDecrypt(buf, []byte(cryptKey)); err != nil {
			return
		}
	} else {
		rb = buf
	}
	copy(b, rb)
	n = len(rb)
	if s.rate != nil {
		s.rate.Get(int64(n))
	}
	return
}

type SnappyConn struct {
	w     *snappy.Writer
	r     *snappy.Reader
	crypt bool
	rate  *Rate
}

func NewSnappyConn(conn net.Conn, crypt bool, rate *Rate) *SnappyConn {
	c := new(SnappyConn)
	c.w = snappy.NewBufferedWriter(conn)
	c.r = snappy.NewReader(conn)
	c.crypt = crypt
	c.rate = rate
	return c
}

//snappy压缩写 包含加密
func (s *SnappyConn) Write(b []byte) (n int, err error) {
	n = len(b)
	if s.crypt {
		if b, err = AesEncrypt(b, []byte(cryptKey)); err != nil {
			log.Println("encode crypt error:", err)
			return
		}
	}
	if _, err = s.w.Write(b); err != nil {
		return
	}
	if err = s.w.Flush(); err != nil {
		return
	}
	if s.rate != nil {
		s.rate.Get(int64(n))
	}
	return
}

//snappy压缩读 包含解密
func (s *SnappyConn) Read(b []byte) (n int, err error) {
	buf := bufPool.Get().([]byte)
	defer func() {
		if err == nil && n == len(IO_EOF) && string(b[:n]) == IO_EOF {
			err = io.EOF
			n = 0
		}
		bufPool.Put(buf)
	}()
	if n, err = s.r.Read(buf); err != nil {
		return
	}
	var bs []byte
	if s.crypt {
		if bs, err = AesDecrypt(buf[:n], []byte(cryptKey)); err != nil {
			log.Println("decode crypt error:", err)
			return
		}
	} else {
		bs = buf[:n]
	}
	n = len(bs)
	copy(b, bs)
	if s.rate != nil {
		s.rate.Get(int64(n))
	}
	return
}

type Conn struct {
	Conn net.Conn
}

//new conn
func NewConn(conn net.Conn) *Conn {
	c := new(Conn)
	c.Conn = conn
	return c
}

//读取指定长度内容
func (s *Conn) ReadLen(cLen int) ([]byte, error) {
	if cLen > poolSize {
		return nil, errors.New("长度错误" + strconv.Itoa(cLen))
	}
	var buf []byte
	if cLen <= poolSizeSmall {
		buf = bufPoolSmall.Get().([]byte)[:cLen]
		defer bufPoolSmall.Put(buf)
	} else {
		buf = bufPoolMax.Get().([]byte)[:cLen]
		defer bufPoolMax.Put(buf)
	}
	if n, err := io.ReadFull(s, buf); err != nil || n != cLen {
		return buf, errors.New("读取指定长度错误" + err.Error())
	}
	return buf, nil
}

//获取长度
func (s *Conn) GetLen() (int, error) {
	val, err := s.ReadLen(4)
	if err != nil {
		return 0, err
	}
	return GetLenByBytes(val)
}

//写入长度+内容 粘包
func (s *Conn) WriteLen(buf []byte) (int, error) {
	var b []byte
	var err error
	if b, err = GetLenBytes(buf); err != nil {
		return 0, err
	}
	return s.Write(b)
}

//读取flag
func (s *Conn) ReadFlag() (string, error) {
	val, err := s.ReadLen(4)
	if err != nil {
		return "", err
	}
	return string(val), err
}

//读取host 连接地址 压缩类型
func (s *Conn) GetHostFromConn() (typeStr string, host string, en, de int, crypt, mux bool, err error) {
retry:
	lType, err := s.ReadLen(3)
	if err != nil {
		return
	}
	if typeStr = string(lType); typeStr == TEST_FLAG {
		en, de, crypt, mux = s.GetConnInfoFromConn()
		goto retry
	} else if typeStr != CONN_TCP && typeStr != CONN_UDP {
		err = errors.New("unknown conn type")
		return
	}
	cLen, err := s.GetLen()
	if err != nil || cLen > poolSize {
		return
	}
	hostByte, err := s.ReadLen(cLen)
	if err != nil {
		return
	}
	host = string(hostByte)
	return
}

//写连接类型 和 host地址
func (s *Conn) WriteHost(ltype string, host string) (int, error) {
	raw := bytes.NewBuffer([]byte{})
	binary.Write(raw, binary.LittleEndian, []byte(ltype))
	binary.Write(raw, binary.LittleEndian, int32(len([]byte(host))))
	binary.Write(raw, binary.LittleEndian, []byte(host))
	return s.Write(raw.Bytes())
}

//设置连接为长连接
func (s *Conn) SetAlive() {
	conn := s.Conn.(*net.TCPConn)
	conn.SetReadDeadline(time.Time{})
	conn.SetKeepAlive(true)
	conn.SetKeepAlivePeriod(time.Duration(2 * time.Second))
}

func (s *Conn) SetReadDeadline(t time.Duration) {
	s.Conn.(*net.TCPConn).SetReadDeadline(time.Now().Add(time.Duration(t) * time.Second))
}

//从tcp报文中解析出host，连接类型等 TODO 多种情况
func (s *Conn) GetHost() (method, address string, rb []byte, err error, r *http.Request) {
	var b [32 * 1024]byte
	var n int
	if n, err = s.Read(b[:]); err != nil {
		return
	}
	rb = b[:n]
	r, err = http.ReadRequest(bufio.NewReader(bytes.NewReader(rb)))
	if err != nil {
		return
	}
	hostPortURL, err := url.Parse(r.Host)
	if err != nil {
		address = r.Host
		err = nil
		return
	}
	if hostPortURL.Opaque == "443" { //https访问
		address = r.Host + ":443"
	} else { //http访问
		if strings.Index(hostPortURL.Host, ":") == -1 { //host不带端口， 默认80
			address = r.Host + ":80"
		} else {
			address = r.Host
		}
	}
	return
}

//单独读（加密|压缩）
func (s *Conn) ReadFrom(b []byte, compress int, crypt bool, rate *Rate) (int, error) {
	if COMPRESS_SNAPY_DECODE == compress {
		return NewSnappyConn(s.Conn, crypt, rate).Read(b)
	}
	return NewCryptConn(s.Conn, crypt, rate).Read(b)
}

//单独写（加密|压缩）
func (s *Conn) WriteTo(b []byte, compress int, crypt bool, rate *Rate) (n int, err error) {
	if COMPRESS_SNAPY_ENCODE == compress {
		return NewSnappyConn(s.Conn, crypt, rate).Write(b)
	}
	return NewCryptConn(s.Conn, crypt, rate).Write(b)
}

//写压缩方式，加密
func (s *Conn) WriteConnInfo(en, de int, crypt, mux bool) {
	s.Write([]byte(strconv.Itoa(en) + strconv.Itoa(de) + GetStrByBool(crypt) + GetStrByBool(mux)))
}

//获取压缩方式，是否加密
func (s *Conn) GetConnInfoFromConn() (en, de int, crypt, mux bool) {
	buf, err := s.ReadLen(4)
	if err != nil {
		return
	}
	en, _ = strconv.Atoi(string(buf[0]))
	de, _ = strconv.Atoi(string(buf[1]))
	crypt = GetBoolByStr(string(buf[2]))
	mux = GetBoolByStr(string(buf[3]))
	return
}

//close
func (s *Conn) Close() error {
	return s.Conn.Close()
}

//write
func (s *Conn) Write(b []byte) (int, error) {
	return s.Conn.Write(b)
}

//read
func (s *Conn) Read(b []byte) (int, error) {
	return s.Conn.Read(b)
}

//write error
func (s *Conn) WriteError() (int, error) {
	return s.Write([]byte(RES_MSG))
}

//write sign flag
func (s *Conn) WriteSign() (int, error) {
	return s.Write([]byte(RES_SIGN))
}

//write sign flag
func (s *Conn) WriteClose() (int, error) {
	return s.Write([]byte(RES_CLOSE))
}

//write main
func (s *Conn) WriteMain() (int, error) {
	return s.Write([]byte(WORK_MAIN))
}

//write chan
func (s *Conn) WriteChan() (int, error) {
	return s.Write([]byte(WORK_CHAN))
}

//write test
func (s *Conn) WriteTest() (int, error) {
	return s.Write([]byte(TEST_FLAG))
}

//write test
func (s *Conn) WriteSuccess() (int, error) {
	return s.Write([]byte(CONN_SUCCESS))
}

//write test
func (s *Conn) WriteFail() (int, error) {
	return s.Write([]byte(CONN_ERROR))
}

//获取长度+内容
func GetLenBytes(buf []byte) (b []byte, err error) {
	raw := bytes.NewBuffer([]byte{})
	if err = binary.Write(raw, binary.LittleEndian, int32(len(buf))); err != nil {
		return
	}
	if err = binary.Write(raw, binary.LittleEndian, buf); err != nil {
		return
	}
	b = raw.Bytes()
	return
}

//解析出长度
func GetLenByBytes(buf []byte) (int, error) {
	nlen := binary.LittleEndian.Uint32(buf)
	if nlen <= 0 {
		return 0, errors.New("数据长度错误")
	}
	return int(nlen), nil
}