From 5d8a0ef605e535f838000ce21f0de581f287a133 Mon Sep 17 00:00:00 2001
From: Thor77 <thor77@thor77.org>
Date: Wed, 7 Nov 2018 19:13:27 +0100
Subject: [PATCH] Serve file directly for curl and wget user agents (#145)

* Serve file directly for curl and wget user agents

Fix #127

* Add test for get with wget user agent

* Add -nodirectagents flag

to disable serving files directly for wget/curl user agents

* Fix TestPutAndGetCLI failing for Go 1.5

It failed because it doesn't include the Content-Type header for every
response.
---
 display.go     |  8 ++++++++
 server.go      |  3 +++
 server_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+)

diff --git a/display.go b/display.go
index c6d8470..4220c76 100644
--- a/display.go
+++ b/display.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"net/http"
 	"path/filepath"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
@@ -18,7 +19,14 @@ import (
 
 const maxDisplayFileSizeBytes = 1024 * 512
 
+var cliUserAgentRe = regexp.MustCompile("(?i)(lib)?curl|wget")
+
 func fileDisplayHandler(c web.C, w http.ResponseWriter, r *http.Request) {
+	if !Config.noDirectAgents && cliUserAgentRe.MatchString(r.Header.Get("User-Agent")) {
+		fileServeHandler(c, w, r)
+		return
+	}
+
 	fileName := c.URLParams["name"]
 
 	err := checkFile(fileName)
diff --git a/server.go b/server.go
index daaf89e..ba0efb8 100644
--- a/server.go
+++ b/server.go
@@ -58,6 +58,7 @@ var Config struct {
 	remoteAuthFile            string
 	addHeaders                headerList
 	googleShorterAPIKey       string
+	noDirectAgents            bool
 }
 
 var Templates = make(map[string]*pongo2.Template)
@@ -243,6 +244,8 @@ func main() {
 		"Add an arbitrary header to the response. This option can be used multiple times.")
 	flag.StringVar(&Config.googleShorterAPIKey, "googleapikey", "",
 		"API Key for Google's URL Shortener.")
+	flag.BoolVar(&Config.noDirectAgents, "nodirectagents", false,
+		"disable serving files directly for wget/curl user agents")
 
 	iniflags.Parse()
 
diff --git a/server_test.go b/server_test.go
index 6fe363c..7727f0c 100644
--- a/server_test.go
+++ b/server_test.go
@@ -1121,3 +1121,50 @@ func TestShutdown(t *testing.T) {
 	os.RemoveAll(Config.filesDir)
 	os.RemoveAll(Config.metaDir)
 }
+
+func TestPutAndGetCLI(t *testing.T) {
+	var myjson RespOkJSON
+	mux := setup()
+
+	// upload file
+	w := httptest.NewRecorder()
+	req, err := http.NewRequest("PUT", "/upload", strings.NewReader("File content"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	req.Header.Set("Accept", "application/json")
+	mux.ServeHTTP(w, req)
+
+	err = json.Unmarshal([]byte(w.Body.String()), &myjson)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// request file without wget user agent
+	w = httptest.NewRecorder()
+	req, err = http.NewRequest("GET", myjson.Url, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	mux.ServeHTTP(w, req)
+
+	contentType := w.Header().Get("Content-Type")
+	if strings.HasPrefix(contentType, "text/plain") {
+		t.Fatalf("Didn't receive file display page but %s", contentType)
+	}
+
+	// request file with wget user agent
+	w = httptest.NewRecorder()
+	req, err = http.NewRequest("GET", myjson.Url, nil)
+	req.Header.Set("User-Agent", "wget")
+	if err != nil {
+		t.Fatal(err)
+	}
+	mux.ServeHTTP(w, req)
+
+	contentType = w.Header().Get("Content-Type")
+	if !strings.HasPrefix(contentType, "text/plain") {
+		t.Fatalf("Didn't receive file directly but %s", contentType)
+	}
+
+}