// Copyright (c) 2020 Proton Technologies AG
//
// This file is part of ProtonMail Bridge.
//
// ProtonMail Bridge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// ProtonMail Bridge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.

package pmapi

import (
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

const (
	TestDoHQuery       = "dMFYGSLTQOJXXI33ONVQWS3BOMNUA.protonpro.xyz"
	TestQuad9Provider  = "https://dns11.quad9.net/dns-query"
	TestGoogleProvider = "https://dns.google/dns-query"
)

func TestProxyProvider_FindProxy(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	proxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	p.dohLookup = func(q, p string) ([]string, error) { return []string{proxy.URL}, nil }

	url, err := p.findReachableServer()
	require.NoError(t, err)
	require.Equal(t, proxy.URL, url)
}

func TestProxyProvider_FindProxy_ChooseReachableProxy(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	badProxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	goodProxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))

	// Close the bad proxy first so it isn't reachable; we should then choose the good proxy.
	badProxy.Close()
	defer goodProxy.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	p.dohLookup = func(q, p string) ([]string, error) { return []string{badProxy.URL, goodProxy.URL}, nil }

	url, err := p.findReachableServer()
	require.NoError(t, err)
	require.Equal(t, goodProxy.URL, url)
}

func TestProxyProvider_FindProxy_FailIfNoneReachable(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	badProxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	anotherBadProxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))

	// Close the proxies to simulate them not being reachable.
	badProxy.Close()
	anotherBadProxy.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	p.dohLookup = func(q, p string) ([]string, error) { return []string{badProxy.URL, anotherBadProxy.URL}, nil }

	_, err := p.findReachableServer()
	require.Error(t, err)
}

func TestProxyProvider_FindProxy_LookupTimeout(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	proxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	p.lookupTimeout = time.Second
	p.dohLookup = func(q, p string) ([]string, error) { time.Sleep(2 * time.Second); return nil, nil }

	// The findReachableServer should fail because lookup takes 2 seconds but we only allow 1 second.
	_, err := p.findReachableServer()
	require.Error(t, err)
}

func TestProxyProvider_FindProxy_FindTimeout(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	slowProxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(2 * time.Second)
	}))
	defer slowProxy.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	p.findTimeout = time.Second
	p.dohLookup = func(q, p string) ([]string, error) { return []string{slowProxy.URL}, nil }

	// The findReachableServer should fail because lookup takes 2 seconds but we only allow 1 second.
	_, err := p.findReachableServer()
	require.Error(t, err)
}

func TestProxyProvider_UseProxy(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	cm := NewClientManager(testClientConfig)

	proxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	cm.proxyProvider = p

	p.dohLookup = func(q, p string) ([]string, error) { return []string{proxy.URL}, nil }
	url, err := cm.switchToReachableServer()
	require.NoError(t, err)
	require.Equal(t, proxy.URL, url)
	require.Equal(t, proxy.URL, cm.getHost())
}

func TestProxyProvider_UseProxy_MultipleTimes(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	cm := NewClientManager(testClientConfig)

	proxy1 := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy1.Close()
	proxy2 := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy2.Close()
	proxy3 := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy3.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	cm.proxyProvider = p

	p.dohLookup = func(q, p string) ([]string, error) { return []string{proxy1.URL}, nil }
	url, err := cm.switchToReachableServer()
	require.NoError(t, err)
	require.Equal(t, proxy1.URL, url)
	require.Equal(t, proxy1.URL, cm.getHost())

	// Have to wait so as to not get rejected.
	time.Sleep(proxyLookupWait)

	p.dohLookup = func(q, p string) ([]string, error) { return []string{proxy2.URL}, nil }
	url, err = cm.switchToReachableServer()
	require.NoError(t, err)
	require.Equal(t, proxy2.URL, url)
	require.Equal(t, proxy2.URL, cm.getHost())

	// Have to wait so as to not get rejected.
	time.Sleep(proxyLookupWait)

	p.dohLookup = func(q, p string) ([]string, error) { return []string{proxy3.URL}, nil }
	url, err = cm.switchToReachableServer()
	require.NoError(t, err)
	require.Equal(t, proxy3.URL, url)
	require.Equal(t, proxy3.URL, cm.getHost())
}

