add snappy compression to connection

xtaci · Jun 27, 2016 · d97ba4c · d97ba4c
1 parent 230d1f4
commit d97ba4c
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 5 deletions.
diff --git a/client/main.go b/client/main.go
@@ -11,6 +11,7 @@ import (
 
 	"golang.org/x/crypto/pbkdf2"
 
+	"github.com/golang/snappy"
 	"github.com/hashicorp/yamux"
 	"github.com/urfave/cli"
 	"github.com/xtaci/kcp-go"
@@ -21,7 +22,35 @@ var (
 	SALT    = "kcp-go"
 )
 
-func handleClient(p1, p2 net.Conn) {
+type compStream struct {
+	conn net.Conn
+	w    *snappy.Writer
+	r    *snappy.Reader
+}
+
+func (c *compStream) Read(p []byte) (n int, err error) {
+	return c.r.Read(p)
+}
+
+func (c *compStream) Write(p []byte) (n int, err error) {
+	n, err = c.w.Write(p)
+	err = c.w.Flush()
+	return n, err
+}
+
+func (c *compStream) Close() error {
+	return c.conn.Close()
+}
+
+func newCompStream(conn net.Conn) *compStream {
+	c := new(compStream)
+	c.conn = conn
+	c.w = snappy.NewBufferedWriter(conn)
+	c.r = snappy.NewReader(conn)
+	return c
+}
+
+func handleClient(p1, p2 io.ReadWriteCloser) {
 	log.Println("stream opened")
 	defer log.Println("stream closed")
 	defer p1.Close()
@@ -107,6 +136,10 @@ func main() {
 			Value: 1024,
 			Usage: "set receive window size(num of packets)",
 		},
+		cli.BoolFlag{
+			Name:  "nocomp",
+			Usage: "disable compression",
+		},
 		cli.IntFlag{
 			Name:  "datashard",
 			Value: 10,
@@ -209,7 +242,12 @@ func main() {
 				MaxStreamWindowSize:    16777216,
 				LogOutput:              os.Stderr,
 			}
-			session, err := yamux.Client(kcpconn, config)
+			var session *yamux.Session
+			if c.Bool("nocomp") {
+				session, err = yamux.Client(kcpconn, config)
+			} else {
+				session, err = yamux.Client(newCompStream(kcpconn), config)
+			}
 			checkError(err)
 			return session
 		}

diff --git a/server/main.go b/server/main.go
@@ -11,6 +11,7 @@ import (
 
 	"golang.org/x/crypto/pbkdf2"
 
+	"github.com/golang/snappy"
 	"github.com/hashicorp/yamux"
 	"github.com/urfave/cli"
 	"github.com/xtaci/kcp-go"
@@ -21,8 +22,36 @@ var (
 	SALT    = "kcp-go"
 )
 
+type compStream struct {
+	conn net.Conn
+	w    *snappy.Writer
+	r    *snappy.Reader
+}
+
+func (c *compStream) Read(p []byte) (n int, err error) {
+	return c.r.Read(p)
+}
+
+func (c *compStream) Write(p []byte) (n int, err error) {
+	n, err = c.w.Write(p)
+	err = c.w.Flush()
+	return n, err
+}
+
+func (c *compStream) Close() error {
+	return c.conn.Close()
+}
+
+func newCompStream(conn net.Conn) *compStream {
+	c := new(compStream)
+	c.conn = conn
+	c.w = snappy.NewBufferedWriter(conn)
+	c.r = snappy.NewReader(conn)
+	return c
+}
+
 // handle multiplex-ed connection
-func handleMux(conn *kcp.UDPSession, target string) {
+func handleMux(conn io.ReadWriteCloser, target string) {
 	// stream multiplex
 	var mux *yamux.Session
 	config := &yamux.Config{
@@ -56,7 +85,7 @@ func handleMux(conn *kcp.UDPSession, target string) {
 	}
 }
 
-func handleClient(p1, p2 net.Conn) {
+func handleClient(p1, p2 io.ReadWriteCloser) {
 	log.Println("stream opened")
 	defer log.Println("stream closed")
 	defer p1.Close()
@@ -130,6 +159,10 @@ func main() {
 			Value: 1024,
 			Usage: "set receive window size(num of packets)",
 		},
+		cli.BoolFlag{
+			Name:  "nocomp",
+			Usage: "disable compression",
+		},
 		cli.IntFlag{
 			Name:  "datashard",
 			Value: 10,
@@ -217,7 +250,12 @@ func main() {
 				conn.SetWindowSize(c.Int("sndwnd"), c.Int("rcvwnd"))
 				conn.SetACKNoDelay(c.Bool("acknodelay"))
 				conn.SetDSCP(c.Int("dscp"))
-				go handleMux(conn, c.String("target"))
+
+				if c.Bool("nocomp") {
+					go handleMux(conn, c.String("target"))
+				} else {
+					go handleMux(newCompStream(conn), c.String("target"))
+				}
 			} else {
 				log.Println(err)
 			}