<commit_before>package modules

import (
	"github.com/Seklfreak/Robyul2/modules/plugins"
	"github.com/Seklfreak/Robyul2/modules/plugins/biasgame"
	"github.com/Seklfreak/Robyul2/modules/plugins/idols"
	"github.com/Seklfreak/Robyul2/modules/plugins/levels"
	"github.com/Seklfreak/Robyul2/modules/plugins/mod"
	"github.com/Seklfreak/Robyul2/modules/plugins/notifications"
	"github.com/Seklfreak/Robyul2/modules/plugins/nugugame"
	"github.com/Seklfreak/Robyul2/modules/plugins/youtube"
)

var (
	pluginCache         map[string]*Plugin
	extendedPluginCache map[string]*ExtendedPlugin

	PluginList = []Plugin{
		&notifications.Handler{},
		&plugins.About{},
		&plugins.Stats{},
		&plugins.Uptime{},
		&plugins.Translator{},
		&plugins.UrbanDict{},
		&plugins.Weather{},
		&plugins.VLive{},
		&plugins.WolframAlpha{},
		&plugins.LastFm{},
		&plugins.Twitch{},
		&plugins.Charts{},
		&plugins.Choice{},
		&plugins.Osu{},
		&plugins.Reminders{},
		&plugins.Ratelimit{},
		&plugins.Gfycat{},
		&plugins.RandomPictures{},
		&youtube.Handler{},
		&plugins.Spoiler{},
		&plugins.RandomCat{},
		&plugins.RPS{},
		&plugins.Nuke{},
		&plugins.Dig{},
		&plugins.Streamable{},
		&plugins.Lyrics{},
		&plugins.Friend{},
		&plugins.Names{},
		&plugins.Reddit{},
		&plugins.Color{},
		&plugins.Dog{},
		&plugins.Debug{},
		&plugins.Donators{},
		&plugins.Ping{},
		&plugins.BotStatus{},
		&plugins.VanityInvite{},
		&plugins.DiscordMoney{},
		&plugins.Whois{},
		&plugins.Isup{},
		&plugins.ModulePermissions{},
		&plugins.M8ball{},
		&plugins.Feedback{},
		&plugins.DM{},
		&plugins.EmbedPost{},
		&plugins.Useruploads{},
		&plugins.Move{},
		&plugins.Crypto{},
		&plugins.Imgur{},
		&plugins.Steam{},
		&plugins.Config{},
		&plugins.Storage{},
		&plugins.Mirror{},

		// &instagram.Handler{},
		// &plugins.Facebook{},
		// &google.Handler{},
	}

	PluginExtendedList = []ExtendedPlugin{
		&plugins.Bias{},
		&plugins.GuildAnnouncements{},
		&levels.Levels{},
		&plugins.Gallery{},
		&plugins.CustomCommands{},
		&plugins.ReactionPolls{},
		&mod.Mod{},
		&plugins.AutoRoles{},
		&plugins.Starboard{},
		&plugins.Autoleaver{},
		&plugins.Persistency{},
		&biasgame.Module{},
		&nugugame.Module{},
		&idols.Module{},

		// &plugins.Twitter{},
		// &eventlog.Handler{},
		// &plugins.Perspective{},
	}
)
<commit_msg>disables names plugin<commit_after>package modules

import (
	"github.com/Seklfreak/Robyul2/modules/plugins"
	"github.com/Seklfreak/Robyul2/modules/plugins/biasgame"
	"github.com/Seklfreak/Robyul2/modules/plugins/idols"
	"github.com/Seklfreak/Robyul2/modules/plugins/levels"
	"github.com/Seklfreak/Robyul2/modules/plugins/mod"
	"github.com/Seklfreak/Robyul2/modules/plugins/notifications"
	"github.com/Seklfreak/Robyul2/modules/plugins/nugugame"
	"github.com/Seklfreak/Robyul2/modules/plugins/youtube"
)

var (
	pluginCache         map[string]*Plugin
	extendedPluginCache map[string]*ExtendedPlugin

	PluginList = []Plugin{
		&notifications.Handler{},
		&plugins.About{},
		&plugins.Stats{},
		&plugins.Uptime{},
		&plugins.Translator{},
		&plugins.UrbanDict{},
		&plugins.Weather{},
		&plugins.VLive{},
		&plugins.WolframAlpha{},
		&plugins.LastFm{},
		&plugins.Twitch{},
		&plugins.Charts{},
		&plugins.Choice{},
		&plugins.Osu{},
		&plugins.Reminders{},
		&plugins.Ratelimit{},
		&plugins.Gfycat{},
		&plugins.RandomPictures{},
		&youtube.Handler{},
		&plugins.Spoiler{},
		&plugins.RandomCat{},
		&plugins.RPS{},
		&plugins.Nuke{},
		&plugins.Dig{},
		&plugins.Streamable{},
		&plugins.Lyrics{},
		&plugins.Friend{},
		&plugins.Reddit{},
		&plugins.Color{},
		&plugins.Dog{},
		&plugins.Debug{},
		&plugins.Donators{},
		&plugins.Ping{},
		&plugins.BotStatus{},
		&plugins.VanityInvite{},
		&plugins.DiscordMoney{},
		&plugins.Whois{},
		&plugins.Isup{},
		&plugins.ModulePermissions{},
		&plugins.M8ball{},
		&plugins.Feedback{},
		&plugins.DM{},
		&plugins.EmbedPost{},
		&plugins.Useruploads{},
		&plugins.Move{},
		&plugins.Crypto{},
		&plugins.Imgur{},
		&plugins.Steam{},
		&plugins.Config{},
		&plugins.Storage{},
		&plugins.Mirror{},

		// &plugins.Names{},
		// &instagram.Handler{},
		// &plugins.Facebook{},
		// &google.Handler{},
	}

	PluginExtendedList = []ExtendedPlugin{
		&plugins.Bias{},
		&plugins.GuildAnnouncements{},
		&levels.Levels{},
		&plugins.Gallery{},
		&plugins.CustomCommands{},
		&plugins.ReactionPolls{},
		&mod.Mod{},
		&plugins.AutoRoles{},
		&plugins.Starboard{},
		&plugins.Autoleaver{},
		&plugins.Persistency{},
		&biasgame.Module{},
		&nugugame.Module{},
		&idols.Module{},

		// &plugins.Twitter{},
		// &eventlog.Handler{},
		// &plugins.Perspective{},
	}
)
<|endoftext|>


<commit_before>// Package modules contains definitions for all of the major modules of Sia, as
// well as some helper functions for performing actions that are common to
// multiple modules.
package modules

import (
	"time"

	"github.com/NebulousLabs/Sia/build"
)

var (
	// SafeMutexDelay is the recommended timeout for the deadlock detecting
	// mutex. This value is DEPRECATED, as safe mutexes are no longer
	// recommended. Instead, the locking conventions should be followed and a
	// traditional mutex or a demote mutex should be used.
	SafeMutexDelay time.Duration
)

func init() {
	if build.Release == "dev" {
		SafeMutexDelay = 40 * time.Second
	} else if build.Release == "standard" {
		SafeMutexDelay = 60 * time.Second
	} else if build.Release == "testing" {
		SafeMutexDelay = 20 * time.Second
	}
}
<commit_msg>extend SafeLocks timeouts<commit_after>// Package modules contains definitions for all of the major modules of Sia, as
// well as some helper functions for performing actions that are common to
// multiple modules.
package modules

import (
	"time"

	"github.com/NebulousLabs/Sia/build"
)

var (
	// SafeMutexDelay is the recommended timeout for the deadlock detecting
	// mutex. This value is DEPRECATED, as safe mutexes are no longer
	// recommended. Instead, the locking conventions should be followed and a
	// traditional mutex or a demote mutex should be used.
	SafeMutexDelay time.Duration
)

func init() {
	if build.Release == "dev" {
		SafeMutexDelay = 60 * time.Second
	} else if build.Release == "standard" {
		SafeMutexDelay = 90 * time.Second
	} else if build.Release == "testing" {
		SafeMutexDelay = 30 * time.Second
	}
}
<|endoftext|>


<commit_before>package session

import (
	"errors"
	"github.com/danjac/photoshare/api/models"
	"github.com/danjac/photoshare/api/settings"
	jwt "github.com/dgrijalva/jwt-go"
	"io/ioutil"
	"net/http"
	"strconv"
	"time"
)

var (
	verifyKey, signKey []byte
	MissingLoginFields = errors.New("Missing login fields")
	userMgr            = models.NewUserManager()
)

func init() {
	var err error
	signKey, err = ioutil.ReadFile(settings.PrivKeyFile)
	if err != nil {
		panic(err)
	}
	verifyKey, err = ioutil.ReadFile(settings.PubKeyFile)
	if err != nil {
		panic(err)
	}

}

// Basic user session info
type SessionInfo struct {
	ID       int64  `json:"id"`
	Name     string `json:"name"`
	IsAdmin  bool   `json:"isAdmin"`
	LoggedIn bool   `json:"loggedIn"`
	Token    string `json:"token"`
}

func NewSessionInfo(user *models.User, token string) *SessionInfo {
	if user == nil || user.ID == 0 {
		return &SessionInfo{}
	}

	return &SessionInfo{user.ID, user.Name, user.IsAdmin, true, token}
}

// Handles user authentication
type Authenticator struct {
	Identifier string `json:"identifier"`
	Password   string `json:"password"`
}

func (auth *Authenticator) Identify() (*models.User, error) {

	if auth.Identifier == "" || auth.Password == "" {
		return nil, MissingLoginFields
	}
	return userMgr.Authenticate(auth.Identifier, auth.Password)
}

func GetCurrentUser(r *http.Request) (*models.User, error) {

	userID, err := readToken(r)
	if err != nil {
		return nil, err
	}
	if userID == "" {
		return nil, nil
	}
	return userMgr.GetActive(userID)
}

func Login(w http.ResponseWriter, user *models.User) (string, error) {
	return createToken(w, strconv.FormatInt(user.ID, 10))
}

func Logout(w http.ResponseWriter) (string, error) {
	return createToken(w, "")
}

func readToken(r *http.Request) (string, error) {
	tokenString := r.Header.Get("X-Auth-Token")
	if tokenString == "" {
		return "", nil
	}
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) ([]byte, error) {
		return verifyKey, nil
	})
	switch err.(type) {
	case nil:
		if !token.Valid {
			return "", nil
		}
		return token.Claims["uid"].(string), nil
	case *jwt.ValidationError:
		return "", nil
	default:
		return "", err
	}
}

func createToken(w http.ResponseWriter, userID string) (string, error) {
	token := jwt.New(jwt.GetSigningMethod("RS256"))
	token.Claims["uid"] = userID
	token.Claims["exp"] = time.Now().Add(time.Minute * 1).Unix()
	tokenString, err := token.SignedString(signKey)
	if err != nil {
		return "", err
	}
	w.Header().Set("X-Auth-Token", tokenString)
	return tokenString, nil
}
<commit_msg>Using a const<commit_after>package session

import (
	"errors"
	"github.com/danjac/photoshare/api/models"
	"github.com/danjac/photoshare/api/settings"
	jwt "github.com/dgrijalva/jwt-go"
	"io/ioutil"
	"net/http"
	"strconv"
	"time"
)

const tokenHeader = "X-Auth-Token"

var (
	verifyKey, signKey []byte
	MissingLoginFields = errors.New("Missing login fields")
	userMgr            = models.NewUserManager()
)

func init() {
	var err error
	signKey, err = ioutil.ReadFile(settings.PrivKeyFile)
	if err != nil {
		panic(err)
	}
	verifyKey, err = ioutil.ReadFile(settings.PubKeyFile)
	if err != nil {
		panic(err)
	}

}

