reverseproxy - redirect.go

     1  package reverseproxy
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"sync/atomic"
     7  )
     8  
     9  type addrService struct {
    10  	copying uint64
    11  	MatchServiceName
    12  	net.Addr
    13  }
    14  
    15  func (a *addrService) Transfer(buf []byte, conn *net.TCPConn) {
    16  	p, err := net.Dial(a.Network(), a.String())
    17  	if err == nil {
    18  		if _, err = p.Write(buf); err == nil {
    19  			atomic.AddUint64(&a.copying, 2)
    20  			go copyConn(p, conn, &a.copying)
    21  			go copyConn(conn, p, &a.copying)
    22  		}
    23  	}
    24  }
    25  
    26  func (a *addrService) Active() bool {
    27  	return atomic.LoadUint64(&a.copying) > 0
    28  }
    29  
    30  func copyConn(a, b net.Conn, c *uint64) {
    31  	io.Copy(a, b)
    32  	a.Close()
    33  	b.Close()
    34  	atomic.AddUint64(c, ^uint64(0))
    35  }
    36  
    37  // AddRedirect sets a port to be redirected to an external service
    38  func AddRedirect(serviceName MatchServiceName, port uint16, to net.Addr) (*Port, error) {
    39  	return addPort(port, &addrService{
    40  		MatchServiceName: serviceName,
    41  		Addr:             to,
    42  	})
    43  }
    44