github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/e2etesting/balancer/proxyproto/proxyproto.go (about)

     1  package proxyproto
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"net"
     7  	"strconv"
     8  
     9  	proxyproto "github.com/pires/go-proxyproto"
    10  )
    11  
    12  func splitHostPort(addr string) (net.IP, uint16, error) {
    13  	ta, err := net.ResolveTCPAddr("tcp", addr)
    14  	if err == nil {
    15  		return ta.IP, uint16(ta.Port), nil
    16  	}
    17  	hostStr, portStr, err := net.SplitHostPort(addr)
    18  	if err != nil {
    19  		return nil, 0, err
    20  	}
    21  	host := net.ParseIP(hostStr)
    22  	if host == nil {
    23  		return nil, 0, fmt.Errorf("Failed to parse IP")
    24  	}
    25  	port, err := strconv.Atoi(portStr)
    26  	if err != nil {
    27  		return nil, 0, fmt.Errorf("Failed to parse port: %v", err)
    28  	}
    29  	return host, uint16(port), nil
    30  }
    31  
    32  func WriteFirstProxyMessage(w io.Writer, srcAddr, dstAddr string) error {
    33  	srcHost, srcPort, err := splitHostPort(srcAddr)
    34  	if err != nil {
    35  		return fmt.Errorf("Failed to parse source (client) address: %v", err)
    36  	}
    37  	dstHost, dstPort, err := splitHostPort(dstAddr)
    38  	if err != nil {
    39  		return fmt.Errorf("Failed to parse destination (server) address: %v", err)
    40  	}
    41  	header := proxyproto.Header{
    42  		Version:           1,
    43  		TransportProtocol: proxyproto.TCPv4,
    44  		SourceAddr:        &net.TCPAddr{IP: srcHost, Port: int(srcPort)},
    45  		DestinationAddr:   &net.TCPAddr{IP: dstHost, Port: int(dstPort)},
    46  	}
    47  	_, err = header.WriteTo(w)
    48  	if err != nil {
    49  		return fmt.Errorf("Failed to write Proxy header: %v", err)
    50  	}
    51  	return nil
    52  }