netbird/client/internal/auth/oauth_test.go

package auth

import (
	"context"
	"errors"
	"fmt"
	"net/url"
	"runtime"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	mgm "github.com/netbirdio/netbird/shared/management/client"
)

// stubFlow is a minimal OAuthFlow returned by the fake initializers below. requestErr, when set,
// is what its RequestAuthInfo returns, standing in for an IdP that rejects the flow.
type stubFlow struct {
	name       string
	hint       string
	requestErr error
}

func (s *stubFlow) RequestAuthInfo(context.Context) (AuthFlowInfo, error) {
	if s.requestErr != nil {
		return AuthFlowInfo{}, s.requestErr
	}
	return AuthFlowInfo{UserCode: s.name}, nil
}

func (s *stubFlow) WaitToken(context.Context, AuthFlowInfo) (TokenInfo, error) {
	return TokenInfo{}, nil
}

func (s *stubFlow) GetClientID(context.Context) string {
	return ""
}

func (s *stubFlow) SetLoginHint(hint string) {
	s.hint = hint
}

// stubInit returns a flow initializer that yields a named stub flow, or err when err is non-nil.
func stubInit(name string, err error) oauthFlowInit {
	return stubInitFlow(name, err, nil)
}

// stubInitFlow is stubInit with control over what the resulting flow's RequestAuthInfo returns.
func stubInitFlow(name string, initErr, requestErr error) oauthFlowInit {
	return oauthFlowInit{
		name: name,
		init: func(_ *Auth, _ *mgm.GrpcClient, hint string) (OAuthFlow, error) {
			if initErr != nil {
				return nil, initErr
			}
			return &stubFlow{name: name, hint: hint, requestErr: requestErr}, nil
		},
	}
}

// stubAuthFactory hands out an Auth without a management connection, which the stub
// initializers above never touch.
func stubAuthFactory(a *Auth) authFactory {
	return func(context.Context) (*Auth, func(), error) {
		return a, func() {}, nil
	}
}

func TestOAuthFlowWithFallback(t *testing.T) {
	notFound := status.Error(codes.NotFound, "no device authorization flow information available")
	unimplemented := status.Error(codes.Unimplemented, "unknown method")
	incompleteConfig := fmt.Errorf("%w: Client ID value is empty", errFlowNotConfigured)
	unreachable := status.Error(codes.Unavailable, "connection refused")

	tests := []struct {
		name          string
		flows         []oauthFlowInit
		expectedFlow  string
		expectedErr   string
		expectedNoSSO bool
	}{
		{
			name:         "preferred flow is used",
			flows:        []oauthFlowInit{stubInit("device", nil), stubInit("pkce", nil)},
			expectedFlow: "device",
		},
		{
			// the RedHat case: device code flow disabled on management, PKCE configured
			name:         "falls back when preferred flow is not configured",
			flows:        []oauthFlowInit{stubInit("device", notFound), stubInit("pkce", nil)},
			expectedFlow: "pkce",
		},
		{
			name:         "falls back on an incomplete flow configuration",
			flows:        []oauthFlowInit{stubInit("pkce", incompleteConfig), stubInit("device", nil)},
			expectedFlow: "device",
		},
		{
			name:        "does not fall back when the preferred flow fails for another reason",
			flows:       []oauthFlowInit{stubInit("pkce", unreachable), stubInit("device", nil)},
			expectedErr: "connection refused",
		},
		{
			// stays neutral about the remedy: --extend and SSH auth cannot use a setup key
			name:          "neither flow configured reports no SSO provider",
			flows:         []oauthFlowInit{stubInit("device", notFound), stubInit("pkce", notFound)},
			expectedErr:   "no SSO provider configured",
			expectedNoSSO: true,
		},
		{
			name:          "old server without the flow RPCs asks for an update",
			flows:         []oauthFlowInit{stubInit("device", unimplemented), stubInit("pkce", unimplemented)},
			expectedErr:   "does not support SSO providers",
			expectedNoSSO: true,
		},
	}

	mgmURL, err := url.Parse("https://api.netbird.io:443")
	require.NoError(t, err)

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			a := &Auth{mgmURL: mgmURL}
			flow, err := oauthFlowWithFallback(a, nil, tt.flows, "user@example.com", stubAuthFactory(a))

			if tt.expectedErr != "" {
				require.Error(t, err)
				assert.Contains(t, err.Error(), tt.expectedErr)

				var ssoUnavailable *ssoUnavailableError
				assert.Equal(t, tt.expectedNoSSO, errors.As(err, &ssoUnavailable),
					"terminal SSO-unavailable classification mismatch for %v", err)
				return
			}

			require.NoError(t, err)
			stub := activeStub(t, flow)
			assert.Equal(t, tt.expectedFlow, stub.name)
			assert.Equal(t, "user@example.com", stub.hint, "login hint must be passed to the flow")
		})
	}
}

