Added ReadAtMost function for reading request data (refs #2292)

ReadAtMost will read up to the specified number of bytes, returning
an error if the request is larger than that limit.

diff --git a/services/keep/keep.go b/services/keep/keep.go
index 7cca5344d8a2abc074cfb0f0f4dbd66d4fd2f0af..c4e1705d08dd6e20b221f9d2e1b877dba483fb4a 100644 (file)
--- a/services/keep/keep.go
+++ b/services/keep/keep.go
@@ -7,6 +7,7 @@ import (
        "errors"
        "fmt"
        "github.com/gorilla/mux"
+       "io"
        "io/ioutil"
        "log"
        "net/http"
@@ -128,12 +129,10 @@ func PutBlockHandler(w http.ResponseWriter, req *http.Request) {
        hash := mux.Vars(req)["hash"]
 
        // Read the block data to be stored.
-       // TODO(twp): decide what to do when the input stream contains
-       // more than BLOCKSIZE bytes.
+       // If the request exceeds BLOCKSIZE bytes, return 500 Request Too Large.
        //
-       buf := make([]byte, BLOCKSIZE)
-       if nread, err := req.Body.Read(buf); err == nil {
-               if err := PutBlock(buf[:nread], hash); err == nil {
+       if buf, err := ReadAtMost(req.Body, BLOCKSIZE); err == nil {
+               if err := PutBlock(buf, hash); err == nil {
                        w.WriteHeader(http.StatusOK)
                } else {
                        ke := err.(*KeepError)
@@ -331,3 +330,17 @@ func FreeDiskSpace(volume string) (free uint64, err error) {
 
        return
 }
+
+// ReadAtMost
+//     Returns a byte slice containing at most N bytes read
+//     from the specified io.Reader.
+//
+func ReadAtMost(r io.Reader, limit int) ([]byte, error) {
+       // Attempt to read one more byte than limit.
+       lr := io.LimitReader(r, int64(limit+1))
+       buf, err := ioutil.ReadAll(lr)
+       if len(buf) > limit {
+               return buf[:limit], errors.New("Request Too Large")
+       }
+       return buf, err
+}