/*
 * Copyright (c) 2023, Psiphon Inc.
 * All rights reserved.
 *
 * This program 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.
 *
 * This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package psiphon

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"regexp"
	"testing"

	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/errors"
	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/inproxy"
	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/parameters"
	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/prng"
	"github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon/common/resolver"
)

func TestInproxyComponents(t *testing.T) {

	// This is a unit test of the in-proxy components internals, such as
	// replay; actual in-proxy broker round trips are exercised in the
	// psiphon/server end-to-end tests.

	err := runInproxyBrokerDialParametersTest()
	if err != nil {
		t.Fatalf(errors.Trace(err).Error())
	}

	err = runInproxySTUNDialParametersTest()
	if err != nil {
		t.Fatalf(errors.Trace(err).Error())
	}

	err = runInproxyNATStateTest()
	if err != nil {
		t.Fatalf(errors.Trace(err).Error())
	}

	// TODO: test inproxyUDPConn multiplexed IPv6Synthesizer
}

func runInproxyBrokerDialParametersTest() error {

	testDataDirName, err := ioutil.TempDir("", "psiphon-inproxy-broker-test")
	if err != nil {
		return errors.Trace(err)
	}
	defer os.RemoveAll(testDataDirName)

	isProxy := false
	propagationChannelID := prng.HexString(8)
	sponsorID := prng.HexString(8)
	networkID := "NETWORK1"
	addressRegex := `[a-z0-9]{5,10}\.example\.org`
	commonCompartmentID, _ := inproxy.MakeID()
	personalCompartmentID, _ := inproxy.MakeID()
	commonCompartmentIDs := []string{commonCompartmentID.String()}
	personalCompartmentIDs := []string{personalCompartmentID.String()}
	privateKey, _ := inproxy.GenerateSessionPrivateKey()
	publicKey, _ := privateKey.GetPublicKey()
	obfuscationSecret, _ := inproxy.GenerateRootObfuscationSecret()
	brokerSpecs := []*parameters.InproxyBrokerSpec{
		{
			BrokerPublicKey:             publicKey.String(),
			BrokerRootObfuscationSecret: obfuscationSecret.String(),
			BrokerFrontingSpecs: []*parameters.FrontingSpec{
				{
					FrontingProviderID: prng.HexString(8),
					Addresses:          []string{addressRegex},
					VerifyServerName:   "example.org",
					Host:               "example.org",
				},
			},
		},
	}
	retainFailed := float64(0.0)

	config := &Config{
		DataRootDirectory:    testDataDirName,
		PropagationChannelId: propagationChannelID,
		SponsorId:            sponsorID,
		NetworkID:            networkID,
	}
	err = config.Commit(false)
	if err != nil {
		return errors.Trace(err)
	}

	err = OpenDataStore(config)
	if err != nil {
		return errors.Trace(err)
	}
	defer CloseDataStore()

	manager := NewInproxyBrokerClientManager(config, isProxy)

	// Test: no broker specs

	_, _, err = manager.GetBrokerClient(networkID)
	if err == nil {
		return errors.TraceNew("unexpected success")
	}

	// Test: select broker and common compartment IDs

	config = &Config{
		DataRootDirectory:           testDataDirName,
		PropagationChannelId:        propagationChannelID,
		SponsorId:                   sponsorID,
		NetworkID:                   networkID,
		InproxyBrokerSpecs:          brokerSpecs,
		InproxyCommonCompartmentIDs: commonCompartmentIDs,
		InproxyReplayBrokerRetainFailedProbability: &retainFailed,
	}
	err = config.Commit(false)
	if err != nil {
		return errors.Trace(err)
	}
	config.SetResolver(resolver.NewResolver(&resolver.NetworkConfig{}, networkID))

	manager = NewInproxyBrokerClientManager(config, isProxy)

	brokerClient, brokerDialParams, err := manager.GetBrokerClient(networkID)
	if err != nil {
		return errors.Trace(err)
	}

	if !regexp.MustCompile(addressRegex).Copy().Match(
		[]byte(brokerDialParams.FrontingDialAddress)) {
		return errors.TraceNew("unexpected FrontingDialAddress")
	}

	if len(brokerClient.GetBrokerDialCoordinator().CommonCompartmentIDs()) != 1 ||
		brokerClient.GetBrokerDialCoordinator().CommonCompartmentIDs()[0].String() !=
			commonCompartmentID.String() {
		return errors.TraceNew("unexpected compartment IDs")
	}

	_ = brokerDialParams.GetMetrics()

	// Test: replay on success

	previousFrontingDialAddress := brokerDialParams.FrontingDialAddress
	previousTLSProfile := brokerDialParams.TLSProfile

	roundTripper, err := brokerClient.GetBrokerDialCoordinator().BrokerClientRoundTripper()
	if err != nil {
		return errors.Trace(err)
	}

	brokerClient.GetBrokerDialCoordinator().BrokerClientRoundTripperSucceeded(roundTripper)

	manager = NewInproxyBrokerClientManager(config, isProxy)

	brokerClient, brokerDialParams, err = manager.GetBrokerClient(networkID)
	if err != nil {
		return errors.Trace(err)
	}

	if !brokerDialParams.isReplay {
		return errors.TraceNew("unexpected non-replay")
	}

	if brokerDialParams.FrontingDialAddress != previousFrontingDialAddress {
		return errors.TraceNew("unexpected replayed FrontingDialAddress")
	}

	if brokerDialParams.TLSProfile != previousTLSProfile {
		return errors.TraceNew("unexpected replayed TLSProfile")
	}

	_ = brokerDialParams.GetMetrics()

	// Test: manager's broker client and dial parameters reinitialized after
	// network ID change

	previousBrokerClient := brokerClient
	previousNetworkID := networkID
	networkID = "NETWORK2"
	config.networkIDGetter = newStaticNetworkGetter(networkID)
	config.SetResolver(resolver.NewResolver(&resolver.NetworkConfig{}, networkID))

	brokerClient, brokerDialParams, err = manager.GetBrokerClient(networkID)
	if err != nil {
		return errors.Trace(err)
	}

	if brokerClient == previousBrokerClient {
		return errors.TraceNew("unexpected brokerClient")
	}

	if brokerDialParams.isReplay {
		return errors.TraceNew("unexpected replay")
	}

	if brokerDialParams.FrontingDialAddress == previousFrontingDialAddress {
		return errors.TraceNew("unexpected non-replayed FrontingDialAddress")
	}

	_ = brokerDialParams.GetMetrics()

	// Test: another replay after switch back to previous network ID

	networkID = previousNetworkID
	config.networkIDGetter = newStaticNetworkGetter(networkID)

	brokerClient, brokerDialParams, err = manager.GetBrokerClient(networkID)
	if err != nil {
		return errors.Trace(err)
	}

	if !brokerDialParams.isReplay {
		return errors.TraceNew("unexpected non-replay")
	}

	if brokerDialParams.FrontingDialAddress != previousFrontingDialAddress {
		return errors.TraceNew("unexpected replayed FrontingDialAddress")
	}

	if brokerDialParams.TLSProfile != previousTLSProfile {
		return errors.TraceNew("unexpected replayed TLSProfile")
	}

	_ = brokerDialParams.GetMetrics()

	// Test: clear replay

	roundTripper, err = brokerClient.GetBrokerDialCoordinator().BrokerClientRoundTripper()
	if err != nil {
		return errors.Trace(err)
	}

	brokerClient.GetBrokerDialCoordinator().BrokerClientRoundTripperFailed(roundTripper)

	manager = NewInproxyBrokerClientManager(config, isProxy)

	brokerClient, brokerDialParams, err = manager.GetBrokerClient(networkID)
	if err != nil {
		return errors.Trace(err)
	}

	if brokerDialParams.isReplay {
		return errors.TraceNew("unexpected replay")
	}

	if brokerDialParams.FrontingDialAddress == previousFrontingDialAddress {
		return errors.TraceNew("unexpected non-replayed FrontingDialAddress")
	}

	_ = brokerDialParams.GetMetrics()

	// Test: no common compartment IDs sent when personal ID is set

	config.InproxyClientPersonalCompartmentIDs = personalCompartmentIDs
	config.InproxyProxyPersonalCompartmentIDs = personalCompartmentIDs

	manager = NewInproxyBrokerClientManager(config, isProxy)

	brokerClient, brokerDialParams, err = manager.GetBrokerClient(networkID)
	if err != nil {
		return errors.Trace(err)
	}

	if len(brokerClient.GetBrokerDialCoordinator().CommonCompartmentIDs()) != 0 ||
		len(brokerClient.GetBrokerDialCoordinator().PersonalCompartmentIDs()) != 1 ||
		brokerClient.GetBrokerDialCoordinator().PersonalCompartmentIDs()[0].String() !=
			personalCompartmentID.String() {
		return errors.TraceNew("unexpected compartment IDs")
	}

	// Test: use persisted common compartment IDs

	config = &Config{
		PropagationChannelId: propagationChannelID,
		SponsorId:            sponsorID,
		NetworkID:            networkID,
	}
	config.InproxyBrokerSpecs = brokerSpecs
	config.InproxyCommonCompartmentIDs = nil
	err = config.Commit(false)
	if err != nil {
		return errors.Trace(err)
	}
	config.SetResolver(resolver.NewResolver(&resolver.NetworkConfig{}, networkID))

	manager = NewInproxyBrokerClientManager(config, isProxy)

	brokerClient, brokerDialParams, err = manager.GetBrokerClient(networkID)
	if err != nil {
		return errors.Trace(err)
	}

	if len(brokerClient.GetBrokerDialCoordinator().CommonCompartmentIDs()) != 1 ||
		brokerClient.GetBrokerDialCoordinator().CommonCompartmentIDs()[0].String() !=
			commonCompartmentID.String() {
		return errors.TraceNew("unexpected compartment IDs")
	}

	_ = brokerDialParams.GetMetrics()

	return nil
}

func runInproxySTUNDialParametersTest() error {

	testDataDirName, err := ioutil.TempDir("", "psiphon-inproxy-stun-test")
	if err != nil {
		return errors.Trace(err)
	}
	defer os.RemoveAll(testDataDirName)

	propagationChannelID := prng.HexString(8)
	sponsorID := prng.HexString(8)
	networkID := "NETWORK1"
	stunServerAddresses := []string{"example.org"}

	config := &Config{
		DataRootDirectory:                 testDataDirName,
		PropagationChannelId:              propagationChannelID,
		SponsorId:                         sponsorID,
		NetworkID:                         networkID,
		InproxySTUNServerAddresses:        stunServerAddresses,
		InproxySTUNServerAddressesRFC5780: stunServerAddresses,
	}
	err = config.Commit(false)
	if err != nil {
		return errors.Trace(err)
	}
	config.SetResolver(resolver.NewResolver(&resolver.NetworkConfig{}, networkID))

	p := config.GetParameters().Get()
	defer p.Close()

	dialParams, err := MakeInproxySTUNDialParameters(config, p, false)
	if err != nil {
		return errors.Trace(err)
	}

	_ = dialParams.GetMetrics()

	dialParamsJSON, err := json.Marshal(dialParams)
	if err != nil {
		return errors.Trace(err)
	}

	var replayDialParams *InproxySTUNDialParameters
	err = json.Unmarshal(dialParamsJSON, &replayDialParams)
	if err != nil {
		return errors.Trace(err)
	}

	replayDialParams.Prepare()

	_ = replayDialParams.GetMetrics()

	return nil
}

func runInproxyNATStateTest() error {

	propagationChannelID := prng.HexString(8)
	sponsorID := prng.HexString(8)
	networkID := "NETWORK1"

	config := &Config{
		PropagationChannelId: propagationChannelID,
		SponsorId:            sponsorID,
		NetworkID:            networkID,
	}
	err := config.Commit(false)
	if err != nil {
		return errors.Trace(err)
	}

	manager := NewInproxyNATStateManager(config)

	// Test: set values stored and cached

	manager.setNATType(networkID, inproxy.NATTypeSymmetric)
	manager.setPortMappingTypes(networkID, inproxy.PortMappingTypes{inproxy.PortMappingTypeUPnP})

	if manager.getNATType(networkID) != inproxy.NATTypeSymmetric {
		return errors.TraceNew("unexpected NAT type")
	}

	portMappingTypes := manager.getPortMappingTypes(networkID)
	if len(portMappingTypes) != 1 || portMappingTypes[0] != inproxy.PortMappingTypeUPnP {
		return errors.TraceNew("unexpected port mapping types")
	}

	// Test: set values ignored when network ID is changing

	otherNetworkID := "NETWORK2"

	manager.setNATType(otherNetworkID, inproxy.NATTypePortRestrictedCone)
	manager.setPortMappingTypes(otherNetworkID, inproxy.PortMappingTypes{inproxy.PortMappingTypePMP})

	if manager.getNATType(networkID) != inproxy.NATTypeSymmetric {
		return errors.TraceNew("unexpected NAT type")
	}

	portMappingTypes = manager.getPortMappingTypes(networkID)
	if len(portMappingTypes) != 1 || portMappingTypes[0] != inproxy.PortMappingTypeUPnP {
		return errors.TraceNew("unexpected port mapping types")
	}

	// Test: reset

	networkID = "NETWORK2"
	config.networkIDGetter = newStaticNetworkGetter(networkID)

	manager.reset()

	if manager.networkID != networkID {
		return errors.TraceNew("unexpected network ID")
	}

	if manager.getNATType(networkID) != inproxy.NATTypeUnknown {
		return errors.TraceNew("unexpected NAT type")
	}

	if len(manager.getPortMappingTypes(networkID)) != 0 {
		return errors.TraceNew("unexpected port mapping types")
	}

	return nil
}