diff --git a/README.md b/README.md index 7eb45c0..0c641e9 100644 --- a/README.md +++ b/README.md @@ -88,6 +88,10 @@ Usage of ./endlessh-go The address for prometheus (default "0.0.0.0") -prometheus_port string The port for prometheus (default "2112") + -proxy_protocol_enabled + Enable PROXY protocol support. This causes the server to expect PROXY protocol headers on incoming connections. + -proxy_protocol_read_header_timeout_ms int + Timeout for reading the PROXY protocol header in milliseconds. If the connection does not send a valid PROXY protocol header in this time, the header is ignored. (default 200) -stderrthreshold value logs at or above this threshold go to stderr (default 2) -v value diff --git a/go.mod b/go.mod index b97d881..f059fd9 100644 --- a/go.mod +++ b/go.mod @@ -15,6 +15,7 @@ require ( github.com/kr/text v0.2.0 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/oschwald/maxminddb-golang v1.13.0 // indirect + github.com/pires/go-proxyproto v0.8.1 // indirect github.com/prometheus/client_model v0.6.2 // indirect github.com/prometheus/common v0.66.1 // indirect github.com/prometheus/procfs v0.16.1 // indirect diff --git a/go.sum b/go.sum index 9d2090b..3e3265c 100644 --- a/go.sum +++ b/go.sum @@ -43,6 +43,8 @@ github.com/pierrre/go-libs v0.17.0 h1:bjxd9unioV/YDkUW7obETp2IFct0kO9HePURn81UL8 github.com/pierrre/go-libs v0.17.0/go.mod h1:920odOqc5mZREW9GFWg056mjQ2prNVRGUZO7HRS2Jlc= github.com/pierrre/pretty v0.14.3 h1:I100hHs1C/MCd3M0D/hIV7J2OXl7amLD0uP2jnB7mRw= github.com/pierrre/pretty v0.14.3/go.mod h1:HTaFDNtT9ELVK5pODLfXRLiEiyIx3MmQUL5UadrR3/0= +github.com/pires/go-proxyproto v0.8.1 h1:9KEixbdJfhrbtjpz/ZwCdWDD2Xem0NZ38qMYaASJgp0= +github.com/pires/go-proxyproto v0.8.1/go.mod h1:ZKAAyp3cgy5Y5Mo4n9AlScrkCZwUy0g3Jf+slqQVcuU= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= diff --git a/main.go b/main.go index c9d0353..94e165c 100644 --- a/main.go +++ b/main.go @@ -28,6 +28,7 @@ import ( "time" "github.com/golang/glog" + proxyproto "github.com/pires/go-proxyproto" ) func startSending(maxClients int64, bannerMaxLength int64, records chan<- metrics.RecordEntry) chan *client.Client { @@ -67,15 +68,22 @@ func startSending(maxClients int64, bannerMaxLength int64, records chan<- metric return clients } -func startAccepting(maxClients int64, connType, connHost, connPort string, interval time.Duration, clients chan<- *client.Client, records chan<- metrics.RecordEntry) { +func startAccepting(maxClients int64, connType, connHost, connPort string, interval time.Duration, clients chan<- *client.Client, records chan<- metrics.RecordEntry, proxyProtocolEnabled bool, proxyProtocolReadHeaderTimeout int) { go func() { l, err := net.Listen(connType, connHost+":"+connPort) if err != nil { glog.Errorf("Error listening: %v", err) os.Exit(1) } + + // Wrap the listener in a proxy protocol listener + if proxyProtocolEnabled { + l = &proxyproto.Listener{Listener: l, ReadHeaderTimeout: time.Duration(proxyProtocolReadHeaderTimeout) * time.Millisecond} + } + // Close the listener when the application closes. defer l.Close() + glog.Infof("Listening on %v:%v", connHost, connPort) for { // Listen for an incoming connection. @@ -125,6 +133,8 @@ func main() { prometheusCleanUnseenSeconds := flag.Int("prometheus_clean_unseen_seconds", 0, "Remove series if the IP is not seen for the given time. Set to 0 to disable. (default 0)") geoipSupplier := flag.String("geoip_supplier", "off", "Supplier to obtain Geohash of IPs. Possible values are \"off\", \"ip-api\", \"max-mind-db\"") maxMindDbFileName := flag.String("max_mind_db", "", "Path to the MaxMind DB file.") + proxyProtocolEnabled := flag.Bool("proxy_protocol_enabled", false, "Enable PROXY protocol support. This causes the server to expect PROXY protocol headers on incoming connections.") + proxyProtocolReadHeaderTimeout := flag.Int("proxy_protocol_read_header_timeout_ms", 200, "Timeout for reading the PROXY protocol header in milliseconds. If the connection does not send a valid PROXY protocol header in this time, the header is ignored.") flag.Usage = func() { fmt.Fprintf(flag.CommandLine.Output(), "Usage of %v \n", os.Args[0]) @@ -155,7 +165,7 @@ func main() { connPorts = append(connPorts, defaultPort) } for _, connPort := range connPorts { - startAccepting(*maxClients, *connType, *connHost, connPort, interval, clients, records) + startAccepting(*maxClients, *connType, *connHost, connPort, interval, clients, records, *proxyProtocolEnabled, *proxyProtocolReadHeaderTimeout) } for { if *prometheusCleanUnseenSeconds <= 0 {