// Basic user session info
type SessionInfo struct {
	ID       int64  `json:"id"`
	Name     string `json:"name"`
	IsAdmin  bool   `json:"isAdmin"`
	LoggedIn bool   `json:"loggedIn"`
	Token    string `json:"token"`
}

func NewSessionInfo(user *models.User, token string) *SessionInfo {
	if user == nil || user.ID == 0 {
		return &SessionInfo{}
	}

	return &SessionInfo{user.ID, user.Name, user.IsAdmin, true, token}
}

// Handles user authentication
type Authenticator struct {
	Identifier string `json:"identifier"`
	Password   string `json:"password"`
}

func (auth *Authenticator) Identify() (*models.User, error) {

	if auth.Identifier == "" || auth.Password == "" {
		return nil, MissingLoginFields
	}
	return userMgr.Authenticate(auth.Identifier, auth.Password)
}

func GetCurrentUser(r *http.Request) (*models.User, error) {

	userID, err := readToken(r)
	if err != nil {
		return nil, err
	}
	if userID == "" {
		return nil, nil
	}
	return userMgr.GetActive(userID)
}

func Login(w http.ResponseWriter, user *models.User) (string, error) {
	return createToken(w, strconv.FormatInt(user.ID, 10))
}

func Logout(w http.ResponseWriter) (string, error) {
	return createToken(w, "")
}

func readToken(r *http.Request) (string, error) {
	tokenString := r.Header.Get(tokenHeader)
	if tokenString == "" {
		return "", nil
	}
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) ([]byte, error) {
		return verifyKey, nil
	})
	switch err.(type) {
	case nil:
		if !token.Valid {
			return "", nil
		}
		return token.Claims["uid"].(string), nil
	case *jwt.ValidationError:
		return "", nil
	default:
		return "", err
	}
}

func createToken(w http.ResponseWriter, userID string) (string, error) {
	token := jwt.New(jwt.GetSigningMethod("RS256"))
	token.Claims["uid"] = userID
	token.Claims["exp"] = time.Now().Add(time.Minute * 1).Unix()
	tokenString, err := token.SignedString(signKey)
	if err != nil {
		return "", err
	}
	w.Header().Set(tokenHeader, tokenString)
	return tokenString, nil
}
<|endoftext|>


<commit_before>package session

import (
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/kleister/kleister-api/store"
)

// SetSolder injects the client and key into the context.
func SetSolder() gin.HandlerFunc {
	return func(c *gin.Context) {
		if c.Query("k") != "" {
			record, res := store.GetKeyByValue(
				c,
				c.Query("k"),
			)

			if res.Error != nil {
				c.JSON(
					http.StatusNotFound,
					gin.H{
						"status":  http.StatusNotFound,
						"message": "Failed to find key",
					},
				)

				c.Abort()
				return
			}

			if !res.RecordNotFound() {
				c.Set(KeyContextKey, record)
			}
		}

		if c.Query("cid") != "" {
			record, res := store.GetClientByValue(
				c,
				c.Query("cid"),
			)

			if res.Error != nil {
				c.JSON(
					http.StatusNotFound,
					gin.H{
						"status":  http.StatusNotFound,
						"message": "Failed to find client",
					},
				)

				c.Abort()
				return
			}

			if !res.RecordNotFound() {
				c.Set(ClientContextKey, record)
			}
		}

		c.Next()
	}
}
<commit_msg>Stop erroring on invalid key/client on solder<commit_after>package session

import (
	"github.com/gin-gonic/gin"
	"github.com/kleister/kleister-api/store"
)

// SetSolder injects the client and key into the context.
func SetSolder() gin.HandlerFunc {
	return func(c *gin.Context) {
		if c.Query("k") != "" {
			record, res := store.GetKeyByValue(
				c,
				c.Query("k"),
			)

			if !res.RecordNotFound() {
				c.Set(KeyContextKey, record)
			}
		}

		if c.Query("cid") != "" {
			record, res := store.GetClientByValue(
				c,
				c.Query("cid"),
			)

			if !res.RecordNotFound() {
				c.Set(ClientContextKey, record)
			}
		}

		c.Next()
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcscaching_test

import (
	"testing"
	"time"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	. "github.com/jacobsa/ogletest"
)

func TestStatCache(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Invariant-checking cache
////////////////////////////////////////////////////////////////////////

type invariantsCache struct {
	wrapped gcscaching.StatCache
}

func (c *invariantsCache) Insert(
	o *gcs.Object,
	expiration time.Time) {
	c.wrapped.CheckInvariants()
	defer c.wrapped.CheckInvariants()

	c.wrapped.Insert(o, expiration)
	return
}

func (c *invariantsCache) AddNegativeEntry(
	name string,
	expiration time.Time) {
	c.wrapped.CheckInvariants()
	defer c.wrapped.CheckInvariants()

	c.wrapped.AddNegativeEntry(name, expiration)
	return
}

func (c *invariantsCache) Erase(name string) {
	c.wrapped.CheckInvariants()
	defer c.wrapped.CheckInvariants()

	c.wrapped.Erase(name)
	return
}

func (c *invariantsCache) LookUp(
	name string,
	now time.Time) (hit bool, o *gcs.Object) {
	c.wrapped.CheckInvariants()
	defer c.wrapped.CheckInvariants()

	hit, o = c.wrapped.LookUp(name, now)
	return
}

func (c *invariantsCache) LookUpOrNil(
	name string,
	now time.Time) (o *gcs.Object) {
	_, o = c.LookUp(name, now)
	return
}

func (c *invariantsCache) Hit(
	name string,
	now time.Time) (hit bool) {
	hit, _ = c.LookUp(name, now)
	return
}

func (c *invariantsCache) NegativeEntry(
	name string,
	now time.Time) (negative bool) {
	hit, o := c.LookUp(name, now)
	negative = hit && o == nil
	return
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

const capacity = 3

var someTime = time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local)
var expiration = someTime.Add(time.Second)

type StatCacheTest struct {
	cache invariantsCache
}

func init() { RegisterTestSuite(&StatCacheTest{}) }

func (t *StatCacheTest) SetUp(ti *TestInfo) {
	t.cache.wrapped = gcscaching.NewStatCache(capacity)
}

////////////////////////////////////////////////////////////////////////
// Test functions
////////////////////////////////////////////////////////////////////////

func (t *StatCacheTest) LookUpInEmptyCache() {
	ExpectFalse(t.cache.Hit("", someTime))
	ExpectFalse(t.cache.Hit("taco", someTime))
}

func (t *StatCacheTest) LookUpUnknownKey() {
	o0 := &gcs.Object{Name: "burrito"}
	o1 := &gcs.Object{Name: "taco"}

	t.cache.Insert(o0, someTime.Add(time.Second))
	t.cache.Insert(o1, someTime.Add(time.Second))

	ExpectFalse(t.cache.Hit("", someTime))
	ExpectFalse(t.cache.Hit("enchilada", someTime))
}

func (t *StatCacheTest) KeysPresentButEverythingIsExpired() {
	o0 := &gcs.Object{Name: "burrito"}
	o1 := &gcs.Object{Name: "taco"}

	t.cache.Insert(o0, someTime.Add(-time.Second))
	t.cache.Insert(o1, someTime.Add(-time.Second))

	ExpectFalse(t.cache.Hit("burrito", someTime))
	ExpectFalse(t.cache.Hit("taco", someTime))
}

func (t *StatCacheTest) FillUpToCapacity() {
	AssertEq(3, capacity)

	o0 := &gcs.Object{Name: "burrito"}
	o1 := &gcs.Object{Name: "taco"}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)
	t.cache.AddNegativeEntry("enchilada", expiration)

	// Before expiration
	justBefore := expiration.Add(-time.Nanosecond)
	ExpectEq(o0, t.cache.LookUpOrNil("burrito", justBefore))
	ExpectEq(o1, t.cache.LookUpOrNil("taco", justBefore))
	ExpectTrue(t.cache.NegativeEntry("enchilada", justBefore))

	// At expiration
	ExpectEq(o0, t.cache.LookUpOrNil("burrito", expiration))
	ExpectEq(o1, t.cache.LookUpOrNil("taco", expiration))
	ExpectTrue(t.cache.NegativeEntry("enchilada", justBefore))

	// After expiration
	justAfter := expiration.Add(time.Nanosecond)
	ExpectFalse(t.cache.Hit("burrito", justAfter))
	ExpectFalse(t.cache.Hit("taco", justAfter))
	ExpectFalse(t.cache.Hit("enchilada", justAfter))
}

func (t *StatCacheTest) ExpiresLeastRecentlyUsed() {
	AssertEq(3, capacity)

	o0 := &gcs.Object{Name: "burrito"}
	o1 := &gcs.Object{Name: "taco"}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)                         // Least recent
	t.cache.AddNegativeEntry("enchilada", expiration)      // Second most recent
	AssertEq(o0, t.cache.LookUpOrNil("burrito", someTime)) // Most recent

	// Insert another.
	o3 := &gcs.Object{Name: "queso"}
	t.cache.Insert(o3, expiration)

	// See what's left.
	ExpectFalse(t.cache.Hit("taco", someTime))
	ExpectEq(o0, t.cache.LookUpOrNil("burrito", someTime))
	ExpectTrue(t.cache.NegativeEntry("enchilada", someTime))
	ExpectEq(o3, t.cache.LookUpOrNil("queso", someTime))
}

func (t *StatCacheTest) Overwrite_NewerGeneration() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 19, MetaGeneration: 1}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o1, t.cache.LookUpOrNil("taco", someTime))

	// The overwritten entry shouldn't count toward capacity.
	AssertEq(3, capacity)

	t.cache.Insert(&gcs.Object{Name: "burrito"}, expiration)
	t.cache.Insert(&gcs.Object{Name: "enchilada"}, expiration)

	ExpectNe(nil, t.cache.LookUpOrNil("taco", someTime))
	ExpectNe(nil, t.cache.LookUpOrNil("burrito", someTime))
	ExpectNe(nil, t.cache.LookUpOrNil("enchilada", someTime))
}

func (t *StatCacheTest) Overwrite_SameGeneration_NewerMetadataGen() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 7}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o1, t.cache.LookUpOrNil("taco", someTime))

	// The overwritten entry shouldn't count toward capacity.
	AssertEq(3, capacity)

	t.cache.Insert(&gcs.Object{Name: "burrito"}, expiration)
	t.cache.Insert(&gcs.Object{Name: "enchilada"}, expiration)

	ExpectNe(nil, t.cache.LookUpOrNil("taco", someTime))
	ExpectNe(nil, t.cache.LookUpOrNil("burrito", someTime))
	ExpectNe(nil, t.cache.LookUpOrNil("enchilada", someTime))
}

func (t *StatCacheTest) Overwrite_SameGeneration_SameMetadataGen() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o1, t.cache.LookUpOrNil("taco", someTime))
}

func (t *StatCacheTest) Overwrite_SameGeneration_OlderMetadataGen() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 3}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o0, t.cache.LookUpOrNil("taco", someTime))
}

func (t *StatCacheTest) Overwrite_OlderGeneration() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 13, MetaGeneration: 7}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o0, t.cache.LookUpOrNil("taco", someTime))
}

func (t *StatCacheTest) Overwrite_NegativeWithPositive() {
	const name = "taco"
	o1 := &gcs.Object{Name: name, Generation: 13, MetaGeneration: 7}

	t.cache.AddNegativeEntry(name, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o1, t.cache.LookUpOrNil("taco", someTime))
}

func (t *StatCacheTest) Overwrite_PositiveWithNegative() {
	AssertTrue(false, "TODO")
}

func (t *StatCacheTest) Overwrite_NegativeWithNegative() {
	AssertTrue(false, "TODO")
}
<commit_msg>StatCacheTest.Overwrite_PositiveWithNegative<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcscaching_test

