This commit is contained in:
9seconds
2021-11-30 15:19:37 +03:00
parent a5e59d9ef7
commit 7b1f86b75d
2 changed files with 18 additions and 4 deletions
+5 -3
View File
@@ -7,6 +7,9 @@ import (
)
func Relay(ctx context.Context, log Logger, telegramConn, clientConn net.Conn) {
defer telegramConn.Close()
defer clientConn.Close()
ctx, cancel := context.WithCancel(ctx)
defer cancel()
@@ -28,13 +31,12 @@ func Relay(ctx context.Context, log Logger, telegramConn, clientConn net.Conn) {
func pump(log Logger, src, dst net.Conn, wg *sync.WaitGroup, direction string) {
defer wg.Done()
defer src.Close()
defer dst.Close()
syncer := acquireSyncPair(src, dst)
defer releaseSyncPair(syncer)
defer syncer.Flush()
if n, err := syncer.Sync(); err != nil {
log.Printf("cannot pump %s (written %d bytes): %w", direction, n, err)
log.Printf("cannot pump %s (written %d bytes): %v", direction, n, err)
}
}
+13 -1
View File
@@ -7,6 +7,7 @@ import (
"io"
"net"
"os"
"sync"
"time"
)
@@ -14,6 +15,7 @@ type syncPair struct {
writer *bufio.Writer
copyBuf []byte
mutex sync.Mutex
reader net.Conn
}
@@ -25,7 +27,7 @@ func (s *syncPair) Read(p []byte) (int, error) {
n, err := s.readBlocking(p, false)
if errors.Is(err, os.ErrDeadlineExceeded) {
if err := s.writer.Flush(); err != nil {
if err := s.Flush(); err != nil {
return 0, fmt.Errorf("cannot flush writer hand-side: %w", err)
}
@@ -36,9 +38,19 @@ func (s *syncPair) Read(p []byte) (int, error) {
}
func (s *syncPair) Write(p []byte) (int, error) {
s.mutex.Lock()
defer s.mutex.Unlock()
return s.writer.Write(p) // nolint: wrapcheck
}
func (s *syncPair) Flush() error {
s.mutex.Lock()
defer s.mutex.Unlock()
return s.writer.Flush()
}
func (s *syncPair) readBlocking(p []byte, blocking bool) (int, error) {
var deadline time.Time