add support for --response.reject-non-prometheus to force upstream re…

…sponse to be in prometheus metrics format

go.mod

@@ -1,10 +1,13 @@
 module github.com/hoffie/sshified
 
-go 1.17
+go 1.21
+
+toolchain go1.22.5
 
 require (
 	github.com/alecthomas/kingpin/v2 v2.4.0
 	github.com/prometheus/client_golang v1.19.1
+	github.com/prometheus/prometheus v0.53.0
 	github.com/sirupsen/logrus v1.9.3
 	golang.org/x/crypto v0.25.0
 )
@@ -13,6 +16,8 @@ require (
 	github.com/alecthomas/units v0.0.0-20240626203959-61d1e3462e30 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/grafana/regexp v0.0.0-20240518133315-a468a5bfb3bc // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/prometheus/client_model v0.6.1 // indirect
 	github.com/prometheus/common v0.55.0 // indirect

go.sum

@@ -706,6 +706,8 @@ github.com/go-pdf/fpdf v0.6.0/go.mod h1:HzcnA+A23uwogo0tp9yU+l3V+KXhiESpt1PMayhO
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/goccy/go-json v0.9.11/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
+github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
+github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
 github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0/go.mod h1:E/TSTwGwJL78qG/PmXZO1EjYhfJinVAhrmmHX6Z8B9k=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
 github.com/golang/glog v1.0.0/go.mod h1:EWib/APOK0SL3dFbYqvxE3UYd8E6s1ouQ7iEp/0LWV4=
@@ -811,6 +813,8 @@ github.com/googleapis/gax-go/v2 v2.10.0/go.mod h1:4UOEnMCrxsSqQ940WnTiD6qJ63le2e
 github.com/googleapis/gax-go/v2 v2.11.0/go.mod h1:DxmR61SGKkGLa2xigwuZIQpkCI2S5iydzRfb3peWZJI=
 github.com/googleapis/go-type-adapters v1.0.0/go.mod h1:zHW75FOG2aur7gAO2B+MLby+cLsWGBF62rFAi7WjWO4=
 github.com/googleapis/google-cloud-go-testing v0.0.0-20200911160855-bcd43fbb19e8/go.mod h1:dvDLG8qkwmyD9a/MJJN3XJcT3xFxOKAvTZGvuZmac9g=
+github.com/grafana/regexp v0.0.0-20240518133315-a468a5bfb3bc h1:GN2Lv3MGO7AS6PrRoT6yV5+wkrOpcszoIsO4+4ds248=
+github.com/grafana/regexp v0.0.0-20240518133315-a468a5bfb3bc/go.mod h1:+JKpmjMGhpgPL+rXZ5nsZieVzvarn86asRlBg4uNGnk=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.7.0/go.mod h1:hgWBS7lorOAVIJEQMi4ZsPv9hVvWI6+ch50m39Pf2Ks=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.11.3/go.mod h1:o//XUCC/F+yRGJoPO/VU0GSB0f8Nhgmxx0VIRUvaC0w=
@@ -831,6 +835,7 @@ github.com/julienschmidt/httprouter v1.3.0/go.mod h1:JR6WtHb+2LUe8TCKY3cZOxFyyO8
 github.com/jung-kurt/gofpdf v1.0.0/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes=
 github.com/jung-kurt/gofpdf v1.0.3-0.20190309125859-24315acbbda5/go.mod h1:7Id9E/uU8ce6rXgefFLlgrJj/GYY22cpxn+r32jIOes=
 github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8=
+github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE=
 github.com/klauspost/compress v1.15.9/go.mod h1:PhcZ0MbTNciWF3rruxRgKxI5NkcHHrHUDtV4Yw2GlzU=
@@ -925,6 +930,8 @@ github.com/prometheus/procfs v0.12.0 h1:jluTpSng7V9hY0O2R9DzzJHYb2xULk9VTR1V1R/k
 github.com/prometheus/procfs v0.12.0/go.mod h1:pcuDEFsWDnvcgNzo4EEweacyhjeA9Zk3cnaOZAZEfOo=
 github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc=
 github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk=
+github.com/prometheus/prometheus v0.53.0 h1:vOnhpUKrDv954jnVBvhG/ZQJ3kqscnKI+Hbdwo2tAhc=
+github.com/prometheus/prometheus v0.53.0/go.mod h1:RZDkzs+ShMBDkAPQkLEaLBXpjmDcjhNxU2drUVPgKUU=
 github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
@@ -1372,6 +1379,7 @@ golang.org/x/tools v0.0.0-20200501065659-ab2804fb9c9d/go.mod h1:EkVYQZoAsY45+roY
 golang.org/x/tools v0.0.0-20200512131952-2bc93b1c0c88/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/tools v0.0.0-20200515010526-7d3b6ebf133d/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/tools v0.0.0-20200618134242-20370b0cb4b2/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
+golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
 golang.org/x/tools v0.0.0-20200729194436-6467de6f59a7/go.mod h1:njjCfa9FT2d7l9Bc6FUM5FLjQPp3cFF28FI3qnDFljA=
 golang.org/x/tools v0.0.0-20200804011535-6c149bb5ef0d/go.mod h1:njjCfa9FT2d7l9Bc6FUM5FLjQPp3cFF28FI3qnDFljA=
 golang.org/x/tools v0.0.0-20200825202427-b303f430e36d/go.mod h1:njjCfa9FT2d7l9Bc6FUM5FLjQPp3cFF28FI3qnDFljA=
@@ -1381,6 +1389,7 @@ golang.org/x/tools v0.0.0-20201124115921-2c860bdd6e78/go.mod h1:emZCQorbCU4vsT4f
 golang.org/x/tools v0.0.0-20201201161351-ac6f37ff4c2a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.0.0-20201208233053-a543418bbed2/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.0.0-20210105154028-b0ab187a4818/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
+golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.0.0-20210108195828-e2f9c7f1fc8e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0=
 golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=

main.go

@@ -12,17 +12,19 @@ import (
 )
 
 var (
-	verbose                = kingpin.Flag("verbose", "Verbose mode.").Short('v').Bool()
-	trace                  = kingpin.Flag("trace", "Trace mode.").Bool()
-	proxyAddr              = kingpin.Flag("proxy.listen-addr", "address the proxy will listen on").Required().String()
-	nextProxyAddr          = kingpin.Flag("next-proxy.addr", "optional address of another http proxy when cascading usage is required").String()
-	metricsAddr            = kingpin.Flag("metrics.listen-addr", "adress the service will listen on for metrics request about itself").String()
-	sshUser                = kingpin.Flag("ssh.user", "username used for connecting via ssh").Required().String()
-	sshKeyFile             = kingpin.Flag("ssh.key-file", "private key file used for connecting via ssh").Required().String()
-	sshKnownHostsFile      = kingpin.Flag("ssh.known-hosts-file", "known hosts file used for connecting via ssh").Required().String()
-	sshPort                = kingpin.Flag("ssh.port", "port used for connecting via ssh").Default("22").Int()
-	timeout                = kingpin.Flag("timeout", "full roundtrip request timeout in seconds").Default("50").Int()
-	timeoutDurationSeconds time.Duration
+	verbose                     = kingpin.Flag("verbose", "Verbose mode.").Short('v').Bool()
+	trace                       = kingpin.Flag("trace", "Trace mode.").Bool()
+	proxyAddr                   = kingpin.Flag("proxy.listen-addr", "address the proxy will listen on").Required().String()
+	nextProxyAddr               = kingpin.Flag("next-proxy.addr", "optional address of another http proxy when cascading usage is required").String()
+	metricsAddr                 = kingpin.Flag("metrics.listen-addr", "adress the service will listen on for metrics request about itself").String()
+	sshUser                     = kingpin.Flag("ssh.user", "username used for connecting via ssh").Required().String()
+	sshKeyFile                  = kingpin.Flag("ssh.key-file", "private key file used for connecting via ssh").Required().String()
+	sshKnownHostsFile           = kingpin.Flag("ssh.known-hosts-file", "known hosts file used for connecting via ssh").Required().String()
+	sshPort                     = kingpin.Flag("ssh.port", "port used for connecting via ssh").Default("22").Int()
+	timeout                     = kingpin.Flag("timeout", "full roundtrip request timeout in seconds").Default("50").Int()
+	timeoutDurationSeconds      time.Duration
+	responseMaxBytes            = kingpin.Flag("response.max-bytes", "maximum length of upstream response in bytes (0 = no limit)").Default("0").Int64()
+	responseRejectNonPrometheus = kingpin.Flag("response.reject-non-prometheus", "parse upstream response as Prometheus metrics and reject unparsable responses").Bool()
 )
 
 func main() {
@@ -35,6 +37,9 @@ func main() {
 	} else {
 		log.SetLevel(log.InfoLevel)
 	}
+	if *responseMaxBytes <= 0 && *responseRejectNonPrometheus {
+		kingpin.Fatalf("setting --response.reject-non-prometheus also requires setting a --response.max-bytes value due to internal buffering needs")
+	}
 	log.WithFields(log.Fields{"addr": *proxyAddr}).Info("Listening")
 	if *nextProxyAddr != "" {
 		log.WithFields(log.Fields{"nextProxyAddr": *nextProxyAddr}).Info("Running in cascading mode: will ssh to nextProxyAddr and use the http proxy there")

proxy.go

@@ -1,12 +1,16 @@
 package main
 
 import (
+	"bytes"
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"strings"
 
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/model/textparse"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -161,8 +165,53 @@ func (pr *proxyRequest) sendRequest() error {
 	return nil
 }
 
+func parsableAsPrometheus(b []byte, contentType string) error {
+	parser, err := textparse.New(b, contentType, false, labels.NewSymbolTable())
+	if err != nil {
+		return errors.New("failed to create parser for Prometheus metrics format")
+	}
+	for {
+		_, err := parser.Next()
+		if errors.Is(err, io.EOF) {
+			break
+		}
+		if err != nil {
+			return fmt.Errorf("failed to parse as Prometheus metrics format: %v", err)
+		}
+	}
+	return nil
+}
+
 func (pr *proxyRequest) forwardResponse() error {
+	assumeHttpErr := true
+	defer func() {
+		if assumeHttpErr {
+			pr.rw.WriteHeader(http.StatusInternalServerError)
+		}
+	}()
 	respHeader := pr.rw.Header()
+	var reader io.Reader
+	if *responseMaxBytes <= 0 {
+		reader = pr.upstreamResponse.Body
+	} else {
+		lr := io.LimitReader(pr.upstreamResponse.Body, *responseMaxBytes)
+		var buf = bytes.Buffer{}
+		_, err := io.Copy(&buf, lr)
+		if err != nil {
+			return fmt.Errorf("failed to copy response to buffer: %v", err)
+		}
+		reader = &buf
+		if *responseRejectNonPrometheus {
+			log.Trace("parsing response as prometheus metrics")
+			err := parsableAsPrometheus(buf.Bytes(), respHeader.Get("Content-Type"))
+			if err != nil {
+				assumeHttpErr = false
+				pr.rw.WriteHeader(http.StatusBadGateway)
+				return err
+			}
+		}
+	}
+
 	for k, vv := range pr.upstreamResponse.Header {
 		if k == "Content-Length" {
 			continue
@@ -172,9 +221,10 @@ func (pr *proxyRequest) forwardResponse() error {
 			respHeader.Add(k, v)
 		}
 	}
+	assumeHttpErr = false
 	pr.rw.WriteHeader(pr.upstreamResponse.StatusCode)
 	log.Trace("copying response body")
-	length, err := io.Copy(pr.rw, pr.upstreamResponse.Body)
+	length, err := io.Copy(pr.rw, reader)
 	if err != nil {
 		log.WithFields(log.Fields{"err": err}).Debug("failed to forward response body")
 		metricErrorsByType.WithLabelValues("response_body_forwarding").Inc()