proxy.go 1.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. package netutil
  2. import (
  3. "io"
  4. "log"
  5. "net"
  6. "sync"
  7. "time"
  8. )
  9. // ProxyTCP proxies between two TCP connections.
  10. // Because TLS connections don't have CloseRead() and CloseWrite() methods, our
  11. // temporary solution is to use timeouts.
  12. func ProxyTCP(conn1, conn2 net.Conn, tlsWriteDeadline, tlsReadDeadline time.Duration) {
  13. var wg sync.WaitGroup
  14. wg.Add(2)
  15. go copyBytes(conn1, conn2, &wg, tlsWriteDeadline, tlsReadDeadline)
  16. go copyBytes(conn2, conn1, &wg, tlsWriteDeadline, tlsReadDeadline)
  17. wg.Wait()
  18. conn1.Close()
  19. conn2.Close()
  20. }
  21. func copyBytes(dst, src net.Conn, wg *sync.WaitGroup, writeDeadline, readDeadline time.Duration) {
  22. defer wg.Done()
  23. _, err := io.Copy(dst, src)
  24. if err != nil {
  25. log.Printf("proxy i/o error: %v", err)
  26. }
  27. if cr, ok := src.(*net.TCPConn); ok {
  28. cr.CloseRead()
  29. } else {
  30. // For TLS connections.
  31. wto := time.Now().Add(writeDeadline)
  32. src.SetWriteDeadline(wto)
  33. }
  34. if cw, ok := dst.(*net.TCPConn); ok {
  35. cw.CloseWrite()
  36. } else {
  37. // For TLS connections.
  38. rto := time.Now().Add(readDeadline)
  39. dst.SetReadDeadline(rto)
  40. }
  41. }