feat(backend): add kiro account support
This commit is contained in:
nianzs
@@ -0,0 +1,369 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
kiropkg "github.com/Wei-Shaw/sub2api/internal/pkg/kiro"
|
||||
)
|
||||
|
||||
const (
|
||||
// Kiro desktop social auth uses localhost loopback callbacks from a fixed
|
||||
// allowlist. Use one of the bundled ports from the official client.
|
||||
kiroSocialRedirectURI = "http://localhost:49153"
|
||||
// AWS IAM Identity Center native/public clients require an explicit loopback IP redirect URI.
|
||||
kiroIDCRedirectURI = "http://127.0.0.1:9876/oauth/callback"
|
||||
)
|
||||
|
||||
type KiroOAuthService struct {
|
||||
sessionStore *kiropkg.SessionStore
|
||||
proxyRepo ProxyRepository
|
||||
}
|
||||
|
||||
func NewKiroOAuthService(proxyRepo ProxyRepository) *KiroOAuthService {
|
||||
return &KiroOAuthService{
|
||||
sessionStore: kiropkg.NewSessionStore(),
|
||||
proxyRepo: proxyRepo,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *KiroOAuthService) Stop() {}
|
||||
|
||||
type KiroAuthURLResult struct {
|
||||
AuthURL string `json:"auth_url"`
|
||||
SessionID string `json:"session_id"`
|
||||
State string `json:"state"`
|
||||
}
|
||||
|
||||
type KiroIDCAuthURLResult struct {
|
||||
AuthURL string `json:"auth_url"`
|
||||
SessionID string `json:"session_id"`
|
||||
State string `json:"state"`
|
||||
ClientID string `json:"client_id"`
|
||||
Region string `json:"region"`
|
||||
StartURL string `json:"start_url"`
|
||||
}
|
||||
|
||||
type KiroTokenInfo struct {
|
||||
AccessToken string `json:"access_token,omitempty"`
|
||||
RefreshToken string `json:"refresh_token,omitempty"`
|
||||
ProfileArn string `json:"profile_arn,omitempty"`
|
||||
ExpiresAt string `json:"expires_at,omitempty"`
|
||||
AuthMethod string `json:"auth_method,omitempty"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
ClientID string `json:"client_id,omitempty"`
|
||||
ClientSecret string `json:"client_secret,omitempty"`
|
||||
ClientIDHash string `json:"client_id_hash,omitempty"`
|
||||
Email string `json:"email,omitempty"`
|
||||
StartURL string `json:"start_url,omitempty"`
|
||||
Region string `json:"region,omitempty"`
|
||||
}
|
||||
|
||||
type KiroGenerateAuthURLInput struct {
|
||||
ProxyID *int64
|
||||
Provider string
|
||||
}
|
||||
|
||||
type KiroExchangeCodeInput struct {
|
||||
SessionID string
|
||||
State string
|
||||
Code string
|
||||
CallbackPath string
|
||||
LoginOption string
|
||||
ProxyID *int64
|
||||
}
|
||||
|
||||
type KiroGenerateIDCAuthURLInput struct {
|
||||
ProxyID *int64
|
||||
StartURL string
|
||||
Region string
|
||||
}
|
||||
|
||||
type KiroRefreshTokenInput struct {
|
||||
RefreshToken string
|
||||
AuthMethod string
|
||||
Provider string
|
||||
ClientID string
|
||||
ClientSecret string
|
||||
StartURL string
|
||||
Region string
|
||||
ProfileArn string
|
||||
ProxyID *int64
|
||||
}
|
||||
|
||||
type KiroImportTokenInput struct {
|
||||
TokenJSON string
|
||||
DeviceRegistrationJSON string
|
||||
}
|
||||
|
||||
func (s *KiroOAuthService) GenerateAuthURL(ctx context.Context, input *KiroGenerateAuthURLInput) (*KiroAuthURLResult, error) {
|
||||
provider := strings.TrimSpace(input.Provider)
|
||||
if provider == "" {
|
||||
provider = string(kiropkg.SocialProviderGoogle)
|
||||
}
|
||||
if provider != string(kiropkg.SocialProviderGoogle) && provider != string(kiropkg.SocialProviderGitHub) {
|
||||
return nil, fmt.Errorf("unsupported kiro social provider: %s", provider)
|
||||
}
|
||||
state, err := kiropkg.GenerateState()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate state failed: %w", err)
|
||||
}
|
||||
codeVerifier, err := kiropkg.GenerateCodeVerifier()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate code verifier failed: %w", err)
|
||||
}
|
||||
sessionID := kiropkg.GenerateSessionID()
|
||||
proxyURL, _ := s.resolveProxyURL(ctx, input.ProxyID)
|
||||
s.sessionStore.Set(sessionID, &kiropkg.AuthSession{
|
||||
State: state,
|
||||
CodeVerifier: codeVerifier,
|
||||
ProxyURL: proxyURL,
|
||||
CreatedAt: time.Now(),
|
||||
AuthType: "social",
|
||||
Provider: provider,
|
||||
RedirectURI: kiroSocialRedirectURI,
|
||||
})
|
||||
return &KiroAuthURLResult{
|
||||
AuthURL: kiropkg.BuildSocialSignInURL(kiroSocialRedirectURI, kiropkg.GenerateCodeChallenge(codeVerifier), state),
|
||||
SessionID: sessionID,
|
||||
State: state,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *KiroOAuthService) ExchangeCode(ctx context.Context, input *KiroExchangeCodeInput) (*KiroTokenInfo, error) {
|
||||
session, ok := s.sessionStore.Get(input.SessionID)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("session not found or expired")
|
||||
}
|
||||
if strings.TrimSpace(input.State) == "" || input.State != session.State {
|
||||
return nil, fmt.Errorf("state invalid")
|
||||
}
|
||||
proxyURL := session.ProxyURL
|
||||
if input.ProxyID != nil {
|
||||
proxyURL, _ = s.resolveProxyURL(ctx, input.ProxyID)
|
||||
}
|
||||
|
||||
switch session.AuthType {
|
||||
case "social":
|
||||
token, err := kiropkg.CreateSocialToken(
|
||||
ctx,
|
||||
proxyURL,
|
||||
input.Code,
|
||||
session.CodeVerifier,
|
||||
buildKiroSocialExchangeRedirectURI(session.RedirectURI, session.Provider, input.CallbackPath, input.LoginOption),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
token.Provider = session.Provider
|
||||
s.sessionStore.Delete(input.SessionID)
|
||||
return toKiroTokenInfo(token), nil
|
||||
case "idc":
|
||||
token, err := kiropkg.ExchangeIDCAuthCode(ctx, proxyURL, session.ClientID, session.ClientSecret, input.Code, session.CodeVerifier, session.RedirectURI, session.Region, session.StartURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s.sessionStore.Delete(input.SessionID)
|
||||
return toKiroTokenInfo(token), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported auth session type: %s", session.AuthType)
|
||||
}
|
||||
}
|
||||
|
||||
func buildKiroSocialExchangeRedirectURI(baseRedirectURI, provider, callbackPath, loginOption string) string {
|
||||
option := strings.ToLower(strings.TrimSpace(loginOption))
|
||||
if option == "" {
|
||||
switch provider {
|
||||
case string(kiropkg.SocialProviderGitHub):
|
||||
option = "github"
|
||||
case string(kiropkg.SocialProviderGoogle):
|
||||
option = "google"
|
||||
}
|
||||
}
|
||||
return kiropkg.BuildSocialTokenRedirectURI(baseRedirectURI, callbackPath, option)
|
||||
}
|
||||
|
||||
func (s *KiroOAuthService) GenerateIDCAuthURL(ctx context.Context, input *KiroGenerateIDCAuthURLInput) (*KiroIDCAuthURLResult, error) {
|
||||
startURL := strings.TrimSpace(input.StartURL)
|
||||
if startURL == "" {
|
||||
startURL = kiropkg.BuilderIDStartURL
|
||||
}
|
||||
region := strings.TrimSpace(input.Region)
|
||||
if region == "" {
|
||||
region = "us-east-1"
|
||||
}
|
||||
state, err := kiropkg.GenerateState()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate state failed: %w", err)
|
||||
}
|
||||
codeVerifier, err := kiropkg.GenerateCodeVerifier()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate code verifier failed: %w", err)
|
||||
}
|
||||
proxyURL, _ := s.resolveProxyURL(ctx, input.ProxyID)
|
||||
reg, err := kiropkg.RegisterIDCClient(ctx, proxyURL, kiroIDCRedirectURI, startURL, region)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sessionID := kiropkg.GenerateSessionID()
|
||||
s.sessionStore.Set(sessionID, &kiropkg.AuthSession{
|
||||
State: state,
|
||||
CodeVerifier: codeVerifier,
|
||||
ProxyURL: proxyURL,
|
||||
CreatedAt: time.Now(),
|
||||
AuthType: "idc",
|
||||
RedirectURI: kiroIDCRedirectURI,
|
||||
ClientID: reg.ClientID,
|
||||
ClientSecret: reg.ClientSecret,
|
||||
Region: region,
|
||||
StartURL: startURL,
|
||||
})
|
||||
return &KiroIDCAuthURLResult{
|
||||
AuthURL: kiropkg.BuildIDCAuthURL(reg.ClientID, kiroIDCRedirectURI, state, kiropkg.GenerateCodeChallenge(codeVerifier), region),
|
||||
SessionID: sessionID,
|
||||
State: state,
|
||||
ClientID: reg.ClientID,
|
||||
Region: region,
|
||||
StartURL: startURL,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *KiroOAuthService) RefreshToken(ctx context.Context, input *KiroRefreshTokenInput) (*KiroTokenInfo, error) {
|
||||
proxyURL, _ := s.resolveProxyURL(ctx, input.ProxyID)
|
||||
authMethod := strings.ToLower(strings.TrimSpace(input.AuthMethod))
|
||||
if authMethod == "" {
|
||||
authMethod = "social"
|
||||
}
|
||||
|
||||
var token *kiropkg.TokenData
|
||||
var err error
|
||||
switch authMethod {
|
||||
case "idc":
|
||||
token, err = kiropkg.RefreshIDCToken(ctx, proxyURL, input.ClientID, input.ClientSecret, input.RefreshToken, input.Region, input.StartURL)
|
||||
default:
|
||||
token, err = kiropkg.RefreshSocialToken(ctx, proxyURL, input.RefreshToken, input.Provider)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if token.ProfileArn == "" {
|
||||
token.ProfileArn = input.ProfileArn
|
||||
}
|
||||
if token.ClientID == "" {
|
||||
token.ClientID = input.ClientID
|
||||
}
|
||||
if token.ClientSecret == "" {
|
||||
token.ClientSecret = input.ClientSecret
|
||||
}
|
||||
if token.StartURL == "" {
|
||||
token.StartURL = input.StartURL
|
||||
}
|
||||
if token.Region == "" {
|
||||
token.Region = input.Region
|
||||
}
|
||||
return toKiroTokenInfo(token), nil
|
||||
}
|
||||
|
||||
func (s *KiroOAuthService) RefreshAccountToken(ctx context.Context, account *Account) (*KiroTokenInfo, error) {
|
||||
if account.Platform != PlatformKiro || account.Type != AccountTypeOAuth {
|
||||
return nil, fmt.Errorf("not a kiro oauth account")
|
||||
}
|
||||
return s.RefreshToken(ctx, &KiroRefreshTokenInput{
|
||||
RefreshToken: account.GetCredential("refresh_token"),
|
||||
AuthMethod: account.GetCredential("auth_method"),
|
||||
Provider: account.GetCredential("provider"),
|
||||
ClientID: account.GetCredential("client_id"),
|
||||
ClientSecret: account.GetCredential("client_secret"),
|
||||
StartURL: account.GetCredential("start_url"),
|
||||
Region: account.GetCredential("region"),
|
||||
ProfileArn: account.GetCredential("profile_arn"),
|
||||
ProxyID: account.ProxyID,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *KiroOAuthService) ImportToken(input *KiroImportTokenInput) (*KiroTokenInfo, error) {
|
||||
token, err := kiropkg.ParseImportedToken(input.TokenJSON, input.DeviceRegistrationJSON)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return toKiroTokenInfo(token), nil
|
||||
}
|
||||
|
||||
func (s *KiroOAuthService) BuildAccountCredentials(tokenInfo *KiroTokenInfo) map[string]any {
|
||||
if tokenInfo == nil {
|
||||
return map[string]any{}
|
||||
}
|
||||
|
||||
creds := map[string]any{}
|
||||
if tokenInfo.AccessToken != "" {
|
||||
creds["access_token"] = tokenInfo.AccessToken
|
||||
}
|
||||
if tokenInfo.RefreshToken != "" {
|
||||
creds["refresh_token"] = tokenInfo.RefreshToken
|
||||
}
|
||||
if tokenInfo.ProfileArn != "" {
|
||||
creds["profile_arn"] = tokenInfo.ProfileArn
|
||||
}
|
||||
if tokenInfo.ExpiresAt != "" {
|
||||
creds["expires_at"] = tokenInfo.ExpiresAt
|
||||
}
|
||||
if tokenInfo.AuthMethod != "" {
|
||||
creds["auth_method"] = tokenInfo.AuthMethod
|
||||
}
|
||||
if tokenInfo.Provider != "" {
|
||||
creds["provider"] = tokenInfo.Provider
|
||||
}
|
||||
if tokenInfo.ClientID != "" {
|
||||
creds["client_id"] = tokenInfo.ClientID
|
||||
}
|
||||
if tokenInfo.ClientSecret != "" {
|
||||
creds["client_secret"] = tokenInfo.ClientSecret
|
||||
}
|
||||
if tokenInfo.ClientIDHash != "" {
|
||||
creds["client_id_hash"] = tokenInfo.ClientIDHash
|
||||
}
|
||||
if tokenInfo.Email != "" {
|
||||
creds["email"] = tokenInfo.Email
|
||||
}
|
||||
if tokenInfo.StartURL != "" {
|
||||
creds["start_url"] = tokenInfo.StartURL
|
||||
}
|
||||
if tokenInfo.Region != "" {
|
||||
creds["region"] = tokenInfo.Region
|
||||
}
|
||||
|
||||
return creds
|
||||
}
|
||||
|
||||
func toKiroTokenInfo(token *kiropkg.TokenData) *KiroTokenInfo {
|
||||
if token == nil {
|
||||
return nil
|
||||
}
|
||||
return &KiroTokenInfo{
|
||||
AccessToken: token.AccessToken,
|
||||
RefreshToken: token.RefreshToken,
|
||||
ProfileArn: token.ProfileArn,
|
||||
ExpiresAt: token.ExpiresAt,
|
||||
AuthMethod: token.AuthMethod,
|
||||
Provider: token.Provider,
|
||||
ClientID: token.ClientID,
|
||||
ClientSecret: token.ClientSecret,
|
||||
ClientIDHash: token.ClientIDHash,
|
||||
Email: token.Email,
|
||||
StartURL: token.StartURL,
|
||||
Region: token.Region,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *KiroOAuthService) resolveProxyURL(ctx context.Context, proxyID *int64) (string, error) {
|
||||
if proxyID == nil || s.proxyRepo == nil {
|
||||
return "", nil
|
||||
}
|
||||
proxy, err := s.proxyRepo.GetByID(ctx, *proxyID)
|
||||
if err != nil || proxy == nil {
|
||||
return "", err
|
||||
}
|
||||
return proxy.URL(), nil
|
||||
}
|
||||
Reference in New Issue
Block a user