<commit_before>// Package applier handles execution of generating diffs between schemas, and
// appropriate application of the generated DDL.
package applier

import (
	"context"
	"fmt"
	"strings"

	log "github.com/sirupsen/logrus"
	"github.com/skeema/skeema/fs"
	"github.com/skeema/tengo"
)

// Result stores the overall result of all operations the worker has completed.
type Result struct {
	Differences      bool
	SkipCount        int
	UnsupportedCount int
}

// Worker reads TargetGroups from the input channel and performs the appropriate
// diff/push operation on each target per TargetGroup. When there are no more
// TargetGroups to read, it writes its aggregate Result to the output channel.
// If a fatal error occurs, it will be returned immediately; Worker is meant to
// be called via an errgroup (see golang.org/x/sync/errgroup).
func Worker(ctx context.Context, targetGroups <-chan TargetGroup, results chan<- Result, printer *Printer) error {
	var result Result
	for tg := range targetGroups {
	TargetsInGroup:
		for _, t := range tg { // iterate over each Target in the TargetGroup
			// Get schema name from t.SchemaFromDir, NOT t.SchemaFromInstance, since
			// t.SchemaFromInstance will be nil if the schema doesn't exist yet
			schemaName := t.SchemaFromDir.Name
			dryRun := t.Dir.Config.GetBool("dry-run")
			brief := dryRun && t.Dir.Config.GetBool("brief")

			if dryRun {
				log.Infof("Generating diff of %s %s vs %s/*.sql", t.Instance, schemaName, t.Dir)
			} else {
				log.Infof("Pushing changes from %s/*.sql to %s %s", t.Dir, t.Instance, schemaName)
			}

			diff := tengo.NewSchemaDiff(t.SchemaFromInstance, t.SchemaFromDir)
			var targetStmtCount int

			if diff.SchemaDDL != "" {
				printer.syncPrintf(t.Instance, "", "%s;\n", diff.SchemaDDL)
				targetStmtCount++
				result.Differences = true
				if !dryRun {
					var schemaDDLErr error
					if strings.HasPrefix(diff.SchemaDDL, "CREATE DATABASE") && t.SchemaFromInstance == nil {
						_, schemaDDLErr = t.Instance.CreateSchema(schemaName, t.SchemaFromDir.CharSet, t.SchemaFromDir.Collation)
					} else if strings.HasPrefix(diff.SchemaDDL, "ALTER DATABASE") {
						schemaDDLErr = t.Instance.AlterSchema(t.SchemaFromInstance.Name, t.SchemaFromDir.CharSet, t.SchemaFromDir.Collation)
					} else {
						return fmt.Errorf("Refusing to run unexpectedly-generated schema-level DDL: %s", diff.SchemaDDL)
					}
					if schemaDDLErr != nil {
						log.Errorf("Error running DDL on %s %s: %s", t.Instance, schemaName, schemaDDLErr)
						skipped := len(diff.TableDiffs) + 1
						result.SkipCount += skipped
						if skipped > 1 {
							log.Warnf("Skipping %d remaining operations for %s %s due to previous error", skipped-1, t.Instance, schemaName)
						}
						continue
					}
				}
			}

			if t.Dir.Config.GetBool("verify") && len(diff.TableDiffs) > 0 && !brief {
				if err := VerifyDiff(diff, t); err != nil {
					return err
				}
			}

			// Obtain StatementModifiers based on the dir's config
			mods, err := StatementModifiersForDir(t.Dir)
			if err != nil {
				return ConfigError(err.Error())
			}
			if configFlavor := tengo.NewFlavor(t.Dir.Config.Get("flavor")); configFlavor != tengo.FlavorUnknown {
				mods.Flavor = configFlavor
			} else {
				mods.Flavor = t.Instance.Flavor()
			}

			// Build DDLStatements for each TableDiff, handling pre-execution errors
			// accordingly
			ddls := make([]*DDLStatement, 0, len(diff.TableDiffs))
			for i, tableDiff := range diff.TableDiffs {
				ddl, err := NewDDLStatement(tableDiff, mods, t)
				if ddl == nil && err == nil {
					continue // Skip entirely if mods made the statement a noop
				}
				targetStmtCount++
				result.Differences = true

				if unsupportedErr, ok := err.(*tengo.UnsupportedDiffError); ok {
					result.UnsupportedCount++
					log.Warnf("Skipping table %s: unable to generate ALTER TABLE due to use of unsupported features. Use --debug for more information.", unsupportedErr.Name)
					DebugLogUnsupportedDiff(unsupportedErr)
					continue
				} else if err != nil {
					result.SkipCount++
					log.Errorf(err.Error())
					skipped := len(diff.TableDiffs) - i
					result.SkipCount += skipped
					if skipped > 1 {
						log.Warnf("Skipping %d remaining operations for %s %s due to previous error", skipped-1, t.Instance, schemaName)
					}
					continue TargetsInGroup
				}

				if dryRun {
					printer.syncPrintf(t.Instance, schemaName, "%s\n", ddl.String())
				} else {
					// IMPORTANT: only append to ddls if NOT dry-run. Do not move this out of
					// this conditional!
					ddls = append(ddls, ddl)
				}
			}

			// Execute DDL. (ddls will be empty if dryRun, due to conditional above)
			for i, ddl := range ddls {
				printer.syncPrintf(t.Instance, schemaName, "%s\n", ddl.String())
				if err := ddl.Execute(); err != nil {
					log.Errorf("Error running DDL on %s %s: %s", t.Instance, schemaName, err)
					skipped := len(ddls) - i
					result.SkipCount += skipped
					if skipped > 1 {
						log.Warnf("Skipping %d remaining operations for %s %s due to previous error", skipped-1, t.Instance, schemaName)
					}
					break
				}
			}

			if targetStmtCount == 0 {
				log.Infof("%s %s: No differences found\n", t.Instance, schemaName)
			} else {
				verb := "push"
				if dryRun {
					verb = "diff"
				}
				log.Infof("%s %s: %s complete\n", t.Instance, schemaName, verb)
			}

			// Exit early if context cancelled
			select {
			case <-ctx.Done():
				return nil
			default:
			}
		}
	}
	results <- result
	return nil
}

// SumResults adds up the supplied results to return a single combined result.
func SumResults(results []Result) Result {
	var total Result
	for _, r := range results {
		total.Differences = total.Differences || r.Differences
		total.SkipCount += r.SkipCount
		total.UnsupportedCount += r.UnsupportedCount
	}
	return total
}

// StatementModifiersForDir returns a set of DDL modifiers, based on the
// directory's configuration.
func StatementModifiersForDir(dir *fs.Dir) (mods tengo.StatementModifiers, err error) {
	mods.NextAutoInc = tengo.NextAutoIncIfIncreased
	forceAllowUnsafe := dir.Config.GetBool("brief") && dir.Config.GetBool("dry-run")
	mods.AllowUnsafe = forceAllowUnsafe || dir.Config.GetBool("allow-unsafe")
	if dir.Config.GetBool("exact-match") {
		mods.StrictIndexOrder = true
		mods.StrictForeignKeyNaming = true
	}
	if mods.AlgorithmClause, err = dir.Config.GetEnum("alter-algorithm", "INPLACE", "COPY", "INSTANT", "DEFAULT"); err != nil {
		return
	}
	if mods.LockClause, err = dir.Config.GetEnum("alter-lock", "NONE", "SHARED", "EXCLUSIVE", "DEFAULT"); err != nil {
		return
	}
	if mods.IgnoreTable, err = dir.Config.GetRegexp("ignore-table"); err != nil {
		return
	}
	return
}

// DebugLogUnsupportedDiff logs (at Debug level) the reason why a table is
// unsupported for diff/alter operations.
func DebugLogUnsupportedDiff(err *tengo.UnsupportedDiffError) {
	for _, line := range strings.Split(err.ExtendedError(), "\n") {
		if len(line) > 0 {
			log.Debug(line)
		}
	}
}

// ConfigError represents a configuration problem encountered at runtime.
type ConfigError string

// Error satisfies the builtin error interface.
func (ce ConfigError) Error() string {
	return string(ce)
}
<commit_msg>push: fix display bug in final error tally (cosmetic only)<commit_after>// Package applier handles execution of generating diffs between schemas, and
// appropriate application of the generated DDL.
package applier

import (
	"context"
	"fmt"
	"strings"

	log "github.com/sirupsen/logrus"
	"github.com/skeema/skeema/fs"
	"github.com/skeema/tengo"
)

// Result stores the overall result of all operations the worker has completed.
type Result struct {
	Differences      bool
	SkipCount        int
	UnsupportedCount int
}

// Worker reads TargetGroups from the input channel and performs the appropriate
// diff/push operation on each target per TargetGroup. When there are no more
// TargetGroups to read, it writes its aggregate Result to the output channel.
// If a fatal error occurs, it will be returned immediately; Worker is meant to
// be called via an errgroup (see golang.org/x/sync/errgroup).
func Worker(ctx context.Context, targetGroups <-chan TargetGroup, results chan<- Result, printer *Printer) error {
	var result Result
	for tg := range targetGroups {
	TargetsInGroup:
		for _, t := range tg { // iterate over each Target in the TargetGroup
			// Get schema name from t.SchemaFromDir, NOT t.SchemaFromInstance, since
			// t.SchemaFromInstance will be nil if the schema doesn't exist yet
			schemaName := t.SchemaFromDir.Name
			dryRun := t.Dir.Config.GetBool("dry-run")
			brief := dryRun && t.Dir.Config.GetBool("brief")

			if dryRun {
				log.Infof("Generating diff of %s %s vs %s/*.sql", t.Instance, schemaName, t.Dir)
			} else {
				log.Infof("Pushing changes from %s/*.sql to %s %s", t.Dir, t.Instance, schemaName)
			}

			diff := tengo.NewSchemaDiff(t.SchemaFromInstance, t.SchemaFromDir)
			var targetStmtCount int

			if diff.SchemaDDL != "" {
				printer.syncPrintf(t.Instance, "", "%s;\n", diff.SchemaDDL)
				targetStmtCount++
				result.Differences = true
				if !dryRun {
					var schemaDDLErr error
					if strings.HasPrefix(diff.SchemaDDL, "CREATE DATABASE") && t.SchemaFromInstance == nil {
						_, schemaDDLErr = t.Instance.CreateSchema(schemaName, t.SchemaFromDir.CharSet, t.SchemaFromDir.Collation)
					} else if strings.HasPrefix(diff.SchemaDDL, "ALTER DATABASE") {
						schemaDDLErr = t.Instance.AlterSchema(t.SchemaFromInstance.Name, t.SchemaFromDir.CharSet, t.SchemaFromDir.Collation)
					} else {
						return fmt.Errorf("Refusing to run unexpectedly-generated schema-level DDL: %s", diff.SchemaDDL)
					}
					if schemaDDLErr != nil {
						log.Errorf("Error running DDL on %s %s: %s", t.Instance, schemaName, schemaDDLErr)
						skipped := len(diff.TableDiffs) + 1
						result.SkipCount += skipped
						if skipped > 1 {
							log.Warnf("Skipping %d remaining operations for %s %s due to previous error", skipped-1, t.Instance, schemaName)
						}
						continue
					}
				}
			}

			if t.Dir.Config.GetBool("verify") && len(diff.TableDiffs) > 0 && !brief {
				if err := VerifyDiff(diff, t); err != nil {
					return err
				}
			}

			// Obtain StatementModifiers based on the dir's config
			mods, err := StatementModifiersForDir(t.Dir)
			if err != nil {
				return ConfigError(err.Error())
			}
			if configFlavor := tengo.NewFlavor(t.Dir.Config.Get("flavor")); configFlavor != tengo.FlavorUnknown {
				mods.Flavor = configFlavor
			} else {
				mods.Flavor = t.Instance.Flavor()
			}

			// Build DDLStatements for each TableDiff, handling pre-execution errors
			// accordingly
			ddls := make([]*DDLStatement, 0, len(diff.TableDiffs))
			for i, tableDiff := range diff.TableDiffs {
				ddl, err := NewDDLStatement(tableDiff, mods, t)
				if ddl == nil && err == nil {
					continue // Skip entirely if mods made the statement a noop
				}
				targetStmtCount++
				result.Differences = true

				if unsupportedErr, ok := err.(*tengo.UnsupportedDiffError); ok {
					result.UnsupportedCount++
					log.Warnf("Skipping table %s: unable to generate ALTER TABLE due to use of unsupported features. Use --debug for more information.", unsupportedErr.Name)
					DebugLogUnsupportedDiff(unsupportedErr)
					continue
				} else if err != nil {
					skipped := len(diff.TableDiffs) - i
					result.SkipCount += skipped
					log.Errorf(err.Error())
					if skipped > 1 {
						log.Warnf("Skipping %d remaining operations for %s %s due to previous error", skipped-1, t.Instance, schemaName)
					}
					continue TargetsInGroup
				}

				if dryRun {
					printer.syncPrintf(t.Instance, schemaName, "%s\n", ddl.String())
				} else {
					// IMPORTANT: only append to ddls if NOT dry-run. Do not move this out of
					// this conditional!
					ddls = append(ddls, ddl)
				}
			}

			// Execute DDL. (ddls will be empty if dryRun, due to conditional above)
			for i, ddl := range ddls {
				printer.syncPrintf(t.Instance, schemaName, "%s\n", ddl.String())
				if err := ddl.Execute(); err != nil {
					log.Errorf("Error running DDL on %s %s: %s", t.Instance, schemaName, err)
					skipped := len(ddls) - i
					result.SkipCount += skipped
					if skipped > 1 {
						log.Warnf("Skipping %d remaining operations for %s %s due to previous error", skipped-1, t.Instance, schemaName)
					}
					break
				}
			}

			if targetStmtCount == 0 {
				log.Infof("%s %s: No differences found\n", t.Instance, schemaName)
			} else {
				verb := "push"
				if dryRun {
					verb = "diff"
				}
				log.Infof("%s %s: %s complete\n", t.Instance, schemaName, verb)
			}

			// Exit early if context cancelled
			select {
			case <-ctx.Done():
				return nil
			default:
			}
		}
	}
	results <- result
	return nil
}

