Several small changes

Don't panic if content-type header is missing.
Allow downstream server to request HTTPS in urls.
gofmt code

Author: foxcpp

README.md

@@ -41,6 +41,8 @@ POST /filedrop/screenshot.png?max-uses=5&store-secs=3600
 Following request will store file screenshot.png for one hour (3600 seconds)
 and allow it to be downloaded not more than 10 times.
 
+**Note** To get `https` scheme in URLs downstream server should set header
+`X-HTTPS-Downstream` to `1` (or you can also set HTTPSDownstream config option)
 
 ### Authorization
 

config.go

@@ -3,14 +3,14 @@ package filedrop
 import "net/http"
 
 type LimitsConfig struct {
-	MaxUses uint
+	MaxUses      uint
 	MaxStoreSecs uint
-	MaxFileSize uint
+	MaxFileSize  uint
 }
 
 type DBConfig struct {
 	Driver string
-	DSN string
+	DSN    string
 }
 
 type AuthConfig struct {
@@ -21,12 +21,12 @@ type AuthConfig struct {
 }
 
 type Config struct {
-	Limits LimitsConfig
-	DB DBConfig
-	DownloadAuth AuthConfig
-	UploadAuth AuthConfig
-	StorageDir string
-	HTTPSUpstream bool
+	Limits          LimitsConfig
+	DB              DBConfig
+	DownloadAuth    AuthConfig
+	UploadAuth      AuthConfig
+	StorageDir      string
+	HTTPSDownstream bool
 }
 
-var Default Config
+var Default Config

db.go

@@ -13,9 +13,9 @@ type db struct {
 	remFile     *sql.Stmt
 	contentType *sql.Stmt
 
-	addUse       *sql.Stmt
-	shouldDelete *sql.Stmt
-	cleanup      *sql.Stmt
+	addUse         *sql.Stmt
+	shouldDelete   *sql.Stmt
+	cleanup        *sql.Stmt
 	pendingCleanup *sql.Stmt
 }
 
@@ -39,7 +39,6 @@ func openDB(driver, dsn string) (*db, error) {
 		return nil, err
 	}
 
-
 	if driver == "sqlite3" {
 		// Also some optimizations for SQLite to make it FAA-A-A-AST.
 		db.Exec(`PRAGMA foreign_keys = ON`)
@@ -163,8 +162,8 @@ func (db *db) ContentType(tx *sql.Tx, fileUUID string) (string, error) {
 		row = db.contentType.QueryRow(fileUUID)
 	}
 
-	res := ""
-	return res, row.Scan(&res)
+	res := sql.NullString{}
+	return res.String, row.Scan(&res)
 }
 
 func (db *db) UnreachableFiles(tx *sql.Tx) ([]string, error) {
@@ -196,4 +195,4 @@ func (db *db) RemoveUnreachableFiles(tx *sql.Tx) error {
 		_, err := db.cleanup.Exec(time.Now().Unix())
 		return err
 	}
-}
+}

server.go

