Upgrade to v0.2

README.md

@@ -12,18 +12,18 @@ when I first used them, I can't remember these complicated parameters for a long
 
 Besides, I think the net programming logic could be optimized. 
 
-For example, while running `lcx -listen 8888 9999` command, client must connect to `:8888` first, then `:9999`, in `iox`, there's no limit to the order in two ports. And while running `lcx -slave 1.1.1.1 8888 1.1.1.1 9999` command, `lcx` will connect two hosts serially, but it's more efficient to connect in concurrently, as `iox` does.
+For example, while running `lcx -listen 8888 9999` command, client must connect to `:8888` first, then `:9999`, in `iox`, there's no limit to the order in two ports. And while running `lcx -slave 1.1.1.1 8888 1.1.1.1 9999` command, `lcx` will connect two hosts serially, but it's more efficient to connect in concurrent, as `iox` does.
 
 And what's more, `iox` provides traffic encryption feature. Actually, you can use `iox` as a simple ShadowSocks.
 
-Of course, because `iox` is written in Go, the static-link-program is a little big, raw program is 2.2MB (800KB for UPX compression)
+Of course, because `iox` is written in Go, the static-link-program is a little large, raw program is 2.2MB (800KB for UPX compression)
 
 ## Feature
 
 + traffic encryption (optional)
 + humanized CLI option
 + logic optimization
-+ UDP traffic forward (TODO)
++ UDP traffic forward
 
 ## Usage
 
@@ -140,6 +140,22 @@ Using `iox` as a simple ShadowSocks
 ./iox fwd -l 1080 -r *VPS:9999 -k 000102
 ```
 
+#### UDP forward
+
+Only need to add CLI option `-u`
+
+```
+./iox fwd -l 53 -r *127.0.0.1:8888 -k 000102 -u
+./iox fwd -l *8888 -l *9999 -k 000102 -u
+./iox fwd -r *127.0.0.1:9999 -r 8.8.8.8:53 -k 000102 -u
+```
+
+**NOTICE: When you make a multistage connection, the `Remote2Remote-UDP-mode` must be started last, which is the No.3 command in above example**
+
+UDP forwarding may have behavior that is not as you expected, because there are many differences between stream & packet.
+
+You can find why in the source code, if you have any ideas, PR / issue are welcomed
+
 ## License
 
 The MIT license

crypto/aes.go

@@ -3,6 +3,12 @@ package crypto
 import (
 	"crypto/aes"
 	"crypto/cipher"
+	"crypto/rand"
+)
+
+var (
+	SECRET_KEY []byte
+	IV         []byte
 )
 
 func expand32(key []byte) ([]byte, []byte) {
@@ -18,28 +24,51 @@ func expand32(key []byte) ([]byte, []byte) {
 	return key[:0x10], key[0x10:0x20]
 }
 
+func ExpandKey(key []byte) {
+	SECRET_KEY, IV = expand32(key)
+}
+
 type Cipher struct {
 	c cipher.Stream
 }
 
 // AES-128-CTR
-func NewCipherPair(key []byte) (*Cipher, *Cipher, error) {
-	secretKey, iv := expand32(key)
-	blockA, err := aes.NewCipher(secretKey)
+func NewCipherPair() (*Cipher, *Cipher, error) {
+	blockA, err := aes.NewCipher(SECRET_KEY)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	blockB, err := aes.NewCipher(secretKey)
+	blockB, err := aes.NewCipher(SECRET_KEY)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	return &Cipher{cipher.NewCTR(blockA, iv)},
-		&Cipher{cipher.NewCTR(blockB, iv)},
+	return &Cipher{cipher.NewCTR(blockA, IV)},
+		&Cipher{cipher.NewCTR(blockB, IV)},
 		nil
 }
 
+func RandomIV() ([]byte, error) {
+	iv := make([]byte, 0x10)
+	_, err := rand.Read(iv)
+	if err != nil {
+		return nil, err
+	}
+	return iv, nil
+}
+
+func NewCipher(iv []byte) (*Cipher, error) {
+	block, err := aes.NewCipher(SECRET_KEY)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Cipher{
+		c: cipher.NewCTR(block, iv),
+	}, nil
+}
+
 func (c Cipher) StreamXOR(dst []byte, src []byte) {
 	c.c.XORKeyStream(dst, src)
 }

docs/README_CN.md

@@ -23,7 +23,7 @@
 + 流量加密（可选）
 + 友好的命令行参数
 + 部分逻辑优化
-+ UDP流量转发（TODO）
++ UDP流量转发
 
 ## 用法
 
@@ -132,12 +132,30 @@ $ curl google.com -x socks5://127.0.0.1:1000
 你也可以把`iox`当做一个简单的ShadowSocks来用：
 
 ```