// SumResults adds up the supplied results to return a single combined result.
func SumResults(results []Result) Result {
	var total Result
	for _, r := range results {
		total.Differences = total.Differences || r.Differences
		total.SkipCount += r.SkipCount
		total.UnsupportedCount += r.UnsupportedCount
	}
	return total
}

// StatementModifiersForDir returns a set of DDL modifiers, based on the
// directory's configuration.
func StatementModifiersForDir(dir *fs.Dir) (mods tengo.StatementModifiers, err error) {
	mods.NextAutoInc = tengo.NextAutoIncIfIncreased
	forceAllowUnsafe := dir.Config.GetBool("brief") && dir.Config.GetBool("dry-run")
	mods.AllowUnsafe = forceAllowUnsafe || dir.Config.GetBool("allow-unsafe")
	if dir.Config.GetBool("exact-match") {
		mods.StrictIndexOrder = true
		mods.StrictForeignKeyNaming = true
	}
	if mods.AlgorithmClause, err = dir.Config.GetEnum("alter-algorithm", "INPLACE", "COPY", "INSTANT", "DEFAULT"); err != nil {
		return
	}
	if mods.LockClause, err = dir.Config.GetEnum("alter-lock", "NONE", "SHARED", "EXCLUSIVE", "DEFAULT"); err != nil {
		return
	}
	if mods.IgnoreTable, err = dir.Config.GetRegexp("ignore-table"); err != nil {
		return
	}
	return
}

// DebugLogUnsupportedDiff logs (at Debug level) the reason why a table is
// unsupported for diff/alter operations.
func DebugLogUnsupportedDiff(err *tengo.UnsupportedDiffError) {
	for _, line := range strings.Split(err.ExtendedError(), "\n") {
		if len(line) > 0 {
			log.Debug(line)
		}
	}
}

// ConfigError represents a configuration problem encountered at runtime.
type ConfigError string

// Error satisfies the builtin error interface.
func (ce ConfigError) Error() string {
	return string(ce)
}
<|endoftext|>


<commit_before>package osc

/*
* Define application client-id(APP_ID) and client-security(APP_SEC) code here which will be gotten when you register an application at www.oschina.net
* Also the key and commonIV for encryption, they must be equal to the same file in client code, the "sec.java".
 */
const (
	APP_ID  = "APP_ID"
	APP_SEC = "APP_SEC"
)

var KEY []byte = []byte{11, 45}
var IV []byte = []byte{35, 23}
<commit_msg>Add redirect-url to sec.<commit_after>package osc

/*
* Define application client-id(APP_ID) and client-security(APP_SEC) code here which will be gotten when you register an application at www.oschina.net
* Also the key and commonIV for encryption, they must be equal to the same file in client code, the "sec.java".
 */
const (
	APP_ID  = "APP_ID"
	APP_SEC = "APP_SEC"
	REDIRECT_URL = "http://www.oschina.net"
)

var KEY []byte = []byte{}
var IV []byte = []byte{}
<|endoftext|>


<commit_before>// Copyright 2017 aaharu All rights reserved.
// This source code is licensed under the BSD-style license found in
// the LICENSE file in the root directory of this source tree.

package codegen

import (
	"bytes"
	"fmt"
	"go/format"
	"strings"

	"github.com/aaharu/schemarshal/utils"
	"github.com/aaharu/schemarshal/version"
)

var enumList = map[string][]interface{}{} // FIXME

// Generator of Go source code from JSON Schema
type Generator struct {
	name    string // package nage
	command string
	imports []*importSpec
	decls   []*typeSpec
}

// NewGenerator create Generator struct
func NewGenerator(packageName string, command string) *Generator {
	return &Generator{
		name:    packageName,
		command: command,
	}
}

// AddImport add an import statement
func (g *Generator) AddImport(path string, name string) {
	if g.imports == nil {
		g.imports = []*importSpec{}
	}
	g.imports = append(g.imports, &importSpec{
		name: name,
		path: path,
	})
}

// AddType add a type statement
func (g *Generator) AddType(name string, jsonType *JSONType) {
	if g.decls == nil {
		g.decls = []*typeSpec{}
	}
	g.decls = append(g.decls, &typeSpec{
		name:     name,
		jsontype: jsonType,
	})
}

// Generate gofmt-ed Go source code
func (g *Generator) Generate() ([]byte, error) {
	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("// generated by %s `%s`\n", version.String(), g.command))
	buf.WriteString("\n")
	buf.WriteString(fmt.Sprintf("package %s\n\n", g.name))

	if len(g.imports) > 1 {
		buf.WriteString(fmt.Sprintf("import (\n"))
		for i := range g.imports {
			buf.WriteString(fmt.Sprintf("%s %s\n", g.imports[i].name, g.imports[i].path))
		}
		buf.WriteString(fmt.Sprintf(")\n\n"))
	} else if len(g.imports) == 1 {
		buf.WriteString(fmt.Sprintf("import %s %s\n\n", g.imports[0].name, g.imports[0].path))
	}

	if g.decls != nil {
		for i := range g.decls {
			buf.WriteString("type " + g.decls[i].name + " ")
			buf.Write(g.decls[i].jsontype.generate())
			buf.WriteString("\n")
		}
	}

	if len(enumList) > 0 {
		buf.WriteString("\n")
		for typeName, enum := range enumList {
			buf.WriteString("type " + typeName + "Enum int\n")
			buf.WriteString("const (\n")
			for i := range enum {
				if i == 0 {
					buf.WriteString("enum" + utils.UpperCamelCase(fmt.Sprintf("%v", enum[0])) + " " + typeName + "Enum = iota\n")
				} else {
					buf.WriteString("enum" + utils.UpperCamelCase(fmt.Sprintf("%v", enum[i])) + "\n")
				}
			}
			buf.WriteString(")\n")
			buf.WriteString("func (enum " + typeName + "Enum) MarshalJSON() ([]byte, error) {\n")
			buf.WriteString("var var" + typeName + " = []interface{}{")
			for i := range enum {
				switch v := enum[i].(type) {
				case string:
					buf.WriteString(`"`)
					buf.WriteString(strings.Replace(v, `"`, `\"`, -1))
					buf.WriteString(`"`)
				default:
					buf.WriteString(fmt.Sprintf("%v", v))
				}
				buf.WriteString(",")
			}
			buf.WriteString("}\n")
			buf.WriteString("switch v:= var" + typeName + "[enum].(type) {\n")
			buf.WriteString("case string:\n")
			buf.WriteString("return []byte(fmt.Sprintf(\"\\\"%v\\\"\", strings.Replace(v, `\"`, `\\\"`, -1))), nil\n")
			buf.WriteString("default:\n")
			buf.WriteString("return []byte(fmt.Sprintf(\"%v\", v)), nil\n")
			buf.WriteString("}\n")
			buf.WriteString("}\n\n")
		}
	}

	return format.Source(buf.Bytes())
}

type importSpec struct {
	name string // ident
	path string
}

type typeSpec struct {
	name     string // type name
	jsontype *JSONType
}

type jsonFormat int

const (
	formatObject jsonFormat = iota
	formatArray
	formatString
	formatBoolean
	formatNumber
	formatInteger
	formatDatetime
)

// JSONType ...
type JSONType struct {
	format   jsonFormat
	nullable bool
	fields   []*field  // only object
	itemType *JSONType // only array
	enumType string
}

// AddField ...
func (t *JSONType) AddField(f *field) {
	if t.fields == nil {
		t.fields = []*field{}
	}
	t.fields = append(t.fields, f)
}

func (t *JSONType) generate() []byte {
	var buf bytes.Buffer
	if t.nullable {
		buf.WriteString("*")
	}
	if t.format == formatObject {
		if t.fields == nil {
			buf.WriteString("map[string]interface{}")
		} else {
			buf.WriteString("struct {\n")
			for i := range t.fields {
				buf.WriteString(t.fields[i].name)
				buf.WriteString(" ")
				buf.Write(t.fields[i].jsontype.generate())
				buf.WriteString(" ")
				buf.Write(t.fields[i].jsontag.generate())
				buf.WriteString("\n")
			}
			buf.WriteString("}")
		}
	} else if t.format == formatArray {
		buf.WriteString("[]")
		buf.Write(t.itemType.generate())
	} else if t.format == formatString {
		buf.WriteString("string")
	} else if t.format == formatBoolean {
		buf.WriteString("bool")
	} else if t.format == formatNumber {
		buf.WriteString("float64")
	} else if t.format == formatInteger {
		buf.WriteString("int")
	} else if t.format == formatDatetime {
		buf.WriteString("time.Time")
	}
	return buf.Bytes()
}

type field struct {
	name     string
	jsontype *JSONType
	jsontag  *jsonTag
}

type jsonTag struct {
	name      string
	omitEmpty bool
}

// Generate JSON tag code
func (t *jsonTag) generate() []byte {
	var buf bytes.Buffer
	buf.WriteString("`json:\"")
	buf.WriteString(t.name)
	if t.omitEmpty {
		buf.WriteString(",omitempty")
	}
	buf.WriteString("\"`")
	return buf.Bytes()
}
<commit_msg>fix enumName<commit_after>// Copyright 2017 aaharu All rights reserved.
// This source code is licensed under the BSD-style license found in
// the LICENSE file in the root directory of this source tree.

package codegen

import (
	"bytes"
	"fmt"
	"go/format"
	"strings"

	"github.com/aaharu/schemarshal/utils"
	"github.com/aaharu/schemarshal/version"
)

var enumList = map[string][]interface{}{} // FIXME

// Generator of Go source code from JSON Schema
type Generator struct {
	name    string // package nage
	command string
	imports []*importSpec
	decls   []*typeSpec
}

// NewGenerator create Generator struct
func NewGenerator(packageName string, command string) *Generator {
	return &Generator{
		name:    packageName,
		command: command,
	}
}

// AddImport add an import statement
func (g *Generator) AddImport(path string, name string) {
	if g.imports == nil {
		g.imports = []*importSpec{}
	}
	g.imports = append(g.imports, &importSpec{
		name: name,
		path: path,
	})
}

// AddType add a type statement
func (g *Generator) AddType(name string, jsonType *JSONType) {
	if g.decls == nil {
		g.decls = []*typeSpec{}
	}
	g.decls = append(g.decls, &typeSpec{
		name:     name,
		jsontype: jsonType,
	})
}

// Generate gofmt-ed Go source code
func (g *Generator) Generate() ([]byte, error) {
	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("// generated by %s `%s`\n", version.String(), g.command))
	buf.WriteString("\n")
	buf.WriteString(fmt.Sprintf("package %s\n\n", g.name))

	if len(g.imports) > 1 {
		buf.WriteString(fmt.Sprintf("import (\n"))
		for i := range g.imports {
			buf.WriteString(fmt.Sprintf("%s %s\n", g.imports[i].name, g.imports[i].path))
		}
		buf.WriteString(fmt.Sprintf(")\n\n"))
	} else if len(g.imports) == 1 {
		buf.WriteString(fmt.Sprintf("import %s %s\n\n", g.imports[0].name, g.imports[0].path))
	}

	if g.decls != nil {
		for i := range g.decls {
			buf.WriteString("type " + g.decls[i].name + " ")
			buf.Write(g.decls[i].jsontype.generate())
			buf.WriteString("\n")
		}
	}

	if len(enumList) > 0 {
		buf.WriteString("\n")
		for typeName, enum := range enumList {
			enumName := typeName + "Enum"
			buf.WriteString("type " + enumName + " int\n")
			buf.WriteString("const (\n")
			for i := range enum {
				if i == 0 {
					buf.WriteString(enumName + utils.UpperCamelCase(fmt.Sprintf("%v", enum[0])) + " " + enumName + " = iota\n")
				} else {
					buf.WriteString(enumName + utils.UpperCamelCase(fmt.Sprintf("%v", enum[i])) + "\n")
				}
			}
			buf.WriteString(")\n")
			buf.WriteString("func (enum " + enumName + ") MarshalJSON() ([]byte, error) {\n")
			buf.WriteString("var var" + typeName + " = []interface{}{")
			for i := range enum {
				switch v := enum[i].(type) {
				case string:
					buf.WriteString(`"`)
					buf.WriteString(strings.Replace(v, `"`, `\"`, -1))
					buf.WriteString(`"`)
				default:
					buf.WriteString(fmt.Sprintf("%v", v))
				}
				buf.WriteString(",")
			}
			buf.WriteString("}\n")
			buf.WriteString("switch v:= var" + typeName + "[enum].(type) {\n")
			buf.WriteString("case string:\n")
			buf.WriteString("return []byte(fmt.Sprintf(\"\\\"%v\\\"\", strings.Replace(v, `\"`, `\\\"`, -1))), nil\n")
			buf.WriteString("default:\n")
			buf.WriteString("return []byte(fmt.Sprintf(\"%v\", v)), nil\n")
			buf.WriteString("}\n")
			buf.WriteString("}\n\n")
		}
	}

	return format.Source(buf.Bytes())
}

type importSpec struct {
	name string // ident
	path string
}

type typeSpec struct {
	name     string // type name
	jsontype *JSONType
}

type jsonFormat int

const (
	formatObject jsonFormat = iota
	formatArray
	formatString
	formatBoolean
	formatNumber
	formatInteger
	formatDatetime
)

// JSONType ...
type JSONType struct {
	format   jsonFormat
	nullable bool
	fields   []*field  // only object
	itemType *JSONType // only array
	enumType string
}

// AddField ...
func (t *JSONType) AddField(f *field) {
	if t.fields == nil {
		t.fields = []*field{}
	}
	t.fields = append(t.fields, f)
}

