<commit_before>package publish

import (
	"database/sql"
	"fmt"
	"reflect"
	"strings"

	"github.com/jinzhu/gorm"
)

type Resolver struct {
	Records      []interface{}
	Dependencies map[string]*Dependency
	DB           *DB
}

type Dependency struct {
	Type        reflect.Type
	PrimaryKeys []string
}

func IncludeValue(value string, values []string) bool {
	for _, v := range values {
		if v == value {
			return true
		}
	}
	return false
}

func (resolver *Resolver) SupportModel(model interface{}) bool {
	var supportedModels []string
	var reflectType = modelType(model)

	if value, ok := resolver.DB.DB.Get("publish:support_models"); ok {
		supportedModels = value.([]string)
	}

	return IncludeValue(reflectType.String(), supportedModels)
}

func (resolver *Resolver) AddDependency(dependency *Dependency) {
	name := dependency.Type.String()
	var newPrimaryKeys []string

	if dep, ok := resolver.Dependencies[name]; ok {
		for _, primaryKey := range dependency.PrimaryKeys {
			if !IncludeValue(primaryKey, dep.PrimaryKeys) {
				newPrimaryKeys = append(newPrimaryKeys, primaryKey)
				dep.PrimaryKeys = append(dep.PrimaryKeys, primaryKey)
			}
		}
	} else {
		resolver.Dependencies[name] = dependency
		newPrimaryKeys = dependency.PrimaryKeys
	}

	if len(newPrimaryKeys) > 0 {
		resolver.GetDependencies(dependency, newPrimaryKeys)
	}
}

func (resolver *Resolver) GetDependencies(dependency *Dependency, primaryKeys []string) {
	value := reflect.New(dependency.Type)
	fromScope := resolver.DB.DB.NewScope(value.Interface())

	draftDB := resolver.DB.DraftMode().Unscoped()
	for _, field := range fromScope.Fields() {
		if relationship := field.Relationship; relationship != nil {
			if resolver.SupportModel(field.Field.Interface()) {
				toType := modelType(field.Field.Interface())
				toScope := draftDB.NewScope(reflect.New(toType).Interface())
				draftTable := DraftTableName(toScope.TableName())
				var dependencyKeys []string
				var rows *sql.Rows
				var err error

				if relationship.Kind == "belongs_to" || relationship.Kind == "has_many" {
					sql := fmt.Sprintf("%v IN (?) and publish_status = ?", gorm.ToSnake(relationship.ForeignKey))
					rows, err = draftDB.Table(draftTable).Select(toScope.PrimaryKey()).Where(sql, primaryKeys, DIRTY).Rows()
				} else if relationship.Kind == "has_one" {
					fromTable := fromScope.TableName()
					fromPrimaryKey := fromScope.PrimaryKey()
					toTable := toScope.TableName()
					toPrimaryKey := toScope.PrimaryKey()

					sql := fmt.Sprintf("%v.%v IN (select %v.%v from %v where %v.%v IN (?)) and %v.publish_status = ?",
						toTable, toPrimaryKey, fromTable, relationship.ForeignKey, fromTable, fromTable, fromPrimaryKey, toTable)

					rows, err = draftDB.Table(draftTable).Select(toTable+"."+toPrimaryKey).Where(sql, primaryKeys, DIRTY).Rows()
				} else if relationship.Kind == "many_to_many" {
				}

				if rows != nil && err == nil {
					for rows.Next() {
						var primaryKey interface{}
						rows.Scan(&primaryKey)
						dependencyKeys = append(dependencyKeys, fmt.Sprintf("%v", primaryKey))
					}

					dependency := Dependency{Type: toType, PrimaryKeys: dependencyKeys}
					resolver.AddDependency(&dependency)
				}
			}
		}
	}
}

func (resolver *Resolver) Publish() {
	for _, record := range resolver.Records {
		if resolver.SupportModel(record) {
			scope := &gorm.Scope{Value: record}
			dependency := Dependency{Type: modelType(record), PrimaryKeys: []string{fmt.Sprintf("%v", scope.PrimaryKeyValue())}}
			resolver.AddDependency(&dependency)
		}
	}

	for _, dependency := range resolver.Dependencies {
		value := reflect.New(dependency.Type)
		productionScope := resolver.DB.ProductionMode().NewScope(value.Interface())
		productionTable := productionScope.TableName()
		productionPrimaryKey := productionScope.PrimaryKey()
		draftTable := DraftTableName(productionTable)

		var columns []string
		for _, field := range productionScope.Fields() {
			if field.IsNormal {
				columns = append(columns, field.DBName)
			}
		}

		var insertColumns []string
		for _, column := range columns {
			insertColumns = append(insertColumns, fmt.Sprintf("%v.%v", productionTable, column))
		}

		var selectColumns []string
		for _, column := range columns {
			selectColumns = append(selectColumns, fmt.Sprintf("%v.%v", draftTable, column))
		}

		deleteSql := fmt.Sprintf("DELETE FROM %v WHERE %v.%v IN (?);",
			productionTable, productionTable, productionScope.PrimaryKey())
		resolver.DB.DB.Exec(deleteSql, dependency.PrimaryKeys)

		sql := fmt.Sprintf("INSERT INTO %v (%v) SELECT %v from %v where %v.%v in (?);",
			productionTable, strings.Join(insertColumns, " ,"), strings.Join(selectColumns, " ,"),
			draftTable, draftTable, productionPrimaryKey)
		resolver.DB.DB.Exec(sql, dependency.PrimaryKeys)
	}
}
<commit_msg>Change publish status after publish<commit_after>package publish

import (
	"database/sql"
	"fmt"
	"reflect"
	"strings"

	"github.com/jinzhu/gorm"
)

type Resolver struct {
	Records      []interface{}
	Dependencies map[string]*Dependency
	DB           *DB
}

type Dependency struct {
	Type        reflect.Type
	PrimaryKeys []string
}

func IncludeValue(value string, values []string) bool {
	for _, v := range values {
		if v == value {
			return true
		}
	}
	return false
}

func (resolver *Resolver) SupportModel(model interface{}) bool {
	var supportedModels []string
	var reflectType = modelType(model)

	if value, ok := resolver.DB.DB.Get("publish:support_models"); ok {
		supportedModels = value.([]string)
	}

	return IncludeValue(reflectType.String(), supportedModels)
}

func (resolver *Resolver) AddDependency(dependency *Dependency) {
	name := dependency.Type.String()
	var newPrimaryKeys []string

	if dep, ok := resolver.Dependencies[name]; ok {
		for _, primaryKey := range dependency.PrimaryKeys {
			if !IncludeValue(primaryKey, dep.PrimaryKeys) {
				newPrimaryKeys = append(newPrimaryKeys, primaryKey)
				dep.PrimaryKeys = append(dep.PrimaryKeys, primaryKey)
			}
		}
	} else {
		resolver.Dependencies[name] = dependency
		newPrimaryKeys = dependency.PrimaryKeys
	}

	if len(newPrimaryKeys) > 0 {
		resolver.GetDependencies(dependency, newPrimaryKeys)
	}
}

func (resolver *Resolver) GetDependencies(dependency *Dependency, primaryKeys []string) {
	value := reflect.New(dependency.Type)
	fromScope := resolver.DB.DB.NewScope(value.Interface())

	draftDB := resolver.DB.DraftMode().Unscoped()
	for _, field := range fromScope.Fields() {
		if relationship := field.Relationship; relationship != nil {
			if resolver.SupportModel(field.Field.Interface()) {
				toType := modelType(field.Field.Interface())
				toScope := draftDB.NewScope(reflect.New(toType).Interface())
				draftTable := DraftTableName(toScope.TableName())
				var dependencyKeys []string
				var rows *sql.Rows
				var err error

				if relationship.Kind == "belongs_to" || relationship.Kind == "has_many" {
					sql := fmt.Sprintf("%v IN (?) and publish_status = ?", gorm.ToSnake(relationship.ForeignKey))
					rows, err = draftDB.Table(draftTable).Select(toScope.PrimaryKey()).Where(sql, primaryKeys, DIRTY).Rows()
				} else if relationship.Kind == "has_one" {
					fromTable := fromScope.TableName()
					fromPrimaryKey := fromScope.PrimaryKey()
					toTable := toScope.TableName()
					toPrimaryKey := toScope.PrimaryKey()

					sql := fmt.Sprintf("%v.%v IN (select %v.%v from %v where %v.%v IN (?)) and %v.publish_status = ?",
						toTable, toPrimaryKey, fromTable, relationship.ForeignKey, fromTable, fromTable, fromPrimaryKey, toTable)

					rows, err = draftDB.Table(draftTable).Select(toTable+"."+toPrimaryKey).Where(sql, primaryKeys, DIRTY).Rows()
				} else if relationship.Kind == "many_to_many" {
				}

				if rows != nil && err == nil {
					for rows.Next() {
						var primaryKey interface{}
						rows.Scan(&primaryKey)
						dependencyKeys = append(dependencyKeys, fmt.Sprintf("%v", primaryKey))
					}

					dependency := Dependency{Type: toType, PrimaryKeys: dependencyKeys}
					resolver.AddDependency(&dependency)
				}
			}
		}
	}
}

func (resolver *Resolver) Publish() {
	for _, record := range resolver.Records {
		if resolver.SupportModel(record) {
			scope := &gorm.Scope{Value: record}
			dependency := Dependency{Type: modelType(record), PrimaryKeys: []string{fmt.Sprintf("%v", scope.PrimaryKeyValue())}}
			resolver.AddDependency(&dependency)
		}
	}

	for _, dependency := range resolver.Dependencies {
		value := reflect.New(dependency.Type)
		productionScope := resolver.DB.ProductionMode().NewScope(value.Interface())
		productionTable := productionScope.TableName()
		productionPrimaryKey := productionScope.PrimaryKey()
		draftTable := DraftTableName(productionTable)

		var columns []string
		for _, field := range productionScope.Fields() {
			if field.IsNormal {
				columns = append(columns, field.DBName)
			}
		}

		var insertColumns []string
		for _, column := range columns {
			insertColumns = append(insertColumns, fmt.Sprintf("%v.%v", productionTable, column))
		}

		var selectColumns []string
		for _, column := range columns {
			selectColumns = append(selectColumns, fmt.Sprintf("%v.%v", draftTable, column))
		}

		deleteSql := fmt.Sprintf("DELETE FROM %v WHERE %v.%v IN (?);",
			productionTable, productionTable, productionScope.PrimaryKey())
		resolver.DB.DB.Exec(deleteSql, dependency.PrimaryKeys)

		publishSql := fmt.Sprintf("INSERT INTO %v (%v) SELECT %v FROM %v WHERE %v.%v IN (?)",
			productionTable, strings.Join(insertColumns, " ,"), strings.Join(selectColumns, " ,"),
			draftTable, draftTable, productionPrimaryKey)
		resolver.DB.DB.Exec(publishSql, dependency.PrimaryKeys)

		updateStateSql := fmt.Sprintf("UPDATE %v SET publish_status = ? WHERE %v.%v IN (?)", draftTable, draftTable, productionPrimaryKey)
		resolver.DB.DB.Exec(updateStateSql, PUBLISHED, dependency.PrimaryKeys)
	}
}
<|endoftext|>


<commit_before>package helpers

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"image"
	"image/jpeg"
	"strconv"
	"strings"
	"time"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/bwmarrin/discordgo"
	jsoniter "github.com/json-iterator/go"
	"github.com/pkg/errors"

	_ "image/gif"
	_ "image/png"
)

func CanRevert(item models.ElasticEventlog) bool {
	if item.Reverted {
		return false
	}

	if len(item.Changes) <= 0 && len(item.Options) <= 0 {
		return false
	}

	switch item.ActionType {
	case models.EventlogTypeChannelUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"channel_name", "channel_topic", "channel_nsfw", "channel_bitrate", "channel_parentid", "channel_permissionoverwrites"},
			nil,
		) {
			return true
		}
	case models.EventlogTypeRoleUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"role_name", "role_mentionable", "role_hoist", "role_color", "role_permissions"},
			nil,
		) {
			return true
		}
	case models.EventlogTypeMemberUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"member_nick"},
			[]string{"member_roles_added", "member_roles_removed"},
		) {
			return true
		}
	case models.EventlogTypeGuildUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"guild_name", "guild_icon_object", "guild_region", "guild_afkchannelid", "guild_afktimeout", "guild_verificationlevel", "guild_defaultmessagenotifications"},
			nil,
		) {
			return true
		}
	case models.EventlogTypeEmojiDelete:
		return true
	case models.EventlogTypeEmojiUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"emoji_name"},
			nil,
		) {
			return true
		}
	case models.EventlogTypeChannelDelete:
		if containsAllowedChangesOrOptions(
			item,
			nil,
			[]string{"channel_name", "channel_type", "channel_topic", "channel_nsfw", "channel_bitrate", "channel_parentid", "channel_permissionoverwrites"},
		) {
			return true
		}
	}

	return false
}