+// ssserver
 ./iox proxy -l *9999 -k 000102
 
 
+// sslocal
 ./iox fwd -l 1080 -r *VPS:9999 -l 000102
 ```
 
+#### UDP转发
+
+只需要添加命令行参数：`-u`
+
+```
+./iox fwd -l 53 -r *127.0.0.1:8888 -k 000102 -u
+./iox fwd -l *8888 -l *9999 -k 000102 -u
+./iox fwd -r *127.0.0.1:9999 -r 8.8.8.8:53 -k 000102 -u
+```
+
+**注意：当你做多级连接的转发时，`Remote2Remote-UDP-mode`必须最后一个被启动，也就是上面示例中的第三条**
+
+UDP转发可能会有一些不合你预期的行为，因为流和包的转发有很多不同点
+
+你可以在源码里找到答案，如果你有什么想法，欢迎提PR / issue
+
 ## 许可
 
 The MIT license

main.go

@@ -7,7 +7,7 @@ import (
 	"os"
 )
 
-const VERSION = "0.1.1"
+const VERSION = "0.2"
 
 func Usage() {
 	fmt.Printf(
@@ -21,6 +21,8 @@ func Usage() {
 			"      remote host to connect, HOST can be IP or Domain. `*` means encrypted socket\n"+
 			"  -k HEX\n"+
 			"      hexadecimal format key, be used to generate AES Key and IV\n"+
+			"  -u\n"+
+			"      udp forward mode\n"+
 			"  -t TIMEOUT\n"+
 			"      set connection timeout(millisecond), default is 5000\n"+
 			"  -v\n"+

netio/forward.go

@@ -8,29 +8,23 @@ import (
 
 // Memory optimized
 func CipherCopy(dst Ctx, src Ctx) (int64, error) {
-	buffer := make([]byte, option.BUFFER_SIZE)
+	buffer := make([]byte, option.TCP_BUFFER_SIZE)
 	var written int64
 	var err error
 
 	for {
 		var nr int
 		var er error
-		if src.IsEncrypted() {
-			nr, er = src.DecryptRead(buffer)
-		} else {
-			nr, er = src.Read(buffer)
-		}
+
+		nr, er = src.DecryptRead(buffer)
 
 		if nr > 0 {
 			logger.Info(" <== [%d bytes] ==", nr)
 
 			var nw int
 			var ew error
-			if dst.IsEncrypted() {
-				nw, ew = dst.EncryptWrite(buffer[:nr])
-			} else {
-				nw, ew = dst.Write(buffer[:nr])
-			}
+
+			nw, ew = dst.EncryptWrite(buffer[:nr])
 
 			if nw > 0 {
 				logger.Info(" == [%d bytes] ==> ", nw)
@@ -52,6 +46,7 @@ func CipherCopy(dst Ctx, src Ctx) (int64, error) {
 			break
 		}
 	}
+
 	return written, err
 }
 
@@ -70,3 +65,127 @@ func PipeForward(ctxA Ctx, ctxB Ctx) {
 
 	<-signal
 }
+
+// This function will run forever
+// If need to do performance optimization in future,
+// I will consider a go-routine pool here, but
+// this can lead to mutex-lock overhead
+func ForwardUDP(ctxA Ctx, ctxB Ctx) {
+	go func() {
+		buffer := make([]byte, option.UDP_PACKET_MAX_SIZE)
+		for {
+			nr, _ := ctxA.DecryptRead(buffer)
+			if nr > 0 {
+				if nr == 4 &&
+					buffer[0] == 0xCC && buffer[1] == 0xDD &&
+					buffer[2] == 0xEE && buffer[3] == 0xFF {
+					continue
+				}
+				logger.Info(" <== [%d bytes] ==", nr)
+				nw, _ := ctxB.EncryptWrite(buffer[:nr])
+				if nw > 0 {
+					logger.Info(" == [%d bytes] ==>", nw)
+				}
+			}
+		}
+	}()
+
+	go func() {
+		buffer := make([]byte, option.UDP_PACKET_MAX_SIZE)
+		for {
+			nr, _ := ctxB.DecryptRead(buffer)
+			if nr > 0 {
+				if nr == 4 &&
+					buffer[0] == 0xCC && buffer[1] == 0xDD &&
+					buffer[2] == 0xEE && buffer[3] == 0xFF {
+					continue
+				}
+				logger.Info(" <== [%d bytes] ==", nr)
+				nw, _ := ctxA.EncryptWrite(buffer[:nr])
+				if nw > 0 {
+					logger.Info(" == [%d bytes] ==>", nw)
+				}
+			}
+		}
+	}()
+
+	select {}
+}
+
+var UDP_INIT_PACKET = []byte{
+	0xCC, 0xDD, 0xEE, 0xFF,
+}
+
+// Each socket only writes the packet to the address
+// which last sent packet to it recently,
+// instead of boardcasting to all the address.
+// I think it is as expected
+func ForwardUnconnectedUDP(ctxA Ctx, ctxB Ctx) {
+	addrRegistedA := false
+	addrRegistedB := false
+	addrRegistedSignalA := make(chan struct{}, 1)
+	addrRegistedSignalB := make(chan struct{}, 1)
+
+	packetChannelA := make(chan []byte, option.UDP_PACKET_CHANNEL_SIZE)
+	packetChannelB := make(chan []byte, option.UDP_PACKET_CHANNEL_SIZE)
+
+	// A read
+	go func() {
+		for {
+			buffer := make([]byte, option.UDP_PACKET_MAX_SIZE)
+			nr, _ := ctxA.DecryptRead(buffer)
+			if nr > 0 {
+				if !addrRegistedA {
+					addrRegistedA = true
+					addrRegistedSignalA <- struct{}{}
+				}
+
+				if !(nr == 4 &&
+					buffer[0] == 0xCC && buffer[1] == 0xDD &&
+					buffer[2] == 0xEE && buffer[3] == 0xFF) {
+					packetChannelB <- buffer[:nr]
+				}
+			}
+		}
+	}()
+
+	// B read
+	go func() {
+		for {
+			buffer := make([]byte, option.UDP_PACKET_MAX_SIZE)
+			nr, _ := ctxB.DecryptRead(buffer)
+			if nr > 0 {
+				if !addrRegistedB {
+					addrRegistedB = true
+					addrRegistedSignalB <- struct{}{}
+				}
+
+				if !(nr == 4 &&
+					buffer[0] == 0xCC && buffer[1] == 0xDD &&
+					buffer[2] == 0xEE && buffer[3] == 0xFF) {
+					packetChannelA <- buffer[:nr]
+				}
+			}
+		}
+	}()
+
+	// A write
+	go func() {
+		<-addrRegistedSignalA
+		for {
+			packet := <-packetChannelA
+			ctxA.EncryptWrite(packet)
+		}
+	}()
+
+	// B write
+	go func() {
+		<-addrRegistedSignalB
+		for {
+			packet := <-packetChannelB
+			ctxB.EncryptWrite(packet)
+		}
+	}()
+
+	select {}
+}