import (
	"testing"
	"time"

	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	. "github.com/jacobsa/ogletest"
)

func TestStatCache(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Invariant-checking cache
////////////////////////////////////////////////////////////////////////

type invariantsCache struct {
	wrapped gcscaching.StatCache
}

func (c *invariantsCache) Insert(
	o *gcs.Object,
	expiration time.Time) {
	c.wrapped.CheckInvariants()
	defer c.wrapped.CheckInvariants()

	c.wrapped.Insert(o, expiration)
	return
}

func (c *invariantsCache) AddNegativeEntry(
	name string,
	expiration time.Time) {
	c.wrapped.CheckInvariants()
	defer c.wrapped.CheckInvariants()

	c.wrapped.AddNegativeEntry(name, expiration)
	return
}

func (c *invariantsCache) Erase(name string) {
	c.wrapped.CheckInvariants()
	defer c.wrapped.CheckInvariants()

	c.wrapped.Erase(name)
	return
}

func (c *invariantsCache) LookUp(
	name string,
	now time.Time) (hit bool, o *gcs.Object) {
	c.wrapped.CheckInvariants()
	defer c.wrapped.CheckInvariants()

	hit, o = c.wrapped.LookUp(name, now)
	return
}

func (c *invariantsCache) LookUpOrNil(
	name string,
	now time.Time) (o *gcs.Object) {
	_, o = c.LookUp(name, now)
	return
}

func (c *invariantsCache) Hit(
	name string,
	now time.Time) (hit bool) {
	hit, _ = c.LookUp(name, now)
	return
}

func (c *invariantsCache) NegativeEntry(
	name string,
	now time.Time) (negative bool) {
	hit, o := c.LookUp(name, now)
	negative = hit && o == nil
	return
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

const capacity = 3

var someTime = time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local)
var expiration = someTime.Add(time.Second)

type StatCacheTest struct {
	cache invariantsCache
}

func init() { RegisterTestSuite(&StatCacheTest{}) }

func (t *StatCacheTest) SetUp(ti *TestInfo) {
	t.cache.wrapped = gcscaching.NewStatCache(capacity)
}

////////////////////////////////////////////////////////////////////////
// Test functions
////////////////////////////////////////////////////////////////////////

func (t *StatCacheTest) LookUpInEmptyCache() {
	ExpectFalse(t.cache.Hit("", someTime))
	ExpectFalse(t.cache.Hit("taco", someTime))
}

func (t *StatCacheTest) LookUpUnknownKey() {
	o0 := &gcs.Object{Name: "burrito"}
	o1 := &gcs.Object{Name: "taco"}

	t.cache.Insert(o0, someTime.Add(time.Second))
	t.cache.Insert(o1, someTime.Add(time.Second))

	ExpectFalse(t.cache.Hit("", someTime))
	ExpectFalse(t.cache.Hit("enchilada", someTime))
}

func (t *StatCacheTest) KeysPresentButEverythingIsExpired() {
	o0 := &gcs.Object{Name: "burrito"}
	o1 := &gcs.Object{Name: "taco"}

	t.cache.Insert(o0, someTime.Add(-time.Second))
	t.cache.Insert(o1, someTime.Add(-time.Second))

	ExpectFalse(t.cache.Hit("burrito", someTime))
	ExpectFalse(t.cache.Hit("taco", someTime))
}

func (t *StatCacheTest) FillUpToCapacity() {
	AssertEq(3, capacity)

	o0 := &gcs.Object{Name: "burrito"}
	o1 := &gcs.Object{Name: "taco"}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)
	t.cache.AddNegativeEntry("enchilada", expiration)

	// Before expiration
	justBefore := expiration.Add(-time.Nanosecond)
	ExpectEq(o0, t.cache.LookUpOrNil("burrito", justBefore))
	ExpectEq(o1, t.cache.LookUpOrNil("taco", justBefore))
	ExpectTrue(t.cache.NegativeEntry("enchilada", justBefore))

	// At expiration
	ExpectEq(o0, t.cache.LookUpOrNil("burrito", expiration))
	ExpectEq(o1, t.cache.LookUpOrNil("taco", expiration))
	ExpectTrue(t.cache.NegativeEntry("enchilada", justBefore))

	// After expiration
	justAfter := expiration.Add(time.Nanosecond)
	ExpectFalse(t.cache.Hit("burrito", justAfter))
	ExpectFalse(t.cache.Hit("taco", justAfter))
	ExpectFalse(t.cache.Hit("enchilada", justAfter))
}

func (t *StatCacheTest) ExpiresLeastRecentlyUsed() {
	AssertEq(3, capacity)

	o0 := &gcs.Object{Name: "burrito"}
	o1 := &gcs.Object{Name: "taco"}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)                         // Least recent
	t.cache.AddNegativeEntry("enchilada", expiration)      // Second most recent
	AssertEq(o0, t.cache.LookUpOrNil("burrito", someTime)) // Most recent

	// Insert another.
	o3 := &gcs.Object{Name: "queso"}
	t.cache.Insert(o3, expiration)

	// See what's left.
	ExpectFalse(t.cache.Hit("taco", someTime))
	ExpectEq(o0, t.cache.LookUpOrNil("burrito", someTime))
	ExpectTrue(t.cache.NegativeEntry("enchilada", someTime))
	ExpectEq(o3, t.cache.LookUpOrNil("queso", someTime))
}

func (t *StatCacheTest) Overwrite_NewerGeneration() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 19, MetaGeneration: 1}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o1, t.cache.LookUpOrNil("taco", someTime))

	// The overwritten entry shouldn't count toward capacity.
	AssertEq(3, capacity)

	t.cache.Insert(&gcs.Object{Name: "burrito"}, expiration)
	t.cache.Insert(&gcs.Object{Name: "enchilada"}, expiration)

	ExpectNe(nil, t.cache.LookUpOrNil("taco", someTime))
	ExpectNe(nil, t.cache.LookUpOrNil("burrito", someTime))
	ExpectNe(nil, t.cache.LookUpOrNil("enchilada", someTime))
}

func (t *StatCacheTest) Overwrite_SameGeneration_NewerMetadataGen() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 7}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o1, t.cache.LookUpOrNil("taco", someTime))

	// The overwritten entry shouldn't count toward capacity.
	AssertEq(3, capacity)

	t.cache.Insert(&gcs.Object{Name: "burrito"}, expiration)
	t.cache.Insert(&gcs.Object{Name: "enchilada"}, expiration)

	ExpectNe(nil, t.cache.LookUpOrNil("taco", someTime))
	ExpectNe(nil, t.cache.LookUpOrNil("burrito", someTime))
	ExpectNe(nil, t.cache.LookUpOrNil("enchilada", someTime))
}

func (t *StatCacheTest) Overwrite_SameGeneration_SameMetadataGen() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o1, t.cache.LookUpOrNil("taco", someTime))
}

func (t *StatCacheTest) Overwrite_SameGeneration_OlderMetadataGen() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 3}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o0, t.cache.LookUpOrNil("taco", someTime))
}

func (t *StatCacheTest) Overwrite_OlderGeneration() {
	o0 := &gcs.Object{Name: "taco", Generation: 17, MetaGeneration: 5}
	o1 := &gcs.Object{Name: "taco", Generation: 13, MetaGeneration: 7}

	t.cache.Insert(o0, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o0, t.cache.LookUpOrNil("taco", someTime))
}

func (t *StatCacheTest) Overwrite_NegativeWithPositive() {
	const name = "taco"
	o1 := &gcs.Object{Name: name, Generation: 13, MetaGeneration: 7}

	t.cache.AddNegativeEntry(name, expiration)
	t.cache.Insert(o1, expiration)

	ExpectEq(o1, t.cache.LookUpOrNil(name, someTime))
}

func (t *StatCacheTest) Overwrite_PositiveWithNegative() {
	const name = "taco"
	o0 := &gcs.Object{Name: name, Generation: 13, MetaGeneration: 7}

	t.cache.Insert(o0, expiration)
	t.cache.AddNegativeEntry(name, expiration)

	ExpectTrue(t.cache.NegativeEntry(name, someTime))
}

func (t *StatCacheTest) Overwrite_NegativeWithNegative() {
	AssertTrue(false, "TODO")
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Pulcy.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package jobs

import (
	"fmt"
	"strings"

	"github.com/juju/errgo"
)

// LinkName is a name of a link consisting of:
// <job>.<task>[@<instance>] or
// <job>.<taskgroup>.<task>[@instance]
type LinkName string

// NewLinkName assembles a link name from its elements.
func NewLinkName(jn JobName, tgn TaskGroupName, tn TaskName, in InstanceName) LinkName {
	result := fmt.Sprintf("%s.%s.%s", jn, tgn, tn)
	if !in.IsEmpty() {
		result = fmt.Sprintf("%s@%s", result, in)
	}
	return LinkName(result)
}

// String returns a link name in format <job>.<taskgroup>.<task>
func (ln LinkName) String() string {
	jn, tgn, tn, in, err := ln.parse()
	if err != nil {
		return ""
	}
	return string(NewLinkName(jn, tgn, tn, in))
}

// PrivateDomainName returns the DNS name (in the private namespace) for the given link name.
func (ln LinkName) PrivateDomainName() string {
	return fmt.Sprintf("%s.private", strings.Replace(ln.String(), "@", ".", -1))
}

// WeaveDomainName returns the DNS name (in the weave namespace) for the given link name.
func (ln LinkName) WeaveDomainName() string {
	return fmt.Sprintf("%s.weave.local", strings.Replace(ln.String(), "@", ".", -1))
}

// EtcdServiceName returns name of the service as it is used in ETCD.
func (ln LinkName) EtcdServiceName() string {
	return strings.Replace(strings.Replace(ln.String(), ".", "-", -1), "@", "-", -1)
}

// Validate checks if a link name follows a valid format
func (ln LinkName) Validate() error {
	_, _, _, _, err := ln.parse()
	return maskAny(err)
}

// Job returns the job part of the link name.
func (ln LinkName) Job() (JobName, error) {
	jn, _, _, _, err := ln.parse()
	return jn, maskAny(err)
}

// TaskGroup returns the taskgroup part of the link name.
func (ln LinkName) TaskGroup() (TaskGroupName, error) {
	_, tg, _, _, err := ln.parse()
	return tg, maskAny(err)
}

// Task returns the task part of the link name.
func (ln LinkName) Task() (TaskName, error) {
	_, _, tn, _, err := ln.parse()
	return tn, maskAny(err)
}

// Instance returns the instance part of the link name.
func (ln LinkName) Instance() (InstanceName, error) {
	_, _, _, in, err := ln.parse()
	return in, maskAny(err)
}

// HasInstance returns true if there is a specific instance in link name.
func (ln LinkName) HasInstance() bool {
	_, _, _, in, _ := ln.parse()
	return !in.IsEmpty()
}

func (ln LinkName) parse() (JobName, TaskGroupName, TaskName, InstanceName, error) {
	var instanceName InstanceName
	parts := strings.Split(string(ln), "@")
	switch len(parts) {
	case 1:
		// Empty instance name
	case 2:
		instanceName = InstanceName(parts[1])
	default:
		return "", "", "", "", maskAny(errgo.WithCausef(nil, InvalidNameError, "unrecognized link '%s'", string(ln)))
	}

	parts = strings.Split(parts[0], ".")
	var jobName JobName
	var taskGroupName TaskGroupName
	var taskName TaskName
	switch len(parts) {
	case 2:
		jobName = JobName(parts[0])
		taskGroupName = TaskGroupName(parts[1])
		taskName = TaskName(parts[1])
	case 3:
		jobName = JobName(parts[0])
		taskGroupName = TaskGroupName(parts[1])
		taskName = TaskName(parts[2])
	default:
		return jobName, taskGroupName, taskName, instanceName,
			maskAny(errgo.WithCausef(nil, InvalidNameError, "unrecognized link '%s'", string(ln)))
	}
	if err := jobName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := taskGroupName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := taskName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := instanceName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	return jobName, taskGroupName, taskName, instanceName, nil
}

func (ln LinkName) normalize() LinkName {
	jn, tgn, tn, in, err := ln.parse()
	if err != nil {
		return ln
	}
	return NewLinkName(jn, tgn, tn, in)
}
<commit_msg>weave dns name now used dashes to seperate job, group, tak instead of dots.<commit_after>// Copyright (c) 2016 Pulcy.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package jobs

import (
	"fmt"
	"strings"

	"github.com/juju/errgo"
)

// LinkName is a name of a link consisting of:
// <job>.<task>[@<instance>] or
// <job>.<taskgroup>.<task>[@instance]
type LinkName string

// NewLinkName assembles a link name from its elements.
func NewLinkName(jn JobName, tgn TaskGroupName, tn TaskName, in InstanceName) LinkName {
	result := fmt.Sprintf("%s.%s.%s", jn, tgn, tn)
	if !in.IsEmpty() {
		result = fmt.Sprintf("%s@%s", result, in)
	}
	return LinkName(result)
}

// String returns a link name in format <job>.<taskgroup>.<task>
func (ln LinkName) String() string {
	jn, tgn, tn, in, err := ln.parse()
	if err != nil {
		return ""
	}
	return string(NewLinkName(jn, tgn, tn, in))
}

// PrivateDomainName returns the DNS name (in the private namespace) for the given link name.
func (ln LinkName) PrivateDomainName() string {
	return fmt.Sprintf("%s.private", strings.Replace(ln.String(), "@", ".", -1))
}

// WeaveDomainName returns the DNS name (in the weave namespace) for the given link name.
func (ln LinkName) WeaveDomainName() string {
	return fmt.Sprintf("%s.weave.local", ln.EtcdServiceName())
}

// EtcdServiceName returns name of the service as it is used in ETCD.
func (ln LinkName) EtcdServiceName() string {
	return strings.Replace(strings.Replace(ln.String(), ".", "-", -1), "@", "-", -1)
}

// Validate checks if a link name follows a valid format
func (ln LinkName) Validate() error {
	_, _, _, _, err := ln.parse()
	return maskAny(err)
}

// Job returns the job part of the link name.
func (ln LinkName) Job() (JobName, error) {
	jn, _, _, _, err := ln.parse()
	return jn, maskAny(err)
}

// TaskGroup returns the taskgroup part of the link name.
func (ln LinkName) TaskGroup() (TaskGroupName, error) {
	_, tg, _, _, err := ln.parse()
	return tg, maskAny(err)
}

// Task returns the task part of the link name.
func (ln LinkName) Task() (TaskName, error) {
	_, _, tn, _, err := ln.parse()
	return tn, maskAny(err)
}

// Instance returns the instance part of the link name.
func (ln LinkName) Instance() (InstanceName, error) {
	_, _, _, in, err := ln.parse()
	return in, maskAny(err)
}

// HasInstance returns true if there is a specific instance in link name.
func (ln LinkName) HasInstance() bool {
	_, _, _, in, _ := ln.parse()
	return !in.IsEmpty()
}

func (ln LinkName) parse() (JobName, TaskGroupName, TaskName, InstanceName, error) {
	var instanceName InstanceName
	parts := strings.Split(string(ln), "@")
	switch len(parts) {
	case 1:
		// Empty instance name
	case 2:
		instanceName = InstanceName(parts[1])
	default:
		return "", "", "", "", maskAny(errgo.WithCausef(nil, InvalidNameError, "unrecognized link '%s'", string(ln)))
	}

	parts = strings.Split(parts[0], ".")
	var jobName JobName
	var taskGroupName TaskGroupName
	var taskName TaskName
	switch len(parts) {
	case 2:
		jobName = JobName(parts[0])
		taskGroupName = TaskGroupName(parts[1])
		taskName = TaskName(parts[1])
	case 3:
		jobName = JobName(parts[0])
		taskGroupName = TaskGroupName(parts[1])
		taskName = TaskName(parts[2])
	default:
		return jobName, taskGroupName, taskName, instanceName,
			maskAny(errgo.WithCausef(nil, InvalidNameError, "unrecognized link '%s'", string(ln)))
	}
	if err := jobName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := taskGroupName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := taskName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	if err := instanceName.Validate(); err != nil {
		return jobName, taskGroupName, taskName, instanceName, maskAny(err)
	}
	return jobName, taskGroupName, taskName, instanceName, nil
}

func (ln LinkName) normalize() LinkName {
	jn, tgn, tn, in, err := ln.parse()
	if err != nil {
		return ln
	}
	return NewLinkName(jn, tgn, tn, in)
}
<|endoftext|>


<commit_before>// tra
package rpcdb

import (
	"bytes"
	"fmt"
	"io"
	"net/rpc"
	"net/http"

	"github.com/ssbl/trago/db"
)

func Run(localDir, localAddr, remoteDir, remoteAddr string) error {
	var args int 					// unused arg variable

	localClient, err := rpc.DialHTTP("tcp", localAddr)
	if err != nil {
		return err
	}
	remoteClient, err := rpc.DialHTTP("tcp", remoteAddr)
	if err != nil {
		return err
	}

	err = startSrv(localClient, localDir)
	if err != nil {
		return err
	}
	err = startSrv(remoteClient, remoteDir)
	if err != nil {
		return err
	}

	localDb := db.TraDb{}
	err = localClient.Call("TraSrv.GetDb", &args, &localDb)
	if err != nil {
		return err
	}

	fmt.Printf("Reply from local trasrv:\n%v\n", localDb)

	remoteDb := db.TraDb{}
	err = remoteClient.Call("TraSrv.GetDb", &args, &remoteDb)
	if err != nil {
		return err
	}

	fmt.Printf("Reply from remote trasrv:\n%v\n", remoteDb)

	fmt.Println("Comparing local with remote...")
	tags := localDb.Compare(&remoteDb)

	for file, tag := range tags {
		if tag == db.File {
			response, err := http.Get("http://"+remoteAddr+"/files/"+file)
			if err != nil {
				return err
			}
			defer response.Body.Close()

			buf := new(bytes.Buffer)
			_, err = io.Copy(buf, response.Body)
			if err != nil {
				return err
			}

			fileData := db.FileData{Name: file, Data: buf.Bytes()}
			err = localClient.Call("TraSrv.PutFile", &fileData, &reply)
			if err != nil {
				return err
			}

			localDb.Files[file] = remoteDb.Files[file]
		} else if tag == db.Deleted {
			err = localClient.Call("TraSrv.RemoveFile", &file, &args)
			if err != nil {
				return err
			}
			delete(localDb.Files, file)
		} else if tag == db.Conflict {
			err = localClient.Call("TraSrv.ShowConflict", &file, &args)
			if err != nil {
				return err
			}
		}
	}

	fmt.Println("Comparing remote with local...")
	tags = remoteDb.Compare(&localDb)

	for file, tag := range tags {
		if tag == db.File {
			response, err := http.Get("http://"+localAddr+"/files/"+file)
			if err != nil {
				return err
			}
			defer response.Body.Close()

			buf := new(bytes.Buffer)
			_, err = io.Copy(buf, response.Body)
			if err != nil {
				return err
			}

			fileData := db.FileData{Name: file, Data: buf.Bytes()}
			err = remoteClient.Call("TraSrv.PutFile", &fileData, &reply)
			if err != nil {
				return err
			}
		} else if tag == db.Deleted {
			err = remoteClient.Call("TraSrv.RemoveFile", &file, &args)
			if err != nil {
				return err
			}
			delete(remoteDb.Files, file)
		} else if tag == db.Conflict {
			err = remoteClient.Call("TraSrv.ShowConflict", &file, &args)
			if err != nil {
				return err
			}
		}
	}

	db.CombineVectors(localDb.VersionVec, remoteDb.VersionVec)
	remoteDb.VersionVec = localDb.VersionVec
	fmt.Println(localDb.VersionVec, remoteDb.VersionVec)

	if err := localClient.Call("TraSrv.PutDb", &localDb, &args); err != nil {
		return err
	}
	if err := remoteClient.Call("TraSrv.PutDb", &remoteDb, &args); err != nil {
		return err
	}

	return nil
}

func startSrv(client *rpc.Client, dir string) error {
	var reply int

	err := client.Call("TraSrv.InitSrv", &dir, &reply)
	return err
}
<commit_msg>put file sending logic in a separate function<commit_after>// tra
package rpcdb

import (
	"bytes"
	"fmt"
	"io"
	"net/rpc"
	"net/http"

	"github.com/ssbl/trago/db"
)

func Run(localDir, localAddr, remoteDir, remoteAddr string) error {
	var args int 					// unused arg variable

	localClient, err := rpc.DialHTTP("tcp", localAddr)
	if err != nil {
		return err
	}
	remoteClient, err := rpc.DialHTTP("tcp", remoteAddr)
	if err != nil {
		return err
	}

	err = startSrv(localClient, localDir)
	if err != nil {
		return err
	}
	err = startSrv(remoteClient, remoteDir)
	if err != nil {
		return err
	}

	localDb := db.TraDb{}
	err = localClient.Call("TraSrv.GetDb", &args, &localDb)
	if err != nil {
		return err
	}

	fmt.Printf("Reply from local trasrv:\n%v\n", localDb)

	remoteDb := db.TraDb{}
	err = remoteClient.Call("TraSrv.GetDb", &args, &remoteDb)
	if err != nil {
		return err
	}

	fmt.Printf("Reply from remote trasrv:\n%v\n", remoteDb)

	fmt.Println("Comparing local with remote...")
	tags := localDb.Compare(&remoteDb)

	for file, tag := range tags {
		if tag == db.File {
			if err := sendFile(localClient, file, remoteAddr); err != nil {
				return err
			}
			localDb.Files[file] = remoteDb.Files[file]
		} else if tag == db.Deleted {
			err = localClient.Call("TraSrv.RemoveFile", &file, &args)
			if err != nil {
				return err
			}
			delete(localDb.Files, file)
		} else if tag == db.Conflict {
			err = localClient.Call("TraSrv.ShowConflict", &file, &args)
			if err != nil {
				return err
			}
		}
	}

	fmt.Println("Comparing remote with local...")
	tags = remoteDb.Compare(&localDb)

	for file, tag := range tags {
		if tag == db.File {
			if err := sendFile(remoteClient, file, localAddr); err != nil {
				return err
			}
		} else if tag == db.Deleted {
			err = remoteClient.Call("TraSrv.RemoveFile", &file, &args)
			if err != nil {
				return err
			}
			delete(remoteDb.Files, file)
		} else if tag == db.Conflict {
			err = remoteClient.Call("TraSrv.ShowConflict", &file, &args)
			if err != nil {
				return err
			}
		}
	}

	db.CombineVectors(localDb.VersionVec, remoteDb.VersionVec)
	remoteDb.VersionVec = localDb.VersionVec
	fmt.Println(localDb.VersionVec, remoteDb.VersionVec)

	if err := localClient.Call("TraSrv.PutDb", &localDb, &args); err != nil {
		return err
	}
	if err := remoteClient.Call("TraSrv.PutDb", &remoteDb, &args); err != nil {
		return err
	}

	return nil
}

func startSrv(client *rpc.Client, dir string) error {
	var reply int

	err := client.Call("TraSrv.InitSrv", &dir, &reply)
	return err
}

func sendFile(client *rpc.Client, file string, addr string) error {
	var reply int

	response, err := http.Get("http://"+addr+"/files/"+file)
	if err != nil {
		return err
	}
	defer response.Body.Close()

	buf := new(bytes.Buffer)
	_, err = io.Copy(buf, response.Body)
	if err != nil {
		return err
	}

	fileData := db.FileData{Name: file, Data: buf.Bytes()}
	err = client.Call("TraSrv.PutFile", &fileData, &reply)

	return err
}
<|endoftext|>


<commit_before>// tra
package rpcdb

import (
	"bytes"
	"fmt"
	"io"
	"net/rpc"
	"net/http"
	"sync"

	"github.com/ssbl/trago/db"
)

func Run(localDir, localAddr, remoteDir, remoteAddr string) error {
	var s    string
	var args int				// unused arg variables

	localClient, err := rpc.DialHTTP("tcp", localAddr)
	if err != nil {
		return err
	}
	remoteClient, err := rpc.DialHTTP("tcp", remoteAddr)
	if err != nil {
		return err
	}
	defer remoteClient.Call("TraSrv.StopSrv", &s, &args)

	err = startSrv(localClient, localDir)
	if err != nil {
		return err
	}
	err = startSrv(remoteClient, remoteDir)
	if err != nil {
		return err
	}

	localDb := db.TraDb{}
	err = localClient.Call("TraSrv.GetDb", &args, &localDb)
	if err != nil {
		return err
	}

	fmt.Printf("Reply from local trasrv:\n%v\n", localDb)

	remoteDb := db.TraDb{}
	err = remoteClient.Call("TraSrv.GetDb", &args, &remoteDb)
	if err != nil {
		return err
	}

	fmt.Printf("Reply from remote trasrv:\n%v\n", remoteDb)

	tags1 := localDb.Compare(&remoteDb)
	tags2 := remoteDb.Compare(&localDb)
	errch := make(chan error, 1)
	var wg sync.WaitGroup

	wg.Add(2)
	go func() {
		defer wg.Done()
		for file, tag := range tags1 {
			if tag == db.File {
				if err := sendFile(localClient, file, remoteAddr); err != nil {
					errch <- err
					return
				}
				localDb.Files[file] = remoteDb.Files[file]
			} else if tag == db.Deleted {
				err = localClient.Call("TraSrv.RemoveFile", &file, &args)
				if err != nil {
					errch <- err
					return
				}
				delete(localDb.Files, file)
			} else if tag == db.Conflict {
				err = localClient.Call("TraSrv.ShowConflict", &file, &args)
				if err != nil {
					errch <- err
					return
				}
			}
		}
	}()

	go func() {
		defer wg.Done()
		for file, tag := range tags2 {
			if tag == db.File {
				if err := sendFile(remoteClient, file, localAddr); err != nil {
					errch <- err
					return
				}
			} else if tag == db.Deleted {
				err = remoteClient.Call("TraSrv.RemoveFile", &file, &args)
				if err != nil {
					errch <- err
					return
				}
				delete(remoteDb.Files, file)
			} else if tag == db.Conflict {
				err = remoteClient.Call("TraSrv.ShowConflict", &file, &args)
				if err != nil {
					errch <- err
					return
				}
			}
		}
	}()
	wg.Wait()

	select {
		case err := <-errch:
		    return err
		default:
	}

	db.CombineVectors(localDb.VersionVec, remoteDb.VersionVec)
	remoteDb.VersionVec = localDb.VersionVec
	fmt.Println(localDb.VersionVec, remoteDb.VersionVec)

	if err := localClient.Call("TraSrv.PutDb", &localDb, &args); err != nil {
		return err
	}
	if err := remoteClient.Call("TraSrv.PutDb", &remoteDb, &args); err != nil {
		return err
	}

	return nil
}

func startSrv(client *rpc.Client, dir string) error {
	var reply int

	return client.Call("TraSrv.InitDb", &dir, &reply)
}

func sendFile(client *rpc.Client, file string, addr string) error {
	var reply int

	response, err := http.Get("http://"+addr+"/files/"+file)
	if err != nil {
		return err
	}
	defer response.Body.Close()

	buf := new(bytes.Buffer)
	_, err = io.Copy(buf, response.Body)
	if err != nil {
		return err
	}

	fileData := db.FileData{Name: file, Data: buf.Bytes()}
	return client.Call("TraSrv.PutFile", &fileData, &reply)
}
<commit_msg>factor out tag handling loop into separate function<commit_after>// tra
package rpcdb

import (
	"bytes"
	"fmt"
	"io"
	"net/rpc"
	"net/http"
	"sync"

	"github.com/ssbl/trago/db"
)

func Run(localDir, localAddr, remoteDir, remoteAddr string) error {
	var s    string
	var args int				// unused arg variables

	localClient, err := rpc.DialHTTP("tcp", localAddr)
	if err != nil {
		return err
	}
	remoteClient, err := rpc.DialHTTP("tcp", remoteAddr)
	if err != nil {
		return err
	}
	defer remoteClient.Call("TraSrv.StopSrv", &s, &args)

	err = startSrv(localClient, localDir)
	if err != nil {
		return err
	}
	err = startSrv(remoteClient, remoteDir)
	if err != nil {
		return err
	}

	localDb := db.TraDb{}
	err = localClient.Call("TraSrv.GetDb", &args, &localDb)
	if err != nil {
		return err
	}

	fmt.Printf("Reply from local trasrv:\n%v\n", localDb)

	remoteDb := db.TraDb{}
	err = remoteClient.Call("TraSrv.GetDb", &args, &remoteDb)
	if err != nil {
		return err
	}

	fmt.Printf("Reply from remote trasrv:\n%v\n", remoteDb)

	tags1 := localDb.Compare(&remoteDb)
	tags2 := remoteDb.Compare(&localDb)
	errch := make(chan error, 1)
	var wg sync.WaitGroup

	wg.Add(2)
	go func() {
		defer wg.Done()
		processFileTags(localClient, localDb, remoteAddr, tags1, errch)
	}()

	go func() {
		defer wg.Done()
		processFileTags(remoteClient, remoteDb, localAddr, tags2, errch)
	}()
	wg.Wait()

	select {
		case err := <-errch:
		    return err
		default:
	}

	db.CombineVectors(localDb.VersionVec, remoteDb.VersionVec)
	remoteDb.VersionVec = localDb.VersionVec
	fmt.Println(localDb.VersionVec, remoteDb.VersionVec)

	if err := localClient.Call("TraSrv.PutDb", &localDb, &args); err != nil {
		return err
	}
	if err := remoteClient.Call("TraSrv.PutDb", &remoteDb, &args); err != nil {
		return err
	}

	return nil
}

func processFileTags(
	client *rpc.Client,
	tradb db.TraDb,
	dest string,
	tags map[string]db.FileTag,
	errch chan error,
) {
	var args int

	for file, tag := range tags {
		if tag == db.File {
			if err := sendFile(client, file, dest); err != nil {
				errch <- err
				return
			}
		} else if tag == db.Deleted {
			err := client.Call("TraSrv.RemoveFile", &file, &args)
			if err != nil {
				errch <- err
				return
			}
			delete(tradb.Files, file)
		} else if tag == db.Conflict {
			err := client.Call("TraSrv.ShowConflict", &file, &args)
			if err != nil {
				errch <- err
				return
			}
		}
	}
}

func startSrv(client *rpc.Client, dir string) error {
	var reply int

	return client.Call("TraSrv.InitDb", &dir, &reply)
}

func sendFile(client *rpc.Client, file string, addr string) error {
	var reply int

	response, err := http.Get("http://"+addr+"/files/"+file)
	if err != nil {
		return err
	}
	defer response.Body.Close()

	buf := new(bytes.Buffer)
	_, err = io.Copy(buf, response.Body)
	if err != nil {
		return err
	}

	fileData := db.FileData{Name: file, Data: buf.Bytes()}
	return client.Call("TraSrv.PutFile", &fileData, &reply)
}
<|endoftext|>


<commit_before>package gatherer

// The gatherer package contains the logic that figures out which powercycle-
// enabled bots or devices are down.
// See the Gatherer interface for more details.

import (
	"sort"
	"sync"
	"time"

	swarming "go.chromium.org/luci/common/api/swarming/swarming/v1"

	"go.skia.org/infra/go/promalertsclient"
	"go.skia.org/infra/go/sklog"
	skswarming "go.skia.org/infra/go/swarming"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/power/go/decider"
	"go.skia.org/infra/power/go/recorder"
)

const (
	// Alert messages from prometheus
	ALERT_BOT_MISSING     = "BotMissing"
	ALERT_BOT_QUARANTINED = "BotQuarantined"

	// Status messages we report.
	STATUS_HOST_MISSING   = "Host Missing"
	STATUS_DEVICE_MISSING = "Device Missing"
)

// Gatherer is a simple interface around the logic behind obtaining
// a list of bots and devices that are down and could be powercycled.
type Gatherer interface {
	// DownBots returns the current set of down bots. It may be cached.
	DownBots() []DownBot
}

// DownBot represents information about a dead or quarantined bot, as well
// as the alert that is related to it.
type DownBot struct {
	BotID      string                                 `json:"bot_id"`
	HostID     string                                 `json:"host_id"`
	Dimensions []*swarming.SwarmingRpcsStringListPair `json:"dimensions"`
	Status     string                                 `json:"status"`
	// Since represents how long the alert been firing
	Since    time.Time `json:"since"`
	Silenced bool      `json:"silenced"`
}

// The gatherer struct implements the Gatherer interface.
type gatherer struct {
	downBots []DownBot
	mutex    sync.Mutex

	iSwarming skswarming.ApiClient
	eSwarming skswarming.ApiClient
	alerts    promalertsclient.APIClient
	decider   decider.Decider
	hostMap   map[string]string // maps bot id -> jumphost name
	recorder  recorder.Recorder
}

// NewPollingGatherer returns a Gatherer created with the given utilities. all the passed in
// clients should be properly authenticated.
func NewPollingGatherer(external, internal skswarming.ApiClient, alerts promalertsclient.APIClient, decider decider.Decider, recorder recorder.Recorder, hostMap map[string]string, period time.Duration) Gatherer {
	g := &gatherer{
		iSwarming: internal,
		eSwarming: external,
		alerts:    alerts,
		decider:   decider,
		hostMap:   hostMap,
		recorder:  recorder,
	}
	if period > 0 {
		go func() {
			g.update()
			for {
				<-time.Tick(period)
				g.update()
			}
		}()
	}
	return g
}

// See the Gatherer interface for more information.
func (g *gatherer) DownBots() []DownBot {
	g.mutex.Lock()
	defer g.mutex.Unlock()
	return g.downBots
}

// set writes to the underlying downBots slice in a thread safe way.
func (g *gatherer) set(bots []DownBot) {
	g.mutex.Lock()
	defer g.mutex.Unlock()
	g.downBots = bots
}

// downBotsFilter is a function that returns true only for alerts about
// dead and quarantined bots.
func downBotsFilter(a promalertsclient.Alert) bool {
	alertName := string(a.Labels["alertname"])
	return alertName == ALERT_BOT_MISSING || alertName == ALERT_BOT_QUARANTINED
}

// update is the "inner loop" of the gatherer. It polls swarming for a list of
// down bots. It then polls alerts for a list of down bots. It constructs the
// intersect of those lists and sets the result in g.downBots.
func (g *gatherer) update() {
	// Ask Swarming API for list of bots down in the pools we care about
	sklog.Infoln("Polling PromAlerts and Swarming API for down bots")
	bots := []*swarming.SwarmingRpcsBotInfo{}
	for _, pool := range skswarming.POOLS_PRIVATE {
		xb, err := g.iSwarming.ListDownBots(pool)
		if err != nil {
			sklog.Warningf("Could not get down bots from internal pool %s: %s", pool, err)
		}
		bots = append(bots, xb...)
	}

	for _, pool := range skswarming.POOLS_PUBLIC {
		xb, err := g.eSwarming.ListDownBots(pool)
		if err != nil {
			sklog.Warningf("Could not get down bots from external pool %s: %s", pool, err)
		}
		bots = append(bots, xb...)
	}

	if len(bots) == 0 {
		g.set([]DownBot{})
		sklog.Info("Swarming reports no down bots")
		return
	}

	sklog.Infof("Swarming reports %d down bots: %s", len(bots), bots)

	// Ask Prometheus for bot alerts related to quarantined and dead
	alerts, err := g.alerts.GetAlerts(downBotsFilter)
	if err != nil {
		sklog.Warningf("Could not get down bots from alerts %s", err)
		return
	}

	if len(alerts) == 0 {
		g.set([]DownBot{})
		sklog.Info("No bot-related alerts")
		return
	}

	sklog.Infof("Promalerts reports %d bot-related alerts: %s", len(alerts), alerts)

	// join these together to create []DownBot
	botsWithAlerts := util.StringSet{}
	alertMap := map[string]promalertsclient.Alert{}
	for _, a := range alerts {
		id := string(a.Labels["bot"])
		botsWithAlerts[id] = true
		alertMap[id] = a
	}
	botsFromSwarming := util.StringSet{}
	for _, b := range bots {
		botsFromSwarming[b.BotId] = true
	}
	matchingBots := botsWithAlerts.Intersect(botsFromSwarming)

	downBots := []DownBot{}
	for _, b := range bots {
		if unique, ok := matchingBots[b.BotId]; ok && unique {
			alert := alertMap[b.BotId]
			if g.decider.ShouldPowercycleBot(b) {
				downBots = append(downBots, DownBot{
					BotID:      b.BotId,
					HostID:     g.hostMap[b.BotId],
					Dimensions: b.Dimensions,
					Status:     STATUS_HOST_MISSING,
					Since:      alert.StartsAt,
					Silenced:   alert.Silenced,
				})
			} else if g.decider.ShouldPowercycleDevice(b) {
				downBots = append(downBots, DownBot{
					BotID:      b.BotId,
					HostID:     g.hostMap[b.BotId+"-device"],
					Dimensions: b.Dimensions,
					Status:     STATUS_DEVICE_MISSING,
					Since:      alert.StartsAt,
					Silenced:   alert.Silenced,
				})
			}
			// Avoid reporting the same bot down more than once
			// This happens when GetDownBots returns a dead and quarantined bot.
			// This assumes that bot ids are unique, even across pools.
			matchingBots[b.BotId] = false
		}
	}

	// Return sorted based on BotID for determinism and organization.
	sort.Slice(downBots, func(i, j int) bool {
		return downBots[i].BotID < downBots[j].BotID
	})

	fixed, broke := g.identifyChangedBots(downBots)

	g.recorder.NewlyFixedBots(fixed)
	g.recorder.NewlyDownBots(broke)

	g.set(downBots)
	sklog.Infof("Done, found %d bots", len(downBots))
}

// identifyChangedBots enumerates the bots that are new since last time and those
//  that were listed and are now fixed.
func (g *gatherer) identifyChangedBots(currentDownBots []DownBot) (fixed []string, broke []string) {
	getName := func(b DownBot) string {
		if b.Status == STATUS_HOST_MISSING {
			return b.BotID
		} else {
			return b.BotID + "-device"
		}
	}
	prev := util.StringSet{}
	for _, b := range g.downBots {
		prev[getName(b)] = true
	}
	curr := util.StringSet{}
	for _, b := range currentDownBots {
		curr[getName(b)] = true
	}

	fixed = prev.Complement(curr).Keys()
	broke = curr.Complement(prev).Keys()

	sort.Strings(fixed)
	sort.Strings(broke)

	return fixed, broke
}
<commit_msg>[power] Reduce noisy logs.<commit_after>package gatherer

// The gatherer package contains the logic that figures out which powercycle-
// enabled bots or devices are down.
// See the Gatherer interface for more details.

import (
	"sort"
	"sync"
	"time"

	swarming "go.chromium.org/luci/common/api/swarming/swarming/v1"

	"go.skia.org/infra/go/promalertsclient"
	"go.skia.org/infra/go/sklog"
	skswarming "go.skia.org/infra/go/swarming"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/power/go/decider"
	"go.skia.org/infra/power/go/recorder"
)

const (
	// Alert messages from prometheus
	ALERT_BOT_MISSING     = "BotMissing"
	ALERT_BOT_QUARANTINED = "BotQuarantined"

	// Status messages we report.
	STATUS_HOST_MISSING   = "Host Missing"
	STATUS_DEVICE_MISSING = "Device Missing"
)

// Gatherer is a simple interface around the logic behind obtaining
// a list of bots and devices that are down and could be powercycled.
type Gatherer interface {
	// DownBots returns the current set of down bots. It may be cached.
	DownBots() []DownBot
}

// DownBot represents information about a dead or quarantined bot, as well
// as the alert that is related to it.
type DownBot struct {
	BotID      string                                 `json:"bot_id"`
	HostID     string                                 `json:"host_id"`
	Dimensions []*swarming.SwarmingRpcsStringListPair `json:"dimensions"`
	Status     string                                 `json:"status"`
	// Since represents how long the alert been firing
	Since    time.Time `json:"since"`
	Silenced bool      `json:"silenced"`
}

// The gatherer struct implements the Gatherer interface.
type gatherer struct {
	downBots []DownBot
	mutex    sync.Mutex

	iSwarming skswarming.ApiClient
	eSwarming skswarming.ApiClient
	alerts    promalertsclient.APIClient
	decider   decider.Decider
	hostMap   map[string]string // maps bot id -> jumphost name
	recorder  recorder.Recorder
}

// NewPollingGatherer returns a Gatherer created with the given utilities. all the passed in
// clients should be properly authenticated.
func NewPollingGatherer(external, internal skswarming.ApiClient, alerts promalertsclient.APIClient, decider decider.Decider, recorder recorder.Recorder, hostMap map[string]string, period time.Duration) Gatherer {
	g := &gatherer{
		iSwarming: internal,
		eSwarming: external,
		alerts:    alerts,
		decider:   decider,
		hostMap:   hostMap,
		recorder:  recorder,
	}
	if period > 0 {
		go func() {
			g.update()
			for {
				<-time.Tick(period)
				g.update()
			}
		}()
	}
	return g
}

// See the Gatherer interface for more information.
func (g *gatherer) DownBots() []DownBot {
	g.mutex.Lock()
	defer g.mutex.Unlock()
	return g.downBots
}

// set writes to the underlying downBots slice in a thread safe way.
func (g *gatherer) set(bots []DownBot) {
	g.mutex.Lock()
	defer g.mutex.Unlock()
	g.downBots = bots
}

// downBotsFilter is a function that returns true only for alerts about
// dead and quarantined bots.
func downBotsFilter(a promalertsclient.Alert) bool {
	alertName := string(a.Labels["alertname"])
	return alertName == ALERT_BOT_MISSING || alertName == ALERT_BOT_QUARANTINED
}

// update is the "inner loop" of the gatherer. It polls swarming for a list of
// down bots. It then polls alerts for a list of down bots. It constructs the
// intersect of those lists and sets the result in g.downBots.
func (g *gatherer) update() {
	// Ask Swarming API for list of bots down in the pools we care about
	sklog.Infoln("Polling PromAlerts and Swarming API for down bots")
	bots := []*swarming.SwarmingRpcsBotInfo{}
	for _, pool := range skswarming.POOLS_PRIVATE {
		xb, err := g.iSwarming.ListDownBots(pool)
		if err != nil {
			sklog.Warningf("Could not get down bots from internal pool %s: %s", pool, err)
		}
		bots = append(bots, xb...)
	}

	for _, pool := range skswarming.POOLS_PUBLIC {
		xb, err := g.eSwarming.ListDownBots(pool)
		if err != nil {
			sklog.Warningf("Could not get down bots from external pool %s: %s", pool, err)
		}
		bots = append(bots, xb...)
	}

	if len(bots) == 0 {
		g.set([]DownBot{})
		sklog.Info("Swarming reports no down bots")
		return
	}

	sklog.Infof("Swarming reports %d down bots", len(bots))

	// Ask Prometheus for bot alerts related to quarantined and dead
	alerts, err := g.alerts.GetAlerts(downBotsFilter)
	if err != nil {
		sklog.Warningf("Could not get down bots from alerts %s", err)
		return
	}

	if len(alerts) == 0 {
		g.set([]DownBot{})
		sklog.Info("No bot-related alerts")
		return
	}

	sklog.Infof("Promalerts reports %d bot-related alerts: %s", len(alerts), alerts)

	// join these together to create []DownBot
	botsWithAlerts := util.StringSet{}
	alertMap := map[string]promalertsclient.Alert{}
	for _, a := range alerts {
		id := string(a.Labels["bot"])
		botsWithAlerts[id] = true
		alertMap[id] = a
	}
	botsFromSwarming := util.StringSet{}
	for _, b := range bots {
		botsFromSwarming[b.BotId] = true
	}
	matchingBots := botsWithAlerts.Intersect(botsFromSwarming)

	downBots := []DownBot{}
	for _, b := range bots {
		if unique, ok := matchingBots[b.BotId]; ok && unique {
			alert := alertMap[b.BotId]
			if g.decider.ShouldPowercycleBot(b) {
				downBots = append(downBots, DownBot{
					BotID:      b.BotId,
					HostID:     g.hostMap[b.BotId],
					Dimensions: b.Dimensions,
					Status:     STATUS_HOST_MISSING,
					Since:      alert.StartsAt,
					Silenced:   alert.Silenced,
				})
			} else if g.decider.ShouldPowercycleDevice(b) {
				downBots = append(downBots, DownBot{
					BotID:      b.BotId,
					HostID:     g.hostMap[b.BotId+"-device"],
					Dimensions: b.Dimensions,
					Status:     STATUS_DEVICE_MISSING,
					Since:      alert.StartsAt,
					Silenced:   alert.Silenced,
				})
			}
			// Avoid reporting the same bot down more than once
			// This happens when GetDownBots returns a dead and quarantined bot.
			// This assumes that bot ids are unique, even across pools.
			matchingBots[b.BotId] = false
		}
	}

	// Return sorted based on BotID for determinism and organization.
	sort.Slice(downBots, func(i, j int) bool {
		return downBots[i].BotID < downBots[j].BotID
	})

	fixed, broke := g.identifyChangedBots(downBots)

	g.recorder.NewlyFixedBots(fixed)
	g.recorder.NewlyDownBots(broke)

	g.set(downBots)
	sklog.Infof("Done, found %d bots", len(downBots))
}

// identifyChangedBots enumerates the bots that are new since last time and those
//  that were listed and are now fixed.
func (g *gatherer) identifyChangedBots(currentDownBots []DownBot) (fixed []string, broke []string) {
	getName := func(b DownBot) string {
		if b.Status == STATUS_HOST_MISSING {
			return b.BotID
		} else {
			return b.BotID + "-device"
		}
	}
	prev := util.StringSet{}
	for _, b := range g.downBots {
		prev[getName(b)] = true
	}
	curr := util.StringSet{}
	for _, b := range currentDownBots {
		curr[getName(b)] = true
	}

	fixed = prev.Complement(curr).Keys()
	broke = curr.Complement(prev).Keys()

	sort.Strings(fixed)
	sort.Strings(broke)

	return fixed, broke
}
<|endoftext|>


<commit_before>package premailer

import (
	"testing"
	"github.com/stretchr/testify/assert"
)

func TestSpecificitySelectorType(t *testing.T) {
	spec := makeSpecificity(1, 2, 100, "table")
	expected := []int{1, 0, 0, 0, 1, 2, 100}
	assert.Equal(t, expected, spec.importantOrders())
}

func TestSpecificitySelectorClass(t *testing.T) {
	// class
	spec := makeSpecificity(1, 2, 102, "table.red")
	expected := []int{1, 0, 1, 0, 1, 2, 102}
	assert.Equal(t, expected, spec.importantOrders())
}

func TestSpecificitySelectorAttr(t *testing.T) {
	// Attribute
	spec := makeSpecificity(1, 3, 103, "span[lang~=\"en-us\"]")
	expected := []int{1, 0, 0, 1, 1, 3, 103}
	assert.Equal(t, expected, spec.importantOrders())
}

func TestSpecificitySelectorId(t *testing.T) {
	// id
	spec := makeSpecificity(0, 3, 104, "#example")
	expected := []int{0, 1, 0, 0, 0, 3, 104}
	assert.Equal(t, expected, spec.importantOrders())
}<commit_msg>Adding test for specificity sorting<commit_after>package premailer

import (
	"testing"
	"github.com/stretchr/testify/assert"
	"sort"
)

func TestSpecificitySelectorType(t *testing.T) {
	spec := makeSpecificity(1, 2, 100, "table")
	expected := []int{1, 0, 0, 0, 1, 2, 100}
	assert.Equal(t, expected, spec.importantOrders())
}

func TestSpecificitySelectorClass(t *testing.T) {
	// class
	spec := makeSpecificity(1, 2, 102, "table.red")
	expected := []int{1, 0, 1, 0, 1, 2, 102}
	assert.Equal(t, expected, spec.importantOrders())
}

func TestSpecificitySelectorAttr(t *testing.T) {
	// Attribute
	spec := makeSpecificity(1, 3, 103, "span[lang~=\"en-us\"]")
	expected := []int{1, 0, 0, 1, 1, 3, 103}
	assert.Equal(t, expected, spec.importantOrders())
}

func TestSpecificitySelectorId(t *testing.T) {
	// id
	spec := makeSpecificity(0, 3, 104, "#example")
	expected := []int{0, 1, 0, 0, 0, 3, 104}
	assert.Equal(t, expected, spec.importantOrders())
}

func TestSpecificitySort(t *testing.T) {
	undertest := make([]*styleRule, 4)
	for i := 0; i < 4; i++ {
		undertest[i] = &styleRule{}
	}
	specificity0 := makeSpecificity(1, 2, 100, "table")
	undertest[0].specificity = specificity0
	specificity1 := makeSpecificity(1, 2, 102, "table.red")
	undertest[1].specificity = specificity1
	specificity2 := makeSpecificity(1, 3, 103, "span[lang~=\"en-us\"]")
	undertest[2].specificity = specificity2
	specificity3 := makeSpecificity(0, 3, 104, "#example")
	undertest[3].specificity = specificity3

	// expected order
	/*
	expected3 := []int{0, 1, 0, 0, 0, 3, 104}
	expected0 := []int{1, 0, 0, 0, 1, 2, 100}
	expected2 := []int{1, 0, 0, 1, 1, 3, 103}
	expected1 := []int{1, 0, 1, 0, 1, 2, 102}
	*/
	sort.Sort(bySpecificity(undertest))

	assert.Equal(t, specificity3, undertest[0].specificity)
	assert.Equal(t, specificity0, undertest[1].specificity)
	assert.Equal(t, specificity2, undertest[2].specificity)
	assert.Equal(t, specificity1, undertest[3].specificity)
}

func TestSpecificitySortRuleSetIndex(t *testing.T) {
	undertest := make([]*styleRule, 2)
	for i := 0; i < 2; i++ {
		undertest[i] = &styleRule{}
	}
	specificity0 := makeSpecificity(1, 2, 102, "table")
	undertest[0].specificity = specificity0
	specificity1 := makeSpecificity(1, 1, 102, "table")
	undertest[1].specificity = specificity1

	sort.Sort(bySpecificity(undertest))

	assert.Equal(t, specificity1, undertest[0].specificity)
	assert.Equal(t, specificity0, undertest[1].specificity)
}

func TestSpecificitySortRuleIndex(t *testing.T) {
	undertest := make([]*styleRule, 2)
	for i := 0; i < 2; i++ {
		undertest[i] = &styleRule{}
	}
	specificity0 := makeSpecificity(1, 1, 102, "table")
	undertest[0].specificity = specificity0
	specificity1 := makeSpecificity(1, 1, 100, "table")
	undertest[1].specificity = specificity1

	sort.Sort(bySpecificity(undertest))

	assert.Equal(t, specificity1, undertest[0].specificity)
	assert.Equal(t, specificity0, undertest[1].specificity)
}<|endoftext|>


<commit_before>// +build linux,amd64 linux,ppc64le

// Build constraint to use this file for amd64 & ppc64le on Linux

package endpoint

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"sync"
	"time"

	"github.com/bluele/gcache"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	log "github.com/sirupsen/logrus"
)

