Detect and set proper Content-Type for avatars

Detect proper avatar type to return instead of returning `image/*`, which is not valid.
go-pkgz · Sep 22, 2024 · ac5b1b4 · ac5b1b4
1 parent 458e57e
commit ac5b1b4
Show file tree

Hide file tree

Showing 6 changed files with 140 additions and 16 deletions.
diff --git a/avatar/avatar.go b/avatar/avatar.go
@@ -23,6 +23,9 @@ import (
 	"github.com/go-pkgz/auth/token"
 )
 
+// http.sniffLen is 512 bytes which is how much we need to read to detect content type
+const sniffLen = 512
+
 // Proxy provides http handler for avatars from avatar.Store
 // On user login token will call Put and it will retrieve and save picture locally.
 type Proxy struct {
@@ -100,7 +103,6 @@ func (p *Proxy) load(url string, client *http.Client) (rc io.ReadCloser, err err
 
 // Handler returns token routes for given provider
 func (p *Proxy) Handler(w http.ResponseWriter, r *http.Request) {
-
 	if r.Method != "GET" {
 		w.WriteHeader(http.StatusMethodNotAllowed)
 	}
@@ -136,9 +138,25 @@ func (p *Proxy) Handler(w http.ResponseWriter, r *http.Request) {
 		}
 	}()
 
-	w.Header().Set("Content-Type", "image/*")
+	buf := make([]byte, sniffLen)
+	n, err := avReader.Read(buf)
+	if err != nil && err != io.EOF {
+		p.Logf("[WARN] can't read from avatar reader for %s, %s", avatarID, err)
+		rest.SendErrorJSON(w, r, p.L, http.StatusInternalServerError, err, "can't read avatar")
+		return
+	}
 	w.Header().Set("Content-Length", strconv.Itoa(size))
+	contentType := http.DetectContentType(buf)
+	if contentType == "application/octet-stream" {
+		contentType = "image/*"
+	}
+	w.Header().Set("Content-Type", contentType)
 	w.WriteHeader(http.StatusOK)
+	if _, err = w.Write(buf[:n]); err != nil {
+		p.Logf("[WARN] can't write response to %s, %s", r.RemoteAddr, err)
+		return
+	}
+	// write the rest of response size if it's bigger than 512 bytes, or nothing as EOF would be sent right away then
 	if _, err = io.Copy(w, avReader); err != nil {
 		p.Logf("[WARN] can't send response to %s, %s", r.RemoteAddr, err)
 	}

diff --git a/avatar/avatar_test.go b/avatar/avatar_test.go
@@ -107,12 +107,21 @@ func TestAvatar_Routes(t *testing.T) {
 
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path == "/pic.png" {
-			w.Header().Set("Content-Type", "image/*")
+			// set wrong content type, in reality would be image/png based on the first bytes
+			w.Header().Set("Content-Type", "image/jpg")
 			w.Header().Set("Custom-Header", "xyz")
-			_, err := fmt.Fprint(w, "some picture bin data")
+			_, err := fmt.Fprint(w, "\x89PNG\x0D\x0A\x1A\x0A some picture bin data")
 			require.NoError(t, err)
 			return
 		}
+		if r.URL.Path == "/circles.png" {
+			http.ServeFile(w, r, "testdata/circles.png")
+			return
+		}
+		if r.URL.Path == "/circles.jpg" {
+			http.ServeFile(w, r, "testdata/circles.jpg")
+			return
+		}
 		http.Error(w, "not found", http.StatusNotFound)
 	}))
 	defer ts.Close()
