|
1 | 1 | package mongo
|
2 | 2 |
|
3 | 3 | import (
|
4 |
| - "github.com/globalsign/mgo" |
5 |
| - "gopkg.in/oauth2.v3" |
6 |
| - "gopkg.in/oauth2.v3/models" |
| 4 | + "context" |
| 5 | + "errors" |
| 6 | + "log" |
| 7 | + "time" |
| 8 | + |
| 9 | + "github.com/go-oauth2/oauth2/v4" |
| 10 | + "github.com/go-oauth2/oauth2/v4/models" |
| 11 | + "go.mongodb.org/mongo-driver/bson" |
| 12 | + "go.mongodb.org/mongo-driver/mongo" |
| 13 | + "go.mongodb.org/mongo-driver/mongo/options" |
7 | 14 | )
|
8 | 15 |
|
| 16 | +// TODO |
| 17 | +// add retry mecanism on all requests |
| 18 | +// if retry fail ping the db |
| 19 | +// if ping fail crash the service |
| 20 | +// that should be optional as a service's restart mecanism should be implemented |
| 21 | + |
| 22 | +// StoreConfig hold configs common to all Configs(ClientConfig, TokenConfig) |
| 23 | +type StoreConfig struct { |
| 24 | + db string |
| 25 | + service string |
| 26 | + connectionTimeout int |
| 27 | + requestTimeout int |
| 28 | + isReplicaSet bool |
| 29 | +} |
| 30 | + |
| 31 | +func NewStoreConfig(ctout, rtout int) *StoreConfig { |
| 32 | + return &StoreConfig{ |
| 33 | + connectionTimeout: ctout, |
| 34 | + requestTimeout: rtout, |
| 35 | + } |
| 36 | +} |
| 37 | + |
| 38 | +func NewDefaultStoreConfig(db, service string, isReplicasSet bool) *StoreConfig { |
| 39 | + return &StoreConfig{ |
| 40 | + db: db, |
| 41 | + service: service, |
| 42 | + connectionTimeout: 0, |
| 43 | + requestTimeout: 0, |
| 44 | + isReplicaSet: isReplicasSet, |
| 45 | + } |
| 46 | +} |
| 47 | + |
| 48 | +// setRequestContext set a WithTimeout or Background context |
| 49 | +func (sc *StoreConfig) setRequestContext() (context.Context, context.CancelFunc) { |
| 50 | + ctx := context.Background() |
| 51 | + if sc.requestTimeout > 0 { |
| 52 | + log.Println("Request timeout: ", sc.requestTimeout) |
| 53 | + timeout := time.Duration(sc.requestTimeout) * time.Second |
| 54 | + return context.WithTimeout(ctx, timeout) |
| 55 | + } |
| 56 | + return nil, func() {} |
| 57 | +} |
| 58 | + |
| 59 | +// setTransactionCreateContext is specific to the transaction(if not a replicaSet) |
| 60 | +func (sc *StoreConfig) setTransactionCreateContext() (context.Context, context.CancelFunc) { |
| 61 | + ctx := context.Background() |
| 62 | + if sc.requestTimeout > 0 { |
| 63 | + // at max TransactionCreate run 9 requests |
| 64 | + timeout := time.Duration(sc.requestTimeout*9) * time.Second |
| 65 | + return context.WithTimeout(ctx, timeout) |
| 66 | + } |
| 67 | + return nil, func() {} |
| 68 | +} |
| 69 | + |
9 | 70 | // ClientConfig client configuration parameters
|
10 | 71 | type ClientConfig struct {
|
11 | 72 | // store clients data collection name(The default is oauth2_clients)
|
12 | 73 | ClientsCName string
|
| 74 | + storeConfig *StoreConfig |
13 | 75 | }
|
14 | 76 |
|
15 | 77 | // NewDefaultClientConfig create a default client configuration
|
16 |
| -func NewDefaultClientConfig() *ClientConfig { |
| 78 | +func NewDefaultClientConfig(strCfgs *StoreConfig) *ClientConfig { |
17 | 79 | return &ClientConfig{
|
18 | 80 | ClientsCName: "oauth2_clients",
|
| 81 | + storeConfig: strCfgs, |
19 | 82 | }
|
20 | 83 | }
|
21 | 84 |
|
22 | 85 | // NewClientStore create a client store instance based on mongodb
|
23 |
| -func NewClientStore(cfg *Config, ccfgs ...*ClientConfig) *ClientStore { |
24 |
| - session, err := mgo.Dial(cfg.URL) |
| 86 | +func NewClientStore(cfg *Config, scfgs ...*StoreConfig) *ClientStore { |
| 87 | + clientOptions := options.Client().ApplyURI(cfg.URL) |
| 88 | + ctx := context.TODO() |
| 89 | + ctxPing := context.TODO() |
| 90 | + |
| 91 | + if len(scfgs) > 0 && scfgs[0].connectionTimeout > 0 { |
| 92 | + newCtx, cancel := context.WithTimeout(context.Background(), time.Duration(scfgs[0].connectionTimeout)*time.Second) |
| 93 | + ctx = newCtx |
| 94 | + defer cancel() |
| 95 | + clientOptions.SetConnectTimeout(time.Duration(scfgs[0].connectionTimeout) * time.Second) |
| 96 | + } |
| 97 | + |
| 98 | + if len(scfgs) > 0 && scfgs[0].requestTimeout > 0 { |
| 99 | + newCtx, cancel := context.WithTimeout(context.Background(), time.Duration(scfgs[0].requestTimeout)*time.Second) |
| 100 | + ctxPing = newCtx |
| 101 | + defer cancel() |
| 102 | + clientOptions.SetConnectTimeout(time.Duration(scfgs[0].requestTimeout) * time.Second) |
| 103 | + } |
| 104 | + |
| 105 | + if !cfg.IsReplicaSet { |
| 106 | + clientOptions.SetAuth(options.Credential{ |
| 107 | + Username: cfg.Username, |
| 108 | + Password: cfg.Password, |
| 109 | + }) |
| 110 | + } |
| 111 | + |
| 112 | + c, err := mongo.Connect(ctx, clientOptions) |
| 113 | + if err != nil { |
| 114 | + log.Fatal("ClientStore failed to connect mongo: ", err) |
| 115 | + } else { |
| 116 | + log.Println("Connection to mongoDB successful") |
| 117 | + } |
| 118 | + |
| 119 | + err = c.Ping(ctxPing, nil) |
25 | 120 | if err != nil {
|
26 |
| - panic(err) |
| 121 | + log.Fatal("MongoDB ping failed:", err) |
27 | 122 | }
|
28 | 123 |
|
29 |
| - return NewClientStoreWithSession(session, cfg.DB, ccfgs...) |
| 124 | + log.Println("Ping db successfull") |
| 125 | + |
| 126 | + return NewClientStoreWithSession(c, cfg, scfgs...) |
30 | 127 | }
|
31 | 128 |
|
32 | 129 | // NewClientStoreWithSession create a client store instance based on mongodb
|
33 |
| -func NewClientStoreWithSession(session *mgo.Session, dbName string, ccfgs ...*ClientConfig) *ClientStore { |
| 130 | +func NewClientStoreWithSession(client *mongo.Client, cfg *Config, scfgs ...*StoreConfig) *ClientStore { |
| 131 | + strCfgs := NewDefaultStoreConfig(cfg.DB, cfg.Service, cfg.IsReplicaSet) |
| 132 | + |
34 | 133 | cs := &ClientStore{
|
35 |
| - dbName: dbName, |
36 |
| - session: session, |
37 |
| - ccfg: NewDefaultClientConfig(), |
| 134 | + client: client, |
| 135 | + ccfg: NewDefaultClientConfig(strCfgs), |
38 | 136 | }
|
39 |
| - if len(ccfgs) > 0 { |
40 |
| - cs.ccfg = ccfgs[0] |
| 137 | + |
| 138 | + if len(scfgs) > 0 { |
| 139 | + if scfgs[0].connectionTimeout > 0 { |
| 140 | + cs.ccfg.storeConfig.connectionTimeout = scfgs[0].connectionTimeout |
| 141 | + } |
| 142 | + if scfgs[0].requestTimeout > 0 { |
| 143 | + cs.ccfg.storeConfig.requestTimeout = scfgs[0].requestTimeout |
| 144 | + } |
41 | 145 | }
|
42 | 146 |
|
43 | 147 | return cs
|
44 | 148 | }
|
45 | 149 |
|
46 | 150 | // ClientStore MongoDB storage for OAuth 2.0
|
47 | 151 | type ClientStore struct {
|
48 |
| - ccfg *ClientConfig |
49 |
| - dbName string |
50 |
| - session *mgo.Session |
| 152 | + ccfg *ClientConfig |
| 153 | + client *mongo.Client |
51 | 154 | }
|
52 | 155 |
|
53 | 156 | // Close close the mongo session
|
54 | 157 | func (cs *ClientStore) Close() {
|
55 |
| - cs.session.Close() |
| 158 | + if err := cs.client.Disconnect(context.Background()); err != nil { |
| 159 | + log.Fatal(err) |
| 160 | + } |
56 | 161 | }
|
57 | 162 |
|
58 |
| -func (cs *ClientStore) c(name string) *mgo.Collection { |
59 |
| - return cs.session.DB(cs.dbName).C(name) |
| 163 | +func (cs *ClientStore) c(name string) *mongo.Collection { |
| 164 | + return cs.client.Database(cs.ccfg.storeConfig.db).Collection(name) |
60 | 165 | }
|
61 | 166 |
|
62 |
| -func (cs *ClientStore) cHandler(name string, handler func(c *mgo.Collection)) { |
63 |
| - session := cs.session.Clone() |
64 |
| - defer session.Close() |
65 |
| - handler(session.DB(cs.dbName).C(name)) |
66 |
| - return |
67 |
| -} |
| 167 | +// Create create client information |
| 168 | +func (cs *ClientStore) Create(info oauth2.ClientInfo) (err error) { |
| 169 | + ctx := context.Background() |
68 | 170 |
|
69 |
| -// Set set client information |
70 |
| -func (cs *ClientStore) Set(info oauth2.ClientInfo) (err error) { |
71 |
| - cs.cHandler(cs.ccfg.ClientsCName, func(c *mgo.Collection) { |
72 |
| - entity := &client{ |
73 |
| - ID: info.GetID(), |
74 |
| - Secret: info.GetSecret(), |
75 |
| - Domain: info.GetDomain(), |
76 |
| - UserID: info.GetUserID(), |
77 |
| - } |
| 171 | + ctxReq, cancel := cs.ccfg.storeConfig.setRequestContext() |
| 172 | + defer cancel() |
| 173 | + if ctxReq != nil { |
| 174 | + ctx = ctxReq |
| 175 | + } |
| 176 | + |
| 177 | + entity := &client{ |
| 178 | + ID: info.GetID(), |
| 179 | + Secret: info.GetSecret(), |
| 180 | + Domain: info.GetDomain(), |
| 181 | + UserID: info.GetUserID(), |
| 182 | + } |
| 183 | + |
| 184 | + collection := cs.c(cs.ccfg.ClientsCName) |
78 | 185 |
|
79 |
| - if cerr := c.Insert(entity); cerr != nil { |
80 |
| - err = cerr |
81 |
| - return |
| 186 | + _, err = collection.InsertOne(ctx, entity) |
| 187 | + if err != nil { |
| 188 | + if !mongo.IsDuplicateKeyError(err) { |
| 189 | + log.Fatal(err) |
| 190 | + } else { |
| 191 | + return nil |
82 | 192 | }
|
83 |
| - }) |
| 193 | + } |
84 | 194 |
|
85 | 195 | return
|
86 | 196 | }
|
87 | 197 |
|
88 | 198 | // GetByID according to the ID for the client information
|
89 |
| -func (cs *ClientStore) GetByID(id string) (info oauth2.ClientInfo, err error) { |
90 |
| - cs.cHandler(cs.ccfg.ClientsCName, func(c *mgo.Collection) { |
91 |
| - entity := new(client) |
| 199 | +func (cs *ClientStore) GetByID(ctx context.Context, id string) (info oauth2.ClientInfo, err error) { |
| 200 | + ctxReq, cancel := cs.ccfg.storeConfig.setRequestContext() |
| 201 | + defer cancel() |
| 202 | + if ctxReq != nil { |
| 203 | + ctx = ctxReq |
| 204 | + } |
92 | 205 |
|
93 |
| - if cerr := c.FindId(id).One(entity); cerr != nil { |
94 |
| - err = cerr |
95 |
| - return |
| 206 | + filter := bson.M{"_id": id} |
| 207 | + result := cs.c(cs.ccfg.ClientsCName).FindOne(ctx, filter) |
| 208 | + if err := result.Err(); err != nil { |
| 209 | + if err == mongo.ErrNoDocuments { |
| 210 | + return nil, err |
96 | 211 | }
|
| 212 | + return nil, errors.New("Internal server error, no client found for this ID") |
97 | 213 |
|
98 |
| - info = &models.Client{ |
99 |
| - ID: entity.ID, |
100 |
| - Secret: entity.Secret, |
101 |
| - Domain: entity.Domain, |
102 |
| - UserID: entity.UserID, |
103 |
| - } |
104 |
| - }) |
| 214 | + } |
| 215 | + |
| 216 | + entity := &client{} |
| 217 | + if err := result.Decode(entity); err != nil { |
| 218 | + log.Println(err) |
| 219 | + } |
| 220 | + |
| 221 | + info = &models.Client{ |
| 222 | + ID: entity.ID, |
| 223 | + Secret: entity.Secret, |
| 224 | + Domain: entity.Domain, |
| 225 | + UserID: entity.UserID, |
| 226 | + } |
105 | 227 |
|
106 | 228 | return
|
107 | 229 | }
|
108 | 230 |
|
109 | 231 | // RemoveByID use the client id to delete the client information
|
110 | 232 | func (cs *ClientStore) RemoveByID(id string) (err error) {
|
111 |
| - cs.cHandler(cs.ccfg.ClientsCName, func(c *mgo.Collection) { |
112 |
| - if cerr := c.RemoveId(id); cerr != nil { |
113 |
| - err = cerr |
114 |
| - return |
115 |
| - } |
116 |
| - }) |
| 233 | + ctx := context.Background() |
117 | 234 |
|
118 |
| - return |
| 235 | + ctxReq, cancel := cs.ccfg.storeConfig.setRequestContext() |
| 236 | + defer cancel() |
| 237 | + if ctxReq != nil { |
| 238 | + ctx = ctxReq |
| 239 | + } |
| 240 | + |
| 241 | + filter := bson.M{"_id": id} |
| 242 | + _, err = cs.c(cs.ccfg.ClientsCName).DeleteOne(ctx, filter) |
| 243 | + return err |
119 | 244 | }
|
120 | 245 |
|
121 | 246 | type client struct {
|
|
0 commit comments