func containsAllowedChangesOrOptions(eventlogEntry models.ElasticEventlog, changes []string, options []string) bool {
	if len(eventlogEntry.Changes) > 0 {
		for _, change := range eventlogEntry.Changes {
			for _, key := range changes {
				if change.Key == key {
					return true
				}
			}
		}
	}
	if len(eventlogEntry.Options) > 0 {
		for _, option := range eventlogEntry.Options {
			for _, key := range options {
				if option.Key == key {
					return true
				}
			}
		}
	}
	return false
}

func Revert(eventlogID, userID string, item models.ElasticEventlog) (err error) {
	switch item.ActionType {
	case models.EventlogTypeChannelUpdate:
		channel, err := GetChannel(item.TargetID)
		if err != nil {
			return err
		}

		channelEdit := &discordgo.ChannelEdit{ // restore ints because go
			Position: channel.Position,
			Bitrate:  channel.Bitrate,
		}
		for _, change := range item.Changes {
			switch change.Key {
			case "channel_name":
				channelEdit.Name = change.OldValue
			case "channel_topic":
				channelEdit.Topic = change.OldValue
			case "channel_nsfw":
				channelEdit.NSFW = GetStringAsBool(change.OldValue)
			case "channel_bitrate":
				newBitrate, err := strconv.Atoi(change.OldValue)
				if err == nil {
					channelEdit.Bitrate = newBitrate
				}
			case "channel_parentid":
				channelEdit.ParentID = change.OldValue
			case "channel_permissionoverwrites":
				newOverwrites := make([]*discordgo.PermissionOverwrite, 0)
				oldOverwritesTexts := strings.Split(change.OldValue, ";")
				for _, oldOverwriteText := range oldOverwritesTexts {
					var oldOverwrite *discordgo.PermissionOverwrite
					err = jsoniter.UnmarshalFromString(oldOverwriteText, &oldOverwrite)
					RelaxLog(err)
					if err == nil && oldOverwrite != nil {
						newOverwrites = append(newOverwrites, oldOverwrite)
					}
				}
				channelEdit.PermissionOverwrites = newOverwrites
			}
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).ChannelEditComplex(item.TargetID, channelEdit)
		if err != nil {
			return err
		}

		return logRevert(channel.GuildID, userID, eventlogID)
	case models.EventlogTypeRoleUpdate:
		role, err := cache.GetSession().SessionForGuildS(item.GuildID).State.Role(item.GuildID, item.TargetID)
		if err != nil {
			return err
		}

		newName := role.Name
		newMentionable := role.Mentionable
		newHoist := role.Hoist
		newColor := role.Color
		newPermissions := role.Permissions

		for _, change := range item.Changes {
			switch change.Key {
			case "role_name":
				newName = change.OldValue
			case "role_mentionable":
				newMentionable = GetStringAsBool(change.OldValue)
			case "role_hoist":
				newHoist = GetStringAsBool(change.OldValue)
			case "role_color":
				newColor = GetDiscordColorFromHex(change.OldValue)
			case "role_permissions":
				tempPermissions, err := strconv.Atoi(change.OldValue)
				if err == nil {
					newPermissions = tempPermissions
				}
			}
		}

		role, err = cache.GetSession().SessionForGuildS(item.GuildID).GuildRoleEdit(item.GuildID, item.TargetID, newName, newColor, newHoist, newPermissions, newMentionable)
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeMemberUpdate:
		for _, change := range item.Changes {
			switch change.Key {
			case "member_nick":
				err = cache.GetSession().SessionForGuildS(item.GuildID).GuildMemberNickname(item.GuildID, item.TargetID, change.OldValue)
				if err != nil {
					return err
				}
			}
		}

		for _, option := range item.Options {
			switch option.Key {
			case "member_roles_added":
				for _, roleID := range strings.Split(option.Value, ";") {
					err = cache.GetSession().SessionForGuildS(item.GuildID).GuildMemberRoleRemove(item.GuildID, item.TargetID, roleID)
					if err != nil {
						return err
					}
				}
			case "member_roles_removed":
				for _, roleID := range strings.Split(option.Value, ";") {
					err = cache.GetSession().SessionForGuildS(item.GuildID).GuildMemberRoleAdd(item.GuildID, item.TargetID, roleID)
					if err != nil {
						return err
					}
				}
			}
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeGuildUpdate:
		guild, err := GetGuildWithoutApi(item.TargetID)
		if err != nil {
			return err
		}

		guildParams := discordgo.GuildParams{
			DefaultMessageNotifications: guild.DefaultMessageNotifications,
			AfkTimeout:                  guild.AfkTimeout,
			AfkChannelID:                guild.AfkChannelID,
		}

		for _, change := range item.Changes {
			switch change.Key {
			case "guild_name":
				guildParams.Name = change.OldValue
			case "guild_icon_object":
				// retrieve previous icon
				iconData, err := RetrieveFile(change.OldValue)
				if err != nil {
					return err
				}

				// read icon
				iconImage, _, err := image.Decode(bytes.NewReader(iconData))
				if err != nil {
					return err
				}

				// convert icon to jpeg
				var jpegIconBuffer bytes.Buffer
				err = jpeg.Encode(bufio.NewWriter(&jpegIconBuffer), iconImage, nil)
				if err != nil {
					return err
				}

				// encode jpeg to base64
				iconJpegBase64 := "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(jpegIconBuffer.Bytes())

				guildParams.Icon = iconJpegBase64
			case "guild_region":
				guildParams.Region = change.OldValue
			case "guild_afkchannelid":
				guildParams.AfkChannelID = change.OldValue
			case "guild_afktimeout":
				newTimeout, err := strconv.Atoi(change.OldValue)
				RelaxLog(err)
				if err == nil {
					guildParams.AfkTimeout = newTimeout
				}
			case "guild_verificationlevel":
				newVerificationLevel, err := strconv.Atoi(change.OldValue)
				RelaxLog(err)
				if err == nil {
					level := discordgo.VerificationLevel(newVerificationLevel)
					guildParams.VerificationLevel = &level
				}
			case "guild_defaultmessagenotifications":
				newDefaultMessageNotifications, err := strconv.Atoi(change.OldValue)
				RelaxLog(err)
				if err == nil {
					guildParams.DefaultMessageNotifications = newDefaultMessageNotifications
				}
			}
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).GuildEdit(item.TargetID, guildParams)
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeEmojiDelete:
		var emojiName, emojiImage, emojiURL string
		var emojiRoles []string

		emojiURL = discordgo.EndpointEmoji(item.TargetID)
		for _, option := range item.Options {
			switch option.Key {
			case "emoji_animated":
				if GetStringAsBool(option.Value) {
					emojiURL = strings.Replace(emojiURL, ".png", ".gif", -1)
				}
			}
		}

		// retrieve previous icon
		iconData, err := NetGetUAWithError(emojiURL, DEFAULT_UA)
		if err != nil {
			return err
		}

		// read icon
		filetype, err := SniffMime(iconData)
		if err != nil {
			return err
		}

		// encode jpeg to base64
		emojiImage = "data:" + filetype + ";base64," + base64.StdEncoding.EncodeToString(iconData)

		for _, option := range item.Options {
			switch option.Key {
			case "emoji_name":
				emojiName = option.Value
			case "emoji_roleids":
				if option.Value != "" {
					emojiRoles = strings.Split(option.Value, ";")
				}
			}
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).GuildEmojiCreate(item.GuildID, emojiName, emojiImage, emojiRoles)
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeEmojiUpdate:
		emoji, err := cache.GetSession().SessionForGuildS(item.GuildID).State.Emoji(item.GuildID, item.TargetID)
		if err != nil {
			return err
		}

		var emojiName string

		for _, change := range item.Changes {
			switch change.Key {
			case "emoji_name":
				emojiName = change.OldValue
			}
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).GuildEmojiEdit(item.GuildID, item.TargetID, emojiName, emoji.Roles)
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeChannelDelete:
		var channelName, channelTopic, channelParentID string
		var channelType discordgo.ChannelType
		var channelNSFW bool
		var channelBitrate int
		channelOverwrites := make([]*discordgo.PermissionOverwrite, 0)

		for _, option := range item.Options {
			switch option.Key {
			case "channel_name":
				channelName = option.Value
			case "channel_type":
				level, err := strconv.Atoi(option.Value)
				if err == nil {
					channelType = discordgo.ChannelType(level)
				}
			case "channel_topic":
				channelTopic = option.Value
			case "channel_nsfw":
				if GetStringAsBool(option.Value) {
					channelNSFW = true
				}
			case "channel_bitrate":
				bitrate, err := strconv.Atoi(option.Value)
				if err == nil {
					channelBitrate = bitrate
				}
			case "channel_parentid":
				channelParentID = option.Value
			case "channel_permissionoverwrites":
				overwritesTexts := strings.Split(option.Value, ";")
				for _, overwriteText := range overwritesTexts {
					var overwrite *discordgo.PermissionOverwrite
					err = jsoniter.UnmarshalFromString(overwriteText, &overwrite)
					RelaxLog(err)
					if err == nil && overwrite != nil {
						channelOverwrites = append(channelOverwrites, overwrite)
					}
				}
			}
		}

		channel, err := cache.GetSession().SessionForGuildS(item.GuildID).GuildChannelCreate(item.GuildID, channelName, channelType)
		if err != nil {
			return err
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).ChannelEditComplex(channel.ID, &discordgo.ChannelEdit{
			Name:                 channelName,
			Topic:                channelTopic,
			NSFW:                 channelNSFW,
			Bitrate:              channelBitrate,
			PermissionOverwrites: channelOverwrites,
			ParentID:             channelParentID,
		})
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	}

	return errors.New("eventlog action type not supported")
}

func logRevert(guildID, userID, eventlogID string) error {
	// add new eventlog entry for revert
	_, err := EventlogLog(time.Now(), guildID, eventlogID,
		models.EventlogTargetTypeRobyulEventlogItem, userID,
		models.EventlogTypeRobyulActionRevert, "",
		nil,
		nil,
		false,
	)
	if err != nil {
		return err
	}

	// get issuer user
	user, err := GetUserWithoutAPI(userID)
	if err != nil {
		user = new(discordgo.User)
		user.ID = userID
		user.Username = "N/A"
		user.Discriminator = "N/A"
	}

	// add option to reverted action with information
	err = EventlogLogUpdate(
		eventlogID,
		"",
		[]models.ElasticEventlogOption{{
			Key:   "reverted_by_userid",
			Value: user.ID,
			Type:  models.EventlogTargetTypeUser,
		}},
		nil,
		"",
		false,
		true,
	)
	return err
}
<commit_msg>[eventlog] set guild icon mime type on demand<commit_after>package helpers

import (
	"encoding/base64"
	"strconv"
	"strings"
	"time"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/Seklfreak/Robyul2/models"
	"github.com/bwmarrin/discordgo"
	jsoniter "github.com/json-iterator/go"
	"github.com/pkg/errors"
)

func CanRevert(item models.ElasticEventlog) bool {
	if item.Reverted {
		return false
	}

	if len(item.Changes) <= 0 && len(item.Options) <= 0 {
		return false
	}

	switch item.ActionType {
	case models.EventlogTypeChannelUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"channel_name", "channel_topic", "channel_nsfw", "channel_bitrate", "channel_parentid", "channel_permissionoverwrites"},
			nil,
		) {
			return true
		}
	case models.EventlogTypeRoleUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"role_name", "role_mentionable", "role_hoist", "role_color", "role_permissions"},
			nil,
		) {
			return true
		}
	case models.EventlogTypeMemberUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"member_nick"},
			[]string{"member_roles_added", "member_roles_removed"},
		) {
			return true
		}
	case models.EventlogTypeGuildUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"guild_name", "guild_icon_object", "guild_region", "guild_afkchannelid", "guild_afktimeout", "guild_verificationlevel", "guild_defaultmessagenotifications"},
			nil,
		) {
			return true
		}
	case models.EventlogTypeEmojiDelete:
		return true
	case models.EventlogTypeEmojiUpdate:
		if containsAllowedChangesOrOptions(
			item,
			[]string{"emoji_name"},
			nil,
		) {
			return true
		}
	case models.EventlogTypeChannelDelete:
		if containsAllowedChangesOrOptions(
			item,
			nil,
			[]string{"channel_name", "channel_type", "channel_topic", "channel_nsfw", "channel_bitrate", "channel_parentid", "channel_permissionoverwrites"},
		) {
			return true
		}
	}

	return false
}

func containsAllowedChangesOrOptions(eventlogEntry models.ElasticEventlog, changes []string, options []string) bool {
	if len(eventlogEntry.Changes) > 0 {
		for _, change := range eventlogEntry.Changes {
			for _, key := range changes {
				if change.Key == key {
					return true
				}
			}
		}
	}
	if len(eventlogEntry.Options) > 0 {
		for _, option := range eventlogEntry.Options {
			for _, key := range options {
				if option.Key == key {
					return true
				}
			}
		}
	}
	return false
}

