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) { func Relay(ctx context.Context, log Logger, telegramConn, clientConn net.Conn) {
defer telegramConn.Close()
defer clientConn.Close()
ctx, cancel := context.WithCancel(ctx) ctx, cancel := context.WithCancel(ctx)
defer cancel() 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) { func pump(log Logger, src, dst net.Conn, wg *sync.WaitGroup, direction string) {
defer wg.Done() defer wg.Done()
defer src.Close()
defer dst.Close()
syncer := acquireSyncPair(src, dst) syncer := acquireSyncPair(src, dst)
defer releaseSyncPair(syncer) defer releaseSyncPair(syncer)
defer syncer.Flush()
if n, err := syncer.Sync(); err != nil { 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" "io"
"net" "net"
"os" "os"
"sync"
"time" "time"
) )
@@ -14,6 +15,7 @@ type syncPair struct {
writer *bufio.Writer writer *bufio.Writer
copyBuf []byte copyBuf []byte
mutex sync.Mutex
reader net.Conn reader net.Conn
} }
@@ -25,7 +27,7 @@ func (s *syncPair) Read(p []byte) (int, error) {
n, err := s.readBlocking(p, false) n, err := s.readBlocking(p, false)
if errors.Is(err, os.ErrDeadlineExceeded) { 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) 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) { func (s *syncPair) Write(p []byte) (int, error) {
s.mutex.Lock()
defer s.mutex.Unlock()
return s.writer.Write(p) // nolint: wrapcheck 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) { func (s *syncPair) readBlocking(p []byte, blocking bool) (int, error) {
var deadline time.Time var deadline time.Time