github.com/masterhung0112/hk_server/v5@v5.0.0-20220302090640-ec71aef15e1c/api4/image_test.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package api4
     5  
     6  import (
     7  	"io/ioutil"
     8  	"net/http"
     9  	"net/http/httptest"
    10  	"net/url"
    11  	"strings"
    12  	"testing"
    13  
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  
    17  	"github.com/masterhung0112/hk_server/v5/model"
    18  )
    19  
    20  func TestGetImage(t *testing.T) {
    21  	th := Setup(t)
    22  	defer th.TearDown()
    23  
    24  	// Prevent the test client from following a redirect
    25  	th.Client.HttpClient.CheckRedirect = func(*http.Request, []*http.Request) error {
    26  		return http.ErrUseLastResponse
    27  	}
    28  
    29  	t.Run("proxy disabled", func(t *testing.T) {
    30  		imageURL := "http://foo.bar/baz.gif"
    31  
    32  		th.App.UpdateConfig(func(cfg *model.Config) {
    33  			cfg.ImageProxySettings.Enable = model.NewBool(false)
    34  		})
    35  
    36  		r, err := http.NewRequest("GET", th.Client.ApiUrl+"/image?url="+url.QueryEscape(imageURL), nil)
    37  		require.NoError(t, err)
    38  		r.Header.Set(model.HEADER_AUTH, th.Client.AuthType+" "+th.Client.AuthToken)
    39  
    40  		resp, err := th.Client.HttpClient.Do(r)
    41  		require.NoError(t, err)
    42  		assert.Equal(t, http.StatusFound, resp.StatusCode)
    43  		assert.Equal(t, imageURL, resp.Header.Get("Location"))
    44  	})
    45  
    46  	t.Run("atmos/camo", func(t *testing.T) {
    47  		imageURL := "http://foo.bar/baz.gif"
    48  		proxiedURL := "https://proxy.foo.bar/004afe2ef382eb5f30c4490f793f8a8c5b33d8a2/687474703a2f2f666f6f2e6261722f62617a2e676966"
    49  
    50  		th.App.UpdateConfig(func(cfg *model.Config) {
    51  			cfg.ImageProxySettings.Enable = model.NewBool(true)
    52  			cfg.ImageProxySettings.ImageProxyType = model.NewString("atmos/camo")
    53  			cfg.ImageProxySettings.RemoteImageProxyOptions = model.NewString("foo")
    54  			cfg.ImageProxySettings.RemoteImageProxyURL = model.NewString("https://proxy.foo.bar")
    55  		})
    56  
    57  		r, err := http.NewRequest("GET", th.Client.ApiUrl+"/image?url="+url.QueryEscape(imageURL), nil)
    58  		require.NoError(t, err)
    59  		r.Header.Set(model.HEADER_AUTH, th.Client.AuthType+" "+th.Client.AuthToken)
    60  
    61  		resp, err := th.Client.HttpClient.Do(r)
    62  		require.NoError(t, err)
    63  		assert.Equal(t, http.StatusFound, resp.StatusCode)
    64  		assert.Equal(t, proxiedURL, resp.Header.Get("Location"))
    65  	})
    66  
    67  	t.Run("local", func(t *testing.T) {
    68  		th.App.UpdateConfig(func(cfg *model.Config) {
    69  			cfg.ImageProxySettings.Enable = model.NewBool(true)
    70  			cfg.ImageProxySettings.ImageProxyType = model.NewString("local")
    71  
    72  			// Allow requests to the "remote" image
    73  			cfg.ServiceSettings.AllowedUntrustedInternalConnections = model.NewString("127.0.0.1")
    74  		})
    75  
    76  		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    77  			w.Header().Set("Content-Type", "image/png")
    78  			w.Write([]byte("success"))
    79  		})
    80  
    81  		imageServer := httptest.NewServer(handler)
    82  		defer imageServer.Close()
    83  
    84  		r, err := http.NewRequest("GET", th.Client.ApiUrl+"/image?url="+url.QueryEscape(imageServer.URL+"/image.png"), nil)
    85  		require.NoError(t, err)
    86  		r.Header.Set(model.HEADER_AUTH, th.Client.AuthType+" "+th.Client.AuthToken)
    87  
    88  		resp, err := th.Client.HttpClient.Do(r)
    89  		require.NoError(t, err)
    90  		assert.Equal(t, http.StatusOK, resp.StatusCode)
    91  
    92  		respBody, err := ioutil.ReadAll(resp.Body)
    93  		require.NoError(t, err)
    94  		assert.Equal(t, "success", string(respBody))
    95  
    96  		// local images should not be proxied, but forwarded
    97  		r, err = http.NewRequest("GET", th.Client.ApiUrl+"/image?url=/plugins/test/image.png", nil)
    98  		require.NoError(t, err)
    99  		r.Header.Set(model.HEADER_AUTH, th.Client.AuthType+" "+th.Client.AuthToken)
   100  
   101  		resp, err = th.Client.HttpClient.Do(r)
   102  		require.NoError(t, err)
   103  		assert.Equal(t, http.StatusFound, resp.StatusCode)
   104  
   105  		// protocol relative URLs should be handled by proxy
   106  		th.App.UpdateConfig(func(cfg *model.Config) {
   107  			cfg.ServiceSettings.SiteURL = model.NewString("http://foo.com")
   108  		})
   109  		r, err = http.NewRequest("GET", th.Client.ApiUrl+"/image?url="+strings.TrimPrefix(imageServer.URL, "http:")+"/image.png", nil)
   110  		require.NoError(t, err)
   111  		r.Header.Set(model.HEADER_AUTH, th.Client.AuthType+" "+th.Client.AuthToken)
   112  
   113  		resp, err = th.Client.HttpClient.Do(r)
   114  		require.NoError(t, err)
   115  		assert.Equal(t, http.StatusOK, resp.StatusCode)
   116  
   117  		// opaque URLs are not supported, should return an error
   118  		r, err = http.NewRequest("GET", th.Client.ApiUrl+"/image?url=mailto:test@example.com", nil)
   119  		require.NoError(t, err)
   120  		r.Header.Set(model.HEADER_AUTH, th.Client.AuthType+" "+th.Client.AuthToken)
   121  
   122  		resp, err = th.Client.HttpClient.Do(r)
   123  		require.NoError(t, err)
   124  		assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
   125  	})
   126  }