func Revert(eventlogID, userID string, item models.ElasticEventlog) (err error) {
	switch item.ActionType {
	case models.EventlogTypeChannelUpdate:
		channel, err := GetChannel(item.TargetID)
		if err != nil {
			return err
		}

		channelEdit := &discordgo.ChannelEdit{ // restore ints because go
			Position: channel.Position,
			Bitrate:  channel.Bitrate,
		}
		for _, change := range item.Changes {
			switch change.Key {
			case "channel_name":
				channelEdit.Name = change.OldValue
			case "channel_topic":
				channelEdit.Topic = change.OldValue
			case "channel_nsfw":
				channelEdit.NSFW = GetStringAsBool(change.OldValue)
			case "channel_bitrate":
				newBitrate, err := strconv.Atoi(change.OldValue)
				if err == nil {
					channelEdit.Bitrate = newBitrate
				}
			case "channel_parentid":
				channelEdit.ParentID = change.OldValue
			case "channel_permissionoverwrites":
				newOverwrites := make([]*discordgo.PermissionOverwrite, 0)
				oldOverwritesTexts := strings.Split(change.OldValue, ";")
				for _, oldOverwriteText := range oldOverwritesTexts {
					var oldOverwrite *discordgo.PermissionOverwrite
					err = jsoniter.UnmarshalFromString(oldOverwriteText, &oldOverwrite)
					RelaxLog(err)
					if err == nil && oldOverwrite != nil {
						newOverwrites = append(newOverwrites, oldOverwrite)
					}
				}
				channelEdit.PermissionOverwrites = newOverwrites
			}
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).ChannelEditComplex(item.TargetID, channelEdit)
		if err != nil {
			return err
		}

		return logRevert(channel.GuildID, userID, eventlogID)
	case models.EventlogTypeRoleUpdate:
		role, err := cache.GetSession().SessionForGuildS(item.GuildID).State.Role(item.GuildID, item.TargetID)
		if err != nil {
			return err
		}

		newName := role.Name
		newMentionable := role.Mentionable
		newHoist := role.Hoist
		newColor := role.Color
		newPermissions := role.Permissions

		for _, change := range item.Changes {
			switch change.Key {
			case "role_name":
				newName = change.OldValue
			case "role_mentionable":
				newMentionable = GetStringAsBool(change.OldValue)
			case "role_hoist":
				newHoist = GetStringAsBool(change.OldValue)
			case "role_color":
				newColor = GetDiscordColorFromHex(change.OldValue)
			case "role_permissions":
				tempPermissions, err := strconv.Atoi(change.OldValue)
				if err == nil {
					newPermissions = tempPermissions
				}
			}
		}

		role, err = cache.GetSession().SessionForGuildS(item.GuildID).GuildRoleEdit(item.GuildID, item.TargetID, newName, newColor, newHoist, newPermissions, newMentionable)
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeMemberUpdate:
		for _, change := range item.Changes {
			switch change.Key {
			case "member_nick":
				err = cache.GetSession().SessionForGuildS(item.GuildID).GuildMemberNickname(item.GuildID, item.TargetID, change.OldValue)
				if err != nil {
					return err
				}
			}
		}

		for _, option := range item.Options {
			switch option.Key {
			case "member_roles_added":
				for _, roleID := range strings.Split(option.Value, ";") {
					err = cache.GetSession().SessionForGuildS(item.GuildID).GuildMemberRoleRemove(item.GuildID, item.TargetID, roleID)
					if err != nil {
						return err
					}
				}
			case "member_roles_removed":
				for _, roleID := range strings.Split(option.Value, ";") {
					err = cache.GetSession().SessionForGuildS(item.GuildID).GuildMemberRoleAdd(item.GuildID, item.TargetID, roleID)
					if err != nil {
						return err
					}
				}
			}
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeGuildUpdate:
		guild, err := GetGuildWithoutApi(item.TargetID)
		if err != nil {
			return err
		}

		guildParams := discordgo.GuildParams{
			DefaultMessageNotifications: guild.DefaultMessageNotifications,
			AfkTimeout:                  guild.AfkTimeout,
			AfkChannelID:                guild.AfkChannelID,
		}

		for _, change := range item.Changes {
			switch change.Key {
			case "guild_name":
				guildParams.Name = change.OldValue
			case "guild_icon_object":
				// retrieve previous icon
				iconData, err := RetrieveFile(change.OldValue)
				if err != nil {
					return err
				}

				// convert icon to jpeg
				mimeType, _ := SniffMime(iconData)

				// encode jpeg to base64
				iconJpegBase64 := "data:" + mimeType + ";base64," + base64.StdEncoding.EncodeToString(iconData)

				guildParams.Icon = iconJpegBase64
			case "guild_region":
				guildParams.Region = change.OldValue
			case "guild_afkchannelid":
				guildParams.AfkChannelID = change.OldValue
			case "guild_afktimeout":
				newTimeout, err := strconv.Atoi(change.OldValue)
				RelaxLog(err)
				if err == nil {
					guildParams.AfkTimeout = newTimeout
				}
			case "guild_verificationlevel":
				newVerificationLevel, err := strconv.Atoi(change.OldValue)
				RelaxLog(err)
				if err == nil {
					level := discordgo.VerificationLevel(newVerificationLevel)
					guildParams.VerificationLevel = &level
				}
			case "guild_defaultmessagenotifications":
				newDefaultMessageNotifications, err := strconv.Atoi(change.OldValue)
				RelaxLog(err)
				if err == nil {
					guildParams.DefaultMessageNotifications = newDefaultMessageNotifications
				}
			}
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).GuildEdit(item.TargetID, guildParams)
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeEmojiDelete:
		var emojiName, emojiImage, emojiURL string
		var emojiRoles []string

		emojiURL = discordgo.EndpointEmoji(item.TargetID)
		for _, option := range item.Options {
			switch option.Key {
			case "emoji_animated":
				if GetStringAsBool(option.Value) {
					emojiURL = strings.Replace(emojiURL, ".png", ".gif", -1)
				}
			}
		}

		// retrieve previous icon
		iconData, err := NetGetUAWithError(emojiURL, DEFAULT_UA)
		if err != nil {
			return err
		}

		// read icon
		filetype, err := SniffMime(iconData)
		if err != nil {
			return err
		}

		// encode jpeg to base64
		emojiImage = "data:" + filetype + ";base64," + base64.StdEncoding.EncodeToString(iconData)

		for _, option := range item.Options {
			switch option.Key {
			case "emoji_name":
				emojiName = option.Value
			case "emoji_roleids":
				if option.Value != "" {
					emojiRoles = strings.Split(option.Value, ";")
				}
			}
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).GuildEmojiCreate(item.GuildID, emojiName, emojiImage, emojiRoles)
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeEmojiUpdate:
		emoji, err := cache.GetSession().SessionForGuildS(item.GuildID).State.Emoji(item.GuildID, item.TargetID)
		if err != nil {
			return err
		}

		var emojiName string

		for _, change := range item.Changes {
			switch change.Key {
			case "emoji_name":
				emojiName = change.OldValue
			}
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).GuildEmojiEdit(item.GuildID, item.TargetID, emojiName, emoji.Roles)
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	case models.EventlogTypeChannelDelete:
		var channelName, channelTopic, channelParentID string
		var channelType discordgo.ChannelType
		var channelNSFW bool
		var channelBitrate int
		channelOverwrites := make([]*discordgo.PermissionOverwrite, 0)

		for _, option := range item.Options {
			switch option.Key {
			case "channel_name":
				channelName = option.Value
			case "channel_type":
				level, err := strconv.Atoi(option.Value)
				if err == nil {
					channelType = discordgo.ChannelType(level)
				}
			case "channel_topic":
				channelTopic = option.Value
			case "channel_nsfw":
				if GetStringAsBool(option.Value) {
					channelNSFW = true
				}
			case "channel_bitrate":
				bitrate, err := strconv.Atoi(option.Value)
				if err == nil {
					channelBitrate = bitrate
				}
			case "channel_parentid":
				channelParentID = option.Value
			case "channel_permissionoverwrites":
				overwritesTexts := strings.Split(option.Value, ";")
				for _, overwriteText := range overwritesTexts {
					var overwrite *discordgo.PermissionOverwrite
					err = jsoniter.UnmarshalFromString(overwriteText, &overwrite)
					RelaxLog(err)
					if err == nil && overwrite != nil {
						channelOverwrites = append(channelOverwrites, overwrite)
					}
				}
			}
		}

		channel, err := cache.GetSession().SessionForGuildS(item.GuildID).GuildChannelCreate(item.GuildID, channelName, channelType)
		if err != nil {
			return err
		}

		_, err = cache.GetSession().SessionForGuildS(item.GuildID).ChannelEditComplex(channel.ID, &discordgo.ChannelEdit{
			Name:                 channelName,
			Topic:                channelTopic,
			NSFW:                 channelNSFW,
			Bitrate:              channelBitrate,
			PermissionOverwrites: channelOverwrites,
			ParentID:             channelParentID,
		})
		if err != nil {
			return err
		}

		return logRevert(item.GuildID, userID, eventlogID)
	}

	return errors.New("eventlog action type not supported")
}

func logRevert(guildID, userID, eventlogID string) error {
	// add new eventlog entry for revert
	_, err := EventlogLog(time.Now(), guildID, eventlogID,
		models.EventlogTargetTypeRobyulEventlogItem, userID,
		models.EventlogTypeRobyulActionRevert, "",
		nil,
		nil,
		false,
	)
	if err != nil {
		return err
	}

	// get issuer user
	user, err := GetUserWithoutAPI(userID)
	if err != nil {
		user = new(discordgo.User)
		user.ID = userID
		user.Username = "N/A"
		user.Discriminator = "N/A"
	}

	// add option to reverted action with information
	err = EventlogLogUpdate(
		eventlogID,
		"",
		[]models.ElasticEventlogOption{{
			Key:   "reverted_by_userid",
			Value: user.ID,
			Type:  models.EventlogTargetTypeUser,
		}},
		nil,
		"",
		false,
		true,
	)
	return err
}
<|endoftext|>


<commit_before>package sudoku

import (
	"testing"
)

func TestForcingChains(t *testing.T) {

	//Steps to test this:
	//* In the forcing chain helper, calculate the steps once, then
	//pass them in each time in a list of ~10 calls to solveTechniqueTEstHelper that we know are valid here.
	//* VERIFY MANUALLY that each step that is returned is actually a valid application of forcingchains.

	options := solveTechniqueTestHelperOptions{
		checkAllSteps: true,
		debugPrint:    true,
	}

	grid, solver, steps := humanSolveTechniqueTestHelperStepGenerator(t,
		"forcingchain_test1.sdk", "Forcing Chain", options)

	options.stepsToCheck.grid = grid
	options.stepsToCheck.solver = solver
	options.stepsToCheck.steps = steps

	//OK, now we'll walk through all of the options in a loop and make sure they all show
	//up in the solve steps.

	type loopOptions struct {
		targetCells  []cellRef
		targetNums   IntSlice
		pointerCells []cellRef
		pointerNums  IntSlice
		description  string
	}

	//Tester puzzle: http://www.komoroske.com/sudoku/index.php?puzzle=Q6Ur5iYGINSUFcyocqaY6G91DpttiqYzs

	tests := []loopOptions{
		{
			targetCells:  []cellRef{{0, 1}},
			targetNums:   IntSlice([]int{7}),
			pointerCells: []cellRef{{1, 0}},
			pointerNums:  IntSlice([]int{1, 2}),
			description:  "cell (1,0) only has two options, 1 and 2, and if you put either one in and see the chain of implications it leads to, both ones end up with 7 in cell (0,1), so we can just fill that number in",
		},
		{
			targetCells:  []cellRef{{1, 0}},
			targetNums:   IntSlice([]int{1}),
			pointerCells: []cellRef{{0, 6}},
			pointerNums:  IntSlice([]int{3, 7}),
			//Explicitly don't test description after the first one.
		},
		{
			targetCells:  []cellRef{{5, 1}},
			targetNums:   IntSlice([]int{1}),
			pointerCells: []cellRef{{0, 6}},
			pointerNums:  IntSlice([]int{3, 7}),
		},
		{
			targetCells:  []cellRef{{8, 3}},
			targetNums:   IntSlice([]int{7}),
			pointerCells: []cellRef{{7, 8}},
			pointerNums:  IntSlice([]int{2, 7}),
		},
		//Skipping 0,1 / 7 / 5,7 / 1,3 because I think implications force it wrong.
		{
			targetCells:  []cellRef{{8, 3}},
			targetNums:   IntSlice([]int{7}),
			pointerCells: []cellRef{{5, 7}},
			pointerNums:  IntSlice([]int{1, 3}),
		},
		{
			targetCells:  []cellRef{{1, 8}},
			targetNums:   IntSlice([]int{4}),
			pointerCells: []cellRef{{4, 0}},
			pointerNums:  IntSlice([]int{1, 2}),
		},
	}

	if len(tests) != len(steps) {
		t.Error("We didn't have enough tests for all of the steps that forcing chains returned. Got", len(tests), "expected", len(steps))
	}

	for _, test := range tests {

		options.targetCells = test.targetCells
		options.targetNums = test.targetNums
		options.pointerCells = test.pointerCells
		options.pointerNums = test.pointerNums
		options.description = test.description

		humanSolveTechniqueTestHelper(t, "forcingchain_test1.sdk", "Forcing Chain", options)
	}

	//TODO: test all other valid steps that could be found at this grid state for this technique.

}
<commit_msg>TESTS FAIL. Added another test.<commit_after>package sudoku

import (
	"testing"
)

func TestForcingChains(t *testing.T) {

	//Steps to test this:
	//* In the forcing chain helper, calculate the steps once, then
	//pass them in each time in a list of ~10 calls to solveTechniqueTEstHelper that we know are valid here.
	//* VERIFY MANUALLY that each step that is returned is actually a valid application of forcingchains.

	options := solveTechniqueTestHelperOptions{
		checkAllSteps: true,
		debugPrint:    true,
	}

	grid, solver, steps := humanSolveTechniqueTestHelperStepGenerator(t,
		"forcingchain_test1.sdk", "Forcing Chain", options)

	options.stepsToCheck.grid = grid
	options.stepsToCheck.solver = solver
	options.stepsToCheck.steps = steps

	//OK, now we'll walk through all of the options in a loop and make sure they all show
	//up in the solve steps.

	type loopOptions struct {
		targetCells  []cellRef
		targetNums   IntSlice
		pointerCells []cellRef
		pointerNums  IntSlice
		description  string
	}

	//Tester puzzle: http://www.komoroske.com/sudoku/index.php?puzzle=Q6Ur5iYGINSUFcyocqaY6G91DpttiqYzs

	tests := []loopOptions{
		{
			targetCells:  []cellRef{{0, 1}},
			targetNums:   IntSlice([]int{7}),
			pointerCells: []cellRef{{1, 0}},
			pointerNums:  IntSlice([]int{1, 2}),
			description:  "cell (1,0) only has two options, 1 and 2, and if you put either one in and see the chain of implications it leads to, both ones end up with 7 in cell (0,1), so we can just fill that number in",
		},
		{
			targetCells:  []cellRef{{1, 0}},
			targetNums:   IntSlice([]int{1}),
			pointerCells: []cellRef{{0, 6}},
			pointerNums:  IntSlice([]int{3, 7}),
			//Explicitly don't test description after the first one.
		},
		{
			targetCells:  []cellRef{{5, 1}},
			targetNums:   IntSlice([]int{1}),
			pointerCells: []cellRef{{0, 6}},
			pointerNums:  IntSlice([]int{3, 7}),
		},
		{
			targetCells:  []cellRef{{8, 3}},
			targetNums:   IntSlice([]int{7}),
			pointerCells: []cellRef{{7, 8}},
			pointerNums:  IntSlice([]int{2, 7}),
		},
		//Skipping 0,1 / 7 / 5,7 / 1,3 because I think implications force it wrong.
		{
			targetCells:  []cellRef{{8, 3}},
			targetNums:   IntSlice([]int{7}),
			pointerCells: []cellRef{{5, 7}},
			pointerNums:  IntSlice([]int{1, 3}),
		},
		{
			targetCells:  []cellRef{{1, 8}},
			targetNums:   IntSlice([]int{4}),
			pointerCells: []cellRef{{4, 0}},
			pointerNums:  IntSlice([]int{1, 2}),
		},
		//This next one's particularly long implication chain
		{
			targetCells:  []cellRef{{0, 1}},
			targetNums:   IntSlice([]int{7}),
			pointerCells: []cellRef{{4, 0}},
			pointerNums:  IntSlice([]int{1, 2}),
		},
	}

	if len(tests) != len(steps) {
		t.Error("We didn't have enough tests for all of the steps that forcing chains returned. Got", len(tests), "expected", len(steps))
	}

	for _, test := range tests {

		options.targetCells = test.targetCells
		options.targetNums = test.targetNums
		options.pointerCells = test.pointerCells
		options.pointerNums = test.pointerNums
		options.description = test.description

		humanSolveTechniqueTestHelper(t, "forcingchain_test1.sdk", "Forcing Chain", options)
	}

	//TODO: test all other valid steps that could be found at this grid state for this technique.

}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 logrus

import (
	"testing"

	"github.com/onsi/gomega"
)

func TestListLoggers(t *testing.T) {
	logRegistry := NewLogRegistry()

	gomega.RegisterTestingT(t)
	loggers := logRegistry.ListLoggers()
	gomega.Expect(loggers).NotTo(gomega.BeNil())

	lg, found := loggers[DefaultLoggerName]
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(lg).NotTo(gomega.BeNil())
}

func TestNewLogger(t *testing.T) {
	logRegistry := NewLogRegistry()

	const loggerName = "myLogger"
	gomega.RegisterTestingT(t)
	lg := logRegistry.NewLogger(loggerName)
	gomega.Expect(lg).NotTo(gomega.BeNil())

	loggers := logRegistry.ListLoggers()
	gomega.Expect(loggers).NotTo(gomega.BeNil())

	fromRegistry, found := loggers[loggerName]
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(fromRegistry).NotTo(gomega.BeNil())
}

func TestGetSetLevel(t *testing.T) {
	logRegistry := NewLogRegistry()

	gomega.RegisterTestingT(t)
	const level = "error"
	//existing logger
	err := logRegistry.SetLevel(DefaultLoggerName, level)
	gomega.Expect(err).To(gomega.BeNil())

	loggers := logRegistry.ListLoggers()
	gomega.Expect(loggers).NotTo(gomega.BeNil())

	logger, found := loggers[DefaultLoggerName]
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(logger).NotTo(gomega.BeNil())
	gomega.Expect(loggers[DefaultLoggerName]).To(gomega.BeEquivalentTo(level))

	currentLevel, err := logRegistry.GetLevel(DefaultLoggerName)
	gomega.Expect(err).To(gomega.BeNil())
	gomega.Expect(level).To(gomega.BeEquivalentTo(currentLevel))

	//non-existing logger
	err = logRegistry.SetLevel("unknown", level)
	gomega.Expect(err).NotTo(gomega.BeNil())

	_, err = logRegistry.GetLevel("unknown")
	gomega.Expect(err).NotTo(gomega.BeNil())
}

func TestGetLoggerByName(t *testing.T) {
	logRegistry := NewLogRegistry()

	const (
		loggerA = "myLoggerA"
		loggerB = "myLoggerB"
	)
	lgA := logRegistry.NewLogger(loggerA)
	gomega.Expect(lgA).NotTo(gomega.BeNil())

	lgB := logRegistry.NewLogger(loggerB)
	gomega.Expect(lgB).NotTo(gomega.BeNil())

	returnedA, found := logRegistry.Lookup(loggerA)
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(returnedA).To(gomega.BeEquivalentTo(lgA))

	returnedB, found := logRegistry.Lookup(loggerB)
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(returnedB).To(gomega.BeEquivalentTo(lgB))

	unknown, found := logRegistry.Lookup("unknown")
	gomega.Expect(found).To(gomega.BeFalse())
	gomega.Expect(unknown).To(gomega.BeNil())
}

func TestClearRegistry(t *testing.T) {
	logRegistry := NewLogRegistry()

	const (
		loggerA = "loggerA"
		loggerB = "loggerB"
	)
	lgA := NewLogger(loggerA)
	gomega.Expect(lgA).NotTo(gomega.BeNil())

	lgB := NewLogger(loggerB)
	gomega.Expect(lgB).NotTo(gomega.BeNil())

	logRegistry.ClearRegistry()

	_, found := logRegistry.Lookup(loggerA)
	gomega.Expect(found).To(gomega.BeFalse())

	_, found = logRegistry.Lookup(loggerB)
	gomega.Expect(found).To(gomega.BeFalse())

	_, found = logRegistry.Lookup(DefaultLoggerName)
	gomega.Expect(found).To(gomega.BeTrue())
}
<commit_msg>Fix GetSetLevel test<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// 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 logrus

import (
	"testing"

	"github.com/onsi/gomega"
)

func TestListLoggers(t *testing.T) {
	logRegistry := NewLogRegistry()

	gomega.RegisterTestingT(t)
	loggers := logRegistry.ListLoggers()
	gomega.Expect(loggers).NotTo(gomega.BeNil())

	lg, found := loggers[DefaultLoggerName]
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(lg).NotTo(gomega.BeNil())
}

func TestNewLogger(t *testing.T) {
	logRegistry := NewLogRegistry()

	const loggerName = "myLogger"
	gomega.RegisterTestingT(t)
	lg := logRegistry.NewLogger(loggerName)
	gomega.Expect(lg).NotTo(gomega.BeNil())

	loggers := logRegistry.ListLoggers()
	gomega.Expect(loggers).NotTo(gomega.BeNil())

	fromRegistry, found := loggers[loggerName]
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(fromRegistry).NotTo(gomega.BeNil())
}

func TestGetSetLevel(t *testing.T) {
	logRegistry := NewLogRegistry()

	gomega.RegisterTestingT(t)
	const level = "error"
	//existing logger
	err := logRegistry.SetLevel(DefaultLoggerName, level)
	gomega.Expect(err).To(gomega.BeNil())

	loggers := logRegistry.ListLoggers()
	gomega.Expect(loggers).NotTo(gomega.BeNil())

	logger, found := loggers[DefaultLoggerName]
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(logger).NotTo(gomega.BeNil())
	gomega.Expect(loggers[DefaultLoggerName]).To(gomega.BeEquivalentTo(level))

	currentLevel, err := logRegistry.GetLevel(DefaultLoggerName)
	gomega.Expect(err).To(gomega.BeNil())
	gomega.Expect(level).To(gomega.BeEquivalentTo(currentLevel))

	//non-existing logger
	err = logRegistry.SetLevel("unknown", level)
	gomega.Expect(err).To(gomega.BeNil()) // will be kept in logger level map in registry

	_, err = logRegistry.GetLevel("unknown")
	gomega.Expect(err).NotTo(gomega.BeNil())
}

func TestGetLoggerByName(t *testing.T) {
	logRegistry := NewLogRegistry()

	const (
		loggerA = "myLoggerA"
		loggerB = "myLoggerB"
	)
	lgA := logRegistry.NewLogger(loggerA)
	gomega.Expect(lgA).NotTo(gomega.BeNil())

	lgB := logRegistry.NewLogger(loggerB)
	gomega.Expect(lgB).NotTo(gomega.BeNil())

	returnedA, found := logRegistry.Lookup(loggerA)
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(returnedA).To(gomega.BeEquivalentTo(lgA))

	returnedB, found := logRegistry.Lookup(loggerB)
	gomega.Expect(found).To(gomega.BeTrue())
	gomega.Expect(returnedB).To(gomega.BeEquivalentTo(lgB))

	unknown, found := logRegistry.Lookup("unknown")
	gomega.Expect(found).To(gomega.BeFalse())
	gomega.Expect(unknown).To(gomega.BeNil())
}

func TestClearRegistry(t *testing.T) {
	logRegistry := NewLogRegistry()

	const (
		loggerA = "loggerA"
		loggerB = "loggerB"
	)
	lgA := NewLogger(loggerA)
	gomega.Expect(lgA).NotTo(gomega.BeNil())

	lgB := NewLogger(loggerB)
	gomega.Expect(lgB).NotTo(gomega.BeNil())

	logRegistry.ClearRegistry()

	_, found := logRegistry.Lookup(loggerA)
	gomega.Expect(found).To(gomega.BeFalse())

	_, found = logRegistry.Lookup(loggerB)
	gomega.Expect(found).To(gomega.BeFalse())

	_, found = logRegistry.Lookup(DefaultLoggerName)
	gomega.Expect(found).To(gomega.BeTrue())
}
<|endoftext|>


<commit_before>package testing

import (
	"database/sql"
	"fmt"
	"log"
	"net/url"
	"os"
	"os/exec"
	"syscall"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach-go/testserver"
	// Import postgres driver.
	_ "github.com/lib/pq"
)

// application represents a single instance of an application running an ORM and
// exposing an HTTP REST API.
type application struct {
	language string
	orm      string
}

func (app application) name() string {
	return fmt.Sprintf("%s/%s", app.language, app.orm)
}

func (app application) dir() string {
	return fmt.Sprintf("../%s", app.name())
}

func (app application) dbName() string {
	return fmt.Sprintf("company_%s", app.orm)
}

// customURLSchemes contains custom schemes for database URLs that are needed
// for test apps that rely on a custom ORM dialect.
var customURLSchemes = map[application]string{
	{language: "python", orm: "sqlalchemy"}: "cockroachdb",
}

// initTestDatabase launches a test database as a subprocess.
func initTestDatabase(t *testing.T, app application) (*sql.DB, *url.URL, func()) {
	ts, err := testserver.NewTestServer()
	if err != nil {
		t.Fatal(err)
	}

	if err := ts.Start(); err != nil {
		t.Fatal(err)
	}

	url := ts.PGURL()
	if url == nil {
		t.Fatalf("url not found")
	}
	url.Path = app.dbName()

	db, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	ts.WaitForInit(db)

	// Create the database if it does not exist.
	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS " + app.dbName()); err != nil {
		t.Fatal(err)
	}

	if scheme, ok := customURLSchemes[app]; ok {
		url.Scheme = scheme
	}
	return db, url, func() {
		_ = db.Close()
		ts.Stop()
	}
}

