Merge pull request #15 from VoycerAG/stability_fixes

Ensure no corrupt images will be stored

application.go

@@ -4,9 +4,9 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"net/http"
 
 	"gopkg.in/mgo.v2"
-	"gopkg.in/tylerb/graceful.v1"
 
 	"github.com/VoycerAG/gridfs-image-server/server"
 )
@@ -52,7 +52,7 @@ func main() {
 
 	log.Printf("Server started. Listening on %d database host is %s\n", *serverPort, *host)
 
-	graceful.Run(fmt.Sprintf(":%d", *serverPort), 0, handler)
+	err = http.ListenAndServe(fmt.Sprintf(":%d", *serverPort), handler)
 	if err != nil {
 		log.Fatal(err)
 	}

server/gridfs.go

@@ -174,9 +174,15 @@ func (g GridfsStorage) StoreChildImage(
 	if err != nil {
 		return nil, err
 	}
+
 	defer targetfile.Close()
 
-	io.Copy(targetfile, reader)
+	_, err = io.Copy(targetfile, reader)
+
+	if err != nil {
+		targetfile.Abort()
+		return nil, err
+	}
 
 	metadata := bson.M{
 		"width":            imageWidth,

server/server.go

@@ -210,8 +210,7 @@ func imageHandler(w http.ResponseWriter, r *http.Request, requestConfig *Configu
 
 		var b bytes.Buffer
 		buffer := bufio.NewWriter(&b)
-		writer := io.MultiWriter(w, buffer)
-		controller.Encode(writer)
+		controller.Encode(buffer)
 		buffer.Flush()
 
 		targetfile, err := storage.StoreChildImage(
@@ -230,6 +229,7 @@ func imageHandler(w http.ResponseWriter, r *http.Request, requestConfig *Configu
 			return
 		}
 
+		w.Write(b.Bytes())
 		setCacheHeaders(targetfile, w)
 
 		log.Printf("%d image succesfully resized and returned.\n", http.StatusOK)