@@ -154,16 +163,16 @@ func TestAvatar_Routes(t *testing.T) {
 		handler.ServeHTTP(rr, req)
 
 		assert.Equal(t, http.StatusOK, rr.Code)
-		assert.Equal(t, []string{"image/*"}, rr.Header()["Content-Type"])
-		assert.Equal(t, []string{"21"}, rr.Header()["Content-Length"])
+		assert.Equal(t, []string{"image/png"}, rr.Header()["Content-Type"])
+		assert.Equal(t, []string{"30"}, rr.Header()["Content-Length"])
 		assert.Equal(t, []string(nil), rr.Header()["Custom-Header"], "strip all custom headers")
 		assert.NotNil(t, rr.Header()["Etag"])
 
 		bb := bytes.Buffer{}
 		sz, err := io.Copy(&bb, rr.Body)
 		assert.NoError(t, err)
-		assert.Equal(t, int64(21), sz)
-		assert.Equal(t, "some picture bin data", bb.String())
+		assert.Equal(t, int64(30), sz)
+		assert.Equal(t, "\x89PNG\x0D\x0A\x1A\x0A some picture bin data", bb.String())
 	}
 
 	{
@@ -181,6 +190,39 @@ func TestAvatar_Routes(t *testing.T) {
 		assert.Equal(t, []string{`"` + id + `"`}, rr.Header()["Etag"])
 	}
 
+	u = token.User{ID: "user2", Name: "user2 name", Picture: ts.URL + "/circles.png"}
+	circlesPngURL, err := p.Put(u, client)
+	t.Log("circlesPngURL", circlesPngURL)
+	assert.NoError(t, err)
+
+	{ // full-size png image is served correctly and with the right header
+		req, err := http.NewRequest("GET", circlesPngURL, http.NoBody)
+		require.NoError(t, err)
+		rr := httptest.NewRecorder()
+		handler := http.Handler(http.HandlerFunc(p.Handler))
+		handler.ServeHTTP(rr, req)
+
+		assert.Equal(t, http.StatusOK, rr.Code)
+		assert.Equal(t, []string{"image/png"}, rr.Header()["Content-Type"])
+		assert.Equal(t, []string{"11392"}, rr.Header()["Content-Length"])
+	}
+
+	u = token.User{ID: "user3", Name: "user3 name", Picture: ts.URL + "/circles.jpg"}
+	circlesJpgURL, err := p.Put(u, client)
+	assert.NoError(t, err)
+
+	{ // full-size jpg image is served correctly and with the right header
+		req, err := http.NewRequest("GET", circlesJpgURL, http.NoBody)
+		require.NoError(t, err)
+		rr := httptest.NewRecorder()
+		handler := http.Handler(http.HandlerFunc(p.Handler))
+		handler.ServeHTTP(rr, req)
+
+		assert.Equal(t, http.StatusOK, rr.Code)
+		assert.Equal(t, []string{"image/jpeg"}, rr.Header()["Content-Type"])
+		assert.Equal(t, []string{"23983"}, rr.Header()["Content-Length"])
+	}
+
 }
 
 func TestAvatar_resize(t *testing.T) {

diff --git a/avatar/noop_test.go b/avatar/noop_test.go
@@ -6,6 +6,7 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -39,6 +40,7 @@ func TestNoOp_Get(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, http.StatusOK, resp.StatusCode)
 		require.Zero(t, resp.ContentLength)
+		assert.Equal(t, "image/*", resp.Header.Get("Content-Type"))
 		body, err := io.ReadAll(resp.Body)
 		require.NoError(t, err)
 		require.Empty(t, body)

diff --git a/v2/avatar/avatar.go b/v2/avatar/avatar.go
@@ -23,6 +23,9 @@ import (
 	"github.com/go-pkgz/auth/v2/token"
 )
 
+// http.sniffLen is 512 bytes which is how much we need to read to detect content type
+const sniffLen = 512
+
 // Proxy provides http handler for avatars from avatar.Store
 // On user login token will call Put and it will retrieve and save picture locally.
 type Proxy struct {
@@ -100,7 +103,6 @@ func (p *Proxy) load(url string, client *http.Client) (rc io.ReadCloser, err err
 
 // Handler returns token routes for given provider
 func (p *Proxy) Handler(w http.ResponseWriter, r *http.Request) {
-
 	if r.Method != "GET" {
 		w.WriteHeader(http.StatusMethodNotAllowed)
 	}
@@ -136,9 +138,25 @@ func (p *Proxy) Handler(w http.ResponseWriter, r *http.Request) {
 		}
 	}()
 
-	w.Header().Set("Content-Type", "image/*")
+	buf := make([]byte, sniffLen)
+	n, err := avReader.Read(buf)
+	if err != nil && err != io.EOF {
+		p.Logf("[WARN] can't read from avatar reader for %s, %s", avatarID, err)
+		rest.SendErrorJSON(w, r, p.L, http.StatusInternalServerError, err, "can't read avatar")
+		return
+	}
 	w.Header().Set("Content-Length", strconv.Itoa(size))
+	contentType := http.DetectContentType(buf)
+	if contentType == "application/octet-stream" {
+		contentType = "image/*"
+	}
+	w.Header().Set("Content-Type", contentType)
 	w.WriteHeader(http.StatusOK)
+	if _, err = w.Write(buf[:n]); err != nil {
+		p.Logf("[WARN] can't write response to %s, %s", r.RemoteAddr, err)
+		return
+	}
+	// write the rest of response size if it's bigger than 512 bytes, or nothing as EOF would be sent right away then
 	if _, err = io.Copy(w, avReader); err != nil {
 		p.Logf("[WARN] can't send response to %s, %s", r.RemoteAddr, err)
 	}

diff --git a/v2/avatar/avatar_test.go b/v2/avatar/avatar_test.go
@@ -107,12 +107,21 @@ func TestAvatar_Routes(t *testing.T) {
 
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path == "/pic.png" {
-			w.Header().Set("Content-Type", "image/*")
+			// set wrong content type, in reality would be image/png based on the first bytes
+			w.Header().Set("Content-Type", "image/jpg")
 			w.Header().Set("Custom-Header", "xyz")
-			_, err := fmt.Fprint(w, "some picture bin data")
+			_, err := fmt.Fprint(w, "\x89PNG\x0D\x0A\x1A\x0A some picture bin data")
 			require.NoError(t, err)
 			return
 		}
+		if r.URL.Path == "/circles.png" {
+			http.ServeFile(w, r, "testdata/circles.png")
+			return
+		}
+		if r.URL.Path == "/circles.jpg" {
+			http.ServeFile(w, r, "testdata/circles.jpg")
+			return
+		}
 		http.Error(w, "not found", http.StatusNotFound)
 	}))
 	defer ts.Close()
@@ -154,16 +163,16 @@ func TestAvatar_Routes(t *testing.T) {
 		handler.ServeHTTP(rr, req)
 
 		assert.Equal(t, http.StatusOK, rr.Code)
-		assert.Equal(t, []string{"image/*"}, rr.Header()["Content-Type"])
-		assert.Equal(t, []string{"21"}, rr.Header()["Content-Length"])
+		assert.Equal(t, []string{"image/png"}, rr.Header()["Content-Type"])
+		assert.Equal(t, []string{"30"}, rr.Header()["Content-Length"])
 		assert.Equal(t, []string(nil), rr.Header()["Custom-Header"], "strip all custom headers")
 		assert.NotNil(t, rr.Header()["Etag"])
 
 		bb := bytes.Buffer{}
 		sz, err := io.Copy(&bb, rr.Body)
 		assert.NoError(t, err)
-		assert.Equal(t, int64(21), sz)
-		assert.Equal(t, "some picture bin data", bb.String())
+		assert.Equal(t, int64(30), sz)
+		assert.Equal(t, "\x89PNG\x0D\x0A\x1A\x0A some picture bin data", bb.String())
 	}
 
 	{
@@ -181,6 +190,39 @@ func TestAvatar_Routes(t *testing.T) {
 		assert.Equal(t, []string{`"` + id + `"`}, rr.Header()["Etag"])
 	}
 
+	u = token.User{ID: "user2", Name: "user2 name", Picture: ts.URL + "/circles.png"}
+	circlesPngURL, err := p.Put(u, client)
+	t.Log("circlesPngURL", circlesPngURL)
+	assert.NoError(t, err)
+
+	{ // full-size png image is served correctly and with the right header
+		req, err := http.NewRequest("GET", circlesPngURL, http.NoBody)
+		require.NoError(t, err)
+		rr := httptest.NewRecorder()
+		handler := http.Handler(http.HandlerFunc(p.Handler))
+		handler.ServeHTTP(rr, req)
+
+		assert.Equal(t, http.StatusOK, rr.Code)
+		assert.Equal(t, []string{"image/png"}, rr.Header()["Content-Type"])
+		assert.Equal(t, []string{"11392"}, rr.Header()["Content-Length"])
+	}
+
+	u = token.User{ID: "user3", Name: "user3 name", Picture: ts.URL + "/circles.jpg"}
+	circlesJpgURL, err := p.Put(u, client)
+	assert.NoError(t, err)
+
+	{ // full-size jpg image is served correctly and with the right header
+		req, err := http.NewRequest("GET", circlesJpgURL, http.NoBody)
+		require.NoError(t, err)
+		rr := httptest.NewRecorder()
+		handler := http.Handler(http.HandlerFunc(p.Handler))
+		handler.ServeHTTP(rr, req)
+
+		assert.Equal(t, http.StatusOK, rr.Code)
+		assert.Equal(t, []string{"image/jpeg"}, rr.Header()["Content-Type"])
+		assert.Equal(t, []string{"23983"}, rr.Header()["Content-Length"])
+	}
+
 }
 
 func TestAvatar_resize(t *testing.T) {

diff --git a/v2/avatar/noop_test.go b/v2/avatar/noop_test.go
@@ -6,6 +6,7 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -39,6 +40,7 @@ func TestNoOp_Get(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, http.StatusOK, resp.StatusCode)
 		require.Zero(t, resp.ContentLength)
+		assert.Equal(t, "image/*", resp.Header.Get("Content-Type"))
 		body, err := io.ReadAll(resp.Body)
 		require.NoError(t, err)
 		require.Empty(t, body)