// initORMApp launches an ORM application as a subprocess and returns a
// function that terminates that process.
func initORMApp(app application, dbURL *url.URL) (func() error, error) {
	cmd := exec.Command("make", "start", "-C", app.dir(), "ADDR="+dbURL.String())
	cmd.Stdout = os.Stderr
	cmd.Stderr = os.Stderr

	// make will launch the application in a child process, and this is the most
	// straightforward way to kill all ancestors.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	killCmd := func() error {
		if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
			return err
		}
		// This error is expected.
		if err := cmd.Wait(); err.Error() != "signal: "+syscall.SIGKILL.String() {
			return err
		}

		// Killing a process is not instant. For example, with the Hibernate server,
		// it often takes ~10 seconds for the listen port to become available after
		// this function is called. This is despite the above code that issues a
		// SIGKILL to the process group for the test server.
		for {
			if !(apiHandler{}).canDial() {
				break
			}
			log.Printf("waiting for app server port to become available")
			time.Sleep(time.Second)
		}

		return nil
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("command %s failed to start: args=%s", cmd.Args, err)
	}

	const maxWait = 3 * time.Minute
	const waitDelay = 250 * time.Millisecond

	for waited := time.Duration(0); ; waited += waitDelay {
		if processState := cmd.ProcessState; processState != nil && processState.Exited() {
			return nil, fmt.Errorf("command %s exited: %v", cmd.Args, cmd.Wait())
		}
		if err := (apiHandler{}).ping(app.name()); err != nil {
			if waited > maxWait {
				if err := killCmd(); err != nil {
					log.Printf("failed to kill command %s with PID %d: %s", cmd.Args, cmd.ProcessState.Pid(), err)
				}
				return nil, err
			}
			time.Sleep(waitDelay)
			continue
		}
		return killCmd, nil
	}
}

