vmess: support chunk stream

conf.go

@@ -2,9 +2,11 @@ package main
 
 import (
 	"fmt"
+	"io/ioutil"
 	"log"
 	"os"
 	"path"
+	"strings"
 
 	"github.com/nadoo/conflag"
 )
@@ -84,6 +86,25 @@ func confInit() {
 
 }
 
+func listDir(dirPth string, suffix string) (files []string, err error) {
+	files = make([]string, 0, 10)
+	dir, err := ioutil.ReadDir(dirPth)
+	if err != nil {
+		return nil, err
+	}
+	PthSep := string(os.PathSeparator)
+	suffix = strings.ToUpper(suffix)
+	for _, fi := range dir {
+		if fi.IsDir() {
+			continue
+		}
+		if strings.HasSuffix(strings.ToUpper(fi.Name()), suffix) {
+			files = append(files, dirPth+PthSep+fi.Name())
+		}
+	}
+	return files, nil
+}
+
 // RuleConf , every ruleForwarder points to a rule file
 type RuleConf struct {
 	name string

main.go

@@ -25,7 +25,7 @@ import (
 )
 
 // VERSION .
-const VERSION = "0.6.0"
+const VERSION = "0.6.2"
 
 func dialerFromConf() proxy.Dialer {
 	// global forwarders in xx.conf
@@ -48,31 +48,16 @@ func dialerFromConf() proxy.Dialer {
 func main() {
 
 	confInit()
-
 	log.F = func(f string, v ...interface{}) {
 		if conf.Verbose {
 			stdlog.Printf(f, v...)
 		}
 	}
 
-	sDialer := NewRuleDialer(conf.rules, dialerFromConf())
+	dialer := NewRuleDialer(conf.rules, dialerFromConf())
-
+	ipsetM, _ := NewIPSetManager(conf.IPSet, conf.rules)
-	for _, listen := range conf.Listen {
-		local, err := proxy.ServerFromURL(listen, sDialer)
-		if err != nil {
-			log.Fatal(err)
-		}
-
-		go local.ListenAndServe()
-	}
-
-	ipsetM, err := NewIPSetManager(conf.IPSet, conf.rules)
-	if err != nil {
-		log.F("create ipset manager error: %s", err)
-	}
-
 	if conf.DNS != "" {
-		d, err := dns.NewDNS(conf.DNS, conf.DNSServer[0], sDialer, false)
+		d, err := dns.NewDNS(conf.DNS, conf.DNSServer[0], dialer, false)
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -87,7 +72,7 @@ func main() {
 		}
 
 		// add a handler to update proxy rules when a domain resolved
-		d.AddAnswerHandler(sDialer.AddDomainIP)
+		d.AddAnswerHandler(dialer.AddDomainIP)
 		if ipsetM != nil {
 			d.AddAnswerHandler(ipsetM.AddDomainIP)
 		}
@@ -95,6 +80,15 @@ func main() {
 		go d.ListenAndServe()
 	}
 
+	for _, listen := range conf.Listen {
+		local, err := proxy.ServerFromURL(listen, dialer)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		go local.ListenAndServe()
+	}
+
 	sigCh := make(chan os.Signal, 1)
 	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
 	<-sigCh

proxy/vmess/chunk.go

@@ -1,3 +1,105 @@
 package vmess
 
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+)
+
 // chunk: plain, AES-128-CFB, AES-128-GCM, ChaCha20-Poly1305
+
+const maxChunkSize = 1 << 14     // 16384
+const defaultChunkSize = 1 << 13 // 8192
+
+type chunkedReader struct {
+	io.Reader
+	buf      []byte
+	leftover []byte
+}
+
+func newChunkedReader(r io.Reader) io.Reader {
+	return &chunkedReader{
+		Reader: r,
+		buf:    make([]byte, 2+maxChunkSize),
+	}
+}
+
+func (r *chunkedReader) read() (int, error) {
+	lenBuf := make([]byte, 2)
+	_, err := io.ReadFull(r.Reader, lenBuf)
+	if err != nil {
+		return 0, err
+	}
+
+	len := binary.BigEndian.Uint16(lenBuf)
+
+	buf := r.buf[:len]
+	_, err = io.ReadFull(r.Reader, buf)
+	if err != nil {
+		return 0, err
+	}
+
+	return int(len), nil
+}
+
+func (r *chunkedReader) Read(b []byte) (int, error) {
+	if len(r.leftover) > 0 {
+		n := copy(b, r.leftover)
+		r.leftover = r.leftover[n:]
+		return n, nil
+	}
+
+	n, err := r.read()
+	m := copy(b, r.buf[:n])
+	if m < n {
+		r.leftover = r.buf[m:n]
+	}
+
+	return m, err
+}
+
+type chunkedWriter struct {
+	io.Writer
+	buf []byte
+}
+
+func newChunkedWriter(w io.Writer) io.Writer {
+	return &chunkedWriter{
+		Writer: w,
+		buf:    make([]byte, 2+maxChunkSize),
+	}
+}
+
+func (w *chunkedWriter) Write(b []byte) (int, error) {
+	n, err := w.ReadFrom(bytes.NewBuffer(b))
+	return int(n), err
+}
+
+func (w *chunkedWriter) ReadFrom(r io.Reader) (n int64, err error) {
+	for {
+		buf := w.buf
+		payloadBuf := buf[2 : 2+defaultChunkSize]
+
+		nr, er := r.Read(payloadBuf)
+		if nr > 0 {
+			n += int64(nr)
+			payloadBuf = payloadBuf[:nr]
+			binary.BigEndian.PutUint16(buf[:], uint16(nr))
+
+			_, ew := w.Writer.Write(buf)
+			if ew != nil {
+				err = ew
+				break
+			}
+		}
+
+		if er != nil {
+			if er != io.EOF { // ignore EOF as per io.ReaderFrom contract
+				err = er
+			}
+			break
+		}
+	}
+
+	return n, err
+}

proxy/vmess/client.go

@@ -18,6 +18,7 @@ import (
 
 // Request Options
 const (
+	OptBasicFormat        byte = 0
 	OptChunkStream        byte = 1
 	OptReuseTCPConnection byte = 2
 	OptMetadataObfuscate  byte = 4
@@ -43,12 +44,14 @@ const (
 type Client struct {
 	users    []*User
 	count    int
+	opt      byte
 	security byte
 }
 
 // Conn is a connection to vmess server
 type Conn struct {
 	user     *User
+	opt      byte
 	security byte
 
 	atyp Atyp
@@ -63,6 +66,9 @@ type Conn struct {
 
 	net.Conn
 	connected bool
+
+	dataReader io.Reader
+	dataWriter io.Writer
 }
 
 // NewClient .
@@ -78,6 +84,9 @@ func NewClient(uuidStr, security string, alterID int) (*Client, error) {
 	c.users = append(c.users, user.GenAlterIDUsers(alterID)...)
 	c.count = len(c.users)
 
+	// TODO: config?
+	c.opt = OptBasicFormat
+
 	security = strings.ToLower(security)
 	switch security {
 	case "aes-128-cfb":
@@ -96,7 +105,7 @@ func NewClient(uuidStr, security string, alterID int) (*Client, error) {
 // NewConn .
 func (c *Client) NewConn(rc net.Conn, target string) (*Conn, error) {
 	r := rand.Intn(c.count)
-	conn := &Conn{user: c.users[r], security: c.security}
+	conn := &Conn{user: c.users[r], opt: c.opt, security: c.security}
 
 	var err error
 	conn.atyp, conn.addr, conn.port, err = ParseAddr(target)
@@ -149,7 +158,7 @@ func (c *Conn) EncodeRequest() ([]byte, error) {
 	buf.Write(c.reqBodyIV[:])  // IV
 	buf.Write(c.reqBodyKey[:]) // Key
 	buf.WriteByte(c.reqRespV)  // V
-	buf.WriteByte(0)           // Opt
+	buf.WriteByte(c.opt)       // Opt
 
 	// pLen and Sec
 	paddingLen := rand.Intn(16)
@@ -218,9 +227,24 @@ func (c *Conn) Read(b []byte) (n int, err error) {
 		c.DecodeRespHeader()
 	}
 
+	if c.opt&OptChunkStream != 0 {
+		if c.dataReader == nil {
+			c.dataReader = newChunkedReader(c.Conn)
+		}
+
+		return c.dataReader.Read(b)
+	}
+
 	return c.Conn.Read(b)
 }
 
 func (c *Conn) Write(b []byte) (n int, err error) {
+	if c.opt&OptChunkStream != 0 {
+		if c.dataWriter == nil {
+			c.dataWriter = newChunkedWriter(c.Conn)
+		}
+
+		return c.dataWriter.Write(b)
+	}
 	return c.Conn.Write(b)
 }

utils.go

@@ -1,26 +0,0 @@
-package main
-
-import (
-	"io/ioutil"
-	"os"
-	"strings"
-)
-
-func listDir(dirPth string, suffix string) (files []string, err error) {
-	files = make([]string, 0, 10)
-	dir, err := ioutil.ReadDir(dirPth)
-	if err != nil {
-		return nil, err
-	}
-	PthSep := string(os.PathSeparator)
-	suffix = strings.ToUpper(suffix)
-	for _, fi := range dir {
-		if fi.IsDir() {
-			continue
-		}
-		if strings.HasSuffix(strings.ToUpper(fi.Name()), suffix) {
-			files = append(files, dirPth+PthSep+fi.Name())
-		}
-	}
-	return files, nil
-}