// activeStub unwraps the flow currently in use, which is behind a fallbackFlow whenever an
// untried flow is left.
func activeStub(t *testing.T, flow OAuthFlow) *stubFlow {
	t.Helper()

	if fallback, ok := flow.(*fallbackFlow); ok {
		flow = fallback.current()
	}

	stub, ok := flow.(*stubFlow)
	require.True(t, ok, "unexpected flow type %T", flow)
	return stub
}

// TestFallbackFlowRequestAuthInfo covers the failure the RedHat report hit: management hands out
// a device flow configuration, but the IdP does not serve the grant and only says so when the
// device code is requested.
func TestFallbackFlowRequestAuthInfo(t *testing.T) {
	mgmURL, err := url.Parse("https://api.netbird.io:443")
	require.NoError(t, err)
	a := &Auth{mgmURL: mgmURL}

	idpRejects := fmt.Errorf("%w: request device code returned status 404", errFlowNotConfigured)

	t.Run("swaps in the untried flow", func(t *testing.T) {
		flows := []oauthFlowInit{stubInitFlow("device", nil, idpRejects), stubInit("pkce", nil)}

		flow, err := oauthFlowWithFallback(a, nil, flows, "", stubAuthFactory(a))
		require.NoError(t, err)
		require.Equal(t, "device", activeStub(t, flow).name)

		info, err := flow.RequestAuthInfo(context.Background())
		require.NoError(t, err)
		assert.Equal(t, "pkce", info.UserCode, "the request must be served by the fallback flow")
		assert.Equal(t, "pkce", activeStub(t, flow).name, "the fallback flow must stay active for WaitToken")
	})

	t.Run("keeps the original error when nothing else is configured", func(t *testing.T) {
		flows := []oauthFlowInit{
			stubInitFlow("device", nil, idpRejects),
			stubInit("pkce", status.Error(codes.NotFound, "no pkce authorization flow information available")),
		}

		flow, err := oauthFlowWithFallback(a, nil, flows, "", stubAuthFactory(a))
		require.NoError(t, err)

		_, err = flow.RequestAuthInfo(context.Background())
		require.Error(t, err)
		assert.Contains(t, err.Error(), "status 404")
	})

	t.Run("keeps the original error when the fallback cannot reach management", func(t *testing.T) {
		flows := []oauthFlowInit{stubInitFlow("device", nil, idpRejects), stubInit("pkce", nil)}

		unreachable := func(context.Context) (*Auth, func(), error) {
			return nil, nil, errors.New("connect to management: connection refused")
		}

		flow, err := oauthFlowWithFallback(a, nil, flows, "", unreachable)
		require.NoError(t, err)

		_, err = flow.RequestAuthInfo(context.Background())
		require.Error(t, err)
		assert.Contains(t, err.Error(), "status 404", "the IdP error must survive a failed fallback")
		assert.Equal(t, "device", activeStub(t, flow).name, "a failed fallback must not swap the flow")
	})

	t.Run("does not swap flows on an unrelated failure", func(t *testing.T) {
		flows := []oauthFlowInit{
			stubInitFlow("device", nil, errors.New("timeout talking to the IdP")),
			stubInit("pkce", nil),
		}

		flow, err := oauthFlowWithFallback(a, nil, flows, "", stubAuthFactory(a))
		require.NoError(t, err)

		_, err = flow.RequestAuthInfo(context.Background())
		require.Error(t, err)
		assert.Equal(t, "device", activeStub(t, flow).name, "the preferred flow must stay active")
	})
}