func testORM(
	t *testing.T, language, orm string, tableNames testTableNames, columnNames testColumnNames,
) {
	app := application{
		language: language,
		orm:      orm,
	}

	db, dbURL, stopDB := initTestDatabase(t, app)
	defer stopDB()

	td := testDriver{
		db:          db,
		dbName:      app.dbName(),
		tableNames:  tableNames,
		columnNames: columnNames,
	}

	t.Run("FirstRun", func(t *testing.T) {
		stopApp, err := initORMApp(app, dbURL)
		if err != nil {
			t.Fatal(err)
		}
		defer func() {
			if err := stopApp(); err != nil {
				t.Fatal(err)
			}
		}()

		// Test that the correct tables were generated.
		t.Run("GeneratedTables", td.TestGeneratedTables)

		// Test that the correct columns in those tables were generated.
		t.Run("GeneratedColumns", parallelTestGroup{
			"CustomersTable":     td.TestGeneratedCustomersTableColumns,
			"ProductsTable":      td.TestGeneratedProductsTableColumns,
			"OrdersTable":        td.TestGeneratedOrdersTableColumns,
			"OrderProductsTable": td.TestGeneratedOrderProductsTableColumns,
		}.T)

		// Test that the tables begin empty.
		t.Run("EmptyTables", parallelTestGroup{
			"CustomersTable":     td.TestCustomersEmpty,
			"ProductsTable":      td.TestProductsTableEmpty,
			"OrdersTable":        td.TestOrdersTableEmpty,
			"OrderProductsTable": td.TestOrderProductsTableEmpty,
		}.T)

		// Test that the API returns empty sets for each collection.
		t.Run("RetrieveFromAPIBeforeCreation", parallelTestGroup{
			"Customers": td.TestRetrieveCustomersBeforeCreation,
			"Products":  td.TestRetrieveProductsBeforeCreation,
			"Orders":    td.TestRetrieveOrdersBeforeCreation,
		}.T)

		// Test the creation of initial objects.
		t.Run("CreateCustomer", td.TestCreateCustomer)
		t.Run("CreateProduct", td.TestCreateProduct)

		// Test that the API returns what we just created.
		t.Run("RetrieveFromAPIAfterInitialCreation", parallelTestGroup{
			"Customers": td.TestRetrieveCustomerAfterCreation,
			"Products":  td.TestRetrieveProductAfterCreation,
		}.T)

		// Test the creation of dependent objects.
		t.Run("CreateOrder", td.TestCreateOrder)

		// Test that the API returns what we just created.
		t.Run("RetrieveFromAPIAfterDependentCreation", parallelTestGroup{
			"Order": td.TestRetrieveProductAfterCreation,
		}.T)
	})

	t.Run("SecondRun", func(t *testing.T) {
		stopApp, err := initORMApp(app, dbURL)
		if err != nil {
			t.Fatal(err)
		}
		defer func() {
			if err := stopApp(); err != nil {
				t.Fatal(err)
			}
		}()

		// Test that the API still returns all created objects.
		t.Run("RetrieveFromAPIAfterRestart", parallelTestGroup{
			"Customers": td.TestRetrieveCustomerAfterCreation,
			"Products":  td.TestRetrieveProductAfterCreation,
			"Order":     td.TestRetrieveProductAfterCreation,
		}.T)
	})
}

func TestGORM(t *testing.T) {
	testORM(t, "go", "gorm", defaultTestTableNames, defaultTestColumnNames)
}

func TestHibernate(t *testing.T) {
	testORM(t, "java", "hibernate", defaultTestTableNames, defaultTestColumnNames)
}

func TestSequelize(t *testing.T) {
	testORM(t, "node", "sequelize", defaultTestTableNames, defaultTestColumnNames)
}

func TestSQLAlchemy(t *testing.T) {
	testORM(t, "python", "sqlalchemy", defaultTestTableNames, defaultTestColumnNames)
}

func TestDjango(t *testing.T) {
	testORM(t, "python", "django", djangoTestTableNames, djangoTestColumnNames)
}

func TestActiveRecord(t *testing.T) {
	testORM(t, "ruby", "activerecord", defaultTestTableNames, defaultTestColumnNames)
}

func TestActiveRecord4(t *testing.T) {
	testORM(t, "ruby", "ar4", defaultTestTableNames, defaultTestColumnNames)
}
<commit_msg>Removing the django test until it stops flaking<commit_after>package testing

import (
	"database/sql"
	"fmt"
	"log"
	"net/url"
	"os"
	"os/exec"
	"syscall"
	"testing"
	"time"

	"github.com/cockroachdb/cockroach-go/testserver"
	// Import postgres driver.
	_ "github.com/lib/pq"
)

// application represents a single instance of an application running an ORM and
// exposing an HTTP REST API.
type application struct {
	language string
	orm      string
}

func (app application) name() string {
	return fmt.Sprintf("%s/%s", app.language, app.orm)
}

func (app application) dir() string {
	return fmt.Sprintf("../%s", app.name())
}

func (app application) dbName() string {
	return fmt.Sprintf("company_%s", app.orm)
}

// customURLSchemes contains custom schemes for database URLs that are needed
// for test apps that rely on a custom ORM dialect.
var customURLSchemes = map[application]string{
	{language: "python", orm: "sqlalchemy"}: "cockroachdb",
}

// initTestDatabase launches a test database as a subprocess.
func initTestDatabase(t *testing.T, app application) (*sql.DB, *url.URL, func()) {
	ts, err := testserver.NewTestServer()
	if err != nil {
		t.Fatal(err)
	}

	if err := ts.Start(); err != nil {
		t.Fatal(err)
	}

	url := ts.PGURL()
	if url == nil {
		t.Fatalf("url not found")
	}
	url.Path = app.dbName()

	db, err := sql.Open("postgres", url.String())
	if err != nil {
		t.Fatal(err)
	}

	ts.WaitForInit(db)

	// Create the database if it does not exist.
	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS " + app.dbName()); err != nil {
		t.Fatal(err)
	}

	if scheme, ok := customURLSchemes[app]; ok {
		url.Scheme = scheme
	}
	return db, url, func() {
		_ = db.Close()
		ts.Stop()
	}
}

// initORMApp launches an ORM application as a subprocess and returns a
// function that terminates that process.
func initORMApp(app application, dbURL *url.URL) (func() error, error) {
	cmd := exec.Command("make", "start", "-C", app.dir(), "ADDR="+dbURL.String())
	cmd.Stdout = os.Stderr
	cmd.Stderr = os.Stderr

	// make will launch the application in a child process, and this is the most
	// straightforward way to kill all ancestors.
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	killCmd := func() error {
		if err := syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL); err != nil {
			return err
		}
		// This error is expected.
		if err := cmd.Wait(); err.Error() != "signal: "+syscall.SIGKILL.String() {
			return err
		}

		// Killing a process is not instant. For example, with the Hibernate server,
		// it often takes ~10 seconds for the listen port to become available after
		// this function is called. This is despite the above code that issues a
		// SIGKILL to the process group for the test server.
		for {
			if !(apiHandler{}).canDial() {
				break
			}
			log.Printf("waiting for app server port to become available")
			time.Sleep(time.Second)
		}

		return nil
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("command %s failed to start: args=%s", cmd.Args, err)
	}

	const maxWait = 3 * time.Minute
	const waitDelay = 250 * time.Millisecond

	for waited := time.Duration(0); ; waited += waitDelay {
		if processState := cmd.ProcessState; processState != nil && processState.Exited() {
			return nil, fmt.Errorf("command %s exited: %v", cmd.Args, cmd.Wait())
		}
		if err := (apiHandler{}).ping(app.name()); err != nil {
			if waited > maxWait {
				if err := killCmd(); err != nil {
					log.Printf("failed to kill command %s with PID %d: %s", cmd.Args, cmd.ProcessState.Pid(), err)
				}
				return nil, err
			}
			time.Sleep(waitDelay)
			continue
		}
		return killCmd, nil
	}
}

func testORM(
	t *testing.T, language, orm string, tableNames testTableNames, columnNames testColumnNames,
) {
	app := application{
		language: language,
		orm:      orm,
	}

	db, dbURL, stopDB := initTestDatabase(t, app)
	defer stopDB()

	td := testDriver{
		db:          db,
		dbName:      app.dbName(),
		tableNames:  tableNames,
		columnNames: columnNames,
	}

	t.Run("FirstRun", func(t *testing.T) {
		stopApp, err := initORMApp(app, dbURL)
		if err != nil {
			t.Fatal(err)
		}
		defer func() {
			if err := stopApp(); err != nil {
				t.Fatal(err)
			}
		}()

		// Test that the correct tables were generated.
		t.Run("GeneratedTables", td.TestGeneratedTables)

		// Test that the correct columns in those tables were generated.
		t.Run("GeneratedColumns", parallelTestGroup{
			"CustomersTable":     td.TestGeneratedCustomersTableColumns,
			"ProductsTable":      td.TestGeneratedProductsTableColumns,
			"OrdersTable":        td.TestGeneratedOrdersTableColumns,
			"OrderProductsTable": td.TestGeneratedOrderProductsTableColumns,
		}.T)

		// Test that the tables begin empty.
		t.Run("EmptyTables", parallelTestGroup{
			"CustomersTable":     td.TestCustomersEmpty,
			"ProductsTable":      td.TestProductsTableEmpty,
			"OrdersTable":        td.TestOrdersTableEmpty,
			"OrderProductsTable": td.TestOrderProductsTableEmpty,
		}.T)

		// Test that the API returns empty sets for each collection.
		t.Run("RetrieveFromAPIBeforeCreation", parallelTestGroup{
			"Customers": td.TestRetrieveCustomersBeforeCreation,
			"Products":  td.TestRetrieveProductsBeforeCreation,
			"Orders":    td.TestRetrieveOrdersBeforeCreation,
		}.T)

		// Test the creation of initial objects.
		t.Run("CreateCustomer", td.TestCreateCustomer)
		t.Run("CreateProduct", td.TestCreateProduct)

		// Test that the API returns what we just created.
		t.Run("RetrieveFromAPIAfterInitialCreation", parallelTestGroup{
			"Customers": td.TestRetrieveCustomerAfterCreation,
			"Products":  td.TestRetrieveProductAfterCreation,
		}.T)

		// Test the creation of dependent objects.
		t.Run("CreateOrder", td.TestCreateOrder)

		// Test that the API returns what we just created.
		t.Run("RetrieveFromAPIAfterDependentCreation", parallelTestGroup{
			"Order": td.TestRetrieveProductAfterCreation,
		}.T)
	})

	t.Run("SecondRun", func(t *testing.T) {
		stopApp, err := initORMApp(app, dbURL)
		if err != nil {
			t.Fatal(err)
		}
		defer func() {
			if err := stopApp(); err != nil {
				t.Fatal(err)
			}
		}()

		// Test that the API still returns all created objects.
		t.Run("RetrieveFromAPIAfterRestart", parallelTestGroup{
			"Customers": td.TestRetrieveCustomerAfterCreation,
			"Products":  td.TestRetrieveProductAfterCreation,
			"Order":     td.TestRetrieveProductAfterCreation,
		}.T)
	})
}

