github.com/niedbalski/juju@v0.0.0-20190215020005-8ff100488e47/state/bakerystorage/storage.go (about)

     1  // Copyright 2014-2016 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package bakerystorage
     5  
     6  import (
     7  	"encoding/json"
     8  	"time"
     9  
    10  	"github.com/juju/errors"
    11  	"gopkg.in/macaroon-bakery.v2-unstable/bakery"
    12  	"gopkg.in/macaroon-bakery.v2-unstable/bakery/mgostorage"
    13  	"gopkg.in/mgo.v2"
    14  )
    15  
    16  // MongoIndexes returns the indexes to apply to the MongoDB collection.
    17  func MongoIndexes() []mgo.Index {
    18  	// Note: this second-guesses the underlying document format
    19  	// used by bakery's mgostorage package.
    20  	// TODO change things so that we can use EnsureIndex instead.
    21  	return []mgo.Index{{
    22  		Key: []string{"-created"},
    23  	}, {
    24  		Key:         []string{"expires"},
    25  		ExpireAfter: time.Second,
    26  	}}
    27  }
    28  
    29  type storage struct {
    30  	config      Config
    31  	expireAfter time.Duration
    32  	rootKeys    *mgostorage.RootKeys
    33  }
    34  
    35  type storageDoc struct {
    36  	Location string    `bson:"_id"`
    37  	Item     string    `bson:"item"`
    38  	ExpireAt time.Time `bson:"expire-at,omitempty"`
    39  }
    40  
    41  type legacyRootKey struct {
    42  	RootKey []byte
    43  }
    44  
    45  // ExpireAfter implements ExpirableStorage.ExpireAfter.
    46  func (s *storage) ExpireAfter(expireAfter time.Duration) ExpirableStorage {
    47  	newStorage := *s
    48  	newStorage.expireAfter = expireAfter
    49  	return &newStorage
    50  }
    51  
    52  // RootKey implements Storage.RootKey
    53  func (s *storage) RootKey() ([]byte, []byte, error) {
    54  	storage, closer := s.getStorage()
    55  	defer closer()
    56  	return storage.RootKey()
    57  }
    58  
    59  func (s *storage) getStorage() (bakery.Storage, func()) {
    60  	coll, closer := s.config.GetCollection()
    61  	return s.config.GetStorage(s.rootKeys, coll, s.expireAfter), closer
    62  }
    63  
    64  // Get implements Storage.Get
    65  func (s *storage) Get(id []byte) ([]byte, error) {
    66  	storage, closer := s.getStorage()
    67  	defer closer()
    68  	i, err := storage.Get(id)
    69  	if err != nil {
    70  		if err == bakery.ErrNotFound {
    71  			return s.legacyGet(id)
    72  		}
    73  		return nil, err
    74  	}
    75  	return i, nil
    76  }
    77  
    78  // legacyGet is attempted as the id we're looking for was created in a previous
    79  // version of Juju while using v1 versions of the macaroon-bakery.
    80  func (s *storage) legacyGet(location []byte) ([]byte, error) {
    81  	coll, closer := s.config.GetCollection()
    82  	defer closer()
    83  
    84  	var i storageDoc
    85  	err := coll.FindId(string(location)).One(&i)
    86  	if err != nil {
    87  		if err == mgo.ErrNotFound {
    88  			return nil, bakery.ErrNotFound
    89  		}
    90  		return nil, errors.Annotatef(err, "cannot get item for location %q", location)
    91  	}
    92  	var rootKey legacyRootKey
    93  	err = json.Unmarshal([]byte(i.Item), &rootKey)
    94  	if err != nil {
    95  		return nil, errors.Annotate(err, "was unable to unmarshal found rootkey")
    96  	}
    97  	return rootKey.RootKey, nil
    98  }