const (
	bufSize                     = 8 * 1024 * 1024 // 8MB
	maxReverseDNSrecords        = 10000
	maxLogsPerDecodingError     = 4
	maxDecodingErrorCardinality = 1000
)

// DNSSnooper is a snopper of DNS queries
type DNSSnooper struct {
	stop       chan struct{}
	pcapHandle *pcap.Handle
	// gcache is goroutine-safe, but the cached values aren't
	reverseDNSMutex     sync.RWMutex
	reverseDNSCache     gcache.Cache
	decodingErrorCounts map[string]uint64 // for limiting
}

// NewDNSSnooper creates a new snooper of DNS queries
func NewDNSSnooper() (*DNSSnooper, error) {
	pcapHandle, err := newPcapHandle()
	if err != nil {
		return nil, err
	}
	reverseDNSCache := gcache.New(maxReverseDNSrecords).LRU().Build()

	s := &DNSSnooper{
		stop:                make(chan struct{}),
		pcapHandle:          pcapHandle,
		reverseDNSCache:     reverseDNSCache,
		decodingErrorCounts: map[string]uint64{},
	}
	go s.run()
	return s, nil
}

func newPcapHandle() (*pcap.Handle, error) {
	inactive, err := pcap.NewInactiveHandle("any")
	if err != nil {
		return nil, err
	}
	defer inactive.CleanUp()
	// pcap timeout blackmagic copied from Weave Net to reduce CPU consumption
	// see https://github.com/weaveworks/weave/commit/025315363d5ea8b8265f1b3ea800f24df2be51a4
	if err = inactive.SetTimeout(time.Duration(math.MaxInt64)); err != nil {
		return nil, err
	}
	if err = inactive.SetImmediateMode(true); err != nil {
		// If gopacket is compiled against an older pcap.h that
		// doesn't have pcap_set_immediate_mode, it supplies a dummy
		// definition that always returns PCAP_ERROR.  That becomes
		// "Generic error", which is not very helpful.  The real
		// pcap_set_immediate_mode never returns PCAP_ERROR, so this
		// turns it into a more informative message.
		if fmt.Sprint(err) == "Generic error" {
			return nil, fmt.Errorf("compiled against an old version of libpcap; please compile against libpcap-1.5.0 or later")
		}

		return nil, err
	}
	if err = inactive.SetBufferSize(bufSize); err != nil {
		return nil, err
	}
	pcapHandle, err := inactive.Activate()
	if err != nil {
		return nil, err
	}
	if err := pcapHandle.SetDirection(pcap.DirectionIn); err != nil {
		pcapHandle.Close()
		return nil, err
	}
	if err := pcapHandle.SetBPFFilter("inbound and port 53"); err != nil {
		pcapHandle.Close()
		return nil, err
	}

	return pcapHandle, nil
}

