mirror of https://github.com/fatedier/frp.git
proto/tcp: fix unexpected close function, fix #332
This commit is contained in:
parent
3a4d769bb3
commit
04f4fd0a81
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue