From daeecfb244acbaac128498e813f4442c7e162d13 Mon Sep 17 00:00:00 2001 From: Zachary Wasserman Date: Fri, 14 Oct 2016 09:15:04 -0700 Subject: [PATCH] Fix bug in inmem ID generation (#312) Fixes #308 --- server/datastore/inmem.go | 19 ++++++++++++++++--- server/datastore/inmem_hosts.go | 5 ++--- server/datastore/inmem_invites.go | 2 +- server/datastore/inmem_labels.go | 4 ++-- server/datastore/inmem_packs.go | 2 +- server/datastore/inmem_password_reset.go | 2 +- server/datastore/inmem_queries.go | 2 +- server/datastore/inmem_sessions.go | 2 +- server/datastore/inmem_users.go | 2 +- 9 files changed, 26 insertions(+), 14 deletions(-) diff --git a/server/datastore/inmem.go b/server/datastore/inmem.go index 3c1c9bf84b..2aefc685a8 100644 --- a/server/datastore/inmem.go +++ b/server/datastore/inmem.go @@ -1,6 +1,7 @@ package datastore import ( + "reflect" "sync" "github.com/kolide/kolide-ose/server/kolide" @@ -8,8 +9,10 @@ import ( type inmem struct { kolide.Datastore - Driver string - mtx sync.RWMutex + Driver string + mtx sync.RWMutex + nextIDs map[interface{}]uint + users map[uint]*kolide.User sessions map[uint]*kolide.Session passwordResets map[uint]*kolide.PasswordResetRequest @@ -19,7 +22,8 @@ type inmem struct { queries map[uint]*kolide.Query packs map[uint]*kolide.Pack hosts map[uint]*kolide.Host - orginfo *kolide.OrgInfo + + orginfo *kolide.OrgInfo } func (orm *inmem) Name() string { @@ -29,6 +33,7 @@ func (orm *inmem) Name() string { func (orm *inmem) Migrate() error { orm.mtx.Lock() defer orm.mtx.Unlock() + orm.nextIDs = make(map[interface{}]uint) orm.users = make(map[uint]*kolide.User) orm.sessions = make(map[uint]*kolide.Session) orm.passwordResets = make(map[uint]*kolide.PasswordResetRequest) @@ -64,3 +69,11 @@ func (orm *inmem) getLimitOffsetSliceBounds(opt kolide.ListOptions, length int) } return offset, max } + +// nextID returns the next ID value that should be used for a struct of the +// given type +func (orm *inmem) nextID(val interface{}) uint { + valType := reflect.TypeOf(val) + orm.nextIDs[valType]++ + return orm.nextIDs[valType] +} diff --git a/server/datastore/inmem_hosts.go b/server/datastore/inmem_hosts.go index 06db8a86aa..459aaf62e9 100644 --- a/server/datastore/inmem_hosts.go +++ b/server/datastore/inmem_hosts.go @@ -18,8 +18,7 @@ func (orm *inmem) NewHost(host *kolide.Host) (*kolide.Host, error) { } } - host.ID = uint(len(orm.hosts) + 1) - orm.hosts[host.ID] = host + host.ID = orm.nextID(host) return host, nil } @@ -122,7 +121,7 @@ func (orm *inmem) EnrollHost(uuid, hostname, ip, platform string, nodeKeySize in } if host.ID == 0 { - host.ID = uint(len(orm.hosts) + 1) + host.ID = orm.nextID(host) } orm.hosts[host.ID] = &host diff --git a/server/datastore/inmem_invites.go b/server/datastore/inmem_invites.go index 6b5339c261..7a08eb0726 100644 --- a/server/datastore/inmem_invites.go +++ b/server/datastore/inmem_invites.go @@ -17,7 +17,7 @@ func (orm *inmem) NewInvite(invite *kolide.Invite) (*kolide.Invite, error) { } } - invite.ID = uint(len(orm.invites) + 1) + invite.ID = orm.nextID(invite) orm.invites[invite.ID] = invite return invite, nil } diff --git a/server/datastore/inmem_labels.go b/server/datastore/inmem_labels.go index 4ab3c93385..bea63b0635 100644 --- a/server/datastore/inmem_labels.go +++ b/server/datastore/inmem_labels.go @@ -20,7 +20,7 @@ func (orm *inmem) NewLabel(label *kolide.Label) (*kolide.Label, error) { } } - newLabel.ID = uint(len(orm.labels) + 1) + newLabel.ID = orm.nextID(label) orm.labels[newLabel.ID] = &newLabel return &newLabel, nil @@ -112,12 +112,12 @@ func (orm *inmem) RecordLabelQueryExecutions(host *kolide.Host, results map[stri if !updated { // Create new execution lqe := kolide.LabelQueryExecution{ - ID: uint(len(orm.labelQueryExecutions) + 1), HostID: host.ID, LabelID: label.ID, UpdatedAt: t, Matches: matches, } + lqe.ID = orm.nextID(lqe) orm.labelQueryExecutions[lqe.ID] = &lqe } } diff --git a/server/datastore/inmem_packs.go b/server/datastore/inmem_packs.go index e6cefecbbb..3949032ce6 100644 --- a/server/datastore/inmem_packs.go +++ b/server/datastore/inmem_packs.go @@ -18,7 +18,7 @@ func (orm *inmem) NewPack(pack *kolide.Pack) error { } } - newPack.ID = uint(len(orm.packs) + 1) + newPack.ID = orm.nextID(pack) orm.packs[newPack.ID] = &newPack return nil diff --git a/server/datastore/inmem_password_reset.go b/server/datastore/inmem_password_reset.go index 94e8c5824e..f8852724a4 100644 --- a/server/datastore/inmem_password_reset.go +++ b/server/datastore/inmem_password_reset.go @@ -6,7 +6,7 @@ func (orm *inmem) NewPasswordResetRequest(req *kolide.PasswordResetRequest) (*ko orm.mtx.Lock() defer orm.mtx.Unlock() - req.ID = uint(len(orm.passwordResets) + 1) + req.ID = orm.nextID(req) orm.passwordResets[req.ID] = req return req, nil } diff --git a/server/datastore/inmem_queries.go b/server/datastore/inmem_queries.go index fc83092e05..bad2766930 100644 --- a/server/datastore/inmem_queries.go +++ b/server/datastore/inmem_queries.go @@ -18,7 +18,7 @@ func (orm *inmem) NewQuery(query *kolide.Query) (*kolide.Query, error) { } } - newQuery.ID = uint(len(orm.queries) + 1) + newQuery.ID = orm.nextID(newQuery) orm.queries[newQuery.ID] = &newQuery return &newQuery, nil diff --git a/server/datastore/inmem_sessions.go b/server/datastore/inmem_sessions.go index adcbf80fd8..796d91e2ca 100644 --- a/server/datastore/inmem_sessions.go +++ b/server/datastore/inmem_sessions.go @@ -48,7 +48,7 @@ func (orm *inmem) NewSession(session *kolide.Session) (*kolide.Session, error) { orm.mtx.Lock() defer orm.mtx.Unlock() - session.ID = uint(len(orm.sessions)) + session.ID = orm.nextID(session) orm.sessions[session.ID] = session if err := orm.MarkSessionAccessed(session); err != nil { return nil, err diff --git a/server/datastore/inmem_users.go b/server/datastore/inmem_users.go index ba1ee76f37..f62d1f0901 100644 --- a/server/datastore/inmem_users.go +++ b/server/datastore/inmem_users.go @@ -16,7 +16,7 @@ func (orm *inmem) NewUser(user *kolide.User) (*kolide.User, error) { } } - user.ID = uint(len(orm.users) + 1) + user.ID = orm.nextID(user) orm.users[user.ID] = user return user, nil