// CachedNamesForIP obtains the domains associated to an IP,
// obtained while snooping A-record queries
func (s *DNSSnooper) CachedNamesForIP(ip string) []string {
	result := []string{}
	if s == nil {
		return result
	}
	domains, err := s.reverseDNSCache.Get(ip)
	if err != nil {
		return result
	}
	s.reverseDNSMutex.RLock()
	for domain := range domains.(map[string]struct{}) {
		result = append(result, domain)
	}
	s.reverseDNSMutex.RUnlock()

	return result
}

// Stop makes the snooper stop inspecting DNS communications
func (s *DNSSnooper) Stop() {
	if s != nil {
		close(s.stop)
	}
}

// Gopacket doesn't provide direct support for DNS over TCP, see https://github.com/google/gopacket/issues/236
type tcpWithDNSSupport struct {
	tcp layers.TCP
}

func (m *tcpWithDNSSupport) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	return m.tcp.DecodeFromBytes(data, df)
}

func (m *tcpWithDNSSupport) CanDecode() gopacket.LayerClass { return m.tcp.CanDecode() }

// Determine if a TCP segment contains a full DNS message (i.e. not fragmented)
func (m *tcpWithDNSSupport) hasSelfContainedDNSPayload() bool {
	payload := m.tcp.LayerPayload()
	if len(payload) < 2 {
		return false
	}

	// Assume it's a self-contained DNS message if the Length field
	// matches the length of the TCP segment
	dnsLengthField := binary.BigEndian.Uint16(payload)
	return int(dnsLengthField) == len(payload)-2
}