@@ -20,8 +20,8 @@ var ErrFileDoesntExists = errors.New("file doesn't exists")
 
 // filedrop server structure, implements http.Handler.
 type Server struct {
-	DB *db
-	Conf Config
+	DB     *db
+	Conf   Config
 	Logger *log.Logger
 
 	fileCleanerStopChan chan bool
@@ -81,7 +81,7 @@ func (s *Server) removeFile(tx *sql.Tx, fileUUID string) error {
 		return errors.Wrap(err, "uuid parse")
 	}
 
-	if err := s.DB.RemoveFile(tx, fileUUID); err !=nil {
+	if err := s.DB.RemoveFile(tx, fileUUID); err != nil {
 		return errors.Wrap(err, "db remove")
 	}
 
@@ -131,7 +131,7 @@ func (s *Server) GetFile(fileUUID string) (r io.Reader, contentType string, err
 
 	if s.DB.ShouldDelete(tx, fileUUID) {
 		if err := s.removeFile(tx, fileUUID); err != nil {
-			s.Logger.Println("Error while trying to remove file", fileUUID + ":", err)
+			s.Logger.Println("Error while trying to remove file", fileUUID+":", err)
 
 		}
 		if err := tx.Commit(); err != nil {
@@ -226,7 +226,11 @@ func (s *Server) acceptFile(w http.ResponseWriter, r *http.Request) {
 
 	// Smart logic to convert request's URL into absolute result URL.
 	resURL := url.URL{}
-	if s.Conf.HTTPSUpstream {
+	if r.Header.Get("X-HTTPS-Downstream") == "1" {
+		resURL.Scheme = "https"
+	} else if r.Header.Get("X-HTTPS-Downstream") == "0" {
+		resURL.Scheme = "http"
+	} else if s.Conf.HTTPSDownstream {
 		resURL.Scheme = "https"
 	} else {
 		resURL.Scheme = "http"
@@ -326,7 +330,7 @@ func (s *Server) cleanupFiles() {
 	}
 
 	for _, fileUUID := range uuids {
-		if err := os.Remove(filepath.Join(s.Conf.StorageDir, fileUUID)); err !=nil {
+		if err := os.Remove(filepath.Join(s.Conf.StorageDir, fileUUID)); err != nil {
 			s.Logger.Println("Failed to remove file during clean-up:", err)
 		}
 	}

server_test.go

@@ -19,7 +19,7 @@ import (
 
 var TestDBConf = filedrop.DBConfig{
 	Driver: "sqlite3",
-	DSN: ":memory:",
+	DSN:    ":memory:",
 }
 
 var file = `Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow Meow 
@@ -43,6 +43,7 @@ func initServ(conf filedrop.Config) *filedrop.Server {
 
 // Test for correct initialization of server.
 func TestNew(t *testing.T) {
+	t.Parallel()
 	conf := filedrop.Default
 	conf.DB = TestDBConf
 	tempDir, err := ioutil.TempDir("", "filedrop-tests")
@@ -74,7 +75,7 @@ func doPOST(t *testing.T, c *http.Client, url string, contentType string, reqBod
 		t.Error("ioutil.ReadAll:", err)
 		t.FailNow()
 	}
-	if resp.StatusCode / 100 != 2 {
+	if resp.StatusCode/100 != 2 {
 		t.Error("POST: HTTP", resp.StatusCode, resp.Status)
 		t.Error("Body:", string(body))
 		t.FailNow()
@@ -89,7 +90,7 @@ func doPOSTFail(t *testing.T, c *http.Client, url string, contentType string, re
 		t.FailNow()
 	}
 	defer resp.Body.Close()
-	if resp.StatusCode / 100 == 2 {
+	if resp.StatusCode/100 == 2 {
 		t.Error("POST: HTTP", resp.StatusCode, resp.Status)
 		t.FailNow()
 	}
@@ -108,7 +109,7 @@ func doGET(t *testing.T, c *http.Client, url string) []byte {
 		t.Error("ioutil.ReadAll:", err)
 		t.FailNow()
 	}
-	if resp.StatusCode / 100 != 2 {
+	if resp.StatusCode/100 != 2 {
 		t.Error("GET: HTTP", resp.Status)
 		t.Error("Body:", string(body))
 		t.FailNow()
@@ -123,7 +124,7 @@ func doGETFail(t *testing.T, c *http.Client, url string) int {
 		t.FailNow()
 	}
 	defer resp.Body.Close()
-	if resp.StatusCode / 100 == 2 {
+	if resp.StatusCode/100 == 2 {
 		t.Error("GET: HTTP", resp.StatusCode, resp.Status)
 		t.FailNow()
 	}
@@ -138,7 +139,7 @@ func TestBasicSubmit(t *testing.T) {
 	defer ts.Close()
 	c := ts.Client()
 
-	url := string(doPOST(t, c, ts.URL + "/filedrop/meow.txt", "text/plain", strings.NewReader(file)))
+	url := string(doPOST(t, c, ts.URL+"/filedrop/meow.txt", "text/plain", strings.NewReader(file)))
 
 	t.Log("File URL:", url)
 	if !strings.HasSuffix(url, "meow.txt") {
@@ -165,7 +166,7 @@ func TestDifferentFilename(t *testing.T) {
 	defer ts.Close()
 	c := ts.Client()
 
-	fileUrl := string(doPOST(t, c, ts.URL + "/filedrop/meow.txt", "text/plain", strings.NewReader(file)))
+	fileUrl := string(doPOST(t, c, ts.URL+"/filedrop/meow.txt", "text/plain", strings.NewReader(file)))
 
 	t.Log("File URL:", fileUrl)
 	if !strings.HasSuffix(fileUrl, "meow.txt") {
@@ -204,7 +205,7 @@ func TestNonExistent(t *testing.T) {
 	defer ts.Close()
 	c := ts.Client()
 
-	code := doGETFail(t, c, ts.URL + "/filedrop/AAAAAAAA-AAAA-AAAA-AAAA-AAAAAAAAAAAA/meow2.txt")
+	code := doGETFail(t, c, ts.URL+"/filedrop/AAAAAAAA-AAAA-AAAA-AAAA-AAAAAAAAAAAA/meow2.txt")
 	if code != 404 {
 		t.Error("GET: HTTP", code)
 		t.FailNow()
@@ -219,7 +220,7 @@ func TestContentTypePreserved(t *testing.T) {
 	defer ts.Close()
 	c := ts.Client()
 
-	url := string(doPOST(t, c, ts.URL + "/filedrop/meow.txt", "text/kitteh", strings.NewReader(file)))
+	url := string(doPOST(t, c, ts.URL+"/filedrop/meow.txt", "text/kitteh", strings.NewReader(file)))
 
 	t.Log("File URL:", url)
 
@@ -234,23 +235,120 @@ func TestContentTypePreserved(t *testing.T) {
 		t.Error("ioutil.ReadAll:", err)
 		t.FailNow()
 	}
-	if resp.StatusCode / 100 != 2 {
+	if resp.StatusCode/100 != 2 {
 		t.Error("GET: HTTP", resp.Status)
 		t.Error("Body:", string(body))
 		t.FailNow()
 	}
 	if resp.Header.Get("Content-Type") != "text/kitteh" {
 		t.Log("Mismatched content type:")
 		t.Log("\tWanted: 'text/kitteh'")
-		t.Log("\tGot:", "'" + resp.Header.Get("Content-Type") + "'")
+		t.Log("\tGot:", "'"+resp.Header.Get("Content-Type")+"'")
 		t.Fail()
 	}
 }
 
+func TestNoContentType(t *testing.T) {
+	serv := initServ(filedrop.Default)
+	ts := httptest.NewServer(serv)
+	defer os.RemoveAll(serv.Conf.StorageDir)
+	defer serv.Close()
+	defer ts.Close()
+	c := ts.Client()
+
+	url := string(doPOST(t, c, ts.URL+"/filedrop/meow.txt", "", strings.NewReader(file)))
+
+	t.Log("File URL:", url)
+
+	resp, err := c.Get(url)
+	if err != nil {
+		t.Error("GET:", err)
+		t.FailNow()
+	}
+	defer resp.Body.Close()
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		t.Error("ioutil.ReadAll:", err)
+		t.FailNow()
+	}
+	if resp.StatusCode/100 != 2 {
+		t.Error("GET: HTTP", resp.Status)
+		t.Error("Body:", string(body))
+		t.FailNow()
+	}
+	t.Log("Got:", "'"+resp.Header.Get("Content-Type")+"'")
+}
+
+func TestHTTPSDownstream(t *testing.T) {
+	serv := initServ(filedrop.Default)
+	ts := httptest.NewServer(serv)
+	defer os.RemoveAll(serv.Conf.StorageDir)
+	defer serv.Close()
+	defer ts.Close()
+	c := ts.Client()
+
+	t.Run("X-HTTPS-Downstream=1", func(t *testing.T) {
+		req, err := http.NewRequest("POST", ts.URL, strings.NewReader(file))
+		if err != nil {
+			t.Error(err)
+			t.FailNow()
+		}
+		req.Header.Set("X-HTTPS-Downstream", "1")
+		resp, err := c.Do(req)
+		if err != nil {
+			t.Error("POST:", err)
+			t.FailNow()
+		}
+		defer resp.Body.Close()
+		body, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			t.Error("ioutil.ReadAll:", err)
+			t.FailNow()
+		}
+		if resp.StatusCode/100 != 2 {
+			t.Error("POST: HTTP", resp.StatusCode, resp.Status)
+			t.Error("Body:", string(body))
+			t.FailNow()
+		}
+		if !strings.HasPrefix(string(body), "https") {
+			t.Error("Got non-HTTPS URl with X-HTTPS-Downstream=1")
+			t.FailNow()
+		}
+	})
+	t.Run("X-HTTPS-Downstream=0", func(t *testing.T) {
+		req, err := http.NewRequest("POST", ts.URL, strings.NewReader(file))
+		if err != nil {
+			t.Error(err)
+			t.FailNow()
+		}
+		req.Header.Set("X-HTTPS-Downstream", "0")
+		resp, err := c.Do(req)
+		if err != nil {
+			t.Error("POST:", err)
+			t.FailNow()
+		}
+		defer resp.Body.Close()
+		body, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			t.Error("ioutil.ReadAll:", err)
+			t.FailNow()
+		}
+		if resp.StatusCode/100 != 2 {
+			t.Error("POST: HTTP", resp.StatusCode, resp.Status)
+			t.Error("Body:", string(body))
+			t.FailNow()
+		}
+		if !strings.HasPrefix(string(body), "http") {
+			t.Error("Got non-HTTP URL with X-HTTPS-Downstream=0")
+			t.FailNow()
+		}
+	})
+}
+
 func testWithPrefix(t *testing.T, ts *httptest.Server, c *http.Client, prefix string) {
 	var URL string
-	t.Run("submit with prefix " + prefix, func(t *testing.T) {
-		URL = string(doPOST(t, c, ts.URL + prefix + "/meow.txt", "text/plain", strings.NewReader(file)))
+	t.Run("submit with prefix "+prefix, func(t *testing.T) {
+		URL = string(doPOST(t, c, ts.URL+prefix+"/meow.txt", "text/plain", strings.NewReader(file)))
 	})
 
 	if !strings.Contains(URL, prefix) {
@@ -259,7 +357,7 @@ func testWithPrefix(t *testing.T, ts *httptest.Server, c *http.Client, prefix st
 	}
 
 	if URL != "" {
-		t.Run("get with " + prefix, func(t *testing.T) {
+		t.Run("get with "+prefix, func(t *testing.T) {
 			body := doGET(t, c, URL)
 			if string(body) != file {
 				t.Error("Got different file!")