github.com/mattermosttest/mattermost-server/v5@v5.0.0-20200917143240-9dfa12e121f9/services/imageproxy/local_test.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package imageproxy
     5  
     6  import (
     7  	"io/ioutil"
     8  	"net/http"
     9  	"net/http/httptest"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/mattermost/mattermost-server/v5/model"
    14  	"github.com/mattermost/mattermost-server/v5/services/httpservice"
    15  	"github.com/mattermost/mattermost-server/v5/utils/testutils"
    16  	"github.com/stretchr/testify/assert"
    17  	"github.com/stretchr/testify/require"
    18  )
    19  
    20  func makeTestLocalProxy() *ImageProxy {
    21  	configService := &testutils.StaticConfigService{
    22  		Cfg: &model.Config{
    23  			ServiceSettings: model.ServiceSettings{
    24  				SiteURL:                             model.NewString("https://mattermost.example.com"),
    25  				AllowedUntrustedInternalConnections: model.NewString("127.0.0.1"),
    26  			},
    27  			ImageProxySettings: model.ImageProxySettings{
    28  				Enable:         model.NewBool(true),
    29  				ImageProxyType: model.NewString(model.IMAGE_PROXY_TYPE_LOCAL),
    30  			},
    31  		},
    32  	}
    33  
    34  	return MakeImageProxy(configService, httpservice.MakeHTTPService(configService), nil)
    35  }
    36  
    37  func TestLocalBackend_GetImage(t *testing.T) {
    38  	t.Run("image", func(t *testing.T) {
    39  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    40  			w.Header().Set("Cache-Control", "max-age=2592000, private")
    41  			w.Header().Set("Content-Type", "image/png")
    42  			w.Header().Set("Content-Length", "10")
    43  
    44  			w.WriteHeader(http.StatusOK)
    45  			w.Write([]byte("1111111111"))
    46  		})
    47  
    48  		mock := httptest.NewServer(handler)
    49  		defer mock.Close()
    50  
    51  		proxy := makeTestLocalProxy()
    52  
    53  		recorder := httptest.NewRecorder()
    54  		request, _ := http.NewRequest(http.MethodGet, "", nil)
    55  		proxy.GetImage(recorder, request, mock.URL+"/image.png")
    56  		resp := recorder.Result()
    57  
    58  		require.Equal(t, http.StatusOK, resp.StatusCode)
    59  		assert.Equal(t, "max-age=2592000, private", resp.Header.Get("Cache-Control"))
    60  		assert.Equal(t, "10", resp.Header.Get("Content-Length"))
    61  
    62  		respBody, _ := ioutil.ReadAll(resp.Body)
    63  		assert.Equal(t, []byte("1111111111"), respBody)
    64  	})
    65  
    66  	t.Run("not an image", func(t *testing.T) {
    67  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    68  			w.WriteHeader(http.StatusNotAcceptable)
    69  		})
    70  
    71  		mock := httptest.NewServer(handler)
    72  		defer mock.Close()
    73  
    74  		proxy := makeTestLocalProxy()
    75  
    76  		recorder := httptest.NewRecorder()
    77  		request, _ := http.NewRequest(http.MethodGet, "", nil)
    78  		proxy.GetImage(recorder, request, mock.URL+"/file.pdf")
    79  		resp := recorder.Result()
    80  
    81  		require.Equal(t, http.StatusNotAcceptable, resp.StatusCode)
    82  	})
    83  
    84  	t.Run("not an image, but remote server ignores accept header", func(t *testing.T) {
    85  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    86  			w.Header().Set("Cache-Control", "max-age=2592000, private")
    87  			w.Header().Set("Content-Type", "application/pdf")
    88  			w.Header().Set("Content-Length", "10")
    89  
    90  			w.WriteHeader(http.StatusOK)
    91  			w.Write([]byte("1111111111"))
    92  		})
    93  
    94  		mock := httptest.NewServer(handler)
    95  		defer mock.Close()
    96  
    97  		proxy := makeTestLocalProxy()
    98  
    99  		recorder := httptest.NewRecorder()
   100  		request, _ := http.NewRequest(http.MethodGet, "", nil)
   101  		proxy.GetImage(recorder, request, mock.URL+"/file.pdf")
   102  		resp := recorder.Result()
   103  
   104  		require.Equal(t, http.StatusForbidden, resp.StatusCode)
   105  	})
   106  
   107  	t.Run("not found", func(t *testing.T) {
   108  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   109  			w.WriteHeader(http.StatusNotFound)
   110  		})
   111  
   112  		mock := httptest.NewServer(handler)
   113  		defer mock.Close()
   114  
   115  		proxy := makeTestLocalProxy()
   116  
   117  		recorder := httptest.NewRecorder()
   118  		request, _ := http.NewRequest(http.MethodGet, "", nil)
   119  		proxy.GetImage(recorder, request, mock.URL+"/image.png")
   120  		resp := recorder.Result()
   121  
   122  		require.Equal(t, http.StatusNotFound, resp.StatusCode)
   123  	})
   124  
   125  	t.Run("other server error", func(t *testing.T) {
   126  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   127  			w.WriteHeader(http.StatusInternalServerError)
   128  		})
   129  
   130  		mock := httptest.NewServer(handler)
   131  		defer mock.Close()
   132  
   133  		proxy := makeTestLocalProxy()
   134  
   135  		recorder := httptest.NewRecorder()
   136  		request, _ := http.NewRequest(http.MethodGet, "", nil)
   137  		proxy.GetImage(recorder, request, mock.URL+"/image.png")
   138  		resp := recorder.Result()
   139  
   140  		require.Equal(t, http.StatusInternalServerError, resp.StatusCode)
   141  	})
   142  
   143  	t.Run("timeout", func(t *testing.T) {
   144  		wait := make(chan bool, 1)
   145  
   146  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   147  			<-wait
   148  		})
   149  
   150  		mock := httptest.NewServer(handler)
   151  		defer mock.Close()
   152  
   153  		proxy := makeTestLocalProxy()
   154  
   155  		// Modify the timeout to be much shorter than the default 30 seconds
   156  		proxy.backend.(*LocalBackend).impl.Timeout = time.Millisecond
   157  
   158  		recorder := httptest.NewRecorder()
   159  		request, _ := http.NewRequest(http.MethodGet, "", nil)
   160  		proxy.GetImage(recorder, request, mock.URL+"/image.png")
   161  		resp := recorder.Result()
   162  
   163  		require.Equal(t, http.StatusGatewayTimeout, resp.StatusCode)
   164  
   165  		wait <- true
   166  	})
   167  
   168  	t.Run("SVG attachment", func(t *testing.T) {
   169  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   170  			w.Header().Set("Cache-Control", "max-age=2592000, private")
   171  			w.Header().Set("Content-Type", "image/svg+xml")
   172  			w.Header().Set("Content-Length", "10")
   173  
   174  			w.WriteHeader(http.StatusOK)
   175  			w.Write([]byte("1111111111"))
   176  		})
   177  
   178  		mock := httptest.NewServer(handler)
   179  		defer mock.Close()
   180  
   181  		proxy := makeTestLocalProxy()
   182  
   183  		recorder := httptest.NewRecorder()
   184  		request, err := http.NewRequest(http.MethodGet, "", nil)
   185  		require.NoError(t, err)
   186  		proxy.GetImage(recorder, request, mock.URL+"/test.svg")
   187  		resp := recorder.Result()
   188  
   189  		assert.Equal(t, http.StatusOK, resp.StatusCode)
   190  		assert.Equal(t, "attachment;filename=\"test.svg\"", resp.Header.Get("Content-Disposition"))
   191  
   192  		_, err = ioutil.ReadAll(resp.Body)
   193  		require.NoError(t, err)
   194  	})
   195  }
   196  
   197  func TestLocalBackend_GetImageDirect(t *testing.T) {
   198  	t.Run("image", func(t *testing.T) {
   199  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   200  			w.Header().Set("Cache-Control", "max-age=2592000, private")
   201  			w.Header().Set("Content-Type", "image/png")
   202  			w.Header().Set("Content-Length", "10")
   203  
   204  			w.WriteHeader(http.StatusOK)
   205  			w.Write([]byte("1111111111"))
   206  		})
   207  
   208  		mock := httptest.NewServer(handler)
   209  		defer mock.Close()
   210  
   211  		proxy := makeTestLocalProxy()
   212  
   213  		body, contentType, err := proxy.GetImageDirect(mock.URL + "/image.png")
   214  
   215  		assert.Nil(t, err)
   216  		assert.Equal(t, "image/png", contentType)
   217  
   218  		respBody, _ := ioutil.ReadAll(body)
   219  		assert.Equal(t, []byte("1111111111"), respBody)
   220  	})
   221  
   222  	t.Run("not an image", func(t *testing.T) {
   223  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   224  			w.WriteHeader(http.StatusNotAcceptable)
   225  		})
   226  
   227  		mock := httptest.NewServer(handler)
   228  		defer mock.Close()
   229  
   230  		proxy := makeTestLocalProxy()
   231  
   232  		body, contentType, err := proxy.GetImageDirect(mock.URL + "/file.pdf")
   233  
   234  		assert.NotNil(t, err)
   235  		assert.Equal(t, "", contentType)
   236  		assert.Equal(t, ErrLocalRequestFailed, err)
   237  		assert.Nil(t, body)
   238  	})
   239  
   240  	t.Run("not an image, but remote server ignores accept header", func(t *testing.T) {
   241  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   242  			w.Header().Set("Cache-Control", "max-age=2592000, private")
   243  			w.Header().Set("Content-Type", "application/pdf")
   244  			w.Header().Set("Content-Length", "10")
   245  
   246  			w.WriteHeader(http.StatusOK)
   247  			w.Write([]byte("1111111111"))
   248  		})
   249  
   250  		mock := httptest.NewServer(handler)
   251  		defer mock.Close()
   252  
   253  		proxy := makeTestLocalProxy()
   254  
   255  		body, contentType, err := proxy.GetImageDirect(mock.URL + "/file.pdf")
   256  
   257  		assert.NotNil(t, err)
   258  		assert.Equal(t, "", contentType)
   259  		assert.Equal(t, ErrLocalRequestFailed, err)
   260  		assert.Nil(t, body)
   261  	})
   262  
   263  	t.Run("not found", func(t *testing.T) {
   264  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   265  			w.WriteHeader(http.StatusNotFound)
   266  		})
   267  
   268  		mock := httptest.NewServer(handler)
   269  		defer mock.Close()
   270  
   271  		proxy := makeTestLocalProxy()
   272  
   273  		body, contentType, err := proxy.GetImageDirect(mock.URL + "/image.png")
   274  
   275  		assert.NotNil(t, err)
   276  		assert.Equal(t, "", contentType)
   277  		assert.Equal(t, ErrLocalRequestFailed, err)
   278  		assert.Nil(t, body)
   279  	})
   280  
   281  	t.Run("other server error", func(t *testing.T) {
   282  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   283  			w.WriteHeader(http.StatusInternalServerError)
   284  		})
   285  
   286  		mock := httptest.NewServer(handler)
   287  		defer mock.Close()
   288  
   289  		proxy := makeTestLocalProxy()
   290  
   291  		body, contentType, err := proxy.GetImageDirect(mock.URL + "/image.png")
   292  
   293  		assert.NotNil(t, err)
   294  		assert.Equal(t, "", contentType)
   295  		assert.Equal(t, ErrLocalRequestFailed, err)
   296  		assert.Nil(t, body)
   297  	})
   298  
   299  	t.Run("timeout", func(t *testing.T) {
   300  		wait := make(chan bool, 1)
   301  
   302  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   303  			<-wait
   304  		})
   305  
   306  		mock := httptest.NewServer(handler)
   307  		defer mock.Close()
   308  
   309  		proxy := makeTestLocalProxy()
   310  
   311  		// Modify the timeout to be much shorter than the default 30 seconds
   312  		proxy.backend.(*LocalBackend).impl.Timeout = time.Millisecond
   313  
   314  		body, contentType, err := proxy.GetImageDirect(mock.URL + "/image.png")
   315  
   316  		assert.NotNil(t, err)
   317  		assert.Equal(t, "", contentType)
   318  		assert.Equal(t, ErrLocalRequestFailed, err)
   319  		assert.Nil(t, body)
   320  
   321  		wait <- true
   322  	})
   323  }