first commit

plugin/idp/oauth2/oauth2.go

@@ -0,0 +1,134 @@
+// Package oauth2 is the plugin for OAuth2 Identity Provider.
+package oauth2
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+
+	"github.com/pkg/errors"
+	"golang.org/x/oauth2"
+
+	"github.com/usememos/memos/plugin/idp"
+	storepb "github.com/usememos/memos/proto/gen/store"
+)
+
+// IdentityProvider represents an OAuth2 Identity Provider.
+type IdentityProvider struct {
+	config *storepb.OAuth2Config
+}
+
+// NewIdentityProvider initializes a new OAuth2 Identity Provider with the given configuration.
+func NewIdentityProvider(config *storepb.OAuth2Config) (*IdentityProvider, error) {
+	for v, field := range map[string]string{
+		config.ClientId:                "clientId",
+		config.ClientSecret:            "clientSecret",
+		config.TokenUrl:                "tokenUrl",
+		config.UserInfoUrl:             "userInfoUrl",
+		config.FieldMapping.Identifier: "fieldMapping.identifier",
+	} {
+		if v == "" {
+			return nil, errors.Errorf(`the field "%s" is empty but required`, field)
+		}
+	}
+
+	return &IdentityProvider{
+		config: config,
+	}, nil
+}
+
+// ExchangeToken returns the exchanged OAuth2 token using the given authorization code.
+// If codeVerifier is provided, it will be used for PKCE (Proof Key for Code Exchange) validation.
+func (p *IdentityProvider) ExchangeToken(ctx context.Context, redirectURL, code, codeVerifier string) (string, error) {
+	conf := &oauth2.Config{
+		ClientID:     p.config.ClientId,
+		ClientSecret: p.config.ClientSecret,
+		RedirectURL:  redirectURL,
+		Scopes:       p.config.Scopes,
+		Endpoint: oauth2.Endpoint{
+			AuthURL:   p.config.AuthUrl,
+			TokenURL:  p.config.TokenUrl,
+			AuthStyle: oauth2.AuthStyleInParams,
+		},
+	}
+
+	// Prepare token exchange options
+	opts := []oauth2.AuthCodeOption{}
+
+	// Add PKCE code_verifier if provided
+	if codeVerifier != "" {
+		opts = append(opts, oauth2.SetAuthURLParam("code_verifier", codeVerifier))
+	}
+
+	token, err := conf.Exchange(ctx, code, opts...)
+	if err != nil {
+		return "", errors.Wrap(err, "failed to exchange access token")
+	}
+
+	// Use the standard AccessToken field instead of Extra()
+	// This is more reliable across different OAuth providers
+	if token.AccessToken == "" {
+		return "", errors.New("missing access token from authorization response")
+	}
+
+	return token.AccessToken, nil
+}
+
+// UserInfo returns the parsed user information using the given OAuth2 token.
+func (p *IdentityProvider) UserInfo(token string) (*idp.IdentityProviderUserInfo, error) {
+	client := &http.Client{}
+	req, err := http.NewRequest(http.MethodGet, p.config.UserInfoUrl, nil)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to new http request")
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to get user information")
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to read response body")
+	}
+
+	var claims map[string]any
+	if err := json.Unmarshal(body, &claims); err != nil {
+		return nil, errors.Wrap(err, "failed to unmarshal response body")
+	}
+	slog.Info("user info claims", "claims", claims)
+	userInfo := &idp.IdentityProviderUserInfo{}
+	if v, ok := claims[p.config.FieldMapping.Identifier].(string); ok {
+		userInfo.Identifier = v
+	}
+	if userInfo.Identifier == "" {
+		return nil, errors.Errorf("the field %q is not found in claims or has empty value", p.config.FieldMapping.Identifier)
+	}
+
+	// Best effort to map optional fields
+	if p.config.FieldMapping.DisplayName != "" {
+		if v, ok := claims[p.config.FieldMapping.DisplayName].(string); ok {
+			userInfo.DisplayName = v
+		}
+	}
+	if userInfo.DisplayName == "" {
+		userInfo.DisplayName = userInfo.Identifier
+	}
+	if p.config.FieldMapping.Email != "" {
+		if v, ok := claims[p.config.FieldMapping.Email].(string); ok {
+			userInfo.Email = v
+		}
+	}
+	if p.config.FieldMapping.AvatarUrl != "" {
+		if v, ok := claims[p.config.FieldMapping.AvatarUrl].(string); ok {
+			userInfo.AvatarURL = v
+		}
+	}
+	slog.Info("user info", "userInfo", userInfo)
+	return userInfo, nil
+}

plugin/idp/oauth2/oauth2_test.go

