From 8d1f6cd567c6bddb2fac668363b5f6cbc740fab7 Mon Sep 17 00:00:00 2001
From: Kimbsen <kim@jotta.no>
Date: Tue, 23 Jun 2020 16:24:11 +0200
Subject: [PATCH] added optional md5 verification of uploaded data

---
 weed/storage/needle/needle_parse_upload.go | 41 +++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/weed/storage/needle/needle_parse_upload.go b/weed/storage/needle/needle_parse_upload.go
index 67b798821..7df0069e9 100644
--- a/weed/storage/needle/needle_parse_upload.go
+++ b/weed/storage/needle/needle_parse_upload.go
@@ -1,7 +1,10 @@
 package needle
 
 import (
+	"compress/gzip"
+	"crypto/md5"
 	"fmt"
+	"hash"
 	"io"
 	"io/ioutil"
 	"mime"
@@ -93,6 +96,23 @@ func parsePut(r *http.Request, sizeLimit int64, pu *ParsedUpload) (e error) {
 	return nil
 }
 
+type ChecksumReader struct {
+	h hash.Hash
+	r io.Reader
+}
+
+func (cr *ChecksumReader) Read(p []byte) (int, error) {
+	n, err := cr.r.Read(p)
+	if err == nil {
+		cr.h.Write(p[:n])
+	}
+	return n, err
+}
+
+func (cr *ChecksumReader) Checksum() string {
+	return fmt.Sprintf("%x", cr.h.Sum(nil))
+}
+
 func parseMultipart(r *http.Request, sizeLimit int64, pu *ParsedUpload) (e error) {
 	defer func() {
 		if e != nil && r.Body != nil {
@@ -120,7 +140,26 @@ func parseMultipart(r *http.Request, sizeLimit int64, pu *ParsedUpload) (e error
 		pu.FileName = path.Base(pu.FileName)
 	}
 
-	pu.Data, e = ioutil.ReadAll(io.LimitReader(part, sizeLimit+1))
+	reader := io.LimitReader(part, sizeLimit+1)
+	if expectedChecksum := r.Header.Get("Content-MD5"); expectedChecksum != "" {
+		if part.Header.Get("Content-Encoding") == "gzip" {
+			gr, err := gzip.NewReader(reader)
+			if err != nil {
+				e = fmt.Errorf("Content-Encoding == gzip but content was not gzipped: %s", err)
+				return
+			}
+			reader = gr
+		}
+		cr := &ChecksumReader{md5.New(), reader}
+		pu.Data, e = ioutil.ReadAll(cr)
+		if expectedChecksum != cr.Checksum() {
+			e = fmt.Errorf("Content-MD5 did not match md5 of file data [%s] != [%s]", expectedChecksum, cr.Checksum())
+			return
+		}
+	} else {
+		pu.Data, e = ioutil.ReadAll(reader)
+	}
+
 	if e != nil {
 		glog.V(0).Infoln("Reading Content [ERROR]", e)
 		return