seperate compStream

This commit is contained in:
xtaci
2019-05-10 21:44:02 +08:00
parent 86cc46f437
commit 116c1e7668
3 changed files with 43 additions and 71 deletions
+1 -35
View File
@@ -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()")
+41
View File
@@ -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
}
+1 -36
View File
@@ -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)