/*
Copyright 2021 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package webclient

import (
	"context"
	"encoding/json"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/stretchr/testify/require"

	"github.com/gravitational/teleport/api/defaults"
)

func newPingHandler(path string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		if req.RequestURI != path {
			w.WriteHeader(http.StatusNotFound)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(PingResponse{ServerVersion: "test"})
	})
}

func TestPlainHttpFallback(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		desc            string
		handler         http.Handler
		actionUnderTest func(addr string, insecure bool) error
	}{
		{
			desc:    "Ping",
			handler: newPingHandler("/webapi/ping"),
			actionUnderTest: func(addr string, insecure bool) error {
				_, err := Ping(
					&Config{Context: context.Background(), ProxyAddr: addr, Insecure: insecure})
				return err
			},
		}, {
			desc:    "Find",
			handler: newPingHandler("/webapi/find"),
			actionUnderTest: func(addr string, insecure bool) error {
				_, err := Find(&Config{Context: context.Background(), ProxyAddr: addr, Insecure: insecure})
				return err
			},
		},
	}

	for _, testCase := range testCases {
		t.Run(testCase.desc, func(t *testing.T) {
			t.Run("Allowed on insecure & loopback", func(t *testing.T) {
				httpSvr := httptest.NewServer(testCase.handler)
				defer httpSvr.Close()

				err := testCase.actionUnderTest(httpSvr.Listener.Addr().String(), true /* insecure */)
				require.NoError(t, err)
			})

			t.Run("Denied on secure", func(t *testing.T) {
				httpSvr := httptest.NewServer(testCase.handler)
				defer httpSvr.Close()

				err := testCase.actionUnderTest(httpSvr.Listener.Addr().String(), false /* secure */)
				require.Error(t, err)
			})

			t.Run("Denied on non-loopback", func(t *testing.T) {
				nonLoopbackSvr := httptest.NewUnstartedServer(testCase.handler)

				// replace the test-supplied loopback listener with the first available
				// non-loopback address
				nonLoopbackSvr.Listener.Close()
				l, err := net.Listen("tcp", "0.0.0.0:0")
				require.NoError(t, err)
				nonLoopbackSvr.Listener = l
				nonLoopbackSvr.Start()
				defer nonLoopbackSvr.Close()

				err = testCase.actionUnderTest(nonLoopbackSvr.Listener.Addr().String(), true /* insecure */)
				require.Error(t, err)
			})
		})
	}
}

func TestTunnelAddr(t *testing.T) {
	cases := []struct {
		name               string
		settings           ProxySettings
		expectedTunnelAddr string
		setup              func(t *testing.T)
	}{
		{
			name: "should use TunnelPublicAddr",
			settings: ProxySettings{
				SSH: SSHProxySettings{
					TunnelPublicAddr: "tunnel.example.com:4024",
					PublicAddr:       "public.example.com",
					SSHPublicAddr:    "ssh.example.com",
					TunnelListenAddr: "[::]:5024",
					WebListenAddr:    "proxy.example.com",
				},
			},
			expectedTunnelAddr: "tunnel.example.com:4024",
		},
		{
			name: "should use SSHPublicAddr and TunnelListenAddr",
			settings: ProxySettings{
				SSH: SSHProxySettings{
					SSHPublicAddr:    "ssh.example.com",
					PublicAddr:       "public.example.com",
					TunnelListenAddr: "[::]:5024",
					WebListenAddr:    "proxy.example.com",
				},
			},
			expectedTunnelAddr: "ssh.example.com:5024",
		},
		{
			name: "should use PublicAddr and TunnelListenAddr",
			settings: ProxySettings{
				SSH: SSHProxySettings{
					PublicAddr:       "public.example.com",
					TunnelListenAddr: "[::]:5024",
					WebListenAddr:    "proxy.example.com",
				},
			},
			expectedTunnelAddr: "public.example.com:5024",
		},
		{
			name: "should use PublicAddr and SSHProxyTunnelListenPort",
			settings: ProxySettings{
				SSH: SSHProxySettings{
					PublicAddr:    "public.example.com",
					WebListenAddr: "proxy.example.com",
				},
			},
			expectedTunnelAddr: "public.example.com:3024",
		},
		{
			name: "should use WebListenAddr and SSHProxyTunnelListenPort",
			settings: ProxySettings{
				SSH: SSHProxySettings{
					WebListenAddr: "proxy.example.com",
				},
			},
			expectedTunnelAddr: "proxy.example.com:3024",
		},
		{
			name: "should use PublicAddr with ProxyWebPort if TLSRoutingEnabled was enabled",
			settings: ProxySettings{
				SSH: SSHProxySettings{
					PublicAddr:       "public.example.com",
					TunnelListenAddr: "[::]:5024",
					TunnelPublicAddr: "tpa.example.com:3032",
					WebListenAddr:    "proxy.example.com:443",
				},
				TLSRoutingEnabled: true,
			},
			expectedTunnelAddr: "public.example.com:443",
		},
		{
			name: "should use PublicAddr with custom port if TLSRoutingEnabled was enabled",
			settings: ProxySettings{
				SSH: SSHProxySettings{
					PublicAddr:       "public.example.com:443",
					TunnelListenAddr: "[::]:5024",
					TunnelPublicAddr: "tpa.example.com:3032",
					WebListenAddr:    "proxy.example.com:443",
				},
				TLSRoutingEnabled: true,
			},
			expectedTunnelAddr: "public.example.com:443",
		},
		{
			name: "should use WebListenAddr with custom ProxyWebPort if TLSRoutingEnabled was enabled",
			settings: ProxySettings{
				SSH: SSHProxySettings{
					TunnelListenAddr: "[::]:5024",
					TunnelPublicAddr: "tpa.example.com:3032",
					WebListenAddr:    "proxy.example.com:443",
				},
				TLSRoutingEnabled: true,
			},
			expectedTunnelAddr: "proxy.example.com:443",
		},
		{
			name: "should use WebListenAddr with default https port if TLSRoutingEnabled was enabled",
			settings: ProxySettings{
				SSH: SSHProxySettings{
					TunnelListenAddr: "[::]:5024",
					TunnelPublicAddr: "tpa.example.com:3032",
					WebListenAddr:    "proxy.example.com",
				},
				TLSRoutingEnabled: true,
			},
			expectedTunnelAddr: "proxy.example.com:443",
		},
		{
			name:               "TELEPORT_TUNNEL_PUBLIC_ADDR overrides tunnel address",
			settings:           ProxySettings{},
			expectedTunnelAddr: "tunnel.example.com:4024",
			setup: func(t *testing.T) {
				t.Setenv(defaults.TunnelPublicAddrEnvar, "tunnel.example.com:4024")
			},
		},
	}

	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			if tt.setup != nil {
				tt.setup(t)
			}
			tunnelAddr, err := tt.settings.TunnelAddr()
			require.NoError(t, err)
			require.Equal(t, tt.expectedTunnelAddr, tunnelAddr)
		})
	}
}

