1 files changed, 15 insertions(+), 22 deletions(-)

M gemini.go

@@ 6,10 6,11 @@ package main
 
 import (
 	"bufio"
+	"bytes"
 	"crypto/tls"
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"log"
 	"net/http"
 	"net/url"


@@ 236,29 237,21 @@ func proxyGemini(w http.ResponseWriter, r *http.Request, u *url.URL) (*url.URL, 
 	return u, err
 }
 
-// serveFile saves a temporary file with the contents of rd, then serves it to w.
+// serveFile creates a buffer with the contents of rd, then serves it to w.
 func serveFile(w http.ResponseWriter, r *http.Request, u *url.URL, rd *bufio.Reader) error {
-	var err error
-	fileName := u.String()[strings.LastIndex(u.String(), "/")+1:]
-
-	f, err := ioutil.TempFile("", "kindleto*-"+fileName)
-	if err != nil {
-		err = fmt.Errorf("serveFile: failed to create temp file: %v", err)
+	const maxSizeBytes int64 = 10 * 1024 * 1024
+	var buf bytes.Buffer
+	_, err := io.CopyN(&buf, rd, maxSizeBytes + 1)
+	if err == nil {
+		return fmt.Errorf("serveFile: file exceeding maximum size of %v MiB", maxSizeBytes / 1024 / 1024)
+	} else if err.Error() != "EOF" {
+		return fmt.Errorf("serveFile: failed to read file: %v", err)
 	}
-	defer os.Remove(f.Name()) // clean up
 
-	if _, err := f.ReadFrom(rd); err != nil {
-		err = fmt.Errorf("serveFile: failed to write to temp file: %v", err)
-	}
-
-	// Note: If we ever want to serve images inline, we'll have to revisit
-	// this content disposition header value.
-	w.Header().Set("Content-Disposition", "attachment; filename="+fileName)
-	http.ServeContent(w, r, fileName, time.Time{}, f)
-
-	if err := f.Close(); err != nil {
-		err = fmt.Errorf("serveFile: failed to close temp file: %v", err)
-	}
+	bufReader := bytes.NewReader(buf.Bytes())
+	fileName := u.String()[strings.LastIndex(u.String(), "/")+1:]
+	w.Header().Set("Content-Disposition", `attachment; filename="`+fileName+`"`)
+	http.ServeContent(w, r, fileName, time.Time{}, bufReader)
 
-	return err
+	return nil
 }