Skip to content

Commit 3b45d11

Browse files
vademecumukvademecumuk
committed
Added AWS cognito provider, based on the OKTA provider.
1 parent b5cb5a6 commit 3b45d11

File tree

4 files changed

+381
-0
lines changed

4 files changed

+381
-0
lines changed

‎providers/cognito/cognito.go‎

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
package cognito
2+
3+
import (
4+
"bytes"
5+
"encoding/json"
6+
"fmt"
7+
"io"
8+
"net/http"
9+
"net/url"
10+
11+
"github.com/markbates/goth"
12+
"golang.org/x/oauth2"
13+
)
14+
15+
// Provider is the implementation of `goth.Provider` for accessing AWS Cognito.
16+
// This is based upon the implementation for okta
17+
type Provider struct {
18+
ClientKey string
19+
Secret string
20+
CallbackURL string
21+
HTTPClient *http.Client
22+
config *oauth2.Config
23+
providerName string
24+
issuerURL string
25+
profileURL string
26+
}
27+
28+
// New creates a new AWS Cognito provider and sets up important connection details.
29+
// You should always call `cognito.New` to get a new provider. Never try to
30+
// create one manually.
31+
func New(clientID, secret, baseUrl, callbackURL string, scopes ...string) *Provider {
32+
issuerURL, _ := url.JoinPath(baseUrl, "/oauth2/default")
33+
authURL, _ := url.JoinPath(baseUrl, "/oauth2/authorize")
34+
tokenURL, _ := url.JoinPath(baseUrl, "/oauth2/token")
35+
profileURL, _ := url.JoinPath(baseUrl, "/oauth2/userInfo")
36+
return NewCustomisedURL(clientID, secret, callbackURL, authURL, tokenURL, issuerURL, profileURL, scopes...)
37+
}
38+
39+
// NewCustomisedURL is similar to New(...) but can be used to set custom URLs to connect to
40+
func NewCustomisedURL(clientID, secret, callbackURL, authURL, tokenURL, issuerURL, profileURL string, scopes ...string) *Provider {
41+
p := &Provider{
42+
ClientKey: clientID,
43+
Secret: secret,
44+
CallbackURL: callbackURL,
45+
providerName: "aws",
46+
issuerURL: issuerURL,
47+
profileURL: profileURL,
48+
}
49+
p.config = newConfig(p, authURL, tokenURL, scopes)
50+
return p
51+
}
52+
53+
// Name is the name used to retrieve this provider later.
54+
func (p *Provider) Name() string {
55+
return p.providerName
56+
}
57+
58+
// SetName is to update the name of the provider (needed in case of multiple providers of 1 type)
59+
func (p *Provider) SetName(name string) {
60+
p.providerName = name
61+
}
62+
63+
func (p *Provider) Client() *http.Client {
64+
return goth.HTTPClientWithFallBack(p.HTTPClient)
65+
}
66+
67+
// Debug is a no-op for the aws package.
68+
func (p *Provider) Debug(debug bool) {
69+
if debug {
70+
fmt.Println("WARNING: Debug request for goth/providers/cognito but no debug is available")
71+
}
72+
}
73+
74+
// BeginAuth asks AWS for an authentication end-point.
75+
func (p *Provider) BeginAuth(state string) (goth.Session, error) {
76+
return &Session{
77+
AuthURL: p.config.AuthCodeURL(state),
78+
}, nil
79+
}
80+
81+
// FetchUser will go to aws and access basic information about the user.
82+
func (p *Provider) FetchUser(session goth.Session) (goth.User, error) {
83+
sess := session.(*Session)
84+
user := goth.User{
85+
AccessToken: sess.AccessToken,
86+
Provider: p.Name(),
87+
RefreshToken: sess.RefreshToken,
88+
ExpiresAt: sess.ExpiresAt,
89+
UserID: sess.UserID,
90+
}
91+
92+
if user.AccessToken == "" {
93+
// data is not yet retrieved since accessToken is still empty
94+
return user, fmt.Errorf("%s cannot get user information without accessToken", p.providerName)
95+
}
96+
97+
req, err := http.NewRequest("GET", p.profileURL, nil)
98+
if err != nil {
99+
return user, err
100+
}
101+
req.Header.Set("Authorization", "Bearer "+sess.AccessToken)
102+
response, err := p.Client().Do(req)
103+
if err != nil {
104+
if response != nil {
105+
_ = response.Body.Close()
106+
}
107+
return user, err
108+
}
109+
defer func(Body io.ReadCloser) {
110+
_ = Body.Close()
111+
}(response.Body)
112+
113+
if response.StatusCode != http.StatusOK {
114+
return user, fmt.Errorf("%s responded with a %d trying to fetch user information", p.providerName, response.StatusCode)
115+
}
116+
117+
bits, err := io.ReadAll(response.Body)
118+
if err != nil {
119+
return user, err
120+
}
121+
122+
err = json.NewDecoder(bytes.NewReader(bits)).Decode(&user.RawData)
123+
if err != nil {
124+
return user, err
125+
}
126+
127+
err = userFromReader(bytes.NewReader(bits), &user)
128+
129+
return user, err
130+
}
131+
132+
func newConfig(provider *Provider, authURL, tokenURL string, scopes []string) *oauth2.Config {
133+
c := &oauth2.Config{
134+
ClientID: provider.ClientKey,
135+
ClientSecret: provider.Secret,
136+
RedirectURL: provider.CallbackURL,
137+
Endpoint: oauth2.Endpoint{
138+
AuthURL: authURL,
139+
TokenURL: tokenURL,
140+
},
141+
Scopes: []string{},
142+
}
143+
144+
if len(scopes) > 0 {
145+
for _, scope := range scopes {
146+
c.Scopes = append(c.Scopes, scope)
147+
}
148+
}
149+
return c
150+
}
151+
152+
func userFromReader(r io.Reader, user *goth.User) error {
153+
u := struct {
154+
Name string `json:"name"`
155+
Email string `json:"email"`
156+
FirstName string `json:"given_name"`
157+
LastName string `json:"family_name"`
158+
NickName string `json:"nickname"`
159+
ID string `json:"sub"`
160+
Locale string `json:"locale"`
161+
ProfileURL string `json:"profile"`
162+
Username string `json:"preferred_username"`
163+
Zoneinfo string `json:"zoneinfo"`
164+
}{}
165+
166+
err := json.NewDecoder(r).Decode(&u)
167+
if err != nil {
168+
return err
169+
}
170+
171+
rd := make(map[string]interface{})
172+
rd["ProfileURL"] = u.ProfileURL
173+
rd["Locale"] = u.Locale
174+
rd["Username"] = u.Username
175+
rd["Zoneinfo"] = u.Zoneinfo
176+
177+
user.UserID = u.ID
178+
user.Email = u.Email
179+
user.Name = u.Name
180+
user.NickName = u.NickName
181+
user.FirstName = u.FirstName
182+
user.LastName = u.LastName
183+
184+
user.RawData = rd
185+
186+
return nil
187+
}
188+
189+
// RefreshTokenAvailable refresh token is provided by auth provider or not
190+
func (p *Provider) RefreshTokenAvailable() bool {
191+
return true
192+
}
193+
194+
// RefreshToken get new access token based on the refresh token
195+
func (p *Provider) RefreshToken(refreshToken string) (*oauth2.Token, error) {
196+
token := &oauth2.Token{RefreshToken: refreshToken}
197+
ts := p.config.TokenSource(goth.ContextForClient(p.Client()), token)
198+
newToken, err := ts.Token()
199+
if err != nil {
200+
return nil, err
201+
}
202+
return newToken, err
203+
}