// TestForcedDeviceFlowHasNoFallback covers Android TV and tvOS: a browserless device must get the
// device code error rather than a PKCE flow it can never complete.
func TestForcedDeviceFlowHasNoFallback(t *testing.T) {
	mgmURL, err := url.Parse("https://api.netbird.io:443")
	require.NoError(t, err)
	a := &Auth{mgmURL: mgmURL}

	notFound := status.Error(codes.NotFound, "no device authorization flow information available")

	t.Run("no wrapper when the device flow works", func(t *testing.T) {
		// flowOrder(force) yields this single-entry list, see TestFlowOrder
		forced := []oauthFlowInit{stubInit("device", nil)}

		flow, err := oauthFlowWithFallback(a, nil, forced, "", stubAuthFactory(a))
		require.NoError(t, err)

		_, wrapped := flow.(*fallbackFlow)
		assert.False(t, wrapped, "nothing may swap the flow later on a browserless device")
	})

	t.Run("reports the device flow error instead of falling back", func(t *testing.T) {
		forced := []oauthFlowInit{stubInit("device", notFound)}

		_, err := oauthFlowWithFallback(a, nil, forced, "", stubAuthFactory(a))
		require.Error(t, err)
		assert.True(t, IsSSOUnavailable(err), "the caller must see that SSO is unavailable here")
	})
}

// TestFallbackFlowSetLoginHint covers the Android SDK's pattern: it sets the login hint after the
// flow is built, through a type assertion that the wrapper must satisfy.
func TestFallbackFlowSetLoginHint(t *testing.T) {
	mgmURL, err := url.Parse("https://api.netbird.io:443")
	require.NoError(t, err)
	a := &Auth{mgmURL: mgmURL}

	idpRejects := fmt.Errorf("%w: request device code returned status 404", errFlowNotConfigured)
	flows := []oauthFlowInit{stubInitFlow("device", nil, idpRejects), stubInit("pkce", nil)}

	flow, err := oauthFlowWithFallback(a, nil, flows, "", stubAuthFactory(a))
	require.NoError(t, err)

	setter, ok := flow.(loginHintSetter)
	require.True(t, ok, "the wrapper must accept a login hint like the concrete flows do")
	setter.SetLoginHint("user@example.com")
	assert.Equal(t, "user@example.com", activeStub(t, flow).hint, "the active flow must get the hint")

	// the device flow is rejected by the IdP here, so the hint has to survive into the fallback
	_, err = flow.RequestAuthInfo(context.Background())
	require.NoError(t, err)
	assert.Equal(t, "pkce", activeStub(t, flow).name)
	assert.Equal(t, "user@example.com", activeStub(t, flow).hint, "the fallback flow must get the hint too")
}

func TestWithSetupKeyAdvice(t *testing.T) {
	other := errors.New("connection refused")
	assert.Equal(t, other, WithSetupKeyAdvice(other), "only an SSO-unavailable error gets advice")

	advised := WithSetupKeyAdvice(&ssoUnavailableError{msg: "no SSO provider configured"})
	assert.Contains(t, advised.Error(), "no SSO provider configured", "the original message must survive")
	assert.Contains(t, advised.Error(), "setup key")
	// a setup key cannot re-enrol a peer whose SSO session expired, and the login paths cannot
	// tell that peer apart from an unregistered one, so the advice must state its condition
	assert.Contains(t, advised.Error(), "not enrolled yet")
	assert.True(t, IsSSOUnavailable(advised), "advice must keep the error classifiable")
}

func TestFlowOrder(t *testing.T) {
	graphical := flowOrder(false, true)
	require.Len(t, graphical, 2, "both flows must be attempted when the device has a browser")
	assert.Equal(t, "pkce authorization flow", graphical[0].name)

	headless := flowOrder(false, false)
	require.Len(t, headless, 2)
	if runtime.GOOS == "linux" || runtime.GOOS == "freebsd" {
		assert.Equal(t, "device code flow", headless[0].name, "a headless unix host prefers the device flow")
	}

	// Android TV and tvOS have no browser, so PKCE cannot complete there even from another
	// device: the redirect must reach the loopback listener of the device being enrolled.
	forced := flowOrder(true, false)
	require.Len(t, forced, 1, "a forced device code flow must not fall back to PKCE")
	assert.Equal(t, "device code flow", forced[0].name)
	assert.Len(t, flowOrder(true, true), 1, "force wins over a reported graphical session")
}

func TestPreferDeviceFlow(t *testing.T) {
	isUnix := runtime.GOOS == "linux" || runtime.GOOS == "freebsd"

	assert.Equal(t, isUnix, preferDeviceFlow(false), "headless unix hosts prefer the device flow")
	assert.False(t, preferDeviceFlow(true), "clients with a graphical session prefer PKCE")
}