package mongodbstore import ( "context" "log" "time" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" ) type item struct { Token string `json:"token"` Object []byte `json:"object"` Expiration int64 `json:"expiration"` } // MongoDBStore represents the session store. type MongoDBStore struct { collection *mongo.Collection stopCleanup chan bool } // New returns a new MongoDBStore instance, with a background cleanup goroutine that // runs every minute to remove expired session data. func New(db *mongo.Database) *MongoDBStore { return NewWithCleanupInterval(db, time.Minute) } // NewWithCleanupInterval returns a new MongoDBStore instance. The cleanupInterval // parameter controls how frequently expired session data is removed by the // background cleanup goroutine. Setting it to 0 prevents the cleanup goroutine // from running (i.e. expired sessions will not be removed). func NewWithCleanupInterval(db *mongo.Database, cleanupInterval time.Duration) *MongoDBStore { collection := db.Collection("sessions") m := &MongoDBStore{ collection: collection, } if cleanupInterval > 0 { go m.startCleanup(cleanupInterval) } return m } // Find returns the data for a given session token from the MongoDBStore instance. // If the session token is not found or is expired, the returned exists flag will // be set to false. func (m *MongoDBStore) Find(token string) ([]byte, bool, error) { filter := bson.M{"token": token} result := m.collection.FindOne(context.Background(), filter) err := result.Err() if err != nil { if err == mongo.ErrNoDocuments { return nil, false, nil } return nil, false, err } var i item err = result.Decode(&i) if err != nil { return nil, false, err } if time.Now().UnixNano() > i.Expiration { return nil, false, nil } return i.Object, true, nil } // Commit adds a session token and data to the MongoDBStore instance with the // given expiry time. If the session token already exists, then the data and expiry // time are updated. func (m *MongoDBStore) Commit(token string, b []byte, expiry time.Time) error { item := item{ Token: token, Object: b, Expiration: expiry.UnixNano(), } // Create or replace the existing item // https://docs.mongodb.com/drivers/node/fundamentals/crud/write-operations/upsert/ filter := bson.M{"token": token} update := bson.M{"$set": item} opts := options.Update().SetUpsert(true) _, err := m.collection.UpdateOne(context.Background(), filter, update, opts) return err } // Delete removes a session token and corresponding data from the MongoDBStore // instance. func (m *MongoDBStore) Delete(token string) error { filter := bson.M{"token": token} _, err := m.collection.DeleteOne(context.Background(), filter) return err } // All returns a map containing the token and data for all active (i.e. // not expired) sessions in the MongoDBStore instance. func (m *MongoDBStore) All() (map[string][]byte, error) { filter := bson.D{} cursor, err := m.collection.Find(context.Background(), filter) if err != nil { if err == mongo.ErrNoDocuments { return nil, nil } return nil, err } sessions := make(map[string][]byte) for cursor.Next(context.Background()) { var i item err := cursor.Decode(&i) if err != nil { return nil, err } if i.Expiration > time.Now().UnixNano() { sessions[i.Token] = i.Object } } if err := cursor.Err(); err != nil { return nil, err } cursor.Close(context.Background()) return sessions, nil } func (m *MongoDBStore) startCleanup(cleanupInterval time.Duration) { m.stopCleanup = make(chan bool) ticker := time.NewTicker(cleanupInterval) for { select { case <-ticker.C: err := m.deleteExpired() if err != nil { log.Println(err) } case <-m.stopCleanup: ticker.Stop() return } } } // StopCleanup terminates the background cleanup goroutine for the MongoDBStore // instance. It's rare to terminate this; generally MongoDBStore instances and // their cleanup goroutines are intended to be long-lived and run for the lifetime // of your application. // // There may be occasions though when your use of the MongoDBStore is transient. // An example is creating a new MongoDBStore instance in a test function. In this // scenario, the cleanup goroutine (which will run forever) will prevent the // MongoDBStore object from being garbage collected even after the test function // has finished. You can prevent this by manually calling StopCleanup. func (m *MongoDBStore) StopCleanup() { if m.stopCleanup != nil { m.stopCleanup <- true } } func (m *MongoDBStore) deleteExpired() error { now := time.Now().UnixNano() filter := bson.M{"expiration": bson.M{"$lt": now}} _, err := m.collection.DeleteMany(context.Background(), filter, nil) if err != nil { return err } return nil }