‎providers/cognito/cognito_test.go‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
package cognito
2+
3+
import (
4+
"os"
5+
"testing"
6+
7+
"github.com/markbates/goth"
8+
"github.com/markbates/goth/providers/okta"
9+
"github.com/stretchr/testify/assert"
10+
)
11+
12+
func Test_New(t *testing.T) {
13+
t.Parallel()
14+
a := assert.New(t)
15+
p := provider()
16+
17+
a.Equal(p.ClientKey, os.Getenv("COGNITO_ID"))
18+
a.Equal(p.Secret, os.Getenv("COGNITO_SECRET"))
19+
a.Equal(p.CallbackURL, "/foo")
20+
}
21+
22+
func Test_NewCustomisedURL(t *testing.T) {
23+
t.Parallel()
24+
a := assert.New(t)
25+
p := urlCustomisedURLProvider()
26+
session, err := p.BeginAuth("test_state")
27+
s := session.(*okta.Session)
28+
a.NoError(err)
29+
a.Contains(s.AuthURL, "http://authURL")
30+
}
31+
32+
func Test_Implements_Provider(t *testing.T) {
33+
t.Parallel()
34+
a := assert.New(t)
35+
a.Implements((*goth.Provider)(nil), provider())
36+
}
37+
38+
func Test_BeginAuth(t *testing.T) {
39+
t.Parallel()
40+
a := assert.New(t)
41+
p := provider()
42+
session, err := p.BeginAuth("test_state")
43+
s := session.(*okta.Session)
44+
a.NoError(err)
45+
a.Contains(s.AuthURL, os.Getenv("COGNITO_ISSUER_URL"))
46+
}
47+
48+
func Test_SessionFromJSON(t *testing.T) {
49+
t.Parallel()
50+
a := assert.New(t)
51+
52+
p := provider()
53+
session, err := p.UnmarshalSession(`{"AuthURL":"` + os.Getenv("COGNITO_ISSUER_URL") + `/oauth2/authorize", "AccessToken":"1234567890"}`)
54+
a.NoError(err)
55+
56+
s := session.(*okta.Session)
57+
a.Equal(s.AuthURL, os.Getenv("COGNITO_ISSUER_URL")+"/oauth2/authorize")
58+
a.Equal(s.AccessToken, "1234567890")
59+
}
60+
61+
func provider() *okta.Provider {
62+
return okta.New(os.Getenv("COGNITO_ID"), os.Getenv("COGNITO_SECRET"), os.Getenv("COGNITO_ISSUER_URL"), "/foo")
63+
}
64+
65+
func urlCustomisedURLProvider() *okta.Provider {
66+
return okta.NewCustomisedURL(os.Getenv("CLIENT_ID"), os.Getenv("CLIENT_SECRET"), "/foo", "http://authURL", "http://tokenURL", "http://issuerURL", "http://profileURL")
67+
}

