OrgGo
/
frp
mirror of https://github.com/fatedier/frp.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

proto/tcp: fix unexpected close function, fix #332

This commit is contained in:
fatedier 2017-06-05 23:52:24 +08:00
parent 3a4d769bb3
commit 04f4fd0a81
1 changed files with 24 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
models/proto/tcp/tcp.go
View File

@ -27,23 +27,29 @@ func WithEncryption(rwc io.ReadWriteCloser, key []byte) (io.ReadWriteCloser, err
if err != nil {
return nil, err
}
return WrapReadWriteCloser(crypto.NewReader(rwc, key), w), nil
return WrapReadWriteCloser(crypto.NewReader(rwc, key), w, func() error {
return rwc.Close()
}), nil
}
func WithCompression(rwc io.ReadWriteCloser) io.ReadWriteCloser {
return WrapReadWriteCloser(snappy.NewReader(rwc), snappy.NewWriter(rwc))
}
func WrapReadWriteCloser(r io.Reader, w io.Writer) io.ReadWriteCloser {
return &ReadWriteCloser{
r: r,
w: w,
}
return WrapReadWriteCloser(snappy.NewReader(rwc), snappy.NewWriter(rwc), func() error {
return rwc.Close()
})
}
type ReadWriteCloser struct {
r io.Reader
w io.Writer
r io.Reader
w io.Writer
closeFn func() error
}
func WrapReadWriteCloser(r io.Reader, w io.Writer, closeFn func() error) io.ReadWriteCloser {
return &ReadWriteCloser{
r: r,
w: w,
closeFn: closeFn,
}
}
func (rwc *ReadWriteCloser) Read(p []byte) (n int, err error) {
@ -69,5 +75,12 @@ func (rwc *ReadWriteCloser) Close() (errRet error) {
errRet = err
}
}
if rwc.closeFn != nil {
err = rwc.closeFn()
if err != nil {
errRet = err
}
}
return
}