guoswork/frp
1
0
Fork 0
mirror of https://gitee.com/IrisVega/frp.git synced 2024-11-01 22:31:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
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 { if err != nil {
return nil, err 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 { func WithCompression(rwc io.ReadWriteCloser) io.ReadWriteCloser {
return WrapReadWriteCloser(snappy.NewReader(rwc), snappy.NewWriter(rwc)) return WrapReadWriteCloser(snappy.NewReader(rwc), snappy.NewWriter(rwc), func() error {
} return rwc.Close()
})
func WrapReadWriteCloser(r io.Reader, w io.Writer) io.ReadWriteCloser {
return &ReadWriteCloser{
r: r,
w: w,
}
} }
type ReadWriteCloser struct { type ReadWriteCloser struct {
r io.Reader r io.Reader
w io.Writer 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) { func (rwc *ReadWriteCloser) Read(p []byte) (n int, err error) {
@ -69,5 +75,12 @@ func (rwc *ReadWriteCloser) Close() (errRet error) {
errRet = err errRet = err
} }
} }
if rwc.closeFn != nil {
err = rwc.closeFn()
if err != nil {
errRet = err
}
}
return return
} }