ce7cb702218693d08ea6d83464d2457c9250027c

services/oci-l4-relay/proxy.go

@@ -0,0 +1,208 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"strconv"
+	"sync"
+	"time"
+)
+
+var proxySignature = []byte("\r\n\r\n\x00\r\nQUIT\n")
+
+type mailProxy struct {
+	listener net.Listener
+	upstream string
+	source   net.IP
+	failures chan<- error
+	slots    chan struct{}
+	mu       sync.Mutex
+	closed   bool
+	conns    map[net.Conn]struct{}
+	workers  sync.WaitGroup
+	done     chan struct{}
+}
+
+func newMailProxy(listener net.Listener, upstream string, source net.IP, failures chan<- error, slots chan struct{}) *mailProxy {
+	return &mailProxy{listener: listener, upstream: upstream, source: source, failures: failures, slots: slots, conns: make(map[net.Conn]struct{}), done: make(chan struct{})}
+}
+
+func (r *relay) startProxies() error {
+	if r.config.mailTarget == nil {
+		return nil
+	}
+	slots := make(chan struct{}, mailConnectionLimit)
+	for _, port := range r.config.mailTarget.ports {
+		listener, err := net.ListenTCP("tcp4", &net.TCPAddr{IP: r.config.mailTarget.address, Port: int(port)})
+		if err != nil {
+			r.stopProxies()
+			return fmt.Errorf("listen on %s:%d: %w", r.config.mailTarget.address, port, err)
+		}
+		proxy := newMailProxy(listener, net.JoinHostPort(r.config.mailUpstream, strconv.Itoa(int(port))), net.ParseIP(r.config.podIP).To4(), r.proxyFailures, slots)
+		r.proxies = append(r.proxies, proxy)
+		go proxy.serve()
+	}
+	return nil
+}
+
+func (r *relay) stopProxies() {
+	for _, proxy := range r.proxies {
+		proxy.close()
+	}
+	r.proxies = nil
+}
+
+func (p *mailProxy) serve() {
+	defer close(p.done)
+	for {
+		client, err := p.listener.Accept()
+		if err != nil {
+			if !p.isClosed() {
+				p.report(fmt.Errorf("accept connection: %w", err))
+			}
+			return
+		}
+		select {
+		case p.slots <- struct{}{}:
+		default:
+			_ = client.Close()
+			continue
+		}
+		if !p.track(client) {
+			<-p.slots
+			_ = client.Close()
+			return
+		}
+		p.workers.Add(1)
+		go p.forward(client)
+	}
+}
+
+func (p *mailProxy) forward(client net.Conn) {
+	defer p.workers.Done()
+	defer func() { <-p.slots }()
+	defer p.untrack(client)
+	defer client.Close()
+	dialer := net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
+	if p.source != nil {
+		dialer.LocalAddr = &net.TCPAddr{IP: p.source}
+	}
+	server, err := dialer.DialContext(context.Background(), "tcp4", p.upstream)
+	if err != nil {
+		return
+	}
+	defer server.Close()
+	if !p.track(server) {
+		return
+	}
+	defer p.untrack(server)
+	if err := writeProxyHeader(server, client.RemoteAddr(), client.LocalAddr()); err != nil {
+		return
+	}
+	clientTCP, clientOK := client.(*net.TCPConn)
+	serverTCP, serverOK := server.(*net.TCPConn)
+	if !clientOK || !serverOK {
+		return
+	}
+	results := make(chan error, 2)
+	go copyStream(results, serverTCP, clientTCP)
+	go copyStream(results, clientTCP, serverTCP)
+	<-results
+	<-results
+}
+
+func (p *mailProxy) close() {
+	p.mu.Lock()
+	if p.closed {
+		p.mu.Unlock()
+		return
+	}
+	p.closed = true
+	connections := make([]net.Conn, 0, len(p.conns))
+	for connection := range p.conns {
+		connections = append(connections, connection)
+	}
+	p.mu.Unlock()
+	_ = p.listener.Close()
+	<-p.done
+	for _, connection := range connections {
+		_ = connection.Close()
+	}
+	p.workers.Wait()
+}
+
+func (p *mailProxy) track(connection net.Conn) bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.closed {
+		return false
+	}
+	p.conns[connection] = struct{}{}
+	return true
+}
+
+func (p *mailProxy) untrack(connection net.Conn) {
+	p.mu.Lock()
+	delete(p.conns, connection)
+	p.mu.Unlock()
+}
+
+func (p *mailProxy) isClosed() bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.closed
+}
+
+func (p *mailProxy) report(err error) {
+	select {
+	case p.failures <- err:
+	default:
+	}
+}
+
+func copyStream(results chan<- error, destination, source *net.TCPConn) {
+	_, err := io.Copy(destination, source)
+	if closeErr := destination.CloseWrite(); err == nil {
+		err = closeErr
+	}
+	results <- err
+}
+
+func writeProxyHeader(writer io.Writer, sourceAddress, destinationAddress net.Addr) error {
+	source, sourceOK := sourceAddress.(*net.TCPAddr)
+	destination, destinationOK := destinationAddress.(*net.TCPAddr)
+	if !sourceOK || !destinationOK {
+		return errors.New("PROXY addresses are not TCP")
+	}
+	sourceIP := source.IP.To4()
+	destinationIP := destination.IP.To4()
+	if sourceIP == nil || destinationIP == nil || source.Port < 1 || source.Port > 65535 || destination.Port < 1 || destination.Port > 65535 {
+		return errors.New("PROXY addresses are not IPv4")
+	}
+	header := make([]byte, 28)
+	copy(header, proxySignature)
+	header[12] = 0x21
+	header[13] = 0x11
+	binary.BigEndian.PutUint16(header[14:16], 12)
+	copy(header[16:20], sourceIP)
+	copy(header[20:24], destinationIP)
+	binary.BigEndian.PutUint16(header[24:26], uint16(source.Port))
+	binary.BigEndian.PutUint16(header[26:28], uint16(destination.Port))
+	_, err := io.Copy(writer, bytes.NewReader(header))
+	return err
+}
+
+func drainFailures(failures <-chan error) {
+	for {
+		select {
+		case <-failures:
+		default:
+			return
+		}
+	}
+}