From 0595a9bb62da3c4c0c9cb008af750d316b8353dd Mon Sep 17 00:00:00 2001
From: Troels Liebe Bentsen <tlb@nversion.dk>
Date: Mon, 27 Jul 2020 13:40:17 +0200
Subject: [PATCH] Limit max request and response size

---
 cmd/authwrapper/setup.go | 4 ++--
 cmd/authwrapper/utils.go | 7 +++++--
 server/http.go           | 7 ++++++-
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/cmd/authwrapper/setup.go b/cmd/authwrapper/setup.go
index 7cdbc50..3dd08f0 100644
--- a/cmd/authwrapper/setup.go
+++ b/cmd/authwrapper/setup.go
@@ -194,7 +194,7 @@ func setupKeyring(config *Config) (agent.ExtendedAgent, error) {
 func fetchUserCert(signingServerURL string, signer ssh.AlgorithmSigner, command string, args []string, principals []string) (*ssh.Certificate, error) {
 	// GET /certificate/challenge # { value: "{ \"timestamp\": \"2020-01-01T10:00:00.000Z\" \"random\": \"...\"}", signature: "signed by CA key" }
 	var challenge server.Challenge
-	err := httpJSONRequest("GET", signingServerURL+"/certificate/challenge", nil, &challenge)
+	err := httpJSONRequest("GET", signingServerURL+"/certificate/challenge", nil, &challenge, 1*1024*1024)
 	if err != nil {
 		return nil, err
 	}
@@ -213,7 +213,7 @@ func fetchUserCert(signingServerURL string, signer ssh.AlgorithmSigner, command
 
 	// get back { certificate: "base64 encoded cert" }
 	var certResponse server.CertificateResponse
-	err = httpJSONRequest("POST", signingServerURL+"/certificate", certRequest, &certResponse)
+	err = httpJSONRequest("POST", signingServerURL+"/certificate", certRequest, &certResponse, 1*1024*1024)
 	if err != nil {
 		return nil, err
 	}
diff --git a/cmd/authwrapper/utils.go b/cmd/authwrapper/utils.go
index 4154c3c..30c525c 100644
--- a/cmd/authwrapper/utils.go
+++ b/cmd/authwrapper/utils.go
@@ -101,7 +101,7 @@ func runCommand(command string, args []string) (exitCode int, err error) {
 	return 0, nil
 }
 
-func httpJSONRequest(method string, url string, requestData interface{}, responseData interface{}) error {
+func httpJSONRequest(method string, url string, requestData interface{}, responseData interface{}, maxResponseSize int64) error {
 	// Convert request to JSON and wrap in io.Reader
 	var requestBody io.Reader
 	if requestData != nil {
@@ -122,7 +122,10 @@ func httpJSONRequest(method string, url string, requestData interface{}, respons
 		return err
 	}
 	defer httpResponse.Body.Close()
-	responseBody, err := ioutil.ReadAll(httpResponse.Body)
+
+	// Limit size of response body we read into memory
+	limitedReader := &io.LimitedReader{R: httpResponse.Body, N: maxResponseSize}
+	responseBody, err := ioutil.ReadAll(limitedReader)
 	if err != nil {
 		return err
 	}
diff --git a/server/http.go b/server/http.go
index 42c3e1f..f11c31c 100644
--- a/server/http.go
+++ b/server/http.go
@@ -3,6 +3,7 @@ package server
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"path/filepath"
@@ -72,7 +73,11 @@ func (s *HTTPSigningServer) getCertificateChallenge(w http.ResponseWriter, r *ht
 
 func (s *HTTPSigningServer) postCertificate(w http.ResponseWriter, r *http.Request) (jsonResponse interface{}, statusError *StatusError) {
 	defer r.Body.Close()
-	body, err := ioutil.ReadAll(r.Body)
+
+	// Limit how much of the body we read in a request
+	limitedReader := &io.LimitedReader{R: r.Body, N: 1 * 1024 * 1024}
+
+	body, err := ioutil.ReadAll(limitedReader)
 	if err != nil {
 		return nil, &StatusError{500, err}
 	}