‎providers/cognito/session.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
package cognito
2+
3+
import (
4+
"encoding/json"
5+
"errors"
6+
"strings"
7+
"time"
8+
9+
"github.com/markbates/goth"
10+
)
11+
12+
// Session stores data during the auth process with AWS Cognito.
13+
type Session struct {
14+
AuthURL string
15+
AccessToken string
16+
RefreshToken string
17+
ExpiresAt time.Time
18+
UserID string
19+
}
20+
21+
var _ goth.Session = &Session{}
22+
23+
// GetAuthURL will return the URL set by calling the `BeginAuth` function on the AWS Cognito provider.
24+
func (s Session) GetAuthURL() (string, error) {
25+
if s.AuthURL == "" {
26+
return "", errors.New(goth.NoAuthUrlErrorMessage)
27+
}
28+
return s.AuthURL, nil
29+
}
30+
31+
// Authorize the session with cognito and return the access token to be stored for future use.
32+
func (s *Session) Authorize(provider goth.Provider, params goth.Params) (string, error) {
33+
p := provider.(*Provider)
34+
token, err := p.config.Exchange(goth.ContextForClient(p.Client()), params.Get("code"))
35+
if err != nil {
36+
return "", err
37+
}
38+
39+
if !token.Valid() {
40+
return "", errors.New("invalid token received from provider")
41+
}
42+
43+
s.AccessToken = token.AccessToken
44+
s.RefreshToken = token.RefreshToken
45+
s.ExpiresAt = token.Expiry
46+
return token.AccessToken, err
47+
}
48+
49+
// Marshal the session into a string
50+
func (s Session) Marshal() string {
51+
b, _ := json.Marshal(s)
52+
return string(b)
53+
}
54+
55+
func (s Session) String() string {
56+
return s.Marshal()
57+
}
58+
59+
// UnmarshalSession wil unmarshal a JSON string into a session.
60+
func (p *Provider) UnmarshalSession(data string) (goth.Session, error) {
61+
s := &Session{}
62+
err := json.NewDecoder(strings.NewReader(data)).Decode(s)
63+
return s, err
64+
}

‎providers/cognito/session_test.go‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
package cognito
2+
3+
import (
4+
"testing"
5+
6+
"github.com/markbates/goth"
7+
"github.com/stretchr/testify/assert"
8+
)
9+
10+
func Test_Implements_Session(t *testing.T) {
11+
t.Parallel()
12+
a := assert.New(t)
13+
s := &Session{}
14+
15+
a.Implements((*goth.Session)(nil), s)
16+
}
17+
18+
func Test_GetAuthURL(t *testing.T) {
19+
t.Parallel()
20+
a := assert.New(t)
21+
s := &Session{}
22+
23+
_, err := s.GetAuthURL()
24+
a.Error(err)
25+
26+
s.AuthURL = "/foo"
27+
28+
url, _ := s.GetAuthURL()
29+
a.Equal(url, "/foo")
30+
}
31+
32+
func Test_ToJSON(t *testing.T) {
33+
t.Parallel()
34+
a := assert.New(t)
35+
s := &Session{}
36+
37+
data := s.Marshal()
38+
a.Equal(data, `{"AuthURL":"","AccessToken":"","RefreshToken":"","ExpiresAt":"0001-01-01T00:00:00Z","UserID":""}`)
39+
}
40+
41+
func Test_String(t *testing.T) {
42+
t.Parallel()
43+
a := assert.New(t)
44+
s := &Session{}
45+
46+
a.Equal(s.String(), s.Marshal())
47+
}

0 commit comments

Comments
 (0)