diff --git a/pkg/connector/connector.go b/pkg/connector/connector.go new file mode 100644 index 000000000..fde3d43f0 --- /dev/null +++ b/pkg/connector/connector.go @@ -0,0 +1,68 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package connector + +import ( + "context" + "encoding/json" + "fmt" + "net/http" +) + +type ( + ProtocolType string + + Connector interface { + Initiate(ctx context.Context, connectorID string, organizationID string, r *http.Request) (string, error) + Complete(ctx context.Context, connectorID string, organizationID string, r *http.Request) (Connection, error) + } + + Connection interface { + Type() ProtocolType + Client(ctx context.Context) (*http.Client, error) + + json.Unmarshaler + json.Marshaler + } +) + +const ( + ProtocolOAuth2 ProtocolType = "oauth2" +) + +func UnmarshalConnection(data []byte) (Connection, error) { + var typeContainer struct { + Type string `json:"type"` + } + + if err := json.Unmarshal(data, &typeContainer); err != nil { + return nil, fmt.Errorf("failed to unmarshal connection type: %w", err) + } + + var conn Connection + + switch ProtocolType(typeContainer.Type) { + case ProtocolOAuth2: + conn = &OAuth2Connection{} + default: + return nil, fmt.Errorf("unknown connection type: %s", typeContainer.Type) + } + + if err := conn.UnmarshalJSON(data); err != nil { + return nil, err + } + + return conn, nil +} diff --git a/pkg/connector/oauth2.go b/pkg/connector/oauth2.go new file mode 100644 index 000000000..bf38b02f2 --- /dev/null +++ b/pkg/connector/oauth2.go @@ -0,0 +1,226 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package connector + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/url" + "strings" + "time" + + "github.com/getprobo/probo/pkg/statelesstoken" +) + +type ( + OAuth2Connector struct { + ClientID string + ClientSecret string + RedirectURI string + Scopes []string + AuthURL string + TokenURL string + } + + OAuth2State struct { + OrganizationID string `json:"oid"` + ConnectorID string `json:"cid"` + } + + OAuth2Connection struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token,omitempty"` + ExpiresAt time.Time `json:"expires_at"` + TokenType string `json:"token_type"` + Scope string `json:"scope,omitempty"` + } +) + +var ( + _ Connector = (*OAuth2Connector)(nil) + _ Connection = (*OAuth2Connection)(nil) + + OAuth2TokenType = "probo/connector/oauth2" + OAuth2TokenTTL = 10 * time.Minute +) + +func (c *OAuth2Connector) Initiate(ctx context.Context, connectorID string, organizationID string, r *http.Request) (string, error) { + stateData := OAuth2State{OrganizationID: organizationID, ConnectorID: connectorID} + state, err := statelesstoken.NewToken("", OAuth2TokenType, OAuth2TokenTTL, stateData) + if err != nil { + return "", fmt.Errorf("cannot create state token: %w", err) + } + + redirectURI, err := url.Parse(c.RedirectURI) + if err != nil { + return "", fmt.Errorf("cannot parse redirect URI: %w", err) + } + + redirectQuery := url.Values{} + redirectQuery.Set("organization_id", organizationID) + redirectQuery.Set("connector_id", connectorID) + + redirectURI.RawQuery = redirectQuery.Encode() + + authCodeQuery := url.Values{} + authCodeQuery.Set("state", state) + authCodeQuery.Set("client_id", c.ClientID) + authCodeQuery.Set("redirect_uri", redirectURI.String()) + authCodeQuery.Set("response_type", "code") + authCodeQuery.Set("scope", strings.Join(c.Scopes, " ")) + + u, err := url.Parse(c.AuthURL) + if err != nil { + return "", fmt.Errorf("cannot parse auth URL: %w", err) + } + + u.RawQuery = authCodeQuery.Encode() + + return u.String(), nil +} + +func (c *OAuth2Connector) Complete(ctx context.Context, connectorID string, organizationID string, r *http.Request) (Connection, error) { + code := r.URL.Query().Get("code") + if code == "" { + return nil, fmt.Errorf("no code in request") + } + + state := r.URL.Query().Get("state") + if state == "" { + return nil, fmt.Errorf("no state in request") + } + + payload, err := statelesstoken.ValidateToken[OAuth2State]("", OAuth2TokenType, state) + if err != nil { + return nil, fmt.Errorf("cannot validate state token: %w", err) + } + + if payload.Data.OrganizationID != organizationID { + return nil, fmt.Errorf("invalid organization ID") + } + + if payload.Data.ConnectorID != connectorID { + return nil, fmt.Errorf("invalid connector ID") + } + + redirectURI, err := url.Parse(c.RedirectURI) + if err != nil { + return nil, fmt.Errorf("cannot parse redirect URI: %w", err) + } + + redirectQuery := url.Values{} + redirectQuery.Set("organization_id", organizationID) + redirectQuery.Set("connector_id", connectorID) + + redirectURI.RawQuery = redirectQuery.Encode() + + tokenRequestData := url.Values{} + tokenRequestData.Set("client_id", c.ClientID) + tokenRequestData.Set("client_secret", c.ClientSecret) + tokenRequestData.Set("code", code) + tokenRequestData.Set("redirect_uri", redirectURI.String()) + tokenRequestData.Set("grant_type", "authorization_code") + + tokenRequest, err := http.NewRequestWithContext(ctx, "POST", c.TokenURL, strings.NewReader(tokenRequestData.Encode())) + if err != nil { + return nil, fmt.Errorf("cannot create token request: %w", err) + } + + tokenRequest.Header.Set("Content-Type", "application/x-www-form-urlencoded; charset=utf-8") + tokenRequest.Header.Set("Accept", "application/json") + tokenRequest.Header.Set("User-Agent", "Probo Connector") + + tokenResp, err := http.DefaultClient.Do(tokenRequest) + if err != nil { + return nil, fmt.Errorf("cannot post token URL: %w", err) + } + + defer tokenResp.Body.Close() + + type tokenResponse struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token"` + ExpiresAt time.Time `json:"expires_at"` + Scope string `json:"scope"` + TokenType string `json:"token_type"` + } + + var token tokenResponse + err = json.NewDecoder(tokenResp.Body).Decode(&token) + if err != nil { + return nil, fmt.Errorf("cannot decode token response: %w", err) + } + + fmt.Printf("\n\n\ntoken: %+v\n\n\n", token) + + return &OAuth2Connection{ + AccessToken: token.AccessToken, + RefreshToken: token.RefreshToken, + ExpiresAt: token.ExpiresAt, + Scope: token.Scope, + TokenType: token.TokenType, + }, nil +} + +func (c *OAuth2Connection) Type() ProtocolType { + return ProtocolOAuth2 +} + +func (c OAuth2Connection) Client(ctx context.Context) (*http.Client, error) { + client := &http.Client{ + Transport: &oauth2Transport{ + token: c.AccessToken, + tokenType: c.TokenType, + underlying: http.DefaultTransport, + }, + } + return client, nil +} + +func (c OAuth2Connection) MarshalJSON() ([]byte, error) { + type Alias OAuth2Connection + return json.Marshal(&struct { + Type string `json:"type"` + Alias + }{ + Type: string(ProtocolOAuth2), + Alias: Alias(c), + }) +} + +func (c *OAuth2Connection) UnmarshalJSON(data []byte) error { + type Alias OAuth2Connection + aux := &struct { + *Alias + }{ + Alias: (*Alias)(c), + } + return json.Unmarshal(data, &aux) +} + +// OAuth transport for adding authorization header +type oauth2Transport struct { + token string + tokenType string + underlying http.RoundTripper +} + +func (t *oauth2Transport) RoundTrip(req *http.Request) (*http.Response, error) { + req2 := req.Clone(req.Context()) + req2.Header.Set("Authorization", t.tokenType+" "+t.token) + return t.underlying.RoundTrip(req2) +} diff --git a/pkg/connector/registry.go b/pkg/connector/registry.go new file mode 100644 index 000000000..28b6982ef --- /dev/null +++ b/pkg/connector/registry.go @@ -0,0 +1,69 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package connector + +import ( + "context" + "fmt" + "net/http" + "sync" +) + +type ( + ConnectorRegistry struct { + sync.RWMutex + connectors map[string]Connector + } +) + +func NewConnectorRegistry() *ConnectorRegistry { + return &ConnectorRegistry{ + connectors: make(map[string]Connector), + } +} + +func (cr *ConnectorRegistry) Register(connectorID string, connector Connector) { + cr.Lock() + defer cr.Unlock() + cr.connectors[connectorID] = connector +} + +func (cr *ConnectorRegistry) Get(connectorID string) (Connector, error) { + cr.RLock() + defer cr.RUnlock() + connector, ok := cr.connectors[connectorID] + if !ok { + return nil, fmt.Errorf("connector %q not found", connectorID) + } + return connector, nil +} + +func (cr *ConnectorRegistry) Initiate(ctx context.Context, connectorID string, organizationID string, r *http.Request) (string, error) { + connector, err := cr.Get(connectorID) + if err != nil { + return "", fmt.Errorf("cannot initiate connector: %w", err) + } + + return connector.Initiate(ctx, connectorID, organizationID, r) +} + +func (cr *ConnectorRegistry) Complete(ctx context.Context, connectorID string, organizationID string, r *http.Request) (Connection, error) { + connector, err := cr.Get(connectorID) + if err != nil { + return nil, fmt.Errorf("cannot complete connector: %w", err) + } + + return connector.Complete(ctx, connectorID, organizationID, r) +}