func (t *JSONType) generate() []byte {
	var buf bytes.Buffer
	if t.nullable {
		buf.WriteString("*")
	}
	if t.format == formatObject {
		if t.fields == nil {
			buf.WriteString("map[string]interface{}")
		} else {
			buf.WriteString("struct {\n")
			for i := range t.fields {
				buf.WriteString(t.fields[i].name)
				buf.WriteString(" ")
				buf.Write(t.fields[i].jsontype.generate())
				buf.WriteString(" ")
				buf.Write(t.fields[i].jsontag.generate())
				buf.WriteString("\n")
			}
			buf.WriteString("}")
		}
	} else if t.format == formatArray {
		buf.WriteString("[]")
		buf.Write(t.itemType.generate())
	} else if t.format == formatString {
		buf.WriteString("string")
	} else if t.format == formatBoolean {
		buf.WriteString("bool")
	} else if t.format == formatNumber {
		buf.WriteString("float64")
	} else if t.format == formatInteger {
		buf.WriteString("int")
	} else if t.format == formatDatetime {
		buf.WriteString("time.Time")
	}
	return buf.Bytes()
}

type field struct {
	name     string
	jsontype *JSONType
	jsontag  *jsonTag
}

type jsonTag struct {
	name      string
	omitEmpty bool
}

// Generate JSON tag code
func (t *jsonTag) generate() []byte {
	var buf bytes.Buffer
	buf.WriteString("`json:\"")
	buf.WriteString(t.name)
	if t.omitEmpty {
		buf.WriteString(",omitempty")
	}
	buf.WriteString("\"`")
	return buf.Bytes()
}
<|endoftext|>


<commit_before>package tools

import (
	"bytes"
	"encoding/base64"
	"log"
	"math/rand"
	"net"

	"strings"
	"time"
)

var letters = []rune("0123456789-+@abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

type Tlv struct {
	T string // type of TLV
	P []byte // payload
}

//RandSeq returns a random string
func RandSeq(n int) string {
	rand.Seed(time.Now().UTC().UnixNano())
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}

//GetLocalIp returns back the IP of the interface hosting the default route
func GetLocalIp() string {
	// testing with  198.18.0.0/15 , see https://www.iana.org/assignments/iana-ipv4-special-registry/iana-ipv4-special-registry.xhtml
	conn, err := net.Dial("udp", "198.18.0.30:80")
	if err != nil {
		log.Printf("[TOOLS][UTILS][OS] : cannot use UDP: %s", err.Error())
		return "127.0.0.1" // wanted to use 0.0.0.0 but golang cannot make use of it
	}
	conn.Close()
	torn := strings.Split(conn.LocalAddr().String(), ":")
	return torn[0]
}

func AddrResolve(fqdn string) (addr string) {

	addresses, err := net.LookupIP(fqdn)

	if err != nil {
		log.Printf("[TOOLS][UTILS][DNS] ERROR %s", err)
		// protocol = "tcp4"
		return "127.0.0.1"

	} else {

		addr := addresses[0].String()

		log.Printf("[TOOLS][UTILS][DNS] Resolution ok: %s -> %s", fqdn, addr)

		if strings.Count(addr, ":") > 2 {
			addr = "[" + addr + "]"
			// protocol = "tcp6"
			return addr
		}

		if strings.Contains(addr, ".") {
			// protocol = "tcp4"
			return addr
		}

	}

	return "127.0.0.1"

}

func CreateTLV(typ string, payload []byte) []byte {

	if len(payload) == 0 {
		return nil
	}

	payloadB64 := base64.StdEncoding.EncodeToString(payload)

	return []byte(typ + "|" + payloadB64)

}

func UnPackTLV(n_tlv []byte) (typ string, ln int, payload []byte) {

	if bytes.Count(n_tlv, []byte("|")) != 1 {
		log.Println("[TOOLS][TLV][UNPACK] WTF is that:", string(n_tlv))
		return "Z", 0, nil
	}
	tlv := bytes.Split(n_tlv, []byte("|"))

	tlvBin, err := base64.StdEncoding.DecodeString(string(tlv[1]))

	if err != nil {
		log.Println("[TOOLS][TLV][UNPACK] Error with base64:", err.Error())
		return "Z", 0, nil
	}

	return string(tlv[0]), len(tlvBin), tlvBin

}

func GetFQDN() string {

	myIP := GetLocalIp()

	names, err := net.LookupAddr(myIP)
	if err != nil {
		log.Println("[TOOLS][UTILS][FQDN] Error getting my hostname:", err.Error())
		return "localhost"
	}

	return names[0]

}
<commit_msg>better sanitization of tlv<commit_after>package tools

import (
	"bytes"
	"encoding/base64"
	"log"
	"math/rand"
	"net"
	"regexp"

	"strings"
	"time"
)

var letters = []rune("0123456789-+@abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

//RandSeq returns a random string
func RandSeq(n int) string {
	rand.Seed(time.Now().UTC().UnixNano())
	b := make([]rune, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}

//GetLocalIp returns back the IP of the interface hosting the default route
func GetLocalIp() string {
	// testing with  198.18.0.0/15 , see https://www.iana.org/assignments/iana-ipv4-special-registry/iana-ipv4-special-registry.xhtml
	conn, err := net.Dial("udp", "198.18.0.30:80")
	if err != nil {
		log.Printf("[TOOLS][UTILS][OS] : cannot use UDP: %s", err.Error())
		return "127.0.0.1" // wanted to use 0.0.0.0 but golang cannot make use of it
	}
	conn.Close()
	torn := strings.Split(conn.LocalAddr().String(), ":")
	return torn[0]
}

func AddrResolve(fqdn string) (addr string) {

	addresses, err := net.LookupIP(fqdn)

	if err != nil {
		log.Printf("[TOOLS][UTILS][DNS] ERROR %s", err)
		// protocol = "tcp4"
		return "127.0.0.1"

	} else {

		addr := addresses[0].String()

		log.Printf("[TOOLS][UTILS][DNS] Resolution ok: %s -> %s", fqdn, addr)

		if strings.Count(addr, ":") > 2 {
			addr = "[" + addr + "]"
			// protocol = "tcp6"
			return addr
		}

		if strings.Contains(addr, ".") {
			// protocol = "tcp4"
			return addr
		}

	}

	return "127.0.0.1"

}

func CreateTLV(typ string, payload []byte) []byte {

	if len(payload) == 0 {
		return nil
	}

	payloadB64 := base64.StdEncoding.EncodeToString(payload)

	var ret bytes.Buffer

	ret.WriteString(typ)
	ret.WriteString("|")
	ret.WriteString(payloadB64)

	return ret.Bytes()

}

func UnPackTLV(n_tlv []byte) (typ string, ln int, payload []byte) {

	re := regexp.MustCompile("^([AQFD])[|]([A-Za-z0-9+/=]+)$")

	tlv := re.FindStringSubmatch(string(n_tlv))

	if tlv == nil {
		log.Println("[TOOLS][TLV][UNPACK] WTF is this: ", string(n_tlv))
		return "Z", 0, nil
	}

	// tlv[1] contains the typ , tlv[2] contains the payload

	tlvBin, err := base64.StdEncoding.DecodeString(tlv[2])

	if err != nil {
		log.Println("[TOOLS][TLV][UNPACK] Error with base64:", err.Error())
		return "Z", 0, nil
	}

	return tlv[1], len(tlvBin), tlvBin

}

func GetFQDN() string {

	myIP := GetLocalIp()

	names, err := net.LookupAddr(myIP)
	if err != nil {
		log.Println("[TOOLS][UTILS][FQDN] Error getting my hostname:", err.Error())
		return "localhost"
	}

	return names[0]

}
<|endoftext|>


<commit_before>package bot

import (
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"sort"

	"encoding/json"

	"log"

	"github.com/nlopes/slack"
	"golang.org/x/net/html"
)

type TimeTableResponse struct {
	XMLName  xml.Name `xml:"response"`
	Year     int      `xml:"year,attr"`
	Semester int      `xml:"semester,attr"`
	Subjects []struct {
		Name struct {
			Value string `xml:"value,attr"`
		} `xml:"name"`
		Professor struct {
			Value string `xml:"value,attr"`
		} `xml:"professor"`
		Times []struct {
			Place     string `xml:"place,attr"`
			EndTime   int    `xml:"endtime,attr"`
			StartTime int    `xml:"starttime,attr"`
			Day       int    `xml:"day,attr"`
		} `xml:"time>data"`
	} `xml:"subject"`
}

type Subject struct {
	Name      string
	Professor string
}
type TimeTableEvent struct {
	Id        int
	Place     string
	StartTime int
	EndTime   int
}

type ByStartTime []TimeTableEvent

func (a ByStartTime) Len() int           { return len(a) }
func (a ByStartTime) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByStartTime) Less(i, j int) bool { return a[i].StartTime < a[j].StartTime }

type TimeTableDay struct {
	Events []TimeTableEvent
}

type TimeTable struct {
	Year     int
	Semester int
	Days     []TimeTableDay
	Subjects []Subject
}

func (resp *TimeTableResponse) toTimeTable() *TimeTable {
	ret := &TimeTable{}

	ret.Year = resp.Year
	ret.Semester = resp.Semester

	ret.Days = make([]TimeTableDay, 5)
	ret.Subjects = make([]Subject, len(resp.Subjects))
	for i, subject := range resp.Subjects {
		ret.Subjects[i] = Subject{
			subject.Name.Value,
			subject.Professor.Value,
		}
		for _, day := range subject.Times {
			ret.Days[day.Day].Events = append(ret.Days[day.Day].Events, TimeTableEvent{
				i,
				day.Place,
				day.StartTime,
				day.EndTime,
			})
		}
	}

	for _, day := range ret.Days {
		sort.Sort(ByStartTime(day.Events))
	}

	return ret
}

// return value is dummy...
func getEveryTimeTable(bot *Meu, userid string, username string, et_nick string) error {
	ret := &slack.OutgoingMessage{}
	quit := func(msg string, args ...interface{}) error {
		ret.Text = fmt.Sprintf("<@%s> "+msg, userid, args)
		log.Print(ret.Text)
		bot.SendMessage(ret)
		return nil
	}

	url_str := fmt.Sprintf("http://everytime.kr/@%s", et_nick)
	resp, err := http.Get(url_str)
	if err != nil {
		return quit("에러: 시간표 목록을 가져오는데 실패했습니다.")
	}

	latest_info, err := parseTimeTableList(resp.Body)
	if err != nil {
		return quit("에러: 시간표 목록을 파싱하는데 실패했습니다. %q", err)
	}

	now := time.Now()
	if now.Year() != latest_info.Year || (int(now.Month())-1)/6+1 != latest_info.Semester {
		return quit("에러: 현재 학기의 시간표가 없습니다.")
	}

	url_str = "http://everytime.kr/ajax/timetable/wizard/tableload"
	resp, err = http.PostForm(url_str, url.Values{
		"id": {latest_info.Id},
	})
	if err != nil {
		return quit("에러: 시간표 정보를 가져오는데 실패했습니다.")
	}

	schedule_body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return quit("에러: 시간표 읽기 실패")
	}

	timetableResp, err := parseTimeTable(schedule_body)
	if err != nil {
		return quit("에러: 시간표 파싱 에러 - %q", err)
	}

	timetable := timetableResp.toTimeTable()
	bot.timetable[userid] = timetable
	serialized, err := json.Marshal(timetable)
	if err != nil {
		return quit("에러: 시간표 저장 실패")
	}

	bot.rc.Set(TimeTableKeyName(userid), serialized, 0)

	return nil
}

func TimeTableKeyName(userid string) string {
	return fmt.Sprintf("et_tt_%s", userid)
}

type TimeTableInfo struct {
	Year     int
	Semester int
	Id       string
}

func parseTimeTableList(reader io.Reader) (*TimeTableInfo, error) {
	doc, err := html.Parse(reader)
	if err != nil {
		return nil, err
	}
	ret := &TimeTableInfo{}

	var parser func(*html.Node) bool
	parser = func(n *html.Node) bool {
		if n.Type == html.ElementNode {
			switch {
			case n.Data == "input":
				var input_id, input_val string
				for _, attr := range n.Attr {
					if attr.Key == "id" {
						input_id = attr.Val
					} else if attr.Key == "value" {
						input_val = attr.Val
					}
				}
				switch input_id {
				case "year":
					ret.Year, _ = strconv.Atoi(input_val)
				case "semester":
					ret.Semester, _ = strconv.Atoi(input_val)
				case "tableId":
					ret.Id = input_val
				}
				break
			}
		}

		for c := n.FirstChild; c != nil; c = c.NextSibling {
			if !parser(c) {
				return false
			}
		}

		return true
	}
	parser(doc)

	return ret, nil
}

func parseTimeTable(body []byte) (*TimeTableResponse, error) {
	schedule := &TimeTableResponse{}
	err := xml.Unmarshal(body, &schedule)
	if err != nil {
		return nil, err
	}
	return schedule, nil
}

func (nextEvent *TimeTableEvent) toSlackAttachment(subject *Subject) slack.Attachment {
	now := time.Now()
	begin := time.Date(now.Year(), now.Month(), now.Day(), nextEvent.StartTime/12, (nextEvent.StartTime%12)*5, 0, 0, now.Location())
	end := time.Date(now.Year(), now.Month(), now.Day(), nextEvent.EndTime/12, (nextEvent.EndTime%12)*5, 0, 0, now.Location())

	return slack.Attachment{
		Fields: []slack.AttachmentField{
			slack.AttachmentField{
				Title: "수업명",
				Value: subject.Name,
				Short: true,
			},
			slack.AttachmentField{
				Title: "교수",
				Value: subject.Professor,
				Short: true,
			},
			slack.AttachmentField{
				Title: "장소",
				Value: nextEvent.Place,
				Short: true,
			},
			slack.AttachmentField{
				Title: "수업 시각",
				Value: fmt.Sprintf("%s ~ %s", begin.Format(time.Kitchen), end.Format(time.Kitchen)),
				Short: true,
			},
		},
	}
}

func register_et(bot *Meu, e *slack.MessageEvent, matched []string) {
	// 에브리타임 기록
	et_nick := matched[1]
	bot.rc.Set(fmt.Sprintf("et_nick_%s", e.User), et_nick, 0)
	bot.replySimple(e, "기억했다 메우. 시간표를 가져오겠다 메우.\n시간이 좀 걸릴거다 메우.")

	user, _ := bot.GetUserInfo(e.User)
	go getEveryTimeTable(bot, e.User, user.Name, et_nick)
}

func next_et(bot *Meu, e *slack.MessageEvent, matched []string) {
	// 에브리타임 다음 시간
	log.Printf("%q", bot.timetable)
	timetable, exists := bot.timetable[e.User]
	if !exists {
		log.Print("Get from redis")
		result := bot.rc.Get(TimeTableKeyName(e.User))
		if result == nil {
			bot.replySimple(e, "시간표 정보가 없다 메우. 등록부터 해달라 메우.")
			return
		}

		bytes, _ := result.Bytes()
		timetable = &TimeTable{}
		if json.Unmarshal(bytes, timetable) != nil {
			bot.replySimple(e, "저장된 시간표가 이상하다 메우. 새로 등록해달라 메우.")
			return
		}

		bot.timetable[e.User] = timetable
	}

	now := time.Now()
	weekDay := int(now.Weekday() - 1)
	curHour := now.Hour()*12 + now.Minute()/5
	var nextEvent *TimeTableEvent
	if weekDay < 0 || weekDay >= 5 {
		nextEvent = nil
	} else {
		log.Printf("%q", timetable)
		todayEvents := timetable.Days[weekDay].Events
		for _, event := range todayEvents {
			if event.StartTime > curHour {
				nextEvent = &event
				break
			}
		}
	}
	if nextEvent == nil {
		bot.replySimple(e, "오늘은 더 이상 수업이 없다 메우.")
	} else {
		subject := &timetable.Subjects[nextEvent.Id]

		bot.PostMessage(e.Channel, fmt.Sprintf("다음 수업은 \"%s\"다 메우.", subject.Name), slack.PostMessageParameters{
			AsUser:    false,
			IconEmoji: ":meu:",
			Username:  "시간표 알려주는 메우",
			Attachments: []slack.Attachment{
				nextEvent.toSlackAttachment(subject),
			},
		})
	}
}
<commit_msg>update everytime timetable parsing code<commit_after>package bot

import (
	"encoding/xml"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"

	"sort"

	"encoding/json"

	"log"

	"github.com/nlopes/slack"
	"golang.org/x/net/html"
)

type TimeTableResponseBase struct {
	Id       string `xml:"id,attr"`
	Year     int    `xml:"year,attr"`
	Semester int    `xml:"semester,attr"`
}

type TimeTableResponse struct {
	XMLName xml.Name                `xml:"response"`
	Tables  []TimeTableResponseBase `xml:"table"`
}

type TimeTableDetailResponse struct {
	XMLName xml.Name `xml:"response"`
	TimeTableResponseBase
	Subjects []struct {
		Name struct {
			Value string `xml:"value,attr"`
		} `xml:"name"`
		Professor struct {
			Value string `xml:"value,attr"`
		} `xml:"professor"`
		Times []struct {
			Place     string `xml:"place,attr"`
			EndTime   int    `xml:"endtime,attr"`
			StartTime int    `xml:"starttime,attr"`
			Day       int    `xml:"day,attr"`
		} `xml:"time>data"`
	} `xml:"subject"`
}

type Subject struct {
	Name      string
	Professor string
}
type TimeTableEvent struct {
	Id        int
	Place     string
	StartTime int
	EndTime   int
}

type ByStartTime []TimeTableEvent

func (a ByStartTime) Len() int           { return len(a) }
func (a ByStartTime) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByStartTime) Less(i, j int) bool { return a[i].StartTime < a[j].StartTime }

type TimeTableDay struct {
	Events []TimeTableEvent
}

type TimeTable struct {
	Year     int
	Semester int
	Days     []TimeTableDay
	Subjects []Subject
}

func (resp *TimeTableDetailResponse) toTimeTable() *TimeTable {
	ret := &TimeTable{}

	ret.Year = resp.Year
	ret.Semester = resp.Semester

	ret.Days = make([]TimeTableDay, 5)
	ret.Subjects = make([]Subject, len(resp.Subjects))
	for i, subject := range resp.Subjects {
		ret.Subjects[i] = Subject{
			subject.Name.Value,
			subject.Professor.Value,
		}
		for _, day := range subject.Times {
			ret.Days[day.Day].Events = append(ret.Days[day.Day].Events, TimeTableEvent{
				i,
				day.Place,
				day.StartTime,
				day.EndTime,
			})
		}
	}

	for _, day := range ret.Days {
		sort.Sort(ByStartTime(day.Events))
	}

	return ret
}

// return value is dummy...
func getEveryTimeTable(bot *Meu, userid string, username string, et_nick string) error {
	ret := &slack.OutgoingMessage{}
	quit := func(msg string, args ...interface{}) error {
		ret.Text = fmt.Sprintf("<@%s> "+msg, userid, args)
		log.Print(ret.Text)
		bot.SendMessage(ret)
		return nil
	}

	url_str := fmt.Sprintf("http://everytime.kr/@%s", et_nick)
	resp, err := http.Get(url_str)
	if err != nil {
		return quit("에러: friendToken을 가져오는데 실패했습니다.")
	}

	friend_token, err := parseEveryTimeUserId(resp.Body)
	if err != nil {
		return quit("에러: friendToken을 파싱하는데 실패했습니다. %q", err)
	}

	url_str = "http://timetable.everytime.kr/ajax/timetable/wizard/getPrimaryTableList"
	resp, err = http.PostForm(url_str, url.Values{
		"userid": {friend_token},
	})
	if err != nil {
		return quit("에러: 시간표 목록을 가져오는데 실패했습니다.")
	}

	table_list_body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return quit("에러: 시간표 목록 읽기 실패")
	}

	table_list, err := parseTimeTableList(table_list_body)
	if err != nil {
		return quit("에러: 시간표 목록 파싱 에러 - %q", err)
	}

	now := time.Now()
	current_year := now.Year()
	current_semester := (int(now.Month())-1)/6 + 1
	var current_table *TimeTableResponseBase = nil
	for _, table := range table_list.Tables {
		if table.Year == current_year && table.Semester == current_semester {
			fmt.Println("%d %d", table.Year, table.Semester)
			current_table = &table
			break
		}
	}
	if current_table == nil {
		return quit("에러: 현재 학기의 시간표가 없습니다.")
	}

	url_str = "http://timetable.everytime.kr/ajax/timetable/wizard/getOneTable"
	resp, err = http.PostForm(url_str, url.Values{
		"id":     {current_table.Id},
		"userid": {friend_token},
	})
	if err != nil {
		return quit("에러: 시간표 정보를 가져오는데 실패했습니다.")
	}

	schedule_body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return quit("에러: 시간표 읽기 실패")
	}

	timetableResp, err := parseTimeTable(schedule_body)
	if err != nil {
		return quit("에러: 시간표 파싱 에러 - %q", err)
	}

	timetable := timetableResp.toTimeTable()
	bot.timetable[userid] = timetable
	serialized, err := json.Marshal(timetable)
	if err != nil {
		return quit("에러: 시간표 저장 실패")
	}

	bot.rc.Set(TimeTableKeyName(userid), serialized, 0)

	return nil
}

func TimeTableKeyName(userid string) string {
	return fmt.Sprintf("et_tt_%s", userid)
}

type TimeTableInfo struct {
	Year     int
	Semester int
	Id       string
}

func parseEveryTimeUserId(reader io.Reader) (string, error) {
	doc, err := html.Parse(reader)
	if err != nil {
		return "", err
	}
	var ret string = ""

	var parser func(*html.Node) bool
	parser = func(n *html.Node) bool {
		if n.Type == html.ElementNode {
			switch {
			case n.Data == "input":
				var input_id, input_val string
				for _, attr := range n.Attr {
					if attr.Key == "id" {
						input_id = attr.Val
					} else if attr.Key == "value" {
						input_val = attr.Val
					}
				}
				switch input_id {
				case "friendToken":
					ret = input_val
				}
				break
			}
		}

		if ret != "" {
			return false
		}
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			if !parser(c) {
				return false
			}
		}

		return true
	}
	parser(doc)

	return ret, nil
}

func parseTimeTableList(body []byte) (*TimeTableResponse, error) {
	list := &TimeTableResponse{}
	err := xml.Unmarshal(body, &list)
	if err != nil {
		return nil, err
	}
	return list, nil
}

func parseTimeTable(body []byte) (*TimeTableDetailResponse, error) {
	schedule := &TimeTableDetailResponse{}
	err := xml.Unmarshal(body, &schedule)
	if err != nil {
		return nil, err
	}
	return schedule, nil
}

func (nextEvent *TimeTableEvent) toSlackAttachment(subject *Subject) slack.Attachment {
	now := time.Now()
	begin := time.Date(now.Year(), now.Month(), now.Day(), nextEvent.StartTime/12, (nextEvent.StartTime%12)*5, 0, 0, now.Location())
	end := time.Date(now.Year(), now.Month(), now.Day(), nextEvent.EndTime/12, (nextEvent.EndTime%12)*5, 0, 0, now.Location())

	return slack.Attachment{
		Fields: []slack.AttachmentField{
			slack.AttachmentField{
				Title: "수업명",
				Value: subject.Name,
				Short: true,
			},
			slack.AttachmentField{
				Title: "교수",
				Value: subject.Professor,
				Short: true,
			},
			slack.AttachmentField{
				Title: "장소",
				Value: nextEvent.Place,
				Short: true,
			},
			slack.AttachmentField{
				Title: "수업 시각",
				Value: fmt.Sprintf("%s ~ %s", begin.Format(time.Kitchen), end.Format(time.Kitchen)),
				Short: true,
			},
		},
	}
}

func register_et(bot *Meu, e *slack.MessageEvent, matched []string) {
	// 에브리타임 기록
	et_nick := matched[1]
	bot.rc.Set(fmt.Sprintf("et_nick_%s", e.User), et_nick, 0)
	bot.replySimple(e, "기억했다 메우. 시간표를 가져오겠다 메우.\n시간이 좀 걸릴거다 메우.")

	user, _ := bot.GetUserInfo(e.User)
	go getEveryTimeTable(bot, e.User, user.Name, et_nick)
}

func next_et(bot *Meu, e *slack.MessageEvent, matched []string) {
	// 에브리타임 다음 시간
	log.Printf("%q", bot.timetable)
	timetable, exists := bot.timetable[e.User]
	if !exists {
		log.Print("Get from redis")
		result := bot.rc.Get(TimeTableKeyName(e.User))
		if result == nil {
			bot.replySimple(e, "시간표 정보가 없다 메우. 등록부터 해달라 메우.")
			return
		}

		bytes, _ := result.Bytes()
		timetable = &TimeTable{}
		if json.Unmarshal(bytes, timetable) != nil {
			bot.replySimple(e, "저장된 시간표가 이상하다 메우. 새로 등록해달라 메우.")
			return
		}

		bot.timetable[e.User] = timetable
	}

	now := time.Now()
	weekDay := int(now.Weekday() - 1)
	curHour := now.Hour()*12 + now.Minute()/5
	var nextEvent *TimeTableEvent
	if weekDay < 0 || weekDay >= 5 {
		nextEvent = nil
	} else {
		log.Printf("%q", timetable)
		todayEvents := timetable.Days[weekDay].Events
		for _, event := range todayEvents {
			if event.StartTime > curHour {
				nextEvent = &event
				break
			}
		}
	}
	if nextEvent == nil {
		bot.replySimple(e, "오늘은 더 이상 수업이 없다 메우.")
	} else {
		subject := &timetable.Subjects[nextEvent.Id]

		bot.PostMessage(e.Channel, fmt.Sprintf("다음 수업은 \"%s\"다 메우.", subject.Name), slack.PostMessageParameters{
			AsUser:    false,
			IconEmoji: ":meu:",
			Username:  "시간표 알려주는 메우",
			Attachments: []slack.Attachment{
				nextEvent.toSlackAttachment(subject),
			},
		})
	}
}
<|endoftext|>


<commit_before>package objectstore

import (
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/Sirupsen/logrus"
	"github.com/rancher/convoy/metadata"
	"github.com/rancher/convoy/util"

	. "github.com/rancher/convoy/logging"
)

type BlockMapping struct {
	Offset        int64
	BlockChecksum string
}

type DeltaBlockBackupOperations interface {
	HasSnapshot(id, volumeID string) bool
	CompareSnapshot(id, compareID, volumeID string) (*metadata.Mappings, error)
	OpenSnapshot(id, volumeID string) error
	ReadSnapshot(id, volumeID string, start int64, data []byte) error
	CloseSnapshot(id, volumeID string) error
}

const (
	DEFAULT_BLOCK_SIZE = 2097152

	BLOCKS_DIRECTORY      = "blocks"
	BLOCK_SEPARATE_LAYER1 = 2
	BLOCK_SEPARATE_LAYER2 = 4
)

func CreateDeltaBlockBackup(volume *Volume, snapshot *Snapshot, destURL string, deltaOps DeltaBlockBackupOperations) (string, error) {
	if deltaOps == nil {
		return "", fmt.Errorf("Missing DeltaBlockBackupOperations")
	}

	bsDriver, err := GetObjectStoreDriver(destURL)
	if err != nil {
		return "", err
	}

	if err := addVolume(volume, bsDriver); err != nil {
		return "", err
	}

	// Update volume from objectstore
	volume, err = loadVolume(volume.Name, bsDriver)
	if err != nil {
		return "", err
	}

	lastBackupName := volume.LastBackupName

	if err := deltaOps.OpenSnapshot(snapshot.Name, volume.Name); err != nil {
		return "", err
	}
	defer deltaOps.CloseSnapshot(snapshot.Name, volume.Name)

	var lastSnapshotName string
	var lastBackup *Backup
	if lastBackupName != "" {
		lastBackup, err = loadBackup(lastBackupName, volume.Name, bsDriver)
		if err != nil {
			return "", err
		}

		lastSnapshotName = lastBackup.SnapshotName
		if lastSnapshotName == snapshot.Name {
			//Generate full snapshot if the snapshot has been backed up last time
			lastSnapshotName = ""
			log.Debug("Would create full snapshot metadata")
		} else if !deltaOps.HasSnapshot(lastSnapshotName, volume.Name) {
			// It's possible that the snapshot in objectstore doesn't exist
			// in local storage
			lastSnapshotName = ""
			log.WithFields(logrus.Fields{
				LOG_FIELD_REASON:   LOG_REASON_FALLBACK,
				LOG_FIELD_OBJECT:   LOG_OBJECT_SNAPSHOT,
				LOG_FIELD_SNAPSHOT: lastSnapshotName,
				LOG_FIELD_VOLUME:   volume.Name,
			}).Debug("Cannot find last snapshot in local storage, would process with full backup")
		}
	}

	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:        LOG_REASON_START,
		LOG_FIELD_OBJECT:        LOG_OBJECT_SNAPSHOT,
		LOG_FIELD_EVENT:         LOG_EVENT_COMPARE,
		LOG_FIELD_SNAPSHOT:      snapshot.Name,
		LOG_FIELD_LAST_SNAPSHOT: lastSnapshotName,
	}).Debug("Generating snapshot changed blocks metadata")

	delta, err := deltaOps.CompareSnapshot(snapshot.Name, lastSnapshotName, volume.Name)
	if err != nil {
		return "", err
	}
	if delta.BlockSize != DEFAULT_BLOCK_SIZE {
		return "", fmt.Errorf("Currently doesn't support different block sizes driver other than %v", DEFAULT_BLOCK_SIZE)
	}
	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:        LOG_REASON_COMPLETE,
		LOG_FIELD_OBJECT:        LOG_OBJECT_SNAPSHOT,
		LOG_FIELD_EVENT:         LOG_EVENT_COMPARE,
		LOG_FIELD_SNAPSHOT:      snapshot.Name,
		LOG_FIELD_LAST_SNAPSHOT: lastSnapshotName,
	}).Debug("Generated snapshot changed blocks metadata")

	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:   LOG_REASON_START,
		LOG_FIELD_EVENT:    LOG_EVENT_BACKUP,
		LOG_FIELD_OBJECT:   LOG_OBJECT_SNAPSHOT,
		LOG_FIELD_SNAPSHOT: snapshot.Name,
	}).Debug("Creating backup")

	deltaBackup := &Backup{
		Name:         util.GenerateName("backup"),
		VolumeName:   volume.Name,
		SnapshotName: snapshot.Name,
		Blocks:       []BlockMapping{},
	}
	mCounts := len(delta.Mappings)
	for m, d := range delta.Mappings {
		if d.Size%delta.BlockSize != 0 {
			return "", fmt.Errorf("Mapping's size %v is not multiples of backup block size %v",
				d.Size, delta.BlockSize)
		}
		block := make([]byte, DEFAULT_BLOCK_SIZE)
		blkCounts := d.Size / delta.BlockSize
		for i := int64(0); i < blkCounts; i++ {
			offset := d.Offset + i*delta.BlockSize
			log.Debugf("Backup for %v: segment %v/%v, blocks %v/%v", snapshot.Name, m+1, mCounts, i+1, blkCounts)
			err := deltaOps.ReadSnapshot(snapshot.Name, volume.Name, offset, block)
			if err != nil {
				return "", err
			}
			checksum := util.GetChecksum(block)
			blkFile := getBlockFilePath(volume.Name, checksum)
			if bsDriver.FileSize(blkFile) >= 0 {
				blockMapping := BlockMapping{
					Offset:        offset,
					BlockChecksum: checksum,
				}
				deltaBackup.Blocks = append(deltaBackup.Blocks, blockMapping)
				log.Debugf("Found existed block match at %v", blkFile)
				continue
			}

			rs, err := util.CompressData(block)
			if err != nil {
				return "", err
			}

			if err := bsDriver.Write(blkFile, rs); err != nil {
				return "", err
			}
			log.Debugf("Created new block file at %v", blkFile)

			blockMapping := BlockMapping{
				Offset:        offset,
				BlockChecksum: checksum,
			}
			deltaBackup.Blocks = append(deltaBackup.Blocks, blockMapping)
		}
	}

	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:   LOG_REASON_COMPLETE,
		LOG_FIELD_EVENT:    LOG_EVENT_BACKUP,
		LOG_FIELD_OBJECT:   LOG_OBJECT_SNAPSHOT,
		LOG_FIELD_SNAPSHOT: snapshot.Name,
	}).Debug("Created snapshot changed blocks")

	backup := mergeSnapshotMap(deltaBackup, lastBackup)
	backup.SnapshotName = snapshot.Name
	backup.SnapshotCreatedAt = snapshot.CreatedTime
	backup.CreatedTime = util.Now()

	if err := saveBackup(backup, bsDriver); err != nil {
		return "", err
	}

	volume.LastBackupName = backup.Name
	if err := saveVolume(volume, bsDriver); err != nil {
		return "", err
	}

	return encodeBackupURL(backup.Name, volume.Name, destURL), nil
}

func mergeSnapshotMap(deltaBackup, lastBackup *Backup) *Backup {
	if lastBackup == nil {
		return deltaBackup
	}
	backup := &Backup{
		Name:         deltaBackup.Name,
		VolumeName:   deltaBackup.VolumeName,
		SnapshotName: deltaBackup.SnapshotName,
		Blocks:       []BlockMapping{},
	}
	var d, l int
	for d, l = 0, 0; d < len(deltaBackup.Blocks) && l < len(lastBackup.Blocks); {
		dB := deltaBackup.Blocks[d]
		lB := lastBackup.Blocks[l]
		if dB.Offset == lB.Offset {
			backup.Blocks = append(backup.Blocks, dB)
			d++
			l++
		} else if dB.Offset < lB.Offset {
			backup.Blocks = append(backup.Blocks, dB)
			d++
		} else {
			//dB.Offset > lB.offset
			backup.Blocks = append(backup.Blocks, lB)
			l++
		}
	}

	if d == len(deltaBackup.Blocks) {
		backup.Blocks = append(backup.Blocks, lastBackup.Blocks[l:]...)
	} else {
		backup.Blocks = append(backup.Blocks, deltaBackup.Blocks[d:]...)
	}

	return backup
}

func RestoreDeltaBlockBackup(backupURL, volDevName string) error {
	bsDriver, err := GetObjectStoreDriver(backupURL)
	if err != nil {
		return err
	}

	srcBackupName, srcVolumeName, err := decodeBackupURL(backupURL)
	if err != nil {
		return err
	}

	if _, err := loadVolume(srcVolumeName, bsDriver); err != nil {
		return generateError(logrus.Fields{
			LOG_FIELD_VOLUME:     srcVolumeName,
			LOG_FIELD_BACKUP_URL: backupURL,
		}, "Volume doesn't exist in objectstore: %v", err)
	}

	volDev, err := os.Create(volDevName)
	if err != nil {
		return err
	}
	defer volDev.Close()

	backup, err := loadBackup(srcBackupName, srcVolumeName, bsDriver)
	if err != nil {
		return err
	}

	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:      LOG_REASON_START,
		LOG_FIELD_EVENT:       LOG_EVENT_RESTORE,
		LOG_FIELD_OBJECT:      LOG_FIELD_SNAPSHOT,
		LOG_FIELD_SNAPSHOT:    srcBackupName,
		LOG_FIELD_ORIN_VOLUME: srcVolumeName,
		LOG_FIELD_VOLUME_DEV:  volDevName,
		LOG_FIELD_BACKUP_URL:  backupURL,
	}).Debug()
	blkCounts := len(backup.Blocks)
	for i, block := range backup.Blocks {
		log.Debugf("Restore for %v: block %v, %v/%v", volDevName, block.BlockChecksum, i+1, blkCounts)
		blkFile := getBlockFilePath(srcVolumeName, block.BlockChecksum)
		rc, err := bsDriver.Read(blkFile)
		if err != nil {
			return err
		}
		r, err := util.DecompressAndVerify(rc, block.BlockChecksum)
		rc.Close()
		if err != nil {
			return err
		}
		if _, err := volDev.Seek(block.Offset, 0); err != nil {
			return err
		}
		if _, err := io.CopyN(volDev, r, DEFAULT_BLOCK_SIZE); err != nil {
			return err
		}
	}

	return nil
}

func DeleteDeltaBlockBackup(backupURL string) error {
	bsDriver, err := GetObjectStoreDriver(backupURL)
	if err != nil {
		return err
	}

	backupName, volumeName, err := decodeBackupURL(backupURL)
	if err != nil {
		return err
	}

	v, err := loadVolume(volumeName, bsDriver)
	if err != nil {
		return fmt.Errorf("Cannot find volume %v in objectstore", volumeName, err)
	}

	backup, err := loadBackup(backupName, volumeName, bsDriver)
	if err != nil {
		return err
	}
	discardBlockSet := make(map[string]bool)
	for _, blk := range backup.Blocks {
		discardBlockSet[blk.BlockChecksum] = true
	}
	discardBlockCounts := len(discardBlockSet)

	if err := removeBackup(backup, bsDriver); err != nil {
		return err
	}

	if backup.Name == v.LastBackupName {
		v.LastBackupName = ""
		if err := saveVolume(v, bsDriver); err != nil {
			return err
		}
	}

	backupNames, err := getBackupNamesForVolume(volumeName, bsDriver)
	if err != nil {
		return err
	}
	if len(backupNames) == 0 {
		log.Debugf("No snapshot existed for the volume %v, removing volume", volumeName)
		if err := removeVolume(volumeName, bsDriver); err != nil {
			log.Warningf("Failed to remove volume %v due to: %v", volumeName, err.Error())
		}
		return nil
	}

	log.Debug("GC started")
	for _, backupName := range backupNames {
		backup, err := loadBackup(backupName, volumeName, bsDriver)
		if err != nil {
			return err
		}
		for _, blk := range backup.Blocks {
			if _, exists := discardBlockSet[blk.BlockChecksum]; exists {
				delete(discardBlockSet, blk.BlockChecksum)
				discardBlockCounts--
				if discardBlockCounts == 0 {
					break
				}
			}
		}
		if discardBlockCounts == 0 {
			break
		}
	}

	var blkFileList []string
	for blk := range discardBlockSet {
		blkFileList = append(blkFileList, getBlockFilePath(volumeName, blk))
		log.Debugf("Found unused blocks %v for volume %v", blk, volumeName)
	}
	if err := bsDriver.Remove(blkFileList...); err != nil {
		return err
	}
	log.Debug("Removed unused blocks for volume ", volumeName)

	log.Debug("GC completed")
	log.Debug("Removed objectstore backup ", backupName)

	return nil
}

func getBlockPath(volumeName string) string {
	return filepath.Join(getVolumePath(volumeName), BLOCKS_DIRECTORY) + "/"
}

func getBlockFilePath(volumeName, checksum string) string {
	blockSubDirLayer1 := checksum[0:BLOCK_SEPARATE_LAYER1]
	blockSubDirLayer2 := checksum[BLOCK_SEPARATE_LAYER1:BLOCK_SEPARATE_LAYER2]
	path := filepath.Join(getBlockPath(volumeName), blockSubDirLayer1, blockSubDirLayer2)
	fileName := checksum + ".blk"

	return filepath.Join(path, fileName)
}
<commit_msg>objectstore: Truncate the file after finish restoration<commit_after>package objectstore

import (
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/Sirupsen/logrus"
	"github.com/rancher/convoy/metadata"
	"github.com/rancher/convoy/util"

	. "github.com/rancher/convoy/logging"
)

type BlockMapping struct {
	Offset        int64
	BlockChecksum string
}

type DeltaBlockBackupOperations interface {
	HasSnapshot(id, volumeID string) bool
	CompareSnapshot(id, compareID, volumeID string) (*metadata.Mappings, error)
	OpenSnapshot(id, volumeID string) error
	ReadSnapshot(id, volumeID string, start int64, data []byte) error
	CloseSnapshot(id, volumeID string) error
}

const (
	DEFAULT_BLOCK_SIZE = 2097152

	BLOCKS_DIRECTORY      = "blocks"
	BLOCK_SEPARATE_LAYER1 = 2
	BLOCK_SEPARATE_LAYER2 = 4
)

func CreateDeltaBlockBackup(volume *Volume, snapshot *Snapshot, destURL string, deltaOps DeltaBlockBackupOperations) (string, error) {
	if deltaOps == nil {
		return "", fmt.Errorf("Missing DeltaBlockBackupOperations")
	}

	bsDriver, err := GetObjectStoreDriver(destURL)
	if err != nil {
		return "", err
	}

	if err := addVolume(volume, bsDriver); err != nil {
		return "", err
	}

	// Update volume from objectstore
	volume, err = loadVolume(volume.Name, bsDriver)
	if err != nil {
		return "", err
	}

	lastBackupName := volume.LastBackupName

	if err := deltaOps.OpenSnapshot(snapshot.Name, volume.Name); err != nil {
		return "", err
	}
	defer deltaOps.CloseSnapshot(snapshot.Name, volume.Name)

	var lastSnapshotName string
	var lastBackup *Backup
	if lastBackupName != "" {
		lastBackup, err = loadBackup(lastBackupName, volume.Name, bsDriver)
		if err != nil {
			return "", err
		}

		lastSnapshotName = lastBackup.SnapshotName
		if lastSnapshotName == snapshot.Name {
			//Generate full snapshot if the snapshot has been backed up last time
			lastSnapshotName = ""
			log.Debug("Would create full snapshot metadata")
		} else if !deltaOps.HasSnapshot(lastSnapshotName, volume.Name) {
			// It's possible that the snapshot in objectstore doesn't exist
			// in local storage
			lastSnapshotName = ""
			log.WithFields(logrus.Fields{
				LOG_FIELD_REASON:   LOG_REASON_FALLBACK,
				LOG_FIELD_OBJECT:   LOG_OBJECT_SNAPSHOT,
				LOG_FIELD_SNAPSHOT: lastSnapshotName,
				LOG_FIELD_VOLUME:   volume.Name,
			}).Debug("Cannot find last snapshot in local storage, would process with full backup")
		}
	}

	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:        LOG_REASON_START,
		LOG_FIELD_OBJECT:        LOG_OBJECT_SNAPSHOT,
		LOG_FIELD_EVENT:         LOG_EVENT_COMPARE,
		LOG_FIELD_SNAPSHOT:      snapshot.Name,
		LOG_FIELD_LAST_SNAPSHOT: lastSnapshotName,
	}).Debug("Generating snapshot changed blocks metadata")

	delta, err := deltaOps.CompareSnapshot(snapshot.Name, lastSnapshotName, volume.Name)
	if err != nil {
		return "", err
	}
	if delta.BlockSize != DEFAULT_BLOCK_SIZE {
		return "", fmt.Errorf("Currently doesn't support different block sizes driver other than %v", DEFAULT_BLOCK_SIZE)
	}
	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:        LOG_REASON_COMPLETE,
		LOG_FIELD_OBJECT:        LOG_OBJECT_SNAPSHOT,
		LOG_FIELD_EVENT:         LOG_EVENT_COMPARE,
		LOG_FIELD_SNAPSHOT:      snapshot.Name,
		LOG_FIELD_LAST_SNAPSHOT: lastSnapshotName,
	}).Debug("Generated snapshot changed blocks metadata")

	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:   LOG_REASON_START,
		LOG_FIELD_EVENT:    LOG_EVENT_BACKUP,
		LOG_FIELD_OBJECT:   LOG_OBJECT_SNAPSHOT,
		LOG_FIELD_SNAPSHOT: snapshot.Name,
	}).Debug("Creating backup")

	deltaBackup := &Backup{
		Name:         util.GenerateName("backup"),
		VolumeName:   volume.Name,
		SnapshotName: snapshot.Name,
		Blocks:       []BlockMapping{},
	}
	mCounts := len(delta.Mappings)
	for m, d := range delta.Mappings {
		if d.Size%delta.BlockSize != 0 {
			return "", fmt.Errorf("Mapping's size %v is not multiples of backup block size %v",
				d.Size, delta.BlockSize)
		}
		block := make([]byte, DEFAULT_BLOCK_SIZE)
		blkCounts := d.Size / delta.BlockSize
		for i := int64(0); i < blkCounts; i++ {
			offset := d.Offset + i*delta.BlockSize
			log.Debugf("Backup for %v: segment %v/%v, blocks %v/%v", snapshot.Name, m+1, mCounts, i+1, blkCounts)
			err := deltaOps.ReadSnapshot(snapshot.Name, volume.Name, offset, block)
			if err != nil {
				return "", err
			}
			checksum := util.GetChecksum(block)
			blkFile := getBlockFilePath(volume.Name, checksum)
			if bsDriver.FileSize(blkFile) >= 0 {
				blockMapping := BlockMapping{
					Offset:        offset,
					BlockChecksum: checksum,
				}
				deltaBackup.Blocks = append(deltaBackup.Blocks, blockMapping)
				log.Debugf("Found existed block match at %v", blkFile)
				continue
			}

			rs, err := util.CompressData(block)
			if err != nil {
				return "", err
			}

			if err := bsDriver.Write(blkFile, rs); err != nil {
				return "", err
			}
			log.Debugf("Created new block file at %v", blkFile)

			blockMapping := BlockMapping{
				Offset:        offset,
				BlockChecksum: checksum,
			}
			deltaBackup.Blocks = append(deltaBackup.Blocks, blockMapping)
		}
	}

	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:   LOG_REASON_COMPLETE,
		LOG_FIELD_EVENT:    LOG_EVENT_BACKUP,
		LOG_FIELD_OBJECT:   LOG_OBJECT_SNAPSHOT,
		LOG_FIELD_SNAPSHOT: snapshot.Name,
	}).Debug("Created snapshot changed blocks")

	backup := mergeSnapshotMap(deltaBackup, lastBackup)
	backup.SnapshotName = snapshot.Name
	backup.SnapshotCreatedAt = snapshot.CreatedTime
	backup.CreatedTime = util.Now()

	if err := saveBackup(backup, bsDriver); err != nil {
		return "", err
	}

	volume.LastBackupName = backup.Name
	if err := saveVolume(volume, bsDriver); err != nil {
		return "", err
	}

	return encodeBackupURL(backup.Name, volume.Name, destURL), nil
}

func mergeSnapshotMap(deltaBackup, lastBackup *Backup) *Backup {
	if lastBackup == nil {
		return deltaBackup
	}
	backup := &Backup{
		Name:         deltaBackup.Name,
		VolumeName:   deltaBackup.VolumeName,
		SnapshotName: deltaBackup.SnapshotName,
		Blocks:       []BlockMapping{},
	}
	var d, l int
	for d, l = 0, 0; d < len(deltaBackup.Blocks) && l < len(lastBackup.Blocks); {
		dB := deltaBackup.Blocks[d]
		lB := lastBackup.Blocks[l]
		if dB.Offset == lB.Offset {
			backup.Blocks = append(backup.Blocks, dB)
			d++
			l++
		} else if dB.Offset < lB.Offset {
			backup.Blocks = append(backup.Blocks, dB)
			d++
		} else {
			//dB.Offset > lB.offset
			backup.Blocks = append(backup.Blocks, lB)
			l++
		}
	}

	if d == len(deltaBackup.Blocks) {
		backup.Blocks = append(backup.Blocks, lastBackup.Blocks[l:]...)
	} else {
		backup.Blocks = append(backup.Blocks, deltaBackup.Blocks[d:]...)
	}

	return backup
}

func RestoreDeltaBlockBackup(backupURL, volDevName string) error {
	bsDriver, err := GetObjectStoreDriver(backupURL)
	if err != nil {
		return err
	}

	srcBackupName, srcVolumeName, err := decodeBackupURL(backupURL)
	if err != nil {
		return err
	}

	vol, err := loadVolume(srcVolumeName, bsDriver)
	if err != nil {
		return generateError(logrus.Fields{
			LOG_FIELD_VOLUME:     srcVolumeName,
			LOG_FIELD_BACKUP_URL: backupURL,
		}, "Volume doesn't exist in objectstore: %v", err)
	}

	if vol.Size == 0 || vol.Size%DEFAULT_BLOCK_SIZE != 0 {
		return fmt.Errorf("Read invalid volume size %v", vol.Size)
	}

	volDev, err := os.Create(volDevName)
	if err != nil {
		return err
	}
	defer volDev.Close()

	stat, err := volDev.Stat()
	if err != nil {
		return err
	}

	backup, err := loadBackup(srcBackupName, srcVolumeName, bsDriver)
	if err != nil {
		return err
	}

	log.WithFields(logrus.Fields{
		LOG_FIELD_REASON:      LOG_REASON_START,
		LOG_FIELD_EVENT:       LOG_EVENT_RESTORE,
		LOG_FIELD_OBJECT:      LOG_FIELD_SNAPSHOT,
		LOG_FIELD_SNAPSHOT:    srcBackupName,
		LOG_FIELD_ORIN_VOLUME: srcVolumeName,
		LOG_FIELD_VOLUME_DEV:  volDevName,
		LOG_FIELD_BACKUP_URL:  backupURL,
	}).Debug()
	blkCounts := len(backup.Blocks)
	for i, block := range backup.Blocks {
		log.Debugf("Restore for %v: block %v, %v/%v", volDevName, block.BlockChecksum, i+1, blkCounts)
		blkFile := getBlockFilePath(srcVolumeName, block.BlockChecksum)
		rc, err := bsDriver.Read(blkFile)
		if err != nil {
			return err
		}
		r, err := util.DecompressAndVerify(rc, block.BlockChecksum)
		rc.Close()
		if err != nil {
			return err
		}
		if _, err := volDev.Seek(block.Offset, 0); err != nil {
			return err
		}
		if _, err := io.CopyN(volDev, r, DEFAULT_BLOCK_SIZE); err != nil {
			return err
		}
	}

	// We want to truncate regular files, but not device
	if stat.Mode()&os.ModeType == 0 {
		log.Debugf("Truncate %v to size %v", volDevName, vol.Size)
		if err := volDev.Truncate(vol.Size); err != nil {
			return err
		}
	}

	return nil
}

func DeleteDeltaBlockBackup(backupURL string) error {
	bsDriver, err := GetObjectStoreDriver(backupURL)
	if err != nil {
		return err
	}

	backupName, volumeName, err := decodeBackupURL(backupURL)
	if err != nil {
		return err
	}

	v, err := loadVolume(volumeName, bsDriver)
	if err != nil {
		return fmt.Errorf("Cannot find volume %v in objectstore", volumeName, err)
	}

	backup, err := loadBackup(backupName, volumeName, bsDriver)
	if err != nil {
		return err
	}
	discardBlockSet := make(map[string]bool)
	for _, blk := range backup.Blocks {
		discardBlockSet[blk.BlockChecksum] = true
	}
	discardBlockCounts := len(discardBlockSet)

	if err := removeBackup(backup, bsDriver); err != nil {
		return err
	}

	if backup.Name == v.LastBackupName {
		v.LastBackupName = ""
		if err := saveVolume(v, bsDriver); err != nil {
			return err
		}
	}

	backupNames, err := getBackupNamesForVolume(volumeName, bsDriver)
	if err != nil {
		return err
	}
	if len(backupNames) == 0 {
		log.Debugf("No snapshot existed for the volume %v, removing volume", volumeName)
		if err := removeVolume(volumeName, bsDriver); err != nil {
			log.Warningf("Failed to remove volume %v due to: %v", volumeName, err.Error())
		}
		return nil
	}

	log.Debug("GC started")
	for _, backupName := range backupNames {
		backup, err := loadBackup(backupName, volumeName, bsDriver)
		if err != nil {
			return err
		}
		for _, blk := range backup.Blocks {
			if _, exists := discardBlockSet[blk.BlockChecksum]; exists {
				delete(discardBlockSet, blk.BlockChecksum)
				discardBlockCounts--
				if discardBlockCounts == 0 {
					break
				}
			}
		}
		if discardBlockCounts == 0 {
			break
		}
	}

	var blkFileList []string
	for blk := range discardBlockSet {
		blkFileList = append(blkFileList, getBlockFilePath(volumeName, blk))
		log.Debugf("Found unused blocks %v for volume %v", blk, volumeName)
	}
	if err := bsDriver.Remove(blkFileList...); err != nil {
		return err
	}
	log.Debug("Removed unused blocks for volume ", volumeName)

	log.Debug("GC completed")
	log.Debug("Removed objectstore backup ", backupName)

	return nil
}

func getBlockPath(volumeName string) string {
	return filepath.Join(getVolumePath(volumeName), BLOCKS_DIRECTORY) + "/"
}

func getBlockFilePath(volumeName, checksum string) string {
	blockSubDirLayer1 := checksum[0:BLOCK_SEPARATE_LAYER1]
	blockSubDirLayer2 := checksum[BLOCK_SEPARATE_LAYER1:BLOCK_SEPARATE_LAYER2]
	path := filepath.Join(getBlockPath(volumeName), blockSubDirLayer1, blockSubDirLayer2)
	fileName := checksum + ".blk"

	return filepath.Join(path, fileName)
}
<|endoftext|>


<commit_before>package bridge

import (
	"errors"
	"net"
	"net/url"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
	"sync"

	dockerapi "github.com/fsouza/go-dockerclient"
)

var serviceIDPattern = regexp.MustCompile(`^(.+?):([a-zA-Z0-9][a-zA-Z0-9_.-]+):[0-9]+(?::udp)?$`)

type Bridge struct {
	sync.Mutex
	registry       RegistryAdapter
	docker         *dockerapi.Client
	services       map[string][]*Service
	deadContainers map[string]*DeadContainer
	config         Config
}

func New(docker *dockerapi.Client, adapterUri string, config Config) (*Bridge, error) {
	uri, err := url.Parse(adapterUri)
	if err != nil {
		return nil, errors.New("bad adapter uri: " + adapterUri)
	}
	factory, found := AdapterFactories.Lookup(uri.Scheme)
	if !found {
		return nil, errors.New("unrecognized adapter: " + adapterUri)
	}

	log.Debug("Using", uri.Scheme, "adapter:", uri)
	return &Bridge{
		docker:         docker,
		config:         config,
		registry:       factory.New(uri),
		services:       make(map[string][]*Service),
		deadContainers: make(map[string]*DeadContainer),
	}, nil
}

func (b *Bridge) Ping() error {
	return b.registry.Ping()
}

func (b *Bridge) Add(containerId string) {
	b.add(containerId, false, "")
}

func (b *Bridge) Remove(containerId string) {
	b.remove(containerId, true)
}

func (b *Bridge) RemoveOnExit(containerId string) {
	b.remove(containerId, b.shouldRemove(containerId))
}

func (b *Bridge) Refresh() {
	for containerId, services := range b.getServicesCopy() {
		for _, service := range services {
			err := b.registry.Refresh(service)
			if err != nil {
				log.Error("refresh failed:", service.ID, err)
				continue
			}
			log.Debug("refreshed:", containerId[:12], service.ID)
		}
	}
}

func (b *Bridge) PruneDeadContainers() {
	b.Lock()
	defer b.Unlock()
	for containerId, deadContainer := range b.deadContainers {
		deadContainer.TTL -= b.config.RefreshInterval
		if deadContainer.TTL <= 0 {
			delete(b.deadContainers, containerId)
		}
	}
}

