Never save attachments with empty content type

attachments.go

@@ -10,6 +10,7 @@
 	"net/http"
 	"net/url"
 	"path/filepath"
+	"strings"
 
 	"github.com/h2non/filetype"
 	"github.com/nyaruka/courier/utils"
@@ -107,45 +108,59 @@
 		return nil, err
 	}
 
-	mimeType := ""
-	extension := filepath.Ext(parsedURL.Path)
-	if extension != "" {
-		extension = extension[1:]
+	mimeType, extension := getAttachmentType(trace)
+
+	storageURL, err := b.SaveAttachment(ctx, channel, mimeType, trace.ResponseBody, extension)
+	if err != nil {
+		return nil, err
 	}
 
+	return &Attachment{ContentType: mimeType, URL: storageURL, Size: len(trace.ResponseBody)}, nil
+}
+
+func getAttachmentType(t *httpx.Trace) (string, string) {
+	var typ string
+
+	// use extension from url path if it exists
+	ext := filepath.Ext(t.Request.URL.Path)
+
 	// prioritize to use the response content type header if provided
-	contentTypeHeader := trace.Response.Header.Get("Content-Type")
-	if contentTypeHeader != "" && contentTypeHeader != "application/octet-stream" {
-		mimeType, _, _ = mime.ParseMediaType(contentTypeHeader)
-		if extension == "" {
-			extensions, err := mime.ExtensionsByType(mimeType)
-			if extensions == nil || err != nil {
-				extension = ""
-			} else {
-				extension = extensions[0][1:]
+	contentTypeHeader := t.Response.Header.Get("Content-Type")
+	if contentTypeHeader != "" {
+		typ, _, _ = mime.ParseMediaType(contentTypeHeader)
+	}
+
+	// if we didn't get a meaningful content type from the header, try to guess it from the body
+	if typ == "" || typ == "application/octet-stream" {
+		fileType, _ := filetype.Match(t.ResponseBody[:300])
+		if fileType != filetype.Unknown {
+			typ = fileType.MIME.Value
+			if ext == "" {
+				ext = fileType.Extension
 			}
 		}
-	} else {
+	}
 
-		// first try getting our mime type from the first 300 bytes of our body
-		fileType, _ := filetype.Match(trace.ResponseBody[:300])
+	// if we still don't have a type but the path has an extension, try to use that
+	if typ == "" && ext != "" {
+		fileType := filetype.GetType(ext)
 		if fileType != filetype.Unknown {
-			mimeType = fileType.MIME.Value
-			extension = fileType.Extension
-		} else {
-			// if that didn't work, try from our extension
-			fileType = filetype.GetType(extension)
-			if fileType != filetype.Unknown {
-				mimeType = fileType.MIME.Value
-				extension = fileType.Extension
-			}
+			typ = fileType.MIME.Value
 		}
 	}
 
-	storageURL, err := b.SaveAttachment(ctx, channel, mimeType, trace.ResponseBody, extension)
-	if err != nil {
-		return nil, err
+	// if we have a type but no extension, try to get one from the type
+	if ext == "" {
+		extensions, err := mime.ExtensionsByType(typ)
+		if len(extensions) > 0 && err == nil {
+			ext = extensions[0][1:]
+		}
 	}
 
-	return &Attachment{ContentType: mimeType, URL: storageURL, Size: len(trace.ResponseBody)}, nil
+	// got to default to something...
+	if typ == "" {
+		typ = "application/octet-stream"
+	}
+
+	return typ, strings.TrimPrefix(ext, ".")
 }

attachments_test.go

@@ -37,6 +37,9 @@ func TestFetchAndStoreAttachment(t *testing.T) {
 		"http://mock.com/media/hello.txt": {
 			httpx.NewMockResponse(200, nil, []byte(`hi`)),
 		},
+		"http://mock.com/media/hello7": {
+			httpx.NewMockResponse(200, nil, []byte(`hello world`)),
+		},
 	}))
 
 	defer uuids.SetGenerator(uuids.DefaultGenerator)
@@ -93,6 +96,12 @@ func TestFetchAndStoreAttachment(t *testing.T) {
 	assert.Equal(t, "https://backend.com/attachments/338ff339-5663-49ed-8ef6-384876655d1b.jpg", att.URL)
 	assert.Equal(t, 17301, att.Size)
 
+	att, err = courier.FetchAndStoreAttachment(ctx, mb, mockChannel, "http://mock.com/media/hello7", clog)
+	assert.NoError(t, err)
+	assert.Equal(t, "application/octet-stream", att.ContentType)
+	assert.Equal(t, "https://backend.com/attachments/9b955e36-ac16-4c6b-8ab6-9b9af5cd042a.", att.URL)
+	assert.Equal(t, 11, att.Size)
+
 	// an actual error on our part should be returned as an error
 	mb.SetStorageError(errors.New("boom"))