func TestProxyProvider_UseProxy_RevertAfterTime(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	cm := NewClientManager(testClientConfig)

	proxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	cm.proxyProvider = p
	cm.proxyUseDuration = time.Second

	p.dohLookup = func(q, p string) ([]string, error) { return []string{proxy.URL}, nil }
	url, err := cm.switchToReachableServer()
	require.NoError(t, err)
	require.Equal(t, proxy.URL, url)
	require.Equal(t, proxy.URL, cm.getHost())

	time.Sleep(2 * time.Second)
	require.Equal(t, RootURL, cm.getHost())
}

func TestProxyProvider_UseProxy_RevertIfProxyStopsWorkingAndOriginalAPIIsReachable(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	cm := NewClientManager(testClientConfig)

	proxy := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	cm.proxyProvider = p

	p.dohLookup = func(q, p string) ([]string, error) { return []string{proxy.URL}, nil }
	url, err := cm.switchToReachableServer()
	require.NoError(t, err)
	require.Equal(t, proxy.URL, url)
	require.Equal(t, proxy.URL, cm.getHost())

	// Simulate that the proxy stops working and that the standard api is reachable again.
	proxy.Close()
	unblockAPI()
	time.Sleep(proxyLookupWait)

	// We should now find the original API URL if it is working again.
	url, err = cm.switchToReachableServer()
	require.NoError(t, err)
	require.Equal(t, RootURL, url)
	require.Equal(t, RootURL, cm.getHost())
}

func TestProxyProvider_UseProxy_FindSecondAlternativeIfFirstFailsAndAPIIsStillBlocked(t *testing.T) {
	blockAPI()
	defer unblockAPI()

	cm := NewClientManager(testClientConfig)

	proxy1 := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy1.Close()
	proxy2 := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
	defer proxy2.Close()

	p := newProxyProvider([]string{"not used"}, "not used")
	cm.proxyProvider = p

	// Find a proxy.
	p.dohLookup = func(q, p string) ([]string, error) { return []string{proxy1.URL, proxy2.URL}, nil }
	url, err := cm.switchToReachableServer()
	require.NoError(t, err)
	require.Equal(t, proxy1.URL, url)
	require.Equal(t, proxy1.URL, cm.getHost())

	// Have to wait so as to not get rejected.
	time.Sleep(proxyLookupWait)

	// The proxy stops working and the protonmail API is still blocked.
	proxy1.Close()

	// Should switch to the second proxy because both the first proxy and the protonmail API are blocked.
	url, err = cm.switchToReachableServer()
	require.NoError(t, err)
	require.Equal(t, proxy2.URL, url)
	require.Equal(t, proxy2.URL, cm.getHost())
}

func TestProxyProvider_DoHLookup_Quad9(t *testing.T) {
	p := newProxyProvider([]string{TestQuad9Provider, TestGoogleProvider}, TestDoHQuery)

	records, err := p.dohLookup(TestDoHQuery, TestQuad9Provider)
	require.NoError(t, err)
	require.NotEmpty(t, records)
}

func TestProxyProvider_DoHLookup_Google(t *testing.T) {
	p := newProxyProvider([]string{TestQuad9Provider, TestGoogleProvider}, TestDoHQuery)

	records, err := p.dohLookup(TestDoHQuery, TestGoogleProvider)
	require.NoError(t, err)
	require.NotEmpty(t, records)
}

func TestProxyProvider_DoHLookup_FindProxy(t *testing.T) {
	p := newProxyProvider([]string{TestQuad9Provider, TestGoogleProvider}, TestDoHQuery)

	url, err := p.findReachableServer()
	require.NoError(t, err)
	require.NotEmpty(t, url)
}

func TestProxyProvider_DoHLookup_FindProxyFirstProviderUnreachable(t *testing.T) {
	p := newProxyProvider([]string{"https://unreachable", TestGoogleProvider}, TestDoHQuery)

	url, err := p.findReachableServer()
	require.NoError(t, err)
	require.NotEmpty(t, url)
}

// testAPIURLBackup is used to hold the globalOriginalURL because we clear it for test purposes and need to restore it.
var testAPIURLBackup = RootURL

// blockAPI prevents tests from reaching the standard API, forcing them to find a proxy.
func blockAPI() {
	RootURL = ""
}

// unblockAPI allow tests to reach the standard API again.
func unblockAPI() {
	RootURL = testAPIURLBackup
}