Support ssz-formatted register_validator calls

.github/workflows/golangci-lint.yml

@@ -19,7 +19,7 @@ jobs:
     steps:
       - uses: actions/setup-go@v3
         with:
-          go-version: 1.19
+          go-version: 1.21
       - uses: actions/checkout@v3
       - name: golangci-lint
         uses: golangci/golangci-lint-action@v3

.github/workflows/test.yaml

@@ -12,7 +12,7 @@ jobs:
       - name: Set up Go
         uses: actions/setup-go@v4
         with:
-          go-version: '1.19'
+          go-version: '1.21'
 
       - name: Test
         run: go test -v ./...

go.mod

@@ -1,9 +1,10 @@
 module github.com/Rocket-Rescue-Node/guarded-beacon-proxy
 
-go 1.19
+go 1.21
 
 require (
 	github.com/ethereum/go-ethereum v1.12.0
+	github.com/ferranbt/fastssz v0.1.4
 	github.com/gorilla/mux v1.8.0
 	github.com/mwitkow/grpc-proxy v0.0.0-20230212185441-f345521cb9c9
 	github.com/prysmaticlabs/prysm/v4 v4.0.6
@@ -18,6 +19,7 @@ require (
 	github.com/btcsuite/btcd/btcec/v2 v2.3.2 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.1.0 // indirect
+	github.com/emicklei/dot v1.6.2 // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/google/go-cmp v0.5.9 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.0.1 // indirect
@@ -33,7 +35,7 @@ require (
 	github.com/prometheus/procfs v0.9.0 // indirect
 	github.com/prysmaticlabs/fastssz v0.0.0-20220628121656-93dfe28febab // indirect
 	github.com/prysmaticlabs/go-bitfield v0.0.0-20210809151128-385d8c5e3fb7 // indirect
-	github.com/prysmaticlabs/gohashtree v0.0.3-alpha // indirect
+	github.com/prysmaticlabs/gohashtree v0.0.4-beta // indirect
 	github.com/thomaso-mirodin/intmath v0.0.0-20160323211736-5dc6d854e46e // indirect
 	golang.org/x/crypto v0.7.0 // indirect
 	golang.org/x/sys v0.7.0 // indirect

guarded-beacon-proxy.go

@@ -3,18 +3,23 @@ package guardedbeaconproxy
 import (
 	"context"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
 	"time"
 
+	"github.com/Rocket-Rescue-Node/guarded-beacon-proxy/jsontypes"
 	"github.com/gorilla/mux"
 	"github.com/mwitkow/grpc-proxy/proxy"
 	"google.golang.org/grpc"
 )
 
+type RegisterValidatorRequest = jsontypes.RegisterValidatorRequest
+type PrepareBeaconProposerRequest = jsontypes.PrepareBeaconProposerRequest
+
 // PrepareBeaconProposerGuard is a function that validates whether or not a PrepareBeaconProposer call
 // should be proxied. The provided Context is whatever was returned by the authenticator.
 type PrepareBeaconProposerGuard func(PrepareBeaconProposerRequest, context.Context) (AuthenticationStatus, error)
@@ -53,6 +58,9 @@ type GuardedBeaconProxy struct {
 	Addr string
 	// Optional GRPC address to listen on
 	GRPCAddr string
+	// Maximum request body size in bytes
+	// If 0, no limit is applied
+	MaxRequestBodySize int64
 	// Pass-through HTTP server settings
 	ReadTimeout       time.Duration
 	ReadHeaderTimeout time.Duration
@@ -113,6 +121,43 @@ func (gbp *GuardedBeaconProxy) init() {
 	gbp.server.ErrorLog = gbp.ErrorLog
 }
 
+func (gbp *GuardedBeaconProxy) limitRequestBodyHandlerFunc(next httpGuard) httpGuard {
+	return func(w http.ResponseWriter, r *http.Request) bool {
+		if gbp.MaxRequestBodySize == 0 {
+			return next(w, r)
+		}
+
+		// Allow 1 extra byte. If it actually gets read, we will return an error.
+		// This lets us detect if the request body is exactly the size of the limit.
+		sizeLimit := gbp.MaxRequestBodySize + 1
+
+		if r.ContentLength >= sizeLimit {
+			gbp.httpError(w, http.StatusRequestEntityTooLarge, fmt.Errorf("request body too large"))
+			return false
+		}
+
+		limited := &io.LimitedReader{
+			R: r.Body,
+			N: sizeLimit,
+		}
+		// According to the docs, http servers don't need to close the body ReadCloser, only Clients do.
+		r.Body = io.NopCloser(limited)
+		shouldProxy := next(w, r)
+		if !shouldProxy {
+			return false
+		}
+		if limited.N == 0 {
+			// The next handler didn't return an error, but we exceeded the limit.
+			// Do not proxy the request, and return StatusRequestEntityTooLarge.
+			gbp.httpError(w, http.StatusRequestEntityTooLarge, fmt.Errorf("request body too large"))
+			return false
+		}
+		// The next handler didn't return an error, and we didn't exceed the limit.
+		// Proxy the request.
+		return true
+	}
+}
+
 // Serve attaches the proxy to the provided listener(s)
 //
 // Serve blocks until Stop is called or an error is encountered.
@@ -124,12 +169,25 @@ func (gbp *GuardedBeaconProxy) Serve(httpListener net.Listener, grpcListener *ne
 	router := mux.NewRouter()
 
 	if gbp.PrepareBeaconProposerGuard != nil {
-		router.Path("/eth/v1/validator/prepare_beacon_proposer").HandlerFunc(gbp.prepareBeaconProposer)
+		router.Path("/eth/v1/validator/prepare_beacon_proposer").HandlerFunc(
+			func(w http.ResponseWriter, r *http.Request) {
+				if gbp.limitRequestBodyHandlerFunc(gbp.prepareBeaconProposer)(w, r) {
+					gbp.proxy.ServeHTTP(w, r)
+				}
+			},
+		)
 	}
 
 	if gbp.RegisterValidatorGuard != nil {
-		router.Path("/eth/v1/validator/register_validator").HandlerFunc(gbp.registerValidator)
+		router.Path("/eth/v1/validator/register_validator").HandlerFunc(
+			func(w http.ResponseWriter, r *http.Request) {
+				if gbp.limitRequestBodyHandlerFunc(gbp.registerValidator)(w, r) {
+					gbp.proxy.ServeHTTP(w, r)
+				}
+			},
+		)
 	}
+
 	router.PathPrefix("/").Handler(gbp.proxy)
 
 	if gbp.HTTPAuthenticator != nil {

http.go

@@ -6,7 +6,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"mime"
 	"net/http"
+
+	"github.com/Rocket-Rescue-Node/guarded-beacon-proxy/ssz"
 )
 
 // HTTPAuthenticator is a function type which can authenticate HTTP requests.
@@ -22,6 +25,9 @@ import (
 // information.
 type HTTPAuthenticator func(*http.Request) (AuthenticationStatus, context.Context, error)
 
+// If true is returned, the upstream will proxy the request.
+type httpGuard func(w http.ResponseWriter, r *http.Request) bool
+
 func (gbp *GuardedBeaconProxy) authenticationMiddleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		status, context, err := gbp.HTTPAuthenticator(r)
@@ -40,16 +46,13 @@ func (gbp *GuardedBeaconProxy) authenticationMiddleware(next http.Handler) http.
 }
 
 func cloneRequestBody(r *http.Request) (io.ReadCloser, error) {
-	// Read the body
-	buf, err := io.ReadAll(r.Body)
-	if err != nil {
-		return nil, err
-	}
+	// Use an io.TeeReader to return a reader that re-writes the body to the original request body.
+	buf := bytes.NewBuffer(nil)
+	tee := io.TeeReader(r.Body, buf)
+	out := io.NopCloser(tee)
+	r.Body = io.NopCloser(buf)
 
-	original := io.NopCloser(bytes.NewBuffer(buf))
-	clone := io.NopCloser(bytes.NewBuffer(buf))
-	r.Body = original
-	return clone, nil
+	return out, nil
 }
 
 func (gbp *GuardedBeaconProxy) httpError(w http.ResponseWriter, code int, err error) {
@@ -60,46 +63,64 @@ func (gbp *GuardedBeaconProxy) httpError(w http.ResponseWriter, code int, err er
 	}
 }
 
-func (gbp *GuardedBeaconProxy) prepareBeaconProposer(w http.ResponseWriter, r *http.Request) {
-	buf, err := cloneRequestBody(r)
+func (gbp *GuardedBeaconProxy) prepareBeaconProposer(w http.ResponseWriter, r *http.Request) bool {
+	reader, err := cloneRequestBody(r)
 	if err != nil {
 		gbp.httpError(w, http.StatusInternalServerError, nil)
-		return
+		return false
 	}
 
 	var proposers PrepareBeaconProposerRequest
-	if err := json.NewDecoder(buf).Decode(&proposers); err != nil {
+	if err := json.NewDecoder(reader).Decode(&proposers); err != nil {
 		gbp.httpError(w, http.StatusBadRequest, nil)
-		return
+		return false
 	}
 
 	status, err := gbp.PrepareBeaconProposerGuard(proposers, r.Context())
 	if status != Allowed {
 		gbp.httpError(w, status.httpStatus(), err)
-		return
+		return false
 	}
 
-	gbp.proxy.ServeHTTP(w, r)
+	return true
 }
 
-func (gbp *GuardedBeaconProxy) registerValidator(w http.ResponseWriter, r *http.Request) {
-	buf, err := cloneRequestBody(r)
+func (gbp *GuardedBeaconProxy) registerValidator(w http.ResponseWriter, r *http.Request) bool {
+	reader, err := cloneRequestBody(r)
 	if err != nil {
 		gbp.httpError(w, http.StatusInternalServerError, nil)
-		return
+		return false
+	}
+
+	// Check the content-type header
+	contentType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
+	if err != nil {
+		gbp.httpError(w, http.StatusUnsupportedMediaType, err)
+		return false
 	}
 
 	var validators RegisterValidatorRequest
-	if err := json.NewDecoder(buf).Decode(&validators); err != nil {
-		gbp.httpError(w, http.StatusBadRequest, nil)
-		return
+	switch contentType {
+	case "application/json":
+		if err := json.NewDecoder(reader).Decode(&validators); err != nil {
+			gbp.httpError(w, http.StatusBadRequest, err)
+			return false
+		}
+	case "application/octet-stream":
+		if err, status := ssz.ToRegisterValidatorRequest(&validators, reader, gbp.MaxRequestBodySize); err != nil {
+			gbp.httpError(w, status, err)
+			return false
+		}
+	default:
+		gbp.httpError(w, http.StatusUnsupportedMediaType, fmt.Errorf("unsupported content type: %s", contentType))
+		return false
 	}
 
 	status, err := gbp.RegisterValidatorGuard(validators, r.Context())
 	if status != Allowed {
 		gbp.httpError(w, status.httpStatus(), err)
-		return
+		return false
 	}
 
-	gbp.proxy.ServeHTTP(w, r)
+	return true
 }