diff --git a/client/main.go b/client/main.go index c026442..e551c75 100644 --- a/client/main.go +++ b/client/main.go @@ -14,7 +14,6 @@ import ( "golang.org/x/crypto/pbkdf2" - "github.com/golang/snappy" "github.com/pkg/errors" "github.com/urfave/cli" kcp "github.com/xtaci/kcp-go" @@ -33,39 +32,6 @@ var VERSION = "SELFBUILD" // A pool for stream copying var xmitBuf sync.Pool -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) { - if _, err := c.w.Write(p); err != nil { - return 0, errors.WithStack(err) - } - - if err := c.w.Flush(); err != nil { - return 0, errors.WithStack(err) - } - return len(p), 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(sess *smux.Session, p1 io.ReadWriteCloser, quiet bool) { logln := func(v ...interface{}) { if !quiet { @@ -407,7 +373,7 @@ func main() { if config.NoComp { session, err = smux.Client(kcpconn, smuxConfig) } else { - session, err = smux.Client(newCompStream(kcpconn), smuxConfig) + session, err = smux.Client(generic.NewCompStream(kcpconn), smuxConfig) } if err != nil { return nil, errors.Wrap(err, "createConn()") diff --git a/generic/comp.go b/generic/comp.go new file mode 100644 index 0000000..03f6aea --- /dev/null +++ b/generic/comp.go @@ -0,0 +1,41 @@ +package generic + +import ( + "net" + + "github.com/golang/snappy" + "github.com/pkg/errors" +) + +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) { + if _, err := c.w.Write(p); err != nil { + return 0, errors.WithStack(err) + } + + if err := c.w.Flush(); err != nil { + return 0, errors.WithStack(err) + } + return len(p), 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 +} diff --git a/server/main.go b/server/main.go index 659927b..fdcad03 100644 --- a/server/main.go +++ b/server/main.go @@ -18,8 +18,6 @@ import ( "path/filepath" - "github.com/golang/snappy" - "github.com/pkg/errors" "github.com/urfave/cli" kcp "github.com/xtaci/kcp-go" "github.com/xtaci/kcptun/generic" @@ -35,39 +33,6 @@ var VERSION = "SELFBUILD" // A pool for stream copying var xmitBuf sync.Pool -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) { - if _, err := c.w.Write(p); err != nil { - return 0, errors.WithStack(err) - } - - if err := c.w.Flush(); err != nil { - return 0, errors.WithStack(err) - } - return len(p), 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 io.ReadWriteCloser, config *Config) { // stream multiplex @@ -414,7 +379,7 @@ func main() { if config.NoComp { go handleMux(conn, &config) } else { - go handleMux(newCompStream(conn), &config) + go handleMux(generic.NewCompStream(conn), &config) } } else { log.Printf("%+v", err)