func TestGORM(t *testing.T) {
	testORM(t, "go", "gorm", defaultTestTableNames, defaultTestColumnNames)
}

func TestHibernate(t *testing.T) {
	testORM(t, "java", "hibernate", defaultTestTableNames, defaultTestColumnNames)
}

func TestSequelize(t *testing.T) {
	testORM(t, "node", "sequelize", defaultTestTableNames, defaultTestColumnNames)
}

func TestSQLAlchemy(t *testing.T) {
	testORM(t, "python", "sqlalchemy", defaultTestTableNames, defaultTestColumnNames)
}

// func TestDjango(t *testing.T) {
// 	testORM(t, "python", "django", djangoTestTableNames, djangoTestColumnNames)
// }

func TestActiveRecord(t *testing.T) {
	testORM(t, "ruby", "activerecord", defaultTestTableNames, defaultTestColumnNames)
}

func TestActiveRecord4(t *testing.T) {
	testORM(t, "ruby", "ar4", defaultTestTableNames, defaultTestColumnNames)
}
<|endoftext|>


<commit_before>package levigo

// #cgo LDFLAGS: -lleveldb
// #include "levigo.h"
import "C"

import (
	"unsafe"
)

// WriteBatch is a batching of Puts, and Deletes to be written atomically to a
// database. A WriteBatch is written when passed to DB.Write.
//
// To prevent memory leaks, call Close when the program no longer needs the
// WriteBatch object.
type WriteBatch struct {
	wbatch *C.leveldb_writebatch_t
}

// NewWriteBatch creates a fully allocated WriteBatch.
func NewWriteBatch() *WriteBatch {
	wb := C.leveldb_writebatch_create()
	return &WriteBatch{wb}
}

// Close releases the underlying memory of a WriteBatch.
func (w *WriteBatch) Close() {
	C.leveldb_writebatch_destroy(w.wbatch)
}

// Put places a key-value pair into the WriteBatch for writing later.
//
// Both the key and value byte slices may be reused as WriteBatch takes a copy
// of them before returning.
//
func (w *WriteBatch) Put(key, value []byte) {
	// leveldb_writebatch_put, and _delete call memcpy() (by way of
	// Memtable::Add) when called, so we do not need to worry about these
	// []byte being reclaimed by GC.
	C.leveldb_writebatch_put(w.wbatch,
		(*C.char)(unsafe.Pointer(&key[0])), C.size_t(len(key)),
		(*C.char)(unsafe.Pointer(&value[0])), C.size_t(len(value)))
}

// Delete queues a deletion of the data at key to be deleted later.
//
// The key byte slice may be reused safely. Delete takes a copy of
// them before returning.
func (w *WriteBatch) Delete(key []byte) {
	C.leveldb_writebatch_delete(w.wbatch,
		(*C.char)(unsafe.Pointer(&key[0])), C.size_t(len(key)))
}

// Clear removes all the enqueued Put and Deletes in the WriteBatch.
func (w *WriteBatch) Clear() {
	C.leveldb_writebatch_clear(w.wbatch)
}
<commit_msg>fixed: Method WriteBatch.Put() should have consistent semantic as DB.Put()<commit_after>package levigo

// #cgo LDFLAGS: -lleveldb
// #include "levigo.h"
import "C"

import (
	"unsafe"
)

// WriteBatch is a batching of Puts, and Deletes to be written atomically to a
// database. A WriteBatch is written when passed to DB.Write.
//
// To prevent memory leaks, call Close when the program no longer needs the
// WriteBatch object.
type WriteBatch struct {
	wbatch *C.leveldb_writebatch_t
}

// NewWriteBatch creates a fully allocated WriteBatch.
func NewWriteBatch() *WriteBatch {
	wb := C.leveldb_writebatch_create()
	return &WriteBatch{wb}
}

// Close releases the underlying memory of a WriteBatch.
func (w *WriteBatch) Close() {
	C.leveldb_writebatch_destroy(w.wbatch)
}

// Put places a key-value pair into the WriteBatch for writing later.
//
// Both the key and value byte slices may be reused as WriteBatch takes a copy
// of them before returning.
//
func (w *WriteBatch) Put(key, value []byte) {
	// leveldb_writebatch_put, and _delete call memcpy() (by way of
	// Memtable::Add) when called, so we do not need to worry about these
	// []byte being reclaimed by GC.
	var k, v *C.char
	if len(key) != 0 {
		k = (*C.char)(unsafe.Pointer(&key[0]))
	}
	if len(value) != 0 {
		v = (*C.char)(unsafe.Pointer(&value[0]))
	}

	lenk := len(key)
	lenv := len(value)

	C.leveldb_writebatch_put(w.wbatch, k, C.size_t(lenk), v, C.size_t(lenv))
}

// Delete queues a deletion of the data at key to be deleted later.
//
// The key byte slice may be reused safely. Delete takes a copy of
// them before returning.
func (w *WriteBatch) Delete(key []byte) {
	C.leveldb_writebatch_delete(w.wbatch,
		(*C.char)(unsafe.Pointer(&key[0])), C.size_t(len(key)))
}

// Clear removes all the enqueued Put and Deletes in the WriteBatch.
func (w *WriteBatch) Clear() {
	C.leveldb_writebatch_clear(w.wbatch)
}
<|endoftext|>


<commit_before>package main

import "bytes"

func better(parts []string) []string {
	found := false
	var buffer bytes.Buffer
	var newParts []string

	for _, val := range parts {
		if found {
			newParts = append(newParts, val)
			continue
		}
		if val[0] == '"' {
			if val[len(val)-1] == '"' {
				found = true
				val = val[:len(val)-1]
				buffer.WriteString(val[1:])
				newParts = append(newParts, buffer.String())
			} else {
				buffer.WriteString(val[1:])
			}
		} else if buffer.Len() != 0 {
			if val[len(val)-1] == '"' {
				found = true
				val = val[:len(val)-1]
				buffer.WriteString(" " + val)
				newParts = append(newParts, buffer.String())
			} else {
				buffer.WriteString(" " + val)
			}
		} else {
			newParts = append(newParts, val)
		}
	}

	return newParts
}
<commit_msg>fromquotes: clean better<commit_after>package main

import (
	"bytes"
	"strings"
)

func better(in []string) []string {
	found := false
	var buffer bytes.Buffer
	out := make([]string, 0, len(in))

	const q string = `"`

	var size, lastPos int
	//var first, last byte
	for _, row := range in {
		size = len(row)
		lastPos = size - 1
		//first = row[0]
		//last = row[lastPos]
		if found {
			out = append(out, row)
			continue
		}
		if strings.HasPrefix(row, q) {
			if strings.HasSuffix(row, q) {
				found = true
				row = row[:lastPos]
				buffer.WriteString(row[1:])
				out = append(out, buffer.String())
			} else {
				buffer.WriteString(row[1:])
			}
		} else if buffer.Len() != 0 {
			if strings.HasSuffix(row, q) {
				found = true
				row = row[:lastPos]
				buffer.WriteString(" " + row)
				out = append(out, buffer.String())
			} else {
				buffer.WriteString(" " + row)
			}
		} else {
			out = append(out, row)
		}
	}

	return out
}
<|endoftext|>


<commit_before>// Copyright 2018 Diego Bernardes. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package repository

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/pkg/errors"

	"github.com/diegobernardes/flare"
)

// Subscription implements the data layer for the subscription service.
type Subscription struct {
	mutex              sync.RWMutex
	resourceRepository resourceRepositorier
	documentRepository flare.DocumentRepositorier

	// resourceID -> []subscription
	subscriptions map[string][]flare.Subscription

	// subscriptionID -> documentID -> document revision
	changes map[string]map[string]int64
}

// Find returns a list of subscriptions.
func (s *Subscription) Find(
	_ context.Context, pagination *flare.Pagination, resourceID string,
) ([]flare.Subscription, *flare.Pagination, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	subscriptions, ok := s.subscriptions[resourceID]
	if !ok {
		return []flare.Subscription{}, &flare.Pagination{
			Total:  0,
			Limit:  pagination.Limit,
			Offset: pagination.Offset,
		}, nil
	}

	var resp []flare.Subscription
	if pagination.Offset > len(subscriptions) {
		resp = subscriptions
	} else if pagination.Limit+pagination.Offset > len(subscriptions) {
		resp = subscriptions[pagination.Offset:]
	} else {
		resp = subscriptions[pagination.Offset : pagination.Offset+pagination.Limit]
	}

	return resp, &flare.Pagination{
		Total:  len(subscriptions),
		Limit:  pagination.Limit,
		Offset: pagination.Offset,
	}, nil
}

// FindByID return the Subscription that match the id.
func (s *Subscription) FindByID(
	ctx context.Context, resourceID, id string,
) (*flare.Subscription, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	return s.findByID(ctx, resourceID, id)
}

func (s *Subscription) findByID(
	_ context.Context, resourceID, id string,
) (*flare.Subscription, error) {
	nf := &errMemory{
		message:  fmt.Sprintf("subscription '%s' at resource '%s', not found", id, resourceID),
		notFound: true,
	}
	subscriptions, ok := s.subscriptions[resourceID]
	if !ok {
		return nil, nf
	}

	for _, subscription := range subscriptions {
		if subscription.ID == id {
			return &subscription, nil
		}
	}
	return nil, nf
}

// FindByPartition find all subscriptions that belongs to a given partition.
func (s *Subscription) FindByPartition(
	_ context.Context, resourceID, partition string,
) (<-chan flare.Subscription, <-chan error, error) {
	chanResult := make(chan flare.Subscription)
	chanErr := make(chan error)
	s.mutex.Lock()

	go func() {
		defer func() {
			close(chanResult)
			s.mutex.Unlock()
		}()

		subscriptions, ok := s.subscriptions[resourceID]
		if !ok {
			return
		}

		for _, subscription := range subscriptions {
			if subscription.Partition == partition {
				chanResult <- subscription
			}
		}
	}()

	return chanResult, chanErr, nil
}

// Create a subscription.
func (s *Subscription) Create(ctx context.Context, subscription *flare.Subscription) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	subscriptions, ok := s.subscriptions[subscription.Resource.ID]
	if !ok {
		s.subscriptions[subscription.Resource.ID] = make([]flare.Subscription, 0)
	}

	for _, subs := range subscriptions {
		if subs.Endpoint.URL.String() == subscription.Endpoint.URL.String() {
			return &errMemory{
				alreadyExists: true,
				message: fmt.Sprintf(
					"already exists a subscription '%s' with the endpoint.URL '%s'",
					subscription.ID,
					subscription.Endpoint.URL.String(),
				),
			}
		}
	}

	partition, err := s.resourceRepository.joinPartition(ctx, subscription.Resource.ID)
	if err != nil {
		return errors.Wrap(err, "error during join partition")
	}

	subscription.Partition = partition
	subscription.CreatedAt = time.Now()
	s.subscriptions[subscription.Resource.ID] = append(subscriptions, *subscription)
	return nil
}

// HasSubscription check if a resource has subscriptions.
func (s *Subscription) HasSubscription(ctx context.Context, resourceId string) (bool, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	subscriptions, ok := s.subscriptions[resourceId]
	if !ok {
		return false, nil
	}
	return len(subscriptions) > 0, nil
}

// Delete a given subscription.
func (s *Subscription) Delete(ctx context.Context, resourceId, id string) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	subscriptions := s.subscriptions[resourceId]
	for i, subscription := range subscriptions {
		if subscription.ID == id {
			err := s.resourceRepository.leavePartition(ctx, subscription.Resource.ID, subscription.Partition)
			if err != nil {
				return errors.Wrap(err, "error during leave partition")
			}

			s.subscriptions[resourceId] = append(subscriptions[:i], subscriptions[i+1:]...)
			return nil
		}
	}

	return &errMemory{
		message:  fmt.Sprintf("subscription '%s' at resource '%s', not found", id, resourceId),
		notFound: true,
	}
}