// Get a deep copy of the current services in a thread-safe way
func (b *Bridge) getServicesCopy() map[string][]*Service {
	b.Lock()
	defer b.Unlock()
	svcsCopy := make(map[string][]*Service)

	for id, svc := range b.services {
		var svcPointers []*Service
		for _, s := range svc {
			t := *s
			z := &t
			svcPointers = append(svcPointers, z)
		}
		svcsCopy[id] = svcPointers
	}
	return svcsCopy
}

func (b *Bridge) Sync(quiet bool) {
	serviceSync(b, quiet, "")
}

func (b *Bridge) AllocateNewIPToServices(ip string) {
	serviceSync(b, true, ip)
}

func (b *Bridge) deleteDeadContainer(containerId string) {
	b.Lock()
	defer b.Unlock()

	if d := b.deadContainers[containerId]; d != nil {
		b.services[containerId] = d.Services
		delete(b.deadContainers, containerId)
	}
}

func (b *Bridge) appendService(containerId string, service *Service) {
	b.Lock()
	defer b.Unlock()
	if b.services[containerId] != nil {
		log.Debug("container, ", containerId[:12], ", already exists, will not append.")
		return
	}
	b.services[containerId] = append(b.services[containerId], service)
	log.Debug("added:", containerId[:12], service.ID)
}

func (b *Bridge) add(containerId string, quiet bool, newIP string) {
	b.deleteDeadContainer(containerId)

	b.Lock()
	if b.services[containerId] != nil {
		log.Debug("container, ", containerId[:12], ", already exists, ignoring")
		// Alternatively, remove and readd or resubmit.
		return
	}
	b.Unlock()

	container, err := b.docker.InspectContainer(containerId)
	if err != nil {
		log.Error("unable to inspect container:", containerId[:12], err)
		return
	}

	ports := make(map[string]ServicePort)

	// Extract configured host port mappings, relevant when using --net=host
	for port, _ := range container.Config.ExposedPorts {
		published := []dockerapi.PortBinding{{"0.0.0.0", port.Port()}}
		serviceP := servicePort(container, port, published)
		if newIP != "" {
			serviceP.HostIP = newIP
		}
		ports[string(port)] = serviceP
	}

	// Extract runtime port mappings, relevant when using --net=bridge
	for port, published := range container.NetworkSettings.Ports {
		ports[string(port)] = servicePort(container, port, published)
	}

	if len(ports) == 0 && !quiet {
		log.Debug("ignored:", container.ID[:12], "no published ports")
		return
	}

	servicePorts := make(map[string]ServicePort)
	for key, port := range ports {
		if b.config.Internal != true && port.HostPort == "" {
			if !quiet {
				log.Debug("ignored:", container.ID[:12], "port", port.ExposedPort, "not published on host")
			}
			continue
		}
		servicePorts[key] = port
	}

	isGroup := len(servicePorts) > 1
	for _, port := range servicePorts {
		log.Infof("Creating new service with IP: %s", port.HostIP)
		service := b.newService(port, isGroup)
		if service == nil {
			if !quiet {
				log.Debug("ignored:", container.ID[:12], "service on port", port.ExposedPort)
			}
			continue
		}
		err := b.registry.Register(service)
		if err != nil {
			log.Error("register failed:", service, err)
			continue
		}
		b.appendService(container.ID, service)
	}
}

func (b *Bridge) newService(port ServicePort, isgroup bool) *Service {
	container := port.container
	defaultName := strings.Split(path.Base(container.Config.Image), ":")[0]
	// not sure about this logic. kind of want to remove it.
	hostname := Hostname
	if hostname == "" {
		hostname = port.HostIP
	}
	if port.HostIP == "0.0.0.0" {
		ip, err := net.ResolveIPAddr("ip", hostname)
		if err == nil {
			port.HostIP = ip.String()
		}
	}

	if b.config.HostIp != "" {
		port.HostIP = b.config.HostIp
	}

	metadata, metadataFromPort := serviceMetaData(container.Config, port.ExposedPort)

	ignore := mapDefault(metadata, "ignore", "")
	if ignore != "" {
		return nil
	}

	if b.config.RequireLabel {
		if strings.ToLower(mapDefault(metadata, "register", "false")) == "false" {
			log.Debugf("Did not find label SERVICE_REGISTER on %s - ignoring.", container.Name)
			return nil
		}
	}

	service := new(Service)
	service.Origin = port
	service.ID = hostname + ":" + container.Name[1:] + ":" + port.ExposedPort
	service.Name = mapDefault(metadata, "name", defaultName)
	if isgroup && !metadataFromPort["name"] {
		service.Name += "-" + port.ExposedPort
	}
	var convertedPort int

	if b.config.Internal == true {
		service.IP = port.ExposedIP
		p, err := strconv.Atoi(port.ExposedPort)
		if err != nil {
			log.Error("Unable to parse string ExposedPort to int: %s", port.ExposedPort)
		} else {
			convertedPort = p
		}
	} else {
		service.IP = port.HostIP
		p, err := strconv.Atoi(port.HostPort)
		if err != nil {
			log.Error("Unable to parse string HostPort to int: %s", port.HostPort)
		} else {
			convertedPort = p
		}
	}
	service.Port = convertedPort

	if b.config.UseIpFromLabel != "" {
		containerIp := container.Config.Labels[b.config.UseIpFromLabel]
		if containerIp != "" {
			slashIndex := strings.LastIndex(containerIp, "/")
			if slashIndex > -1 {
				service.IP = containerIp[:slashIndex]
			} else {
				service.IP = containerIp
			}
			log.Debug("using container IP " + service.IP + " from label '" +
				b.config.UseIpFromLabel + "'")
		} else {
			log.Debug("Label '" + b.config.UseIpFromLabel +
				"' not found in container configuration")
		}
	}

	// NetworkMode can point to another container (kuberenetes pods)
	networkMode := container.HostConfig.NetworkMode
	if networkMode != "" {
		if strings.HasPrefix(networkMode, "container:") {
			networkContainerId := strings.Split(networkMode, ":")[1]
			log.Debug(service.Name + ": detected container NetworkMode, linked to: " + networkContainerId[:12])
			networkContainer, err := b.docker.InspectContainer(networkContainerId)
			if err != nil {
				log.Error("unable to inspect network container:", networkContainerId[:12], err)
			} else {
				service.IP = networkContainer.NetworkSettings.IPAddress
				log.Debug(service.Name + ": using network container IP " + service.IP)
			}
		}
	}

	if port.PortType == "udp" {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags, "udp")
		service.ID = service.ID + ":udp"
	} else {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags)
	}

	// Look for ECS labels and add them to metadata if present
	for k, v := range container.Config.Labels {
		if strings.Contains(k, "com.amazonaws.ecs") {
			metadata[k] = v
		}
	}

	id := mapDefault(metadata, "id", "")
	if id != "" {
		service.ID = id
	}

	delete(metadata, "id")
	delete(metadata, "tags")
	delete(metadata, "name")
	service.Attrs = metadata
	service.TTL = b.config.RefreshTtl

	return service
}

func (b *Bridge) remove(containerId string, deregister bool) {
	log.Debugf("container stop detected for: %v", containerId)
	b.Lock()
	defer b.Unlock()

	if deregister {
		deregisterAll := func(services []*Service) {
			for _, service := range services {
				err := b.registry.Deregister(service)
				if err != nil {
					log.Error("deregister failed:", service.ID, err)
					continue
				}
				log.Debug("removed:", containerId[:12], service.ID)
			}
		}
		deregisterAll(b.services[containerId])
		if d := b.deadContainers[containerId]; d != nil {
			deregisterAll(d.Services)
			delete(b.deadContainers, containerId)
		}
	} else if b.config.RefreshTtl != 0 && b.services[containerId] != nil {
		// need to stop the refreshing, but can't delete it yet
		b.deadContainers[containerId] = &DeadContainer{b.config.RefreshTtl, b.services[containerId]}
	}
	delete(b.services, containerId)
}

// bit set on ExitCode if it represents an exit via a signal
var dockerSignaledBit = 128

func (b *Bridge) shouldRemove(containerId string) bool {
	if b.config.DeregisterCheck == "always" {
		return true
	}
	container, err := b.docker.InspectContainer(containerId)
	if _, ok := err.(*dockerapi.NoSuchContainer); ok {
		// the container has already been removed from Docker
		// e.g. probabably run with "--rm" to remove immediately
		// so its exit code is not accessible
		log.Debugf("registrator: container %v was removed, could not fetch exit code", containerId[:12])
		return true
	}

	switch {
	case err != nil:
		log.Errorf("registrator: error fetching status for container %v on \"die\" event: %v\n", containerId[:12], err)
		return false
	case container.State.Running:
		log.Debugf("registrator: not removing container %v, still running", containerId[:12])
		return false
	case container.State.ExitCode == 0:
		return true
	case container.State.ExitCode&dockerSignaledBit == dockerSignaledBit:
		return true
	}
	return false
}

var Hostname string

func init() {
	// It's ok for Hostname to ultimately be an empty string
	// An empty string will fall back to trying to make a best guess
	Hostname, _ = os.Hostname()
}
<commit_msg>Another log addition to bridge<commit_after>package bridge

import (
	"errors"
	"net"
	"net/url"
	"os"
	"path"
	"regexp"
	"strconv"
	"strings"
	"sync"

	dockerapi "github.com/fsouza/go-dockerclient"
)

var serviceIDPattern = regexp.MustCompile(`^(.+?):([a-zA-Z0-9][a-zA-Z0-9_.-]+):[0-9]+(?::udp)?$`)

type Bridge struct {
	sync.Mutex
	registry       RegistryAdapter
	docker         *dockerapi.Client
	services       map[string][]*Service
	deadContainers map[string]*DeadContainer
	config         Config
}

func New(docker *dockerapi.Client, adapterUri string, config Config) (*Bridge, error) {
	uri, err := url.Parse(adapterUri)
	if err != nil {
		return nil, errors.New("bad adapter uri: " + adapterUri)
	}
	factory, found := AdapterFactories.Lookup(uri.Scheme)
	if !found {
		return nil, errors.New("unrecognized adapter: " + adapterUri)
	}

	log.Debug("Using", uri.Scheme, "adapter:", uri)
	return &Bridge{
		docker:         docker,
		config:         config,
		registry:       factory.New(uri),
		services:       make(map[string][]*Service),
		deadContainers: make(map[string]*DeadContainer),
	}, nil
}

func (b *Bridge) Ping() error {
	return b.registry.Ping()
}

func (b *Bridge) Add(containerId string) {
	b.add(containerId, false, "")
}

func (b *Bridge) Remove(containerId string) {
	b.remove(containerId, true)
}

func (b *Bridge) RemoveOnExit(containerId string) {
	b.remove(containerId, b.shouldRemove(containerId))
}

func (b *Bridge) Refresh() {
	for containerId, services := range b.getServicesCopy() {
		for _, service := range services {
			err := b.registry.Refresh(service)
			if err != nil {
				log.Error("refresh failed:", service.ID, err)
				continue
			}
			log.Debug("refreshed:", containerId[:12], service.ID)
		}
	}
}

func (b *Bridge) PruneDeadContainers() {
	b.Lock()
	defer b.Unlock()
	for containerId, deadContainer := range b.deadContainers {
		deadContainer.TTL -= b.config.RefreshInterval
		if deadContainer.TTL <= 0 {
			delete(b.deadContainers, containerId)
		}
	}
}

// Get a deep copy of the current services in a thread-safe way
func (b *Bridge) getServicesCopy() map[string][]*Service {
	b.Lock()
	defer b.Unlock()
	svcsCopy := make(map[string][]*Service)

	for id, svc := range b.services {
		var svcPointers []*Service
		for _, s := range svc {
			t := *s
			z := &t
			svcPointers = append(svcPointers, z)
		}
		svcsCopy[id] = svcPointers
	}
	return svcsCopy
}

func (b *Bridge) Sync(quiet bool) {
	serviceSync(b, quiet, "")
}

func (b *Bridge) AllocateNewIPToServices(ip string) {
	serviceSync(b, true, ip)
}

func (b *Bridge) deleteDeadContainer(containerId string) {
	b.Lock()
	defer b.Unlock()

	if d := b.deadContainers[containerId]; d != nil {
		b.services[containerId] = d.Services
		delete(b.deadContainers, containerId)
	}
}

func (b *Bridge) appendService(containerId string, service *Service) {
	b.Lock()
	defer b.Unlock()
	if b.services[containerId] != nil {
		log.Debug("container, ", containerId[:12], ", already exists, will not append.")
		return
	}
	b.services[containerId] = append(b.services[containerId], service)
	log.Debug("added:", containerId[:12], service.ID)
}

func (b *Bridge) add(containerId string, quiet bool, newIP string) {
	b.deleteDeadContainer(containerId)

	b.Lock()
	if b.services[containerId] != nil {
		log.Debug("container, ", containerId[:12], ", already exists, ignoring")
		// Alternatively, remove and readd or resubmit.
		return
	}
	b.Unlock()

	container, err := b.docker.InspectContainer(containerId)
	if err != nil {
		log.Error("unable to inspect container:", containerId[:12], err)
		return
	}

	ports := make(map[string]ServicePort)

	// Extract configured host port mappings, relevant when using --net=host
	for port, _ := range container.Config.ExposedPorts {
		published := []dockerapi.PortBinding{{"0.0.0.0", port.Port()}}
		serviceP := servicePort(container, port, published)
		if newIP != "" {
			serviceP.HostIP = newIP
		}
		ports[string(port)] = serviceP
	}

	// Extract runtime port mappings, relevant when using --net=bridge
	for port, published := range container.NetworkSettings.Ports {
		ports[string(port)] = servicePort(container, port, published)
	}

	if len(ports) == 0 && !quiet {
		log.Debug("ignored:", container.ID[:12], "no published ports")
		return
	}

	servicePorts := make(map[string]ServicePort)
	for key, port := range ports {
		if b.config.Internal != true && port.HostPort == "" {
			if !quiet {
				log.Debug("ignored:", container.ID[:12], "port", port.ExposedPort, "not published on host")
			}
			continue
		}
		servicePorts[key] = port
	}

	isGroup := len(servicePorts) > 1
	for _, port := range servicePorts {
		log.Infof("Creating new service with IP: %s", port.HostIP)
		service := b.newService(port, isGroup)
		if service == nil {
			if !quiet {
				log.Debug("ignored:", container.ID[:12], "service on port", port.ExposedPort)
			}
			continue
		}
		err := b.registry.Register(service)
		if err != nil {
			log.Error("register failed:", service, err)
			continue
		}
		b.appendService(container.ID, service)
	}
}