func (m *tcpWithDNSSupport) NextLayerType() gopacket.LayerType {
	// TODO: deal with TCP fragmentation and out-of-order segments
	if (m.tcp.SrcPort == 53 || m.tcp.DstPort == 53) && m.hasSelfContainedDNSPayload() {
		return layers.LayerTypeDNS
	}
	return m.tcp.NextLayerType()
}

func (m *tcpWithDNSSupport) LayerPayload() []byte {
	payload := m.tcp.LayerPayload()
	if len(payload) > 1 && (m.tcp.SrcPort == 53 || m.tcp.DstPort == 53) {
		// Omit the DNS length field, only included
		// in TCP, in order to reuse the DNS UDP parser
		payload = payload[2:]
	}
	return payload
}

func (s *DNSSnooper) run() {
	var (
		decodedLayers []gopacket.LayerType
		dns           layers.DNS
		udp           layers.UDP
		tcp           tcpWithDNSSupport
		ip4           layers.IPv4
		ip6           layers.IPv6
		eth           layers.Ethernet
		dot1q         layers.Dot1Q
		sll           layers.LinuxSLL
	)

	// assumes that the "any" interface is being used (see https://wiki.wireshark.org/SLL)
	packetParser := gopacket.NewDecodingLayerParser(layers.LayerTypeLinuxSLL, &sll, &dot1q, &eth, &ip4, &ip6, &udp, &tcp, &dns)

	for {
		select {
		case <-s.stop:
			s.pcapHandle.Close()
			return
		default:
		}

		packet, _, err := s.pcapHandle.ZeroCopyReadPacketData()
		if err != nil {
			// TimeoutExpired is acceptable due to the Timeout black magic
			// on the handle.
			if err != pcap.NextErrorTimeoutExpired {
				log.Errorf("DNSSnooper: error reading packet data: %s", err)
			}
			continue
		}

		if err := packetParser.DecodeLayers(packet, &decodedLayers); err != nil {
			// LayerTypePayload indicates the TCP payload has non-DNS data, which we are not interested in
			if layer, ok := err.(gopacket.UnsupportedLayerType); !ok || gopacket.LayerType(layer) != gopacket.LayerTypePayload {
				s.handleDecodingError(err)
			}
			continue
		}

		for _, layerType := range decodedLayers {
			if layerType == layers.LayerTypeDNS {
				s.processDNSMessage(&dns)
			}
		}
	}
}

// handleDecodeError logs errors up to the maximum allowed count
func (s *DNSSnooper) handleDecodingError(err error) {
	// prevent potential memory leak
	if len(s.decodingErrorCounts) > maxDecodingErrorCardinality {
		return
	}

	str := err.Error()
	count := s.decodingErrorCounts[str]
	count++
	s.decodingErrorCounts[str] = count
	switch {
	case count == maxLogsPerDecodingError:
		log.Errorf("DNSSnooper: error decoding packet: %s (reached %d occurrences, silencing)", str, maxLogsPerDecodingError)
	case count < maxLogsPerDecodingError:
		log.Errorf("DNSSnooper: error decoding packet: %s", str)
	}
}