func TestParse(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		addr     string
		hostPort string
		host     string
		port     int
	}{
		{
			addr:     "example.com",
			hostPort: "example.com",
			host:     "example.com",
			port:     0,
		}, {
			addr:     "example.com:443",
			hostPort: "example.com:443",
			host:     "example.com",
			port:     443,
		}, {
			addr:     "http://example.com:443",
			hostPort: "example.com:443",
			host:     "example.com",
			port:     443,
		}, {
			addr:     "https://example.com:443",
			hostPort: "example.com:443",
			host:     "example.com",
			port:     443,
		}, {
			addr:     "tcp://example.com:443",
			hostPort: "example.com:443",
			host:     "example.com",
			port:     443,
		}, {
			addr:     "file://host/path",
			hostPort: "",
			host:     "",
			port:     0,
		}, {
			addr:     "[::]:443",
			hostPort: "[::]:443",
			host:     "::",
			port:     443,
		}, {
			addr:     "https://example.com:443/path?query=query#fragment",
			hostPort: "example.com:443",
			host:     "example.com",
			port:     443,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.addr, func(t *testing.T) {
			hostPort, err := parseAndJoinHostPort(tc.addr)
			if tc.hostPort == "" {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
				require.Equal(t, tc.hostPort, hostPort)
			}

			host, _, err := ParseHostPort(tc.addr)
			if tc.host == "" {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
				require.Equal(t, tc.host, host)
			}

			port, err := parsePort(tc.addr)
			if tc.port == 0 {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
				require.Equal(t, tc.port, port)
			}
		})
	}
}

func TestNewWebClientRespectHTTPProxy(t *testing.T) {
	t.Setenv("HTTPS_PROXY", "fakeproxy.example.com:9999")
	client, err := newWebClient(&Config{
		Context:   context.Background(),
		ProxyAddr: "localhost:3080",
	})
	require.NoError(t, err)
	//nolint:bodyclose // resp should be nil, so there will be no body to close.
	resp, err := client.Get("https://fakedomain.example.com")
	// Client should try to proxy through nonexistent server at localhost.
	require.Error(t, err, "GET unexpectedly succeeded: %+v", resp)
	require.Contains(t, err.Error(), "proxyconnect")
	require.Contains(t, err.Error(), "lookup fakeproxy.example.com")
	require.Contains(t, err.Error(), "no such host")
}