func (b *Bridge) newService(port ServicePort, isgroup bool) *Service {
	container := port.container
	defaultName := strings.Split(path.Base(container.Config.Image), ":")[0]
	// not sure about this logic. kind of want to remove it.
	hostname := Hostname
	if hostname == "" {
		hostname = port.HostIP
	}
	if port.HostIP == "0.0.0.0" {
		ip, err := net.ResolveIPAddr("ip", hostname)
		if err == nil {
			port.HostIP = ip.String()
		}
	}

	log.Infof("Creation of new service. Config HostIp is %s", b.config.HostIp)
	if b.config.HostIp != "" {
		port.HostIP = b.config.HostIp
	}

	metadata, metadataFromPort := serviceMetaData(container.Config, port.ExposedPort)

	ignore := mapDefault(metadata, "ignore", "")
	if ignore != "" {
		return nil
	}

	if b.config.RequireLabel {
		if strings.ToLower(mapDefault(metadata, "register", "false")) == "false" {
			log.Debugf("Did not find label SERVICE_REGISTER on %s - ignoring.", container.Name)
			return nil
		}
	}

	service := new(Service)
	service.Origin = port
	service.ID = hostname + ":" + container.Name[1:] + ":" + port.ExposedPort
	service.Name = mapDefault(metadata, "name", defaultName)
	if isgroup && !metadataFromPort["name"] {
		service.Name += "-" + port.ExposedPort
	}
	var convertedPort int

	log.Infof("New Service has config: Internal=%s", b.config.Internal)
	if b.config.Internal == true {
		service.IP = port.ExposedIP
		p, err := strconv.Atoi(port.ExposedPort)
		if err != nil {
			log.Error("Unable to parse string ExposedPort to int: %s", port.ExposedPort)
		} else {
			convertedPort = p
		}
	} else {
		service.IP = port.HostIP
		p, err := strconv.Atoi(port.HostPort)
		if err != nil {
			log.Error("Unable to parse string HostPort to int: %s", port.HostPort)
		} else {
			convertedPort = p
		}
	}
	service.Port = convertedPort

	if b.config.UseIpFromLabel != "" {
		containerIp := container.Config.Labels[b.config.UseIpFromLabel]
		if containerIp != "" {
			slashIndex := strings.LastIndex(containerIp, "/")
			if slashIndex > -1 {
				service.IP = containerIp[:slashIndex]
			} else {
				service.IP = containerIp
			}
			log.Debug("using container IP " + service.IP + " from label '" +
				b.config.UseIpFromLabel + "'")
		} else {
			log.Debug("Label '" + b.config.UseIpFromLabel +
				"' not found in container configuration")
		}
	}

	// NetworkMode can point to another container (kuberenetes pods)
	networkMode := container.HostConfig.NetworkMode
	if networkMode != "" {
		if strings.HasPrefix(networkMode, "container:") {
			networkContainerId := strings.Split(networkMode, ":")[1]
			log.Debug(service.Name + ": detected container NetworkMode, linked to: " + networkContainerId[:12])
			networkContainer, err := b.docker.InspectContainer(networkContainerId)
			if err != nil {
				log.Error("unable to inspect network container:", networkContainerId[:12], err)
			} else {
				service.IP = networkContainer.NetworkSettings.IPAddress
				log.Debug(service.Name + ": using network container IP " + service.IP)
			}
		}
	}

	if port.PortType == "udp" {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags, "udp")
		service.ID = service.ID + ":udp"
	} else {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags)
	}

	// Look for ECS labels and add them to metadata if present
	for k, v := range container.Config.Labels {
		if strings.Contains(k, "com.amazonaws.ecs") {
			metadata[k] = v
		}
	}

	id := mapDefault(metadata, "id", "")
	if id != "" {
		service.ID = id
	}

	delete(metadata, "id")
	delete(metadata, "tags")
	delete(metadata, "name")
	service.Attrs = metadata
	service.TTL = b.config.RefreshTtl

	return service
}

func (b *Bridge) remove(containerId string, deregister bool) {
	log.Debugf("container stop detected for: %v", containerId)
	b.Lock()
	defer b.Unlock()

	if deregister {
		deregisterAll := func(services []*Service) {
			for _, service := range services {
				err := b.registry.Deregister(service)
				if err != nil {
					log.Error("deregister failed:", service.ID, err)
					continue
				}
				log.Debug("removed:", containerId[:12], service.ID)
			}
		}
		deregisterAll(b.services[containerId])
		if d := b.deadContainers[containerId]; d != nil {
			deregisterAll(d.Services)
			delete(b.deadContainers, containerId)
		}
	} else if b.config.RefreshTtl != 0 && b.services[containerId] != nil {
		// need to stop the refreshing, but can't delete it yet
		b.deadContainers[containerId] = &DeadContainer{b.config.RefreshTtl, b.services[containerId]}
	}
	delete(b.services, containerId)
}

// bit set on ExitCode if it represents an exit via a signal
var dockerSignaledBit = 128

func (b *Bridge) shouldRemove(containerId string) bool {
	if b.config.DeregisterCheck == "always" {
		return true
	}
	container, err := b.docker.InspectContainer(containerId)
	if _, ok := err.(*dockerapi.NoSuchContainer); ok {
		// the container has already been removed from Docker
		// e.g. probabably run with "--rm" to remove immediately
		// so its exit code is not accessible
		log.Debugf("registrator: container %v was removed, could not fetch exit code", containerId[:12])
		return true
	}

	switch {
	case err != nil:
		log.Errorf("registrator: error fetching status for container %v on \"die\" event: %v\n", containerId[:12], err)
		return false
	case container.State.Running:
		log.Debugf("registrator: not removing container %v, still running", containerId[:12])
		return false
	case container.State.ExitCode == 0:
		return true
	case container.State.ExitCode&dockerSignaledBit == dockerSignaledBit:
		return true
	}
	return false
}

var Hostname string

func init() {
	// It's ok for Hostname to ultimately be an empty string
	// An empty string will fall back to trying to make a best guess
	Hostname, _ = os.Hostname()
}
<|endoftext|>


<commit_before>package broker

import (
	"encoding/json"
	"os"
	"sync"
	"time"

	"github.com/asim/mq/go/client"
)

var (
	Default Broker = newBroker()
)

// internal broker
type broker struct {
	options *Options

	sync.RWMutex
	topics map[string][]chan []byte

	mtx       sync.RWMutex
	persisted map[string]bool
}

// internal message for persistence
type message struct {
	Timestamp int64  `json:"timestamp"`
	Topic     string `json:"topic"`
	Payload   []byte `json:"payload"`
}

// Broker is the message broker
type Broker interface {
	Publish(topic string, payload []byte) error
	Subscribe(topic string) (<-chan []byte, error)
	Unsubscribe(topic string, sub <-chan []byte) error
}

func newBroker(opts ...Option) *broker {
	options := &Options{
		Client: client.New(),
	}

	for _, o := range opts {
		o(options)
	}

	return &broker{
		options:   options,
		topics:    make(map[string][]chan []byte),
		persisted: make(map[string]bool),
	}
}

func (b *broker) persist(topic string) error {
	b.mtx.Lock()
	defer b.mtx.Unlock()

	if b.persisted[topic] {
		return nil
	}

	ch, err := b.Subscribe(topic)
	if err != nil {
		return err
	}

	f, err := os.OpenFile(topic+".mq", os.O_CREATE|os.O_RDWR|os.O_APPEND, 0660)
	if err != nil {
		return err
	}

	go func() {
		for p := range ch {
			b, err := json.Marshal(&message{
				Timestamp: time.Now().UnixNano(),
				Topic:     topic,
				Payload:   p,
			})
			if err != nil {
				continue
			}
			f.Write(b)
			f.Write([]byte{'\n'})
		}
	}()

	b.persisted[topic] = true

	return nil
}

func (b *broker) Publish(topic string, payload []byte) error {
	if b.options.Proxy {
		return b.options.Client.Publish(topic, payload)
	}

	b.RLock()
	subscribers, ok := b.topics[topic]
	b.RUnlock()
	if !ok {
		// persist?
		if !b.options.Persist {
			return nil
		}
		if err := b.persist(topic); err != nil {
			return err
		}
	}

	go func() {
		for _, subscriber := range subscribers {
			select {
			case subscriber <- payload:
			default:
			}
		}
	}()

	return nil
}

func (b *broker) Subscribe(topic string) (<-chan []byte, error) {
	if b.options.Proxy {
		return b.options.Client.Subscribe(topic)
	}

	ch := make(chan []byte, 100)
	b.Lock()
	b.topics[topic] = append(b.topics[topic], ch)
	b.Unlock()
	return ch, nil
}

func (b *broker) Unsubscribe(topic string, sub <-chan []byte) error {
	if b.options.Proxy {
		return b.options.Client.Unsubscribe(sub)
	}

	b.RLock()
	subscribers, ok := b.topics[topic]
	b.RUnlock()

	if !ok {
		return nil
	}

	var subs []chan []byte
	for _, subscriber := range subscribers {
		if subscriber == sub {
			continue
		}
		subs = append(subs, subscriber)
	}

	b.Lock()
	b.topics[topic] = subs
	b.Unlock()
	return nil
}

func Publish(topic string, payload []byte) error {
	return Default.Publish(topic, payload)
}

func Subscribe(topic string) (<-chan []byte, error) {
	return Default.Subscribe(topic)
}

func Unsubscribe(topic string, sub <-chan []byte) error {
	return Default.Unsubscribe(topic, sub)
}

func New(opts ...Option) *broker {
	return newBroker(opts...)
}
<commit_msg>write to file in batches<commit_after>package broker

import (
	"encoding/json"
	"os"
	"sync"
	"time"

	"github.com/asim/mq/go/client"
)

var (
	Default Broker = newBroker()
)

// internal broker
type broker struct {
	options *Options

	sync.RWMutex
	topics map[string][]chan []byte

	mtx       sync.RWMutex
	persisted map[string]bool
}

// internal message for persistence
type message struct {
	Timestamp int64  `json:"timestamp"`
	Topic     string `json:"topic"`
	Payload   []byte `json:"payload"`
}

// Broker is the message broker
type Broker interface {
	Publish(topic string, payload []byte) error
	Subscribe(topic string) (<-chan []byte, error)
	Unsubscribe(topic string, sub <-chan []byte) error
}

func newBroker(opts ...Option) *broker {
	options := &Options{
		Client: client.New(),
	}

	for _, o := range opts {
		o(options)
	}

	return &broker{
		options:   options,
		topics:    make(map[string][]chan []byte),
		persisted: make(map[string]bool),
	}
}

func (b *broker) persist(topic string) error {
	b.mtx.Lock()
	defer b.mtx.Unlock()

	if b.persisted[topic] {
		return nil
	}

	ch, err := b.Subscribe(topic)
	if err != nil {
		return err
	}

	f, err := os.OpenFile(topic+".mq", os.O_CREATE|os.O_RDWR|os.O_APPEND, 0660)
	if err != nil {
		return err
	}

	go func() {
		var pending []byte
		newline := []byte{'\n'}
		t := time.NewTicker(time.Second)
		defer t.Stop()

		for {
			select {
			case p := <-ch:
				b, err := json.Marshal(&message{
					Timestamp: time.Now().UnixNano(),
					Topic:     topic,
					Payload:   p,
				})
				if err != nil {
					continue
				}
				pending = append(pending, b...)
				pending = append(pending, newline...)
			case <-t.C:
				if len(pending) == 0 {
					continue
				}
				f.Write(pending)
				pending = nil
			}
		}
	}()

	b.persisted[topic] = true

	return nil
}

func (b *broker) Publish(topic string, payload []byte) error {
	if b.options.Proxy {
		return b.options.Client.Publish(topic, payload)
	}

	b.RLock()
	subscribers, ok := b.topics[topic]
	b.RUnlock()
	if !ok {
		// persist?
		if !b.options.Persist {
			return nil
		}
		if err := b.persist(topic); err != nil {
			return err
		}
	}

	go func() {
		for _, subscriber := range subscribers {
			select {
			case subscriber <- payload:
			default:
			}
		}
	}()

	return nil
}

func (b *broker) Subscribe(topic string) (<-chan []byte, error) {
	if b.options.Proxy {
		return b.options.Client.Subscribe(topic)
	}

	ch := make(chan []byte, 100)
	b.Lock()
	b.topics[topic] = append(b.topics[topic], ch)
	b.Unlock()
	return ch, nil
}

func (b *broker) Unsubscribe(topic string, sub <-chan []byte) error {
	if b.options.Proxy {
		return b.options.Client.Unsubscribe(sub)
	}

	b.RLock()
	subscribers, ok := b.topics[topic]
	b.RUnlock()

	if !ok {
		return nil
	}

	var subs []chan []byte
	for _, subscriber := range subscribers {
		if subscriber == sub {
			continue
		}
		subs = append(subs, subscriber)
	}

	b.Lock()
	b.topics[topic] = subs
	b.Unlock()
	return nil
}

func Publish(topic string, payload []byte) error {
	return Default.Publish(topic, payload)
}

func Subscribe(topic string) (<-chan []byte, error) {
	return Default.Subscribe(topic)
}

func Unsubscribe(topic string, sub <-chan []byte) error {
	return Default.Unsubscribe(topic, sub)
}

func New(opts ...Option) *broker {
	return newBroker(opts...)
}
<|endoftext|>