@@ -0,0 +1,164 @@
+package oauth2
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/usememos/memos/plugin/idp"
+	storepb "github.com/usememos/memos/proto/gen/store"
+)
+
+func TestNewIdentityProvider(t *testing.T) {
+	tests := []struct {
+		name        string
+		config      *storepb.OAuth2Config
+		containsErr string
+	}{
+		{
+			name: "no tokenUrl",
+			config: &storepb.OAuth2Config{
+				ClientId:     "test-client-id",
+				ClientSecret: "test-client-secret",
+				AuthUrl:      "",
+				TokenUrl:     "",
+				UserInfoUrl:  "https://example.com/api/user",
+				FieldMapping: &storepb.FieldMapping{
+					Identifier: "login",
+				},
+			},
+			containsErr: `the field "tokenUrl" is empty but required`,
+		},
+		{
+			name: "no userInfoUrl",
+			config: &storepb.OAuth2Config{
+				ClientId:     "test-client-id",
+				ClientSecret: "test-client-secret",
+				AuthUrl:      "",
+				TokenUrl:     "https://example.com/token",
+				UserInfoUrl:  "",
+				FieldMapping: &storepb.FieldMapping{
+					Identifier: "login",
+				},
+			},
+			containsErr: `the field "userInfoUrl" is empty but required`,
+		},
+		{
+			name: "no field mapping identifier",
+			config: &storepb.OAuth2Config{
+				ClientId:     "test-client-id",
+				ClientSecret: "test-client-secret",
+				AuthUrl:      "",
+				TokenUrl:     "https://example.com/token",
+				UserInfoUrl:  "https://example.com/api/user",
+				FieldMapping: &storepb.FieldMapping{
+					Identifier: "",
+				},
+			},
+			containsErr: `the field "fieldMapping.identifier" is empty but required`,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(*testing.T) {
+			_, err := NewIdentityProvider(test.config)
+			assert.ErrorContains(t, err, test.containsErr)
+		})
+	}
+}
+
+func newMockServer(t *testing.T, code, accessToken string, userinfo []byte) *httptest.Server {
+	mux := http.NewServeMux()
+
+	var rawIDToken string
+	mux.HandleFunc("/oauth2/token", func(w http.ResponseWriter, r *http.Request) {
+		require.Equal(t, http.MethodPost, r.Method)
+
+		body, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		vals, err := url.ParseQuery(string(body))
+		require.NoError(t, err)
+
+		require.Equal(t, code, vals.Get("code"))
+		require.Equal(t, "authorization_code", vals.Get("grant_type"))
+
+		w.Header().Set("Content-Type", "application/json")
+		err = json.NewEncoder(w).Encode(map[string]any{
+			"access_token": accessToken,
+			"token_type":   "Bearer",
+			"expires_in":   3600,
+			"id_token":     rawIDToken,
+		})
+		require.NoError(t, err)
+	})
+	mux.HandleFunc("/oauth2/userinfo", func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_, err := w.Write(userinfo)
+		require.NoError(t, err)
+	})
+
+	s := httptest.NewServer(mux)
+
+	return s
+}
+
+func TestIdentityProvider(t *testing.T) {
+	ctx := context.Background()
+
+	const (
+		testClientID    = "test-client-id"
+		testCode        = "test-code"
+		testAccessToken = "test-access-token"
+		testSubject     = "123456789"
+		testName        = "John Doe"
+		testEmail       = "john.doe@example.com"
+	)
+	userInfo, err := json.Marshal(
+		map[string]any{
+			"sub":   testSubject,
+			"name":  testName,
+			"email": testEmail,
+		},
+	)
+	require.NoError(t, err)
+
+	s := newMockServer(t, testCode, testAccessToken, userInfo)
+
+	oauth2, err := NewIdentityProvider(
+		&storepb.OAuth2Config{
+			ClientId:     testClientID,
+			ClientSecret: "test-client-secret",
+			TokenUrl:     fmt.Sprintf("%s/oauth2/token", s.URL),
+			UserInfoUrl:  fmt.Sprintf("%s/oauth2/userinfo", s.URL),
+			FieldMapping: &storepb.FieldMapping{
+				Identifier:  "sub",
+				DisplayName: "name",
+				Email:       "email",
+			},
+		},
+	)
+	require.NoError(t, err)
+
+	redirectURL := "https://example.com/oauth/callback"
+	// Test without PKCE (backward compatibility)
+	oauthToken, err := oauth2.ExchangeToken(ctx, redirectURL, testCode, "")
+	require.NoError(t, err)
+	require.Equal(t, testAccessToken, oauthToken)
+
+	userInfoResult, err := oauth2.UserInfo(oauthToken)
+	require.NoError(t, err)
+
+	wantUserInfo := &idp.IdentityProviderUserInfo{
+		Identifier:  testSubject,
+		DisplayName: testName,
+		Email:       testEmail,
+	}
+	assert.Equal(t, wantUserInfo, userInfoResult)
+}