// Trigger process the update on a document.
func (s *Subscription) Trigger(
	ctx context.Context,
	action string,
	rawDocument *flare.Document,
	rawSubscription *flare.Subscription,
	fn func(context.Context, *flare.Document, *flare.Subscription, string) error,
) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	subscription, doc, err := s.triggerDocumentAndSubscription(
		ctx, action, rawDocument, rawSubscription,
	)
	if err != nil {
		return err
	}

	subscriptionMap, ok := s.changes[subscription.ID]
	if !ok {
		if action == flare.SubscriptionTriggerDelete {
			return nil
		}
		return s.triggerProcess(ctx, subscription, doc, flare.SubscriptionTriggerCreate, fn)
	}

	revision, ok := subscriptionMap[doc.ID]
	if !ok {
		if action == flare.SubscriptionTriggerDelete {
			return nil
		}
		return s.triggerProcess(ctx, subscription, doc, flare.SubscriptionTriggerCreate, fn)
	}

	reference, err := s.documentRepository.FindByID(ctx, doc.ID)
	if err != nil {
		return errors.Wrap(
			err,
			"error while loading reference document to process the suscription trigger",
		)
	}

	if action == flare.SubscriptionTriggerDelete {
		return s.triggerProcess(ctx, subscription, doc, flare.SubscriptionTriggerDelete, fn)
	}

	if reference.Revision > revision {
		return s.triggerProcess(ctx, subscription, doc, flare.SubscriptionTriggerUpdate, fn)
	}
	return nil
}

func (s *Subscription) triggerDocumentAndSubscription(
	ctx context.Context,
	action string,
	rawDocument *flare.Document,
	rawSubscription *flare.Subscription,
) (*flare.Subscription, *flare.Document, error) {
	subscription, err := s.findByID(ctx, rawDocument.Resource.ID, rawSubscription.ID)
	if err != nil {
		if repoErr := err.(flare.SubscriptionRepositoryError); repoErr.NotFound() {
			return nil, nil, nil
		}
		return nil, nil, err
	}

	doc, err := s.documentRepository.FindByID(ctx, rawDocument.ID)
	if err != nil {
		if repoErr := err.(flare.DocumentRepositoryError); repoErr.NotFound() {
			return nil, nil, nil
		}
		return nil, nil, err
	}

	return subscription, doc, nil
}

func (s *Subscription) triggerProcess(
	ctx context.Context,
	subs *flare.Subscription,
	doc *flare.Document,
	action string,
	fn func(context.Context, *flare.Document, *flare.Subscription, string) error,
) error {
	if err := fn(ctx, doc, subs, action); err != nil {
		return err
	}

	subscriptionMap, ok := s.changes[subs.ID]
	if !ok {
		subscriptionMap = make(map[string]int64)
		s.changes[subs.ID] = subscriptionMap
	}

	if action == flare.SubscriptionTriggerDelete {
		delete(subscriptionMap, doc.ID)
		return nil
	}

	subscriptionMap[doc.ID] = doc.Revision
	return nil
}

func (s *Subscription) init() {
	s.subscriptions = make(map[string][]flare.Subscription)
	s.changes = make(map[string]map[string]int64)
}
<commit_msg>provider/memory: fix subscription duplication error<commit_after>// Copyright 2018 Diego Bernardes. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package repository

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/pkg/errors"

	"github.com/diegobernardes/flare"
)

// Subscription implements the data layer for the subscription service.
type Subscription struct {
	mutex              sync.RWMutex
	resourceRepository resourceRepositorier
	documentRepository flare.DocumentRepositorier

	// resourceID -> []subscription
	subscriptions map[string][]flare.Subscription

	// subscriptionID -> documentID -> document revision
	changes map[string]map[string]int64
}

// Find returns a list of subscriptions.
func (s *Subscription) Find(
	_ context.Context, pagination *flare.Pagination, resourceID string,
) ([]flare.Subscription, *flare.Pagination, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	subscriptions, ok := s.subscriptions[resourceID]
	if !ok {
		return []flare.Subscription{}, &flare.Pagination{
			Total:  0,
			Limit:  pagination.Limit,
			Offset: pagination.Offset,
		}, nil
	}

	var resp []flare.Subscription
	if pagination.Offset > len(subscriptions) {
		resp = subscriptions
	} else if pagination.Limit+pagination.Offset > len(subscriptions) {
		resp = subscriptions[pagination.Offset:]
	} else {
		resp = subscriptions[pagination.Offset : pagination.Offset+pagination.Limit]
	}

	return resp, &flare.Pagination{
		Total:  len(subscriptions),
		Limit:  pagination.Limit,
		Offset: pagination.Offset,
	}, nil
}

// FindByID return the Subscription that match the id.
func (s *Subscription) FindByID(
	ctx context.Context, resourceID, id string,
) (*flare.Subscription, error) {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	return s.findByID(ctx, resourceID, id)
}

func (s *Subscription) findByID(
	_ context.Context, resourceID, id string,
) (*flare.Subscription, error) {
	nf := &errMemory{
		message:  fmt.Sprintf("subscription '%s' at resource '%s', not found", id, resourceID),
		notFound: true,
	}
	subscriptions, ok := s.subscriptions[resourceID]
	if !ok {
		return nil, nf
	}

	for _, subscription := range subscriptions {
		if subscription.ID == id {
			return &subscription, nil
		}
	}
	return nil, nf
}

// FindByPartition find all subscriptions that belongs to a given partition.
func (s *Subscription) FindByPartition(
	_ context.Context, resourceID, partition string,
) (<-chan flare.Subscription, <-chan error, error) {
	chanResult := make(chan flare.Subscription)
	chanErr := make(chan error)
	s.mutex.Lock()

	go func() {
		defer func() {
			close(chanResult)
			s.mutex.Unlock()
		}()

		subscriptions, ok := s.subscriptions[resourceID]
		if !ok {
			return
		}

		for _, subscription := range subscriptions {
			if subscription.Partition == partition {
				chanResult <- subscription
			}
		}
	}()

	return chanResult, chanErr, nil
}

// Create a subscription.
func (s *Subscription) Create(ctx context.Context, subscription *flare.Subscription) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	subscriptions, ok := s.subscriptions[subscription.Resource.ID]
	if !ok {
		s.subscriptions[subscription.Resource.ID] = make([]flare.Subscription, 0)
	}

	fetchEndpoints := func(endpoint flare.SubscriptionEndpoint) []string {
		var endpoints []string

		if endpoint.URL != nil {
			endpoints = append(endpoints, endpoint.URL.String())
		}

		for _, ea := range endpoint.Action {
			if ea.URL != nil {
				endpoints = append(endpoints, ea.URL.String())
			}
		}

		return endpoints
	}

	newEndpoints := fetchEndpoints(subscription.Endpoint)
	for _, subs := range subscriptions {
		referenceEndpoints := fetchEndpoints(subs.Endpoint)

		for _, ne := range newEndpoints {
			for _, re := range referenceEndpoints {
				if ne == re {
					return &errMemory{
						alreadyExists: true,
						message: fmt.Sprintf(
							"already exists a subscription '%s' with the endpoint.URL '%s'",
							subscription.ID,
							ne,
						),
					}
				}
			}
		}
	}

	partition, err := s.resourceRepository.joinPartition(ctx, subscription.Resource.ID)
	if err != nil {
		return errors.Wrap(err, "error during join partition")
	}

	subscription.Partition = partition
	subscription.CreatedAt = time.Now()
	s.subscriptions[subscription.Resource.ID] = append(subscriptions, *subscription)
	return nil
}

// HasSubscription check if a resource has subscriptions.
func (s *Subscription) HasSubscription(ctx context.Context, resourceId string) (bool, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	subscriptions, ok := s.subscriptions[resourceId]
	if !ok {
		return false, nil
	}
	return len(subscriptions) > 0, nil
}

// Delete a given subscription.
func (s *Subscription) Delete(ctx context.Context, resourceId, id string) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	subscriptions := s.subscriptions[resourceId]
	for i, subscription := range subscriptions {
		if subscription.ID == id {
			err := s.resourceRepository.leavePartition(ctx, subscription.Resource.ID, subscription.Partition)
			if err != nil {
				return errors.Wrap(err, "error during leave partition")
			}

			s.subscriptions[resourceId] = append(subscriptions[:i], subscriptions[i+1:]...)
			return nil
		}
	}

	return &errMemory{
		message:  fmt.Sprintf("subscription '%s' at resource '%s', not found", id, resourceId),
		notFound: true,
	}
}

// Trigger process the update on a document.
func (s *Subscription) Trigger(
	ctx context.Context,
	action string,
	rawDocument *flare.Document,
	rawSubscription *flare.Subscription,
	fn func(context.Context, *flare.Document, *flare.Subscription, string) error,
) error {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	subscription, doc, err := s.triggerDocumentAndSubscription(
		ctx, action, rawDocument, rawSubscription,
	)
	if err != nil {
		return err
	}

	subscriptionMap, ok := s.changes[subscription.ID]
	if !ok {
		if action == flare.SubscriptionTriggerDelete {
			return nil
		}
		return s.triggerProcess(ctx, subscription, doc, flare.SubscriptionTriggerCreate, fn)
	}

	revision, ok := subscriptionMap[doc.ID]
	if !ok {
		if action == flare.SubscriptionTriggerDelete {
			return nil
		}
		return s.triggerProcess(ctx, subscription, doc, flare.SubscriptionTriggerCreate, fn)
	}

	reference, err := s.documentRepository.FindByID(ctx, doc.ID)
	if err != nil {
		return errors.Wrap(
			err,
			"error while loading reference document to process the suscription trigger",
		)
	}

	if action == flare.SubscriptionTriggerDelete {
		return s.triggerProcess(ctx, subscription, doc, flare.SubscriptionTriggerDelete, fn)
	}

	if reference.Revision > revision {
		return s.triggerProcess(ctx, subscription, doc, flare.SubscriptionTriggerUpdate, fn)
	}
	return nil
}

func (s *Subscription) triggerDocumentAndSubscription(
	ctx context.Context,
	action string,
	rawDocument *flare.Document,
	rawSubscription *flare.Subscription,
) (*flare.Subscription, *flare.Document, error) {
	subscription, err := s.findByID(ctx, rawDocument.Resource.ID, rawSubscription.ID)
	if err != nil {
		if repoErr := err.(flare.SubscriptionRepositoryError); repoErr.NotFound() {
			return nil, nil, nil
		}
		return nil, nil, err
	}

	doc, err := s.documentRepository.FindByID(ctx, rawDocument.ID)
	if err != nil {
		if repoErr := err.(flare.DocumentRepositoryError); repoErr.NotFound() {
			return nil, nil, nil
		}
		return nil, nil, err
	}

	return subscription, doc, nil
}

func (s *Subscription) triggerProcess(
	ctx context.Context,
	subs *flare.Subscription,
	doc *flare.Document,
	action string,
	fn func(context.Context, *flare.Document, *flare.Subscription, string) error,
) error {
	if err := fn(ctx, doc, subs, action); err != nil {
		return err
	}

	subscriptionMap, ok := s.changes[subs.ID]
	if !ok {
		subscriptionMap = make(map[string]int64)
		s.changes[subs.ID] = subscriptionMap
	}

	if action == flare.SubscriptionTriggerDelete {
		delete(subscriptionMap, doc.ID)
		return nil
	}

	subscriptionMap[doc.ID] = doc.Revision
	return nil
}

func (s *Subscription) init() {
	s.subscriptions = make(map[string][]flare.Subscription)
	s.changes = make(map[string]map[string]int64)
}
<|endoftext|>


<commit_before>package restart

import (
	"bytes"
	"fmt"

	"log"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
	"github.com/masterzen/winrm"
)

var DefaultRestartCommand = "shutdown /r /f /t 0 /c \"packer restart\""
var DefaultRestartCheckCommand = winrm.Powershell(`echo "${env:COMPUTERNAME} restarted."`)
var retryableSleep = 5 * time.Second
var TryCheckReboot = "shutdown.exe -f -r -t 60"
var AbortReboot = "shutdown.exe -a"

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// The command used to restart the guest machine
	RestartCommand string `mapstructure:"restart_command"`

	// The command used to check if the guest machine has restarted
	// The output of this command will be displayed to the user
	RestartCheckCommand string `mapstructure:"restart_check_command"`

	// The timeout for waiting for the machine to restart
	RestartTimeout time.Duration `mapstructure:"restart_timeout"`

	ctx interpolate.Context
}

type Provisioner struct {
	config     Config
	comm       packer.Communicator
	ui         packer.Ui
	cancel     chan struct{}
	cancelLock sync.Mutex
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.RestartCommand == "" {
		p.config.RestartCommand = DefaultRestartCommand
	}

	if p.config.RestartCheckCommand == "" {
		p.config.RestartCheckCommand = DefaultRestartCheckCommand
	}

	if p.config.RestartTimeout == 0 {
		p.config.RestartTimeout = 5 * time.Minute
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	p.cancelLock.Lock()
	p.cancel = make(chan struct{})
	p.cancelLock.Unlock()

	ui.Say("Restarting Machine")
	p.comm = comm
	p.ui = ui

	var cmd *packer.RemoteCmd
	command := p.config.RestartCommand
	err := p.retryable(func() error {
		cmd = &packer.RemoteCmd{Command: command}
		return cmd.StartWithUi(comm, ui)
	})

	if err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Restart script exited with non-zero exit status: %d", cmd.ExitStatus)
	}

	return waitForRestart(p, comm)
}