func (s *DNSSnooper) processDNSMessage(dns *layers.DNS) {

	// Only consider responses to singleton, A-record questions
	if !dns.QR || dns.ResponseCode != 0 || len(dns.Questions) != 1 {
		return
	}
	question := dns.Questions[0]
	if question.Type != layers.DNSTypeA || question.Class != layers.DNSClassIN {
		return
	}

	var (
		domainQueried = question.Name
		records       = append(dns.Answers, dns.Additionals...)
		ips           = map[string]struct{}{}
		aliases       = [][]byte{}
	)

	// Traverse all the CNAME records and the get the aliases. There are cases when the A record is for only one of the
	// aliases. We traverse CNAME records first because there is no guarantee that the A records will be the first ones
	for _, record := range records {
		if record.Type == layers.DNSTypeCNAME && record.Class == layers.DNSClassIN {
			aliases = append(aliases, record.CNAME)
		}
	}

	// Finally, get the answer
	for _, record := range records {
		if record.Type != layers.DNSTypeA || record.Class != layers.DNSClassIN {
			continue
		}
		if bytes.Equal(domainQueried, record.Name) {
			ips[record.IP.String()] = struct{}{}
			continue
		}
		for _, alias := range aliases {
			if bytes.Equal(alias, record.Name) {
				ips[record.IP.String()] = struct{}{}
				continue
			}
		}
	}

	// Update cache
	newDomain := string(domainQueried)
	log.Debugf("DNSSnooper: caught DNS lookup: %s -> %v", newDomain, ips)
	for ip := range ips {
		if existingDomains, err := s.reverseDNSCache.Get(ip); err != nil {
			s.reverseDNSCache.Set(ip, map[string]struct{}{newDomain: {}})
		} else {
			// TODO: Be smarter about the expiration of entries with pre-existing associated domains
			s.reverseDNSMutex.Lock()
			existingDomains.(map[string]struct{})[newDomain] = struct{}{}
			s.reverseDNSMutex.Unlock()
		}
	}
}
<commit_msg>Use break instead of continue.<commit_after>// +build linux,amd64 linux,ppc64le

// Build constraint to use this file for amd64 & ppc64le on Linux

package endpoint

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"math"
	"sync"
	"time"

	"github.com/bluele/gcache"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	log "github.com/sirupsen/logrus"
)

const (
	bufSize                     = 8 * 1024 * 1024 // 8MB
	maxReverseDNSrecords        = 10000
	maxLogsPerDecodingError     = 4
	maxDecodingErrorCardinality = 1000
)

// DNSSnooper is a snopper of DNS queries
type DNSSnooper struct {
	stop       chan struct{}
	pcapHandle *pcap.Handle
	// gcache is goroutine-safe, but the cached values aren't
	reverseDNSMutex     sync.RWMutex
	reverseDNSCache     gcache.Cache
	decodingErrorCounts map[string]uint64 // for limiting
}

// NewDNSSnooper creates a new snooper of DNS queries
func NewDNSSnooper() (*DNSSnooper, error) {
	pcapHandle, err := newPcapHandle()
	if err != nil {
		return nil, err
	}
	reverseDNSCache := gcache.New(maxReverseDNSrecords).LRU().Build()

	s := &DNSSnooper{
		stop:                make(chan struct{}),
		pcapHandle:          pcapHandle,
		reverseDNSCache:     reverseDNSCache,
		decodingErrorCounts: map[string]uint64{},
	}
	go s.run()
	return s, nil
}

func newPcapHandle() (*pcap.Handle, error) {
	inactive, err := pcap.NewInactiveHandle("any")
	if err != nil {
		return nil, err
	}
	defer inactive.CleanUp()
	// pcap timeout blackmagic copied from Weave Net to reduce CPU consumption
	// see https://github.com/weaveworks/weave/commit/025315363d5ea8b8265f1b3ea800f24df2be51a4
	if err = inactive.SetTimeout(time.Duration(math.MaxInt64)); err != nil {
		return nil, err
	}
	if err = inactive.SetImmediateMode(true); err != nil {
		// If gopacket is compiled against an older pcap.h that
		// doesn't have pcap_set_immediate_mode, it supplies a dummy
		// definition that always returns PCAP_ERROR.  That becomes
		// "Generic error", which is not very helpful.  The real
		// pcap_set_immediate_mode never returns PCAP_ERROR, so this
		// turns it into a more informative message.
		if fmt.Sprint(err) == "Generic error" {
			return nil, fmt.Errorf("compiled against an old version of libpcap; please compile against libpcap-1.5.0 or later")
		}

		return nil, err
	}
	if err = inactive.SetBufferSize(bufSize); err != nil {
		return nil, err
	}
	pcapHandle, err := inactive.Activate()
	if err != nil {
		return nil, err
	}
	if err := pcapHandle.SetDirection(pcap.DirectionIn); err != nil {
		pcapHandle.Close()
		return nil, err
	}
	if err := pcapHandle.SetBPFFilter("inbound and port 53"); err != nil {
		pcapHandle.Close()
		return nil, err
	}

	return pcapHandle, nil
}

// CachedNamesForIP obtains the domains associated to an IP,
// obtained while snooping A-record queries
func (s *DNSSnooper) CachedNamesForIP(ip string) []string {
	result := []string{}
	if s == nil {
		return result
	}
	domains, err := s.reverseDNSCache.Get(ip)
	if err != nil {
		return result
	}
	s.reverseDNSMutex.RLock()
	for domain := range domains.(map[string]struct{}) {
		result = append(result, domain)
	}
	s.reverseDNSMutex.RUnlock()

	return result
}

// Stop makes the snooper stop inspecting DNS communications
func (s *DNSSnooper) Stop() {
	if s != nil {
		close(s.stop)
	}
}

// Gopacket doesn't provide direct support for DNS over TCP, see https://github.com/google/gopacket/issues/236
type tcpWithDNSSupport struct {
	tcp layers.TCP
}

func (m *tcpWithDNSSupport) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	return m.tcp.DecodeFromBytes(data, df)
}

func (m *tcpWithDNSSupport) CanDecode() gopacket.LayerClass { return m.tcp.CanDecode() }

// Determine if a TCP segment contains a full DNS message (i.e. not fragmented)
func (m *tcpWithDNSSupport) hasSelfContainedDNSPayload() bool {
	payload := m.tcp.LayerPayload()
	if len(payload) < 2 {
		return false
	}

	// Assume it's a self-contained DNS message if the Length field
	// matches the length of the TCP segment
	dnsLengthField := binary.BigEndian.Uint16(payload)
	return int(dnsLengthField) == len(payload)-2
}

func (m *tcpWithDNSSupport) NextLayerType() gopacket.LayerType {
	// TODO: deal with TCP fragmentation and out-of-order segments
	if (m.tcp.SrcPort == 53 || m.tcp.DstPort == 53) && m.hasSelfContainedDNSPayload() {
		return layers.LayerTypeDNS
	}
	return m.tcp.NextLayerType()
}

func (m *tcpWithDNSSupport) LayerPayload() []byte {
	payload := m.tcp.LayerPayload()
	if len(payload) > 1 && (m.tcp.SrcPort == 53 || m.tcp.DstPort == 53) {
		// Omit the DNS length field, only included
		// in TCP, in order to reuse the DNS UDP parser
		payload = payload[2:]
	}
	return payload
}

func (s *DNSSnooper) run() {
	var (
		decodedLayers []gopacket.LayerType
		dns           layers.DNS
		udp           layers.UDP
		tcp           tcpWithDNSSupport
		ip4           layers.IPv4
		ip6           layers.IPv6
		eth           layers.Ethernet
		dot1q         layers.Dot1Q
		sll           layers.LinuxSLL
	)

	// assumes that the "any" interface is being used (see https://wiki.wireshark.org/SLL)
	packetParser := gopacket.NewDecodingLayerParser(layers.LayerTypeLinuxSLL, &sll, &dot1q, &eth, &ip4, &ip6, &udp, &tcp, &dns)

	for {
		select {
		case <-s.stop:
			s.pcapHandle.Close()
			return
		default:
		}

		packet, _, err := s.pcapHandle.ZeroCopyReadPacketData()
		if err != nil {
			// TimeoutExpired is acceptable due to the Timeout black magic
			// on the handle.
			if err != pcap.NextErrorTimeoutExpired {
				log.Errorf("DNSSnooper: error reading packet data: %s", err)
			}
			continue
		}

		if err := packetParser.DecodeLayers(packet, &decodedLayers); err != nil {
			// LayerTypePayload indicates the TCP payload has non-DNS data, which we are not interested in
			if layer, ok := err.(gopacket.UnsupportedLayerType); !ok || gopacket.LayerType(layer) != gopacket.LayerTypePayload {
				s.handleDecodingError(err)
			}
			continue
		}

		for _, layerType := range decodedLayers {
			if layerType == layers.LayerTypeDNS {
				s.processDNSMessage(&dns)
			}
		}
	}
}

// handleDecodeError logs errors up to the maximum allowed count
func (s *DNSSnooper) handleDecodingError(err error) {
	// prevent potential memory leak
	if len(s.decodingErrorCounts) > maxDecodingErrorCardinality {
		return
	}

	str := err.Error()
	count := s.decodingErrorCounts[str]
	count++
	s.decodingErrorCounts[str] = count
	switch {
	case count == maxLogsPerDecodingError:
		log.Errorf("DNSSnooper: error decoding packet: %s (reached %d occurrences, silencing)", str, maxLogsPerDecodingError)
	case count < maxLogsPerDecodingError:
		log.Errorf("DNSSnooper: error decoding packet: %s", str)
	}
}

func (s *DNSSnooper) processDNSMessage(dns *layers.DNS) {

	// Only consider responses to singleton, A-record questions
	if !dns.QR || dns.ResponseCode != 0 || len(dns.Questions) != 1 {
		return
	}
	question := dns.Questions[0]
	if question.Type != layers.DNSTypeA || question.Class != layers.DNSClassIN {
		return
	}

	var (
		domainQueried = question.Name
		records       = append(dns.Answers, dns.Additionals...)
		ips           = map[string]struct{}{}
		aliases       = [][]byte{}
	)

	// Traverse all the CNAME records and the get the aliases. There are cases when the A record is for only one of the
	// aliases. We traverse CNAME records first because there is no guarantee that the A records will be the first ones
	for _, record := range records {
		if record.Type == layers.DNSTypeCNAME && record.Class == layers.DNSClassIN {
			aliases = append(aliases, record.CNAME)
		}
	}

	// Finally, get the answer
	for _, record := range records {
		if record.Type != layers.DNSTypeA || record.Class != layers.DNSClassIN {
			continue
		}
		if bytes.Equal(domainQueried, record.Name) {
			ips[record.IP.String()] = struct{}{}
			continue
		}
		for _, alias := range aliases {
			if bytes.Equal(alias, record.Name) {
				ips[record.IP.String()] = struct{}{}
				break
			}
		}
	}

	// Update cache
	newDomain := string(domainQueried)
	log.Debugf("DNSSnooper: caught DNS lookup: %s -> %v", newDomain, ips)
	for ip := range ips {
		if existingDomains, err := s.reverseDNSCache.Get(ip); err != nil {
			s.reverseDNSCache.Set(ip, map[string]struct{}{newDomain: {}})
		} else {
			// TODO: Be smarter about the expiration of entries with pre-existing associated domains
			s.reverseDNSMutex.Lock()
			existingDomains.(map[string]struct{})[newDomain] = struct{}{}
			s.reverseDNSMutex.Unlock()
		}
	}
}
<|endoftext|>