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