func TestNewWebClientNoProxy(t *testing.T) {
	t.Setenv("HTTPS_PROXY", "fakeproxy.example.com:9999")
	t.Setenv("NO_PROXY", "fakedomain.example.com")
	client, err := newWebClient(&Config{
		Context:   context.Background(),
		ProxyAddr: "localhost:3080",
	})
	require.NoError(t, err)
	//nolint:bodyclose // resp should be nil, so there will be no body to close.
	resp, err := client.Get("https://fakedomain.example.com")
	require.Error(t, err, "GET unexpectedly succeeded: %+v", resp)
	require.NotContains(t, err.Error(), "proxyconnect")
	require.Contains(t, err.Error(), "lookup fakedomain.example.com")
	require.Contains(t, err.Error(), "no such host")
}

func TestSSHProxyHostPort(t *testing.T) {
	t.Parallel()

	tests := []struct {
		testName        string
		inProxySettings ProxySettings
		outHost         string
		outPort         string
	}{
		{
			testName: "TLS routing enabled, web public addr",
			inProxySettings: ProxySettings{
				SSH: SSHProxySettings{
					PublicAddr:    "proxy.example.com:443",
					WebListenAddr: "127.0.0.1:3080",
				},
				TLSRoutingEnabled: true,
			},
			outHost: "proxy.example.com",
			outPort: "443",
		},
		{
			testName: "TLS routing enabled, web public addr with listen addr",
			inProxySettings: ProxySettings{
				SSH: SSHProxySettings{
					PublicAddr:    "proxy.example.com",
					WebListenAddr: "127.0.0.1:443",
				},
				TLSRoutingEnabled: true,
			},
			outHost: "proxy.example.com",
			outPort: "443",
		},
		{
			testName: "TLS routing enabled, web listen addr",
			inProxySettings: ProxySettings{
				SSH: SSHProxySettings{
					WebListenAddr: "127.0.0.1:3080",
				},
				TLSRoutingEnabled: true,
			},
			outHost: "127.0.0.1",
			outPort: "3080",
		},
		{
			testName: "TLS routing disabled, SSH public addr",
			inProxySettings: ProxySettings{
				SSH: SSHProxySettings{
					SSHPublicAddr: "ssh.example.com:3023",
					PublicAddr:    "proxy.example.com:443",
					ListenAddr:    "127.0.0.1:3023",
				},
				TLSRoutingEnabled: false,
			},
			outHost: "ssh.example.com",
			outPort: "3023",
		},
		{
			testName: "TLS routing disabled, web public addr",
			inProxySettings: ProxySettings{
				SSH: SSHProxySettings{
					PublicAddr: "proxy.example.com:443",
					ListenAddr: "127.0.0.1:3023",
				},
				TLSRoutingEnabled: false,
			},
			outHost: "proxy.example.com",
			outPort: "3023",
		},
		{
			testName: "TLS routing disabled, SSH listen addr",
			inProxySettings: ProxySettings{
				SSH: SSHProxySettings{
					ListenAddr: "127.0.0.1:3023",
				},
				TLSRoutingEnabled: false,
			},
			outHost: "127.0.0.1",
			outPort: "3023",
		},
	}
	for _, test := range tests {
		t.Run(test.testName, func(t *testing.T) {
			host, port, err := test.inProxySettings.SSHProxyHostPort()
			require.NoError(t, err)
			require.Equal(t, test.outHost, host)
			require.Equal(t, test.outPort, port)
		})
	}
}

// TestWebClientClosesIdleConnections verifies that all http connections
// are closed when the http.Client created by newWebClient is no longer
// being used.
func TestWebClientClosesIdleConnections(t *testing.T) {
	expectedResponse := &PingResponse{
		Proxy: ProxySettings{
			TLSRoutingEnabled: true,
		},
		ServerVersion:    "1.2.3",
		MinClientVersion: "0.1.2",
		ClusterName:      "test",
	}

	expectedStates := []http.ConnState{
		http.StateNew, http.StateActive, http.StateClosed, // the https request will fail and cause us to fallback to http
		http.StateNew, http.StateActive, http.StateIdle, http.StateClosed, // the http request should be processed and closed
	}

	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/webapi/find":
			json.NewEncoder(w).Encode(expectedResponse)
		default:
			w.WriteHeader(http.StatusBadRequest)
		}
	}))

	stateChange := make(chan http.ConnState, len(expectedStates))
	srv.Config.ConnState = func(conn net.Conn, state http.ConnState) {
		stateChange <- state
	}

	srv.Start()
	t.Cleanup(srv.Close)

	resp, err := Find(&Config{
		Context:   context.Background(),
		ProxyAddr: strings.TrimPrefix(srv.URL, "http://"),
		Insecure:  true,
	})
	require.NoError(t, err)
	require.Empty(t, cmp.Diff(expectedResponse, resp))

	for _, expected := range expectedStates {
		select {
		case state := <-stateChange:
			require.Equal(t, expected, state, "expected connection state %s got %s", expected.String(), state.String())
		case <-time.After(3 * time.Second):
			t.Fatalf("timeout waiting for expected connection state %s", expected.String())
		}
	}

}