proto/tcp: fix unexpected close function, fix
@@ -27,25 +27,31 @@ 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))
+ return WrapReadWriteCloser(snappy.NewReader(rwc), snappy.NewWriter(rwc), func() error {
+ return rwc.Close()
+ })
}
-func WrapReadWriteCloser(r io.Reader, w io.Writer) io.ReadWriteCloser {
+type ReadWriteCloser struct {
+ 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,
}
}
-type ReadWriteCloser struct {
- r io.Reader
- w io.Writer
-}
-
func (rwc *ReadWriteCloser) Read(p []byte) (n int, err error) {
return rwc.r.Read(p)
}
@@ -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
}
