新功能+bug修复

utils/conn.go

@@ -18,19 +18,17 @@ import (
 
 const cryptKey = "1234567812345678"
 const poolSize = 64 * 1024
+const poolSizeSmall = 10
 
 type CryptConn struct {
 	conn  net.Conn
 	crypt bool
-	rb    []byte
-	rn    int
 }
 
 func NewCryptConn(conn net.Conn, crypt bool) *CryptConn {
 	c := new(CryptConn)
 	c.conn = conn
 	c.crypt = crypt
-	c.rb = make([]byte, poolSize)
 	return c
 }
 
@@ -51,21 +49,6 @@ func (s *CryptConn) Write(b []byte) (n int, err error) {
 
 //解密读
 func (s *CryptConn) Read(b []byte) (n int, err error) {
-read:
-	if len(s.rb) > 0 {
-		if len(b) >= s.rn {
-			n = s.rn
-			copy(b, s.rb[:s.rn])
-			s.rn = 0
-			s.rb = s.rb[:0]
-		} else {
-			n = len(b)
-			copy(b, s.rb[:len(b)])
-			s.rn = n - len(b)
-			s.rb = s.rb[len(b):]
-		}
-		return
-	}
 	defer func() {
 		if err == nil && n == len(IO_EOF) && string(b[:n]) == IO_EOF {
 			err = io.EOF
@@ -74,6 +57,7 @@ read:
 	}()
 	var lens int
 	var buf []byte
+	var rb []byte
 	c := NewConn(s.conn)
 	if lens, err = c.GetLen(); err != nil {
 		return
@@ -82,14 +66,15 @@ read:
 		return
 	}
 	if s.crypt {
-		if s.rb, err = AesDecrypt(buf, []byte(cryptKey)); err != nil {
+		if rb, err = AesDecrypt(buf, []byte(cryptKey)); err != nil {
 			return
 		}
 	} else {
-		s.rb = buf
+		rb = buf
 	}
-	s.rn = len(s.rb)
-	goto read
+	copy(b, rb)
+	n = len(rb)
+	return
 }
 
 type SnappyConn struct {
@@ -164,7 +149,12 @@ func (s *Conn) ReadLen(cLen int) ([]byte, error) {
 	if cLen > poolSize {
 		return nil, errors.New("长度错误" + strconv.Itoa(cLen))
 	}
-	buf := bufPool.Get().([]byte)[:cLen]
+	var buf []byte
+	if cLen <= poolSizeSmall {
+		buf = bufPoolSmall.Get().([]byte)[:cLen]
+	} else {
+		buf = bufPool.Get().([]byte)[:cLen]
+	}
 	if n, err := io.ReadFull(s, buf); err != nil || n != cLen {
 		return buf, errors.New("读取指定长度错误" + err.Error())
 	}

utils/util.go

@@ -40,21 +40,21 @@ WWW-Authenticate: Basic realm="easyProxy"
 func Relay(in, out net.Conn, compressType int, crypt, mux bool) {
 	switch compressType {
 	case COMPRESS_SNAPY_ENCODE:
-		io.Copy(NewSnappyConn(in, crypt), out)
+		copyBuffer(NewSnappyConn(in, crypt), out)
 		out.Close()
 		NewSnappyConn(in, crypt).Write([]byte(IO_EOF))
 	case COMPRESS_SNAPY_DECODE:
-		io.Copy(in, NewSnappyConn(out, crypt))
+		copyBuffer(in, NewSnappyConn(out, crypt))
 		in.Close()
 		if !mux {
 			out.Close()
 		}
 	case COMPRESS_NONE_ENCODE:
-		io.Copy(NewCryptConn(in, crypt), out)
+		copyBuffer(NewCryptConn(in, crypt), out)
 		out.Close()
 		NewCryptConn(in, crypt).Write([]byte(IO_EOF))
 	case COMPRESS_NONE_DECODE:
-		io.Copy(in, NewCryptConn(out, crypt))
+		copyBuffer(in, NewCryptConn(out, crypt))
 		in.Close()
 		if !mux {
 			out.Close()
@@ -150,11 +150,16 @@ var bufPool = sync.Pool{
 		return make([]byte, poolSize)
 	},
 }
+var bufPoolSmall = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, poolSizeSmall)
+	},
+}
 // io.copy的优化版，读取buffer长度原为32*1024，与snappy不同，导致读取出的内容存在差异，不利于解密，特此修改
 //废除
 func copyBuffer(dst io.Writer, src io.Reader) (written int64, err error) {
 	//TODO 回收问题
-	buf := bufPool.Get().([]byte)
+	buf := bufPool.Get().([]byte)[:32*1024]
 	for {
 		nr, er := src.Read(buf)
 		if nr > 0 {
@@ -197,3 +202,33 @@ func FlushConn(c net.Conn) {
 func Getverifyval(vkey string) string {
 	return Md5(vkey)
 }
+
+//wait replay group
+func ReplayWaitGroup(conn1 net.Conn, conn2 net.Conn, compressEncode, compressDecode int, crypt, mux bool) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		Relay(conn1, conn2, compressEncode, crypt, mux)
+		wg.Done()
+	}()
+	Relay(conn2, conn1, compressDecode, crypt, mux)
+	wg.Wait()
+}
+
+func ChangeHostAndHeader(r *http.Request, host string, header string, addr string) {
+	if host != "" {
+		r.Host = host
+	}
+	if header != "" {
+		h := strings.Split(header, "\n")
+		for _, v := range h {
+			hd := strings.Split(v, ":")
+			if len(hd) == 2 {
+				r.Header.Set(hd[0], hd[1])
+			}
+		}
+	}
+	addr = strings.Split(addr, ":")[0]
+	r.Header.Set("X-Forwarded-For", addr)
+	r.Header.Set("X-Real-IP", addr)
+}