github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/server/components/https/proxy.go (about)

     1  // Copyright 2019 Google Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package https provides generic implementations and utility methods for
    16  // Fleetspeak https communication component type.
    17  package https
    18  
    19  import (
    20  	"bufio"
    21  	"fmt"
    22  	"net"
    23  
    24  	log "github.com/golang/glog"
    25  
    26  	proxyproto "github.com/pires/go-proxyproto"
    27  )
    28  
    29  // ProxyListener wraps a net.Listener and adapts the perceived source of the connection
    30  // to be that provided by the PROXY protocol header string. Should be used only when the
    31  // server is behind a load balancer which implements the PROXY protocol.
    32  //
    33  // https://www.haproxy.org/download/1.8/doc/proxy-protocol.txt
    34  type ProxyListener struct {
    35  	net.Listener
    36  }
    37  
    38  func (l *ProxyListener) Accept() (net.Conn, error) {
    39  	c, err := l.Listener.Accept()
    40  	if err != nil {
    41  		return c, err
    42  	}
    43  	r := bufio.NewReader(c)
    44  	h, err := proxyproto.Read(r)
    45  	if err == proxyproto.ErrNoProxyProtocol {
    46  		log.Warningf("Received connection from %v without proxy header.", c.RemoteAddr())
    47  		return &proxyConn{
    48  			Conn:   c,
    49  			reader: r,
    50  			remote: c.RemoteAddr(),
    51  		}, nil
    52  	}
    53  	if err != nil {
    54  		return nil, fmt.Errorf("error parsing proxy header: %v", err)
    55  	}
    56  	return &proxyConn{
    57  		Conn:   c,
    58  		reader: r,
    59  		remote: h.SourceAddr,
    60  	}, nil
    61  }
    62  
    63  type proxyConn struct {
    64  	net.Conn
    65  	reader *bufio.Reader
    66  	remote net.Addr
    67  }
    68  
    69  func (c *proxyConn) Read(b []byte) (int, error) {
    70  	return c.reader.Read(b)
    71  }
    72  
    73  func (c *proxyConn) RemoteAddr() net.Addr {
    74  	return c.remote
    75  }