proxy.go 1.2 KB

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