package httpfetcher

import (
	"context"
	"errors"
	"net"
	"os"
	"syscall"
	"testing"
	"time"
)

// TestFetchFailedBeforeAnswer checks the error of a fetch whose connection
// fails:
//   - refused: ErrUpstreamRequestFailed.
//   - timed out: ErrUpstreamTimeout, still matching the dial's error, even for
//     a timeout error that does not match context.DeadlineExceeded, such as
//     the TLS handshake timeout.
//   - cancelled by its caller: context.Canceled, not ErrUpstreamRequestFailed.
func TestFetchFailedBeforeAnswer(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name       string
		dial       func(cancel context.CancelFunc) error
		wantErrs   []error
		wantNotErr error
	}{
		{
			name: "connection refused",
			dial: func(context.CancelFunc) error {
				return syscall.ECONNREFUSED
			},
			wantErrs:   []error{ErrUpstreamRequestFailed},
			wantNotErr: context.Canceled,
		},
		{
			name: "timed out while connecting",
			dial: func(context.CancelFunc) error {
				return os.ErrDeadlineExceeded
			},
			wantErrs:   []error{ErrUpstreamTimeout, os.ErrDeadlineExceeded},
			wantNotErr: ErrUpstreamRequestFailed,
		},
		{
			name: "cancelled while connecting",
			dial: func(cancel context.CancelFunc) error {
				cancel()

				return context.Canceled
			},
			wantErrs:   []error{context.Canceled},
			wantNotErr: ErrUpstreamRequestFailed,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			ctx, cancel := context.WithCancel(testContext(t))
			defer cancel()

			f := New(&Config{
				Timeout:               30 * time.Second,
				MaxResponseSize:       50 << 20,
				AllowHTTP:             true,
				MaxConnectionsPerHost: 20,
				MaxConnections:        64,
				DialContext: func(context.Context, string, string) (net.Conn, error) {
					return nil, tt.dial(cancel)
				},
			})

			_, err := f.Fetch(ctx, upstreamSource("/image"), "", "")
			t.Logf("Fetch() error: %v", err)

			for _, wantErr := range tt.wantErrs {
				if !errors.Is(err, wantErr) {
					t.Errorf("Fetch() error = %v, want %v", err, wantErr)
				}
			}

			if errors.Is(err, tt.wantNotErr) {
				t.Errorf("Fetch() error = %v, must not be %v", err, tt.wantNotErr)
			}
		})
	}
}

// TestFetchAddressLookupTimedOut checks that a fetch whose address lookup
// times out fails with ErrUpstreamTimeout, both for the lookup of the source
// host and for the lookup while connecting. The error must match neither
// ErrInvalidHost nor ErrUpstreamRequestFailed, which the image routes answer
// 502. A context whose deadline has passed makes the lookup time out.
func TestFetchAddressLookupTimedOut(t *testing.T) {
	t.Parallel()

	checkTimedOut := func(t *testing.T, err error) {
		t.Helper()
		t.Logf("Fetch() error: %v", err)

		if !errors.Is(err, ErrUpstreamTimeout) {
			t.Errorf("Fetch() error = %v, want ErrUpstreamTimeout", err)
		}

		if errors.Is(err, ErrInvalidHost) || errors.Is(err, ErrUpstreamRequestFailed) {
			t.Errorf("Fetch() error = %v, must match neither ErrInvalidHost "+
				"nor ErrUpstreamRequestFailed", err)
		}
	}

	t.Run("lookup of the source host", func(t *testing.T) {
		t.Parallel()

		ctx, cancel := context.WithTimeout(testContext(t), 0)
		defer cancel()

		f := New(&Config{
			Timeout:               30 * time.Second,
			MaxResponseSize:       50 << 20,
			MaxConnectionsPerHost: 20,
			MaxConnections:        64,
		})

		_, err := f.Fetch(ctx, imageSourceOn("cdn.example.com"), "", "")
		checkTimedOut(t, err)
	})

	t.Run("lookup while connecting", func(t *testing.T) {
		t.Parallel()

		// The source host is an address, which validateURL does not look up;
		// the dialer looks up a name
		f := New(&Config{
			Timeout:               30 * time.Second,
			MaxResponseSize:       50 << 20,
			MaxConnectionsPerHost: 20,
			MaxConnections:        64,
			DialContext: func(ctx context.Context, network, _ string) (net.Conn, error) {
				ctx, cancel := context.WithTimeout(ctx, 0)
				defer cancel()

				return dialSSRFSafe(ctx, network, "img.example.com:443", nil)
			},
		})

		_, err := f.Fetch(testContext(t), upstreamSource("/image"), "", "")
		checkTimedOut(t, err)
	})
}

// TestFetchConnectionTimedOut checks that a fetch whose connection in
// dialSSRFSafe times out fails with ErrUpstreamTimeout, which the image
// routes answer 504. The error must match neither ErrUpstreamRequestFailed
// nor ErrInvalidHost, which they answer 502. A context whose deadline has
// passed makes the connection time out without touching the network.
func TestFetchConnectionTimedOut(t *testing.T) {
	t.Parallel()

	f := New(&Config{
		Timeout:               30 * time.Second,
		MaxResponseSize:       50 << 20,
		MaxConnectionsPerHost: 20,
		MaxConnections:        64,
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			ctx, cancel := context.WithTimeout(ctx, 0)
			defer cancel()

			return dialSSRFSafe(ctx, network, addr, nil)
		},
	})

	_, err := f.Fetch(testContext(t), upstreamSource("/image"), "", "")
	t.Logf("Fetch() error: %v", err)

	if !errors.Is(err, ErrUpstreamTimeout) {
		t.Errorf("Fetch() error = %v, want ErrUpstreamTimeout", err)
	}

	if errors.Is(err, ErrUpstreamRequestFailed) || errors.Is(err, ErrInvalidHost) {
		t.Errorf("Fetch() error = %v, must match neither "+
			"ErrUpstreamRequestFailed nor ErrInvalidHost", err)
	}
}

// TestFetchAddressLookupCancelled checks that a fetch whose address lookup
// its caller cancelled fails with an error matching context.Canceled, not
// ErrInvalidHost, which the image routes answer 502.
func TestFetchAddressLookupCancelled(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithCancel(testContext(t))
	cancel()

	f := New(&Config{
		Timeout:               30 * time.Second,
		MaxResponseSize:       50 << 20,
		MaxConnectionsPerHost: 20,
		MaxConnections:        64,
	})

	_, err := f.Fetch(ctx, imageSourceOn("cdn.example.com"), "", "")
	t.Logf("Fetch() error: %v", err)

	if !errors.Is(err, context.Canceled) {
		t.Errorf("Fetch() error = %v, want context.Canceled", err)
	}

	if errors.Is(err, ErrInvalidHost) {
		t.Errorf("Fetch() error = %v, must not be ErrInvalidHost", err)
	}
}