var waitForRestart = func(p *Provisioner, comm packer.Communicator) error {
	ui := p.ui
	ui.Say("Waiting for machine to restart...")
	waitDone := make(chan bool, 1)
	timeout := time.After(p.config.RestartTimeout)
	var err error

	p.comm = comm
	var cmd *packer.RemoteCmd
	trycommand := TryCheckReboot
	abortcommand := AbortReboot
	// Stolen from Vagrant reboot checker
	for {
		log.Printf("Check if machine is rebooting...")
		cmd = &packer.RemoteCmd{Command: trycommand}
		err = cmd.StartWithUi(comm, ui)
		if err != nil {
			// Couldn't execute, we assume machine is rebooting already
			break
		}

		if cmd.ExitStatus == 1115 || cmd.ExitStatus == 1190 {
			// Reboot already in progress but not completed
			log.Printf("Reboot already in progress, waiting...")
			time.Sleep(10 * time.Second)
		}
		if cmd.ExitStatus == 0 {
			// Cancel reboot we created to test if machine was already rebooting
			cmd = &packer.RemoteCmd{Command: abortcommand}
			cmd.StartWithUi(comm, ui)
			break
		}
	}

	go func() {
		log.Printf("Waiting for machine to become available...")
		err = waitForCommunicator(p)
		waitDone <- true
	}()

	log.Printf("Waiting for machine to reboot with timeout: %s", p.config.RestartTimeout)

WaitLoop:
	for {
		// Wait for either WinRM to become available, a timeout to occur,
		// or an interrupt to come through.
		select {
		case <-waitDone:
			if err != nil {
				ui.Error(fmt.Sprintf("Error waiting for machine to restart: %s", err))
				return err
			}

			ui.Say("Machine successfully restarted, moving on")
			close(p.cancel)
			break WaitLoop
		case <-timeout:
			err := fmt.Errorf("Timeout waiting for machine to restart.")
			ui.Error(err.Error())
			close(p.cancel)
			return err
		case <-p.cancel:
			close(waitDone)
			return fmt.Errorf("Interrupt detected, quitting waiting for machine to restart")
		}
	}

	return nil

}

var waitForCommunicator = func(p *Provisioner) error {
	runCustomRestartCheck := true
	if p.config.RestartCheckCommand == DefaultRestartCheckCommand {
		runCustomRestartCheck = false
	}
	// This command is configurable by the user to make sure that the
	// vm has met their necessary criteria for having restarted. If the
	// user doesn't set a special restart command, we just run the
	// default as cmdModuleLoad below.
	cmdRestartCheck := &packer.RemoteCmd{Command: p.config.RestartCheckCommand}
	log.Printf("Checking that communicator is connected with: '%s'",
		cmdRestartCheck.Command)
	for {
		select {
		case <-p.cancel:
			log.Println("Communicator wait canceled, exiting loop")
			return fmt.Errorf("Communicator wait canceled")
		case <-time.After(retryableSleep):
		}
		if runCustomRestartCheck == true {
			// run user-configured restart check
			err := cmdRestartCheck.StartWithUi(p.comm, p.ui)
			if err != nil {
				log.Printf("Communication connection err: %s", err)
				continue
			}
			log.Printf("Connected to machine")
			runCustomRestartCheck = false
		}

		// This is the non-user-configurable check that powershell
		// modules have loaded.

		// If we catch the restart in just the right place, we will be able
		// to run the restart check but the output will be an error message
		// about how it needs powershell modules to load, and we will start
		// provisioning before powershell is actually ready.
		// In this next check, we parse stdout to make sure that the command is
		// actually running as expected.
		var buf bytes.Buffer
		cmdModuleLoad := &packer.RemoteCmd{
			Command: DefaultRestartCheckCommand,
			Stdin:   nil,
			Stdout:  &buf,
			Stderr:  &buf}

		// cmdModuleLoad.StartWithUi(p.comm, p.ui)
		p.comm.Start(cmdModuleLoad)
		cmdModuleLoad.Wait()

		stdoutToRead := buf.String()
		if !strings.Contains(stdoutToRead, "restarted.") {
			log.Printf("echo didn't succeed; retrying...")
			continue
		}
		break
	}

	return nil
}

func (p *Provisioner) Cancel() {
	log.Printf("Received interrupt Cancel()")

	p.cancelLock.Lock()
	defer p.cancelLock.Unlock()
	if p.cancel != nil {
		close(p.cancel)
	}
}

// retryable will retry the given function over and over until a
// non-error is returned.
func (p *Provisioner) retryable(f func() error) error {
	startTimeout := time.After(p.config.RestartTimeout)
	for {
		var err error
		if err = f(); err == nil {
			return nil
		}

		// Create an error and log it
		err = fmt.Errorf("Retryable error: %s", err)
		log.Print(err.Error())

		// Check if we timed out, otherwise we retry. It is safe to
		// retry since the only error case above is if the command
		// failed to START.
		select {
		case <-startTimeout:
			return err
		default:
			time.Sleep(retryableSleep)
		}
	}
}
<commit_msg>remove unnecessary boolean operator<commit_after>package restart

import (
	"bytes"
	"fmt"

	"log"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/packer/common"
	"github.com/hashicorp/packer/helper/config"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/template/interpolate"
	"github.com/masterzen/winrm"
)

var DefaultRestartCommand = "shutdown /r /f /t 0 /c \"packer restart\""
var DefaultRestartCheckCommand = winrm.Powershell(`echo "${env:COMPUTERNAME} restarted."`)
var retryableSleep = 5 * time.Second
var TryCheckReboot = "shutdown.exe -f -r -t 60"
var AbortReboot = "shutdown.exe -a"

type Config struct {
	common.PackerConfig `mapstructure:",squash"`

	// The command used to restart the guest machine
	RestartCommand string `mapstructure:"restart_command"`

	// The command used to check if the guest machine has restarted
	// The output of this command will be displayed to the user
	RestartCheckCommand string `mapstructure:"restart_check_command"`

	// The timeout for waiting for the machine to restart
	RestartTimeout time.Duration `mapstructure:"restart_timeout"`

	ctx interpolate.Context
}

type Provisioner struct {
	config     Config
	comm       packer.Communicator
	ui         packer.Ui
	cancel     chan struct{}
	cancelLock sync.Mutex
}

func (p *Provisioner) Prepare(raws ...interface{}) error {
	err := config.Decode(&p.config, &config.DecodeOpts{
		Interpolate:        true,
		InterpolateContext: &p.config.ctx,
		InterpolateFilter: &interpolate.RenderFilter{
			Exclude: []string{
				"execute_command",
			},
		},
	}, raws...)
	if err != nil {
		return err
	}

	if p.config.RestartCommand == "" {
		p.config.RestartCommand = DefaultRestartCommand
	}

	if p.config.RestartCheckCommand == "" {
		p.config.RestartCheckCommand = DefaultRestartCheckCommand
	}

	if p.config.RestartTimeout == 0 {
		p.config.RestartTimeout = 5 * time.Minute
	}

	return nil
}

func (p *Provisioner) Provision(ui packer.Ui, comm packer.Communicator) error {
	p.cancelLock.Lock()
	p.cancel = make(chan struct{})
	p.cancelLock.Unlock()

	ui.Say("Restarting Machine")
	p.comm = comm
	p.ui = ui

	var cmd *packer.RemoteCmd
	command := p.config.RestartCommand
	err := p.retryable(func() error {
		cmd = &packer.RemoteCmd{Command: command}
		return cmd.StartWithUi(comm, ui)
	})

	if err != nil {
		return err
	}

	if cmd.ExitStatus != 0 {
		return fmt.Errorf("Restart script exited with non-zero exit status: %d", cmd.ExitStatus)
	}

	return waitForRestart(p, comm)
}

var waitForRestart = func(p *Provisioner, comm packer.Communicator) error {
	ui := p.ui
	ui.Say("Waiting for machine to restart...")
	waitDone := make(chan bool, 1)
	timeout := time.After(p.config.RestartTimeout)
	var err error

	p.comm = comm
	var cmd *packer.RemoteCmd
	trycommand := TryCheckReboot
	abortcommand := AbortReboot
	// Stolen from Vagrant reboot checker
	for {
		log.Printf("Check if machine is rebooting...")
		cmd = &packer.RemoteCmd{Command: trycommand}
		err = cmd.StartWithUi(comm, ui)
		if err != nil {
			// Couldn't execute, we assume machine is rebooting already
			break
		}

		if cmd.ExitStatus == 1115 || cmd.ExitStatus == 1190 {
			// Reboot already in progress but not completed
			log.Printf("Reboot already in progress, waiting...")
			time.Sleep(10 * time.Second)
		}
		if cmd.ExitStatus == 0 {
			// Cancel reboot we created to test if machine was already rebooting
			cmd = &packer.RemoteCmd{Command: abortcommand}
			cmd.StartWithUi(comm, ui)
			break
		}
	}

	go func() {
		log.Printf("Waiting for machine to become available...")
		err = waitForCommunicator(p)
		waitDone <- true
	}()

	log.Printf("Waiting for machine to reboot with timeout: %s", p.config.RestartTimeout)

WaitLoop:
	for {
		// Wait for either WinRM to become available, a timeout to occur,
		// or an interrupt to come through.
		select {
		case <-waitDone:
			if err != nil {
				ui.Error(fmt.Sprintf("Error waiting for machine to restart: %s", err))
				return err
			}

			ui.Say("Machine successfully restarted, moving on")
			close(p.cancel)
			break WaitLoop
		case <-timeout:
			err := fmt.Errorf("Timeout waiting for machine to restart.")
			ui.Error(err.Error())
			close(p.cancel)
			return err
		case <-p.cancel:
			close(waitDone)
			return fmt.Errorf("Interrupt detected, quitting waiting for machine to restart")
		}
	}

	return nil

}

var waitForCommunicator = func(p *Provisioner) error {
	runCustomRestartCheck := true
	if p.config.RestartCheckCommand == DefaultRestartCheckCommand {
		runCustomRestartCheck = false
	}
	// This command is configurable by the user to make sure that the
	// vm has met their necessary criteria for having restarted. If the
	// user doesn't set a special restart command, we just run the
	// default as cmdModuleLoad below.
	cmdRestartCheck := &packer.RemoteCmd{Command: p.config.RestartCheckCommand}
	log.Printf("Checking that communicator is connected with: '%s'",
		cmdRestartCheck.Command)
	for {
		select {
		case <-p.cancel:
			log.Println("Communicator wait canceled, exiting loop")
			return fmt.Errorf("Communicator wait canceled")
		case <-time.After(retryableSleep):
		}
		if runCustomRestartCheck {
			// run user-configured restart check
			err := cmdRestartCheck.StartWithUi(p.comm, p.ui)
			if err != nil {
				log.Printf("Communication connection err: %s", err)
				continue
			}
			log.Printf("Connected to machine")
			runCustomRestartCheck = false
		}

		// This is the non-user-configurable check that powershell
		// modules have loaded.

		// If we catch the restart in just the right place, we will be able
		// to run the restart check but the output will be an error message
		// about how it needs powershell modules to load, and we will start
		// provisioning before powershell is actually ready.
		// In this next check, we parse stdout to make sure that the command is
		// actually running as expected.
		var buf bytes.Buffer
		cmdModuleLoad := &packer.RemoteCmd{
			Command: DefaultRestartCheckCommand,
			Stdin:   nil,
			Stdout:  &buf,
			Stderr:  &buf}

		// cmdModuleLoad.StartWithUi(p.comm, p.ui)
		p.comm.Start(cmdModuleLoad)
		cmdModuleLoad.Wait()

		stdoutToRead := buf.String()
		if !strings.Contains(stdoutToRead, "restarted.") {
			log.Printf("echo didn't succeed; retrying...")
			continue
		}
		break
	}

	return nil
}

func (p *Provisioner) Cancel() {
	log.Printf("Received interrupt Cancel()")

	p.cancelLock.Lock()
	defer p.cancelLock.Unlock()
	if p.cancel != nil {
		close(p.cancel)
	}
}

// retryable will retry the given function over and over until a
// non-error is returned.
func (p *Provisioner) retryable(f func() error) error {
	startTimeout := time.After(p.config.RestartTimeout)
	for {
		var err error
		if err = f(); err == nil {
			return nil
		}

		// Create an error and log it
		err = fmt.Errorf("Retryable error: %s", err)
		log.Print(err.Error())

		// Check if we timed out, otherwise we retry. It is safe to
		// retry since the only error case above is if the command
		// failed to START.
		select {
		case <-startTimeout:
			return err
		default:
			time.Sleep(retryableSleep)
		}
	}
}
<|endoftext|>