diff --git a/mtglib/internal/relay/relay.go b/mtglib/internal/relay/relay.go index 294fa5d..d4abe87 100644 --- a/mtglib/internal/relay/relay.go +++ b/mtglib/internal/relay/relay.go @@ -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) } } diff --git a/mtglib/internal/relay/sync_pair.go b/mtglib/internal/relay/sync_pair.go index e83c58e..68c1e62 100644 --- a/mtglib/internal/relay/sync_pair.go +++ b/mtglib/internal/relay/sync_pair.go @@ -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