<commit_before>package collectors

import (
	"fmt"
	"strings"
	"time"

	"bosun.org/cmd/scollector/conf"
	"bosun.org/metadata"
	"bosun.org/opentsdb"
)

const (
	ifAlias              = ".1.3.6.1.2.1.31.1.1.1.18"
	ifDescr              = ".1.3.6.1.2.1.2.2.1.2"
	ifType               = ".1.3.6.1.2.1.2.2.1.3"
	ifMTU                = ".1.3.6.1.2.1.2.2.1.4"
	ifHighSpeed          = ".1.3.6.1.2.1.31.1.1.1.15"
	ifAdminStatus        = ".1.3.6.1.2.1.2.2.1.7"
	ifOperStatus         = ".1.3.6.1.2.1.2.2.1.8"
	ifHCInBroadcastPkts  = ".1.3.6.1.2.1.31.1.1.1.9"
	ifHCInMulticastPkts  = ".1.3.6.1.2.1.31.1.1.1.8"
	ifHCInUcastPkts      = ".1.3.6.1.2.1.31.1.1.1.7"
	ifHCOutBroadcastPkts = ".1.3.6.1.2.1.31.1.1.1.13"
	ifHCOutMulticastPkts = ".1.3.6.1.2.1.31.1.1.1.12"
	ifHCOutOctets        = ".1.3.6.1.2.1.31.1.1.1.10"
	ifHCOutUcastPkts     = ".1.3.6.1.2.1.31.1.1.1.11"
	ifHCinOctets         = ".1.3.6.1.2.1.31.1.1.1.6"
	ifInDiscards         = ".1.3.6.1.2.1.2.2.1.13"
	ifInErrors           = ".1.3.6.1.2.1.2.2.1.14"
	ifInPauseFrames      = ".1.3.6.1.2.1.10.7.10.1.3"
	ifName               = ".1.3.6.1.2.1.31.1.1.1.1"
	ifOutDiscards        = ".1.3.6.1.2.1.2.2.1.19"
	ifOutErrors          = ".1.3.6.1.2.1.2.2.1.20"
	ifOutPauseFrames     = ".1.3.6.1.2.1.10.7.10.1.4"
)

// SNMPIfaces registers a SNMP Interfaces collector for the given community and host.
func SNMPIfaces(cfg conf.SNMP) {
	collectors = append(collectors, &IntervalCollector{
		F: func() (opentsdb.MultiDataPoint, error) {
			return c_snmp_ifaces(cfg.Community, cfg.Host)
		},
		Interval: time.Second * 30,
		name:     fmt.Sprintf("snmp-ifaces-%s", cfg.Host),
	})
}

const osNet = "os.net"

func switchInterfaceMetric(metric string, iname string, ifType int64) string {
	switch ifType {
	case 6:
		return metric
	case 53, 161:
		return osNet + ".bond" + strings.TrimPrefix(metric, osNet)
	case 135:
		return osNet + ".virtual" + strings.TrimPrefix(metric, osNet)
	case 131:
		return osNet + ".tunnel" + strings.TrimPrefix(metric, osNet)
	default:
		//Cisco ASAs don't mark port channels correctly
		if strings.Contains(iname, "port-channel") {
			return osNet + ".bond" + strings.TrimPrefix(metric, osNet)
		}
		return osNet + ".other" + strings.TrimPrefix(metric, osNet)
	}
}

func c_snmp_ifaces(community, host string) (opentsdb.MultiDataPoint, error) {
	ifNamesRaw, err := snmp_subtree(host, community, ifName)
	if err != nil || len(ifNamesRaw) == 0 {
		ifNamesRaw, err = snmp_subtree(host, community, ifDescr)
		if err != nil {
			return nil, err
		}
	}
	ifAliasesRaw, err := snmp_subtree(host, community, ifAlias)
	if err != nil {
		return nil, err
	}
	ifTypesRaw, err := snmp_subtree(host, community, ifType)
	if err != nil {
		return nil, err
	}
	ifNames := make(map[interface{}]string, len(ifNamesRaw))
	ifAliases := make(map[interface{}]string, len(ifAliasesRaw))
	ifTypes := make(map[interface{}]int64, len(ifTypesRaw))
	for k, v := range ifNamesRaw {
		ifNames[k] = fmt.Sprintf("%s", v)
	}
	for k, v := range ifTypesRaw {
		val, ok := v.(int64)
		if !ok {
			return nil, fmt.Errorf("unexpected type from from MIB::ifType")
		}
		ifTypes[k] = val
	}
	for k, v := range ifAliasesRaw {
		// In case clean would come up empty, prevent the point from being removed
		// by setting our own empty case.
		ifAliases[k], _ = opentsdb.Clean(fmt.Sprintf("%s", v))
		if ifAliases[k] == "" {
			ifAliases[k] = "NA"
		}
	}
	var md opentsdb.MultiDataPoint
	add := func(sA snmpAdd) error {
		m, err := snmp_subtree(host, community, sA.oid)
		if err != nil {
			return err
		}
		var sum int64
		for k, v := range m {
			tags := opentsdb.TagSet{
				"host":  host,
				"iface": fmt.Sprintf("%s", k),
				"iname": ifNames[k],
			}
			if sA.dir != "" {
				tags["direction"] = sA.dir
			}
			if iVal, ok := v.(int64); ok && ifTypes[k] == 6 {
				sum += iVal
			}
			Add(&md, switchInterfaceMetric(sA.metric, ifNames[k], ifTypes[k]), v, tags, sA.rate, sA.unit, sA.desc)
			metadata.AddMeta("", tags, "alias", ifAliases[k], false)
		}
		if sA.metric == osNetBytes {
			tags := opentsdb.TagSet{"host": host, "direction": sA.dir}
			Add(&md, osNetBytes+".total", sum, tags, metadata.Counter, metadata.Bytes, "The total number of bytes transfered through the network device.")
		}
		return nil
	}
	oids := []snmpAdd{
		{ifHCInBroadcastPkts, osNetBroadcast, "in", metadata.Counter, metadata.Packet, osNetBroadcastDesc},
		{ifHCInMulticastPkts, osNetMulticast, "in", metadata.Counter, metadata.Packet, osNetMulticastDesc},
		{ifHCInUcastPkts, osNetUnicast, "in", metadata.Counter, metadata.Packet, osNetUnicastDesc},
		{ifHCOutBroadcastPkts, osNetBroadcast, "out", metadata.Counter, metadata.Packet, osNetBroadcastDesc},
		{ifHCOutMulticastPkts, osNetMulticast, "out", metadata.Counter, metadata.Packet, osNetMulticastDesc},
		{ifHCOutOctets, osNetBytes, "out", metadata.Counter, metadata.Bytes, osNetBytesDesc},
		{ifHCOutUcastPkts, osNetUnicast, "out", metadata.Counter, metadata.Packet, osNetUnicastDesc},
		{ifHCinOctets, osNetBytes, "in", metadata.Counter, metadata.Bytes, osNetBytesDesc},
		{ifInDiscards, osNetDropped, "in", metadata.Counter, metadata.Packet, osNetDroppedDesc},
		{ifInErrors, osNetErrors, "in", metadata.Counter, metadata.Error, osNetErrorsDesc},
		{ifOutDiscards, osNetDropped, "out", metadata.Counter, metadata.Packet, osNetDroppedDesc},
		{ifOutErrors, osNetErrors, "out", metadata.Counter, metadata.Error, osNetErrorsDesc},
		{ifInPauseFrames, osNetPauseFrames, "in", metadata.Counter, metadata.Frame, osNetPauseFrameDesc},
		{ifOutPauseFrames, osNetPauseFrames, "out", metadata.Counter, metadata.Frame, osNetPauseFrameDesc},
		{ifMTU, osNetMTU, "", metadata.Gauge, metadata.Bytes, osNetMTUDesc},
		{ifHighSpeed, osNetIfSpeed, "", metadata.Gauge, metadata.Megabit, osNetIfSpeedDesc},
		{ifAdminStatus, osNetAdminStatus, "", metadata.Gauge, metadata.StatusCode, osNetAdminStatusDesc},
		{ifOperStatus, osNetOperStatus, "", metadata.Gauge, metadata.StatusCode, osNetOperStatusDesc},
	}
	for _, sA := range oids {
		if err := add(sA); err != nil {
			return nil, err
		}
	}
	return md, nil
}

type snmpAdd struct {
	oid    string
	metric string
	dir    string
	rate   metadata.RateType
	unit   metadata.Unit
	desc   string
}
<commit_msg>cmd/scollector: Get physical addresses of snmp ifaces<commit_after>package collectors

import (
	"fmt"
	"strings"
	"time"

	"bosun.org/cmd/scollector/conf"
	"bosun.org/metadata"
	"bosun.org/opentsdb"
)

const (
	ifAlias              = ".1.3.6.1.2.1.31.1.1.1.18"
	ifDescr              = ".1.3.6.1.2.1.2.2.1.2"
	ifType               = ".1.3.6.1.2.1.2.2.1.3"
	ifMTU                = ".1.3.6.1.2.1.2.2.1.4"
	ifPhysAddress        = ".1.3.6.1.2.1.2.2.1.6"
	ifHighSpeed          = ".1.3.6.1.2.1.31.1.1.1.15"
	ifAdminStatus        = ".1.3.6.1.2.1.2.2.1.7"
	ifOperStatus         = ".1.3.6.1.2.1.2.2.1.8"
	ifHCInBroadcastPkts  = ".1.3.6.1.2.1.31.1.1.1.9"
	ifHCInMulticastPkts  = ".1.3.6.1.2.1.31.1.1.1.8"
	ifHCInUcastPkts      = ".1.3.6.1.2.1.31.1.1.1.7"
	ifHCOutBroadcastPkts = ".1.3.6.1.2.1.31.1.1.1.13"
	ifHCOutMulticastPkts = ".1.3.6.1.2.1.31.1.1.1.12"
	ifHCOutOctets        = ".1.3.6.1.2.1.31.1.1.1.10"
	ifHCOutUcastPkts     = ".1.3.6.1.2.1.31.1.1.1.11"
	ifHCinOctets         = ".1.3.6.1.2.1.31.1.1.1.6"
	ifInDiscards         = ".1.3.6.1.2.1.2.2.1.13"
	ifInErrors           = ".1.3.6.1.2.1.2.2.1.14"
	ifInPauseFrames      = ".1.3.6.1.2.1.10.7.10.1.3"
	ifName               = ".1.3.6.1.2.1.31.1.1.1.1"
	ifOutDiscards        = ".1.3.6.1.2.1.2.2.1.19"
	ifOutErrors          = ".1.3.6.1.2.1.2.2.1.20"
	ifOutPauseFrames     = ".1.3.6.1.2.1.10.7.10.1.4"
)

// SNMPIfaces registers a SNMP Interfaces collector for the given community and host.
func SNMPIfaces(cfg conf.SNMP) {
	collectors = append(collectors, &IntervalCollector{
		F: func() (opentsdb.MultiDataPoint, error) {
			return c_snmp_ifaces(cfg.Community, cfg.Host)
		},
		Interval: time.Second * 30,
		name:     fmt.Sprintf("snmp-ifaces-%s", cfg.Host),
	})
}

const osNet = "os.net"

func switchInterfaceMetric(metric string, iname string, ifType int64) string {
	switch ifType {
	case 6:
		return metric
	case 53, 161:
		return osNet + ".bond" + strings.TrimPrefix(metric, osNet)
	case 135:
		return osNet + ".virtual" + strings.TrimPrefix(metric, osNet)
	case 131:
		return osNet + ".tunnel" + strings.TrimPrefix(metric, osNet)
	default:
		//Cisco ASAs don't mark port channels correctly
		if strings.Contains(iname, "port-channel") {
			return osNet + ".bond" + strings.TrimPrefix(metric, osNet)
		}
		return osNet + ".other" + strings.TrimPrefix(metric, osNet)
	}
}

func c_snmp_ifaces(community, host string) (opentsdb.MultiDataPoint, error) {
	ifNamesRaw, err := snmp_subtree(host, community, ifName)
	if err != nil || len(ifNamesRaw) == 0 {
		ifNamesRaw, err = snmp_subtree(host, community, ifDescr)
		if err != nil {
			return nil, err
		}
	}
	ifAliasesRaw, err := snmp_subtree(host, community, ifAlias)
	if err != nil {
		return nil, err
	}
	ifTypesRaw, err := snmp_subtree(host, community, ifType)
	if err != nil {
		return nil, err
	}
	ifPhysAddressRaw, err := snmp_subtree(host, community, ifPhysAddress)
	if err != nil {
		return nil, err
	}
	ifNames := make(map[interface{}]string, len(ifNamesRaw))
	ifAliases := make(map[interface{}]string, len(ifAliasesRaw))
	ifTypes := make(map[interface{}]int64, len(ifTypesRaw))
	ifPhysAddresses := make(map[interface{}]string, len(ifPhysAddressRaw))
	for k, v := range ifNamesRaw {
		ifNames[k] = fmt.Sprintf("%s", v)
	}
	for k, v := range ifTypesRaw {
		val, ok := v.(int64)
		if !ok {
			return nil, fmt.Errorf("unexpected type from from MIB::ifType")
		}
		ifTypes[k] = val
	}
	for k, v := range ifAliasesRaw {
		// In case clean would come up empty, prevent the point from being removed
		// by setting our own empty case.
		ifAliases[k], _ = opentsdb.Clean(fmt.Sprintf("%s", v))
		if ifAliases[k] == "" {
			ifAliases[k] = "NA"
		}
	}
	for k, v := range ifPhysAddressRaw {
		ifPhysAddresses[k] = fmt.Sprintf("%X", v)
	}
	var md opentsdb.MultiDataPoint
	add := func(sA snmpAdd) error {
		m, err := snmp_subtree(host, community, sA.oid)
		if err != nil {
			return err
		}
		var sum int64
		for k, v := range m {
			tags := opentsdb.TagSet{
				"host":  host,
				"iface": fmt.Sprintf("%s", k),
				"iname": ifNames[k],
			}
			if sA.dir != "" {
				tags["direction"] = sA.dir
			}
			if iVal, ok := v.(int64); ok && ifTypes[k] == 6 {
				sum += iVal
			}
			Add(&md, switchInterfaceMetric(sA.metric, ifNames[k], ifTypes[k]), v, tags, sA.rate, sA.unit, sA.desc)
			metadata.AddMeta("", tags, "alias", ifAliases[k], false)
			metadata.AddMeta("", tags, "mac", ifPhysAddresses[k], false)
		}
		if sA.metric == osNetBytes {
			tags := opentsdb.TagSet{"host": host, "direction": sA.dir}
			Add(&md, osNetBytes+".total", sum, tags, metadata.Counter, metadata.Bytes, "The total number of bytes transfered through the network device.")
		}
		return nil
	}
	oids := []snmpAdd{
		{ifHCInBroadcastPkts, osNetBroadcast, "in", metadata.Counter, metadata.Packet, osNetBroadcastDesc},
		{ifHCInMulticastPkts, osNetMulticast, "in", metadata.Counter, metadata.Packet, osNetMulticastDesc},
		{ifHCInUcastPkts, osNetUnicast, "in", metadata.Counter, metadata.Packet, osNetUnicastDesc},
		{ifHCOutBroadcastPkts, osNetBroadcast, "out", metadata.Counter, metadata.Packet, osNetBroadcastDesc},
		{ifHCOutMulticastPkts, osNetMulticast, "out", metadata.Counter, metadata.Packet, osNetMulticastDesc},
		{ifHCOutOctets, osNetBytes, "out", metadata.Counter, metadata.Bytes, osNetBytesDesc},
		{ifHCOutUcastPkts, osNetUnicast, "out", metadata.Counter, metadata.Packet, osNetUnicastDesc},
		{ifHCinOctets, osNetBytes, "in", metadata.Counter, metadata.Bytes, osNetBytesDesc},
		{ifInDiscards, osNetDropped, "in", metadata.Counter, metadata.Packet, osNetDroppedDesc},
		{ifInErrors, osNetErrors, "in", metadata.Counter, metadata.Error, osNetErrorsDesc},
		{ifOutDiscards, osNetDropped, "out", metadata.Counter, metadata.Packet, osNetDroppedDesc},
		{ifOutErrors, osNetErrors, "out", metadata.Counter, metadata.Error, osNetErrorsDesc},
		{ifInPauseFrames, osNetPauseFrames, "in", metadata.Counter, metadata.Frame, osNetPauseFrameDesc},
		{ifOutPauseFrames, osNetPauseFrames, "out", metadata.Counter, metadata.Frame, osNetPauseFrameDesc},
		{ifMTU, osNetMTU, "", metadata.Gauge, metadata.Bytes, osNetMTUDesc},
		{ifHighSpeed, osNetIfSpeed, "", metadata.Gauge, metadata.Megabit, osNetIfSpeedDesc},
		{ifAdminStatus, osNetAdminStatus, "", metadata.Gauge, metadata.StatusCode, osNetAdminStatusDesc},
		{ifOperStatus, osNetOperStatus, "", metadata.Gauge, metadata.StatusCode, osNetOperStatusDesc},
	}
	for _, sA := range oids {
		if err := add(sA); err != nil {
			return nil, err
		}
	}
	return md, nil
}

type snmpAdd struct {
	oid    string
	metric string
	dir    string
	rate   metadata.RateType
	unit   metadata.Unit
	desc   string
}
<|endoftext|>


<commit_before>package model

import (
	"context"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/oinume/lekcije/server/config"
	"github.com/oinume/lekcije/server/util"
)

type TestHelper struct {
	dbURL string
	db    *gorm.DB
	//mCountryService *MCountryService
}

func NewTestHelper() *TestHelper {
	return &TestHelper{}
}

func (h *TestHelper) DB(t *testing.T) *gorm.DB {
	if h.db != nil {
		return h.db
	}
	config.MustProcessDefault()
	h.dbURL = ReplaceToTestDBURL(t, config.DefaultVars.DBURL())
	db, err := OpenDB(h.dbURL, 1, config.DefaultVars.DebugSQL)
	if err != nil {
		e := fmt.Errorf("OpenDB failed: %v", err)
		if t == nil {
			panic(e)
		} else {
			t.Fatal(e)
		}
	}
	h.db = db
	return db
}

func (h *TestHelper) GetDBName(dbURL string) string {
	if index := strings.LastIndex(dbURL, "/"); index != -1 {
		return dbURL[index+1:]
	}
	return ""
}

func (h *TestHelper) LoadAllTables(t *testing.T, db *gorm.DB) []string {
	type Table struct {
		Name string `gorm:"column:table_name"`
	}
	tables := []Table{}
	sql := "SELECT table_name FROM information_schema.tables WHERE table_schema = ?"
	if err := db.Raw(sql, h.GetDBName(h.dbURL)).Scan(&tables).Error; err != nil {
		e := fmt.Errorf("failed to select table names: err=%v", err)
		if t == nil {
			panic(e)
		} else {
			t.Fatal(e)
		}
	}

	tableNames := []string{}
	for _, t := range tables {
		if t.Name == "goose_db_version" {
			continue
		}
		tableNames = append(tableNames, t.Name)
	}
	return tableNames
}

func (h *TestHelper) TruncateAllTables(t *testing.T) {
	//fmt.Printf("TruncateAllTables() called!\n--- stack ---\n%+v\n", errors.NewInternalError().StackTrace())
	db := h.DB(t)
	tables := h.LoadAllTables(t, db)
	for _, table := range tables {
		if strings.HasPrefix(table, "m_") {
			continue
		}
		if err := db.Exec("TRUNCATE TABLE " + table).Error; err != nil {
			e := fmt.Errorf("failed to truncate table: table=%v, err=%v", t, err)
			if t == nil {
				panic(e)
			} else {
				t.Fatal(e)
			}
		}
	}
}

func (h *TestHelper) CreateUser(t *testing.T, name, email string) *User {
	db := h.DB(t)
	user, err := NewUserService(db).Create(name, email)
	if err != nil {
		panic(fmt.Sprintf("Failed to CreateUser(): err=%v", err))
	}
	return user
}

func (h *TestHelper) CreateRandomUser(t *testing.T) *User {
	name := util.RandomString(16)
	return h.CreateUser(t, name, name+"@example.com")
}

func (h *TestHelper) CreateUserGoogle(t *testing.T, googleID string, userID uint32) *UserGoogle {
	userGoogle := &UserGoogle{
		GoogleID: googleID,
		UserID:   userID,
	}
	if err := h.DB(t).Create(userGoogle).Error; err != nil {
		t.Fatal(fmt.Errorf("CreateUserGoogle failed: %v", err))
	}
	return userGoogle
}

func (h *TestHelper) CreateTeacher(t *testing.T, id uint32, name string) *Teacher {
	db := h.DB(t)
	teacher := &Teacher{
		ID:           id,
		Name:         name,
		Gender:       "female",
		Birthday:     time.Date(2001, 1, 1, 0, 0, 0, 0, time.UTC),
		LastLessonAt: time.Now().Add(-1 * 24 * time.Hour), // 1 day ago
	}
	if err := NewTeacherService(db).CreateOrUpdate(teacher); err != nil {
		t.Fatal(fmt.Sprintf("CreateTeacher failed: %v", err))
	}
	return teacher
}

func (h *TestHelper) CreateRandomTeacher(t *testing.T) *Teacher {
	return h.CreateTeacher(t, uint32(util.RandomInt(9999999)), util.RandomString(6))
}

func (h *TestHelper) LoadMCountries(t *testing.T) *MCountries {
	db := h.DB(t)
	// TODO: cache
	mCountries, err := NewMCountryService(db).LoadAll(context.Background())
	if err != nil {
		e := fmt.Errorf("MCountryService.LoadAll failed: %v", err)
		if t == nil {
			panic(e)
		} else {
			t.Fatal(e)
		}
	}
	return mCountries
}

func (h *TestHelper) CreateFollowingTeacher(t *testing.T, userID uint32, teacher *Teacher) *FollowingTeacher {
	now := time.Now()
	ft, err := NewFollowingTeacherService(h.DB(t)).FollowTeacher(userID, teacher, now)
	if err != nil {
		t.Fatal(fmt.Errorf("FollowTeacher failed: %v", err))
	}
	return ft
}
<commit_msg>Fix TRUNCATE TABLE<commit_after>package model

import (
	"context"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/oinume/lekcije/server/config"
	"github.com/oinume/lekcije/server/util"
)

type TestHelper struct {
	dbURL string
	db    *gorm.DB
	//mCountryService *MCountryService
}

func NewTestHelper() *TestHelper {
	return &TestHelper{}
}

func (h *TestHelper) DB(t *testing.T) *gorm.DB {
	if h.db != nil {
		return h.db
	}
	config.MustProcessDefault()
	h.dbURL = ReplaceToTestDBURL(t, config.DefaultVars.DBURL())
	db, err := OpenDB(h.dbURL, 1, config.DefaultVars.DebugSQL)
	if err != nil {
		e := fmt.Errorf("OpenDB failed: %v", err)
		if t == nil {
			panic(e)
		} else {
			t.Fatal(e)
		}
	}
	h.db = db
	return db
}

func (h *TestHelper) GetDBName(dbURL string) string {
	if index := strings.LastIndex(dbURL, "/"); index != -1 {
		return dbURL[index+1:]
	}
	return ""
}

func (h *TestHelper) LoadAllTables(t *testing.T, db *gorm.DB) []string {
	type Table struct {
		Name string `gorm:"column:table_name"`
	}
	tables := []Table{}
	sql := "SELECT table_name AS table_name FROM information_schema.tables WHERE table_schema = ?"
	if err := db.Raw(sql, h.GetDBName(h.dbURL)).Scan(&tables).Error; err != nil {
		e := fmt.Errorf("failed to select table names: err=%v", err)
		if t == nil {
			panic(e)
		} else {
			t.Fatal(e)
		}
	}

	tableNames := []string{}
	for _, t := range tables {
		if t.Name == "goose_db_version" {
			continue
		}
		tableNames = append(tableNames, t.Name)
	}
	return tableNames
}

func (h *TestHelper) TruncateAllTables(t *testing.T) {
	//fmt.Printf("TruncateAllTables() called!\n--- stack ---\n%+v\n", errors.NewInternalError().StackTrace())
	db := h.DB(t)
	tables := h.LoadAllTables(t, db)
	for _, table := range tables {
		if strings.HasPrefix(table, "m_") {
			continue
		}
		if err := db.Exec("TRUNCATE TABLE " + table).Error; err != nil {
			e := fmt.Errorf("failed to truncate table: table=%v, err=%v", table, err)
			if t == nil {
				panic(e)
			} else {
				t.Fatal(e)
			}
		}
	}
}

func (h *TestHelper) CreateUser(t *testing.T, name, email string) *User {
	db := h.DB(t)
	user, err := NewUserService(db).Create(name, email)
	if err != nil {
		panic(fmt.Sprintf("Failed to CreateUser(): err=%v", err))
	}
	return user
}

func (h *TestHelper) CreateRandomUser(t *testing.T) *User {
	name := util.RandomString(16)
	return h.CreateUser(t, name, name+"@example.com")
}

func (h *TestHelper) CreateUserGoogle(t *testing.T, googleID string, userID uint32) *UserGoogle {
	userGoogle := &UserGoogle{
		GoogleID: googleID,
		UserID:   userID,
	}
	if err := h.DB(t).Create(userGoogle).Error; err != nil {
		t.Fatal(fmt.Errorf("CreateUserGoogle failed: %v", err))
	}
	return userGoogle
}

func (h *TestHelper) CreateTeacher(t *testing.T, id uint32, name string) *Teacher {
	db := h.DB(t)
	teacher := &Teacher{
		ID:           id,
		Name:         name,
		Gender:       "female",
		Birthday:     time.Date(2001, 1, 1, 0, 0, 0, 0, time.UTC),
		LastLessonAt: time.Now().Add(-1 * 24 * time.Hour), // 1 day ago
	}
	if err := NewTeacherService(db).CreateOrUpdate(teacher); err != nil {
		t.Fatal(fmt.Sprintf("CreateTeacher failed: %v", err))
	}
	return teacher
}

func (h *TestHelper) CreateRandomTeacher(t *testing.T) *Teacher {
	return h.CreateTeacher(t, uint32(util.RandomInt(9999999)), util.RandomString(6))
}

func (h *TestHelper) LoadMCountries(t *testing.T) *MCountries {
	db := h.DB(t)
	// TODO: cache
	mCountries, err := NewMCountryService(db).LoadAll(context.Background())
	if err != nil {
		e := fmt.Errorf("MCountryService.LoadAll failed: %v", err)
		if t == nil {
			panic(e)
		} else {
			t.Fatal(e)
		}
	}
	return mCountries
}

func (h *TestHelper) CreateFollowingTeacher(t *testing.T, userID uint32, teacher *Teacher) *FollowingTeacher {
	now := time.Now()
	ft, err := NewFollowingTeacherService(h.DB(t)).FollowTeacher(userID, teacher, now)
	if err != nil {
		t.Fatal(fmt.Errorf("FollowTeacher failed: %v", err))
	}
	return ft
}
<|endoftext|>


<commit_before>package notifier

import (
	"bytes"
	"fmt"
	"os"
	"sort"
	"strings"
	"sync"
	"text/template"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/oinume/lekcije/server/config"
	"github.com/oinume/lekcije/server/errors"
	"github.com/oinume/lekcije/server/fetcher"
	"github.com/oinume/lekcije/server/logger"
	"github.com/oinume/lekcije/server/model"
	"github.com/sendgrid/sendgrid-go"
	"github.com/sendgrid/sendgrid-go/helpers/mail"
	"github.com/uber-go/zap"
)

type Notifier struct {
	db            *gorm.DB
	fetcher       *fetcher.TeacherLessonFetcher
	dryRun        bool
	lessonService *model.LessonService
	teachers      map[uint32]*model.Teacher
	fetchedLessons map[uint32][]*model.Lesson
}

func NewNotifier(db *gorm.DB, concurrency int, dryRun bool) *Notifier {
	if concurrency < 1 {
		concurrency = 1
	}
	return &Notifier{
		db:       db,
		fetcher:  fetcher.NewTeacherLessonFetcher(nil, concurrency, logger.App),
		dryRun:   dryRun,
		teachers: make(map[uint32]*model.Teacher, 1000),
	}
}

func (n *Notifier) SendNotification(user *model.User) error {
	followingTeacherService := model.NewFollowingTeacherService(n.db)
	n.lessonService = model.NewLessonService(n.db)

	teacherIDs, err := followingTeacherService.FindTeacherIDsByUserID(user.ID)
	if err != nil {
		return errors.Wrapperf(err, "Failed to FindTeacherIDsByUserID(): userID=%v", user.ID)
	}

	availableLessonsPerTeacher := make(map[uint32][]*model.Lesson, 1000)
	allFetchedLessons := make([]*model.Lesson, 0, 5000)
	wg := &sync.WaitGroup{}
	for _, teacherID := range teacherIDs {
		wg.Add(1)
		go func(teacherID uint32) {
			defer wg.Done()
			teacher, fetchedLessons, newAvailableLessons, err := n.fetchAndExtractNewAvailableLessons(teacherID)
			if err != nil {
				switch err.(type) {
				case *errors.NotFound:
					// TODO: update teacher table flag
					// TODO: Not need to log
					logger.App.Warn("Cannot find teacher", zap.Uint("teacherID", uint(teacherID)))
				default:
					logger.App.Error("Cannot fetch teacher", zap.Uint("teacherID", uint(teacherID)), zap.Error(err))
				}
				return
			}

			allFetchedLessons = append(allFetchedLessons, fetchedLessons...) // TODO: delete
			n.teachers[teacherID] = teacher
			if _, ok := n.fetchedLessons[teacherID]; ok {
				n.fetchedLessons[teacherID] = append(n.fetchedLessons[teacherID], fetchedLessons...)
			} else {
				n.fetchedLessons[teacherID] = make([]*model.Lesson, 0, 5000)
			}
			if len(newAvailableLessons) > 0 {
				availableLessonsPerTeacher[teacherID] = newAvailableLessons
			}
		}(teacherID)

		if err != nil {
			return err
		}
	}
	wg.Wait()

	if err := n.sendNotificationToUser(user, availableLessonsPerTeacher); err != nil {
		return err
	}

	if !n.dryRun {
		n.lessonService.UpdateLessons(allFetchedLessons)
	}

	time.Sleep(500 * time.Millisecond)

	return nil
}

// Returns teacher, fetchedLessons, newAvailableLessons, error
func (n *Notifier) fetchAndExtractNewAvailableLessons(teacherID uint32) (
	*model.Teacher, []*model.Lesson, []*model.Lesson, error,
) {
	teacher, fetchedLessons, err := n.fetcher.Fetch(teacherID)
	if err != nil {
		logger.App.Error(
			"TeacherLessonFetcher.Fetch",
			zap.Uint("teacherID", uint(teacherID)), zap.Error(err),
		)
		return nil, nil, nil, err
	}
	logger.App.Info(
		"TeacherLessonFetcher.Fetch",
		zap.Uint("teacherID", uint(teacher.ID)),
		zap.String("teacherName", teacher.Name),
		zap.Int("fetchedLessons", len(fetchedLessons)),
	)

	//fmt.Printf("fetchedLessons ---\n")
	//for _, l := range fetchedLessons {
	//	fmt.Printf("teacherID=%v, datetime=%v, status=%v\n", l.TeacherId, l.Datetime, l.Status)
	//}

	now := time.Now()
	fromDate := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, config.LocalTimezone())
	toDate := fromDate.Add(24 * 6 * time.Hour)
	lastFetchedLessons, err := n.lessonService.FindLessons(teacher.ID, fromDate, toDate)
	if err != nil {
		return nil, nil, nil, err
	}
	//fmt.Printf("lastFetchedLessons ---\n")
	//for _, l := range lastFetchedLessons {
	//	fmt.Printf("teacherID=%v, datetime=%v, status=%v\n", l.TeacherId, l.Datetime, l.Status)
	//}

	newAvailableLessons := n.lessonService.GetNewAvailableLessons(lastFetchedLessons, fetchedLessons)
	//fmt.Printf("newAvailableLessons ---\n")
	//for _, l := range newAvailableLessons {
	//	fmt.Printf("teacherID=%v, datetime=%v, status=%v\n", l.TeacherId, l.Datetime, l.Status)
	//}
	return teacher, fetchedLessons, newAvailableLessons, nil
}

func (n *Notifier) sendNotificationToUser(
	user *model.User,
	lessonsPerTeacher map[uint32][]*model.Lesson,
) error {
	lessonsCount := 0
	var teacherIDs []int
	for teacherID, lessons := range lessonsPerTeacher {
		teacherIDs = append(teacherIDs, int(teacherID))
		lessonsCount += len(lessons)
	}
	if lessonsCount == 0 {
		// Don't send notification
		return nil
	}

	sort.Ints(teacherIDs)
	var teacherIDs2 []uint32
	var teacherNames []string
	for _, id := range teacherIDs {
		teacherIDs2 = append(teacherIDs2, uint32(id))
		teacherNames = append(teacherNames, n.teachers[uint32(id)].Name)
	}

	t := template.New("email")
	t = template.Must(t.Parse(getEmailTemplateJP()))
	type TemplateData struct {
		TeacherIDs        []uint32
		Teachers          map[uint32]*model.Teacher
		LessonsPerTeacher map[uint32][]*model.Lesson
		WebURL            string
	}
	data := &TemplateData{
		TeacherIDs:        teacherIDs2,
		Teachers:          n.teachers,
		LessonsPerTeacher: lessonsPerTeacher,
		WebURL:            config.WebURL(),
	}

	var body bytes.Buffer
	if err := t.Execute(&body, data); err != nil {
		return errors.InternalWrapf(err, "Failed to execute template.")
	}
	//fmt.Printf("--- mail ---\n%s", body.String())

	logger.App.Info("sendNotificationToUser", zap.String("email", user.Email.Raw()))
	//subject := "Schedule of teacher " + strings.Join(teacherNames, ", ")
	subject := strings.Join(teacherNames, ", ") + "の空きレッスンがあります"
	sender := &EmailNotificationSender{}
	return sender.Send(user, subject, body.String())
}

func getEmailTemplateJP() string {
	return strings.TrimSpace(`
{{- range $teacherID := .TeacherIDs }}
{{- $teacher := index $.Teachers $teacherID -}}
--- {{ $teacher.Name }} ---
  {{- $lessons := index $.LessonsPerTeacher $teacherID }}
  {{- range $lesson := $lessons }}
{{ $lesson.Datetime.Format "2006-01-02 15:04" }}
  {{- end }}

レッスンの予約はこちらから:
PC
Mobile

{{ end }}
空きレッスンの通知の解除はこちら
	`)
}

func getEmailTemplateEN() string {
	return strings.TrimSpace(`
{{- range $teacherID := .TeacherIDs }}
{{- $teacher := index $.Teachers $teacherID -}}
--- {{ $teacher.Name }} ---
  {{- $lessons := index $.LessonsPerTeacher $teacherID }}
  {{- range $lesson := $lessons }}
{{ $lesson.Datetime.Format "2006-01-02 15:04" }}
  {{- end }}

Reserve here:
PC
Mobile
{{ end }}
Click here if you want to stop notification of the teacher.
	`)
}

func (n *Notifier) Close() {
	n.fetcher.Close()
}

type NotificationSender interface {
	Send(user *model.User, subject, body string) error
}

type EmailNotificationSender struct{}

func (s *EmailNotificationSender) Send(user *model.User, subject, body string) error {
	from := mail.NewEmail("lekcije", "lekcije@lekcije.com")
	to := mail.NewEmail(user.Name, user.Email.Raw())
	content := mail.NewContent("text/html", strings.Replace(body, "\n", "
", -1))
	m := mail.NewV3MailInit(from, subject, to, content)

	req := sendgrid.GetRequest(
		os.Getenv("SENDGRID_API_KEY"),
		"/v3/mail/send",
		"https://api.sendgrid.com",
	)
	req.Method = "POST"
	req.Body = mail.GetRequestBody(m)
	resp, err := sendgrid.API(req)
	if err != nil {
		return errors.InternalWrapf(err, "Failed to send email by sendgrid")
	}
	if resp.StatusCode >= 300 {
		message := fmt.Sprintf(
			"Failed to send email by sendgrid: statusCode=%v, body=%v",
			resp.StatusCode, strings.Replace(resp.Body, "\n", "\\n", -1),
		)
		logger.App.Error(message)
		return errors.InternalWrapf(
			err,
			"Failed to send email by sendgrid: statusCode=%v",
			resp.StatusCode,
		)
	}

	return nil
}
<commit_msg>Update lessons before notifier finishes<commit_after>package notifier

import (
	"bytes"
	"fmt"
	"os"
	"sort"
	"strings"
	"sync"
	"text/template"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/oinume/lekcije/server/config"
	"github.com/oinume/lekcije/server/errors"
	"github.com/oinume/lekcije/server/fetcher"
	"github.com/oinume/lekcije/server/logger"
	"github.com/oinume/lekcije/server/model"
	"github.com/sendgrid/sendgrid-go"
	"github.com/sendgrid/sendgrid-go/helpers/mail"
	"github.com/uber-go/zap"
)

type Notifier struct {
	db            *gorm.DB
	fetcher       *fetcher.TeacherLessonFetcher
	dryRun        bool
	lessonService *model.LessonService
	teachers      map[uint32]*model.Teacher
	fetchedLessons map[uint32][]*model.Lesson
}

func NewNotifier(db *gorm.DB, concurrency int, dryRun bool) *Notifier {
	if concurrency < 1 {
		concurrency = 1
	}
	return &Notifier{
		db:       db,
		fetcher:  fetcher.NewTeacherLessonFetcher(nil, concurrency, logger.App),
		dryRun:   dryRun,
		teachers: make(map[uint32]*model.Teacher, 1000),
		fetchedLessons: make(map[uint32][]*model.Lesson, 1000),
	}
}

func (n *Notifier) SendNotification(user *model.User) error {
	followingTeacherService := model.NewFollowingTeacherService(n.db)
	n.lessonService = model.NewLessonService(n.db)

	teacherIDs, err := followingTeacherService.FindTeacherIDsByUserID(user.ID)
	if err != nil {
		return errors.Wrapperf(err, "Failed to FindTeacherIDsByUserID(): userID=%v", user.ID)
	}

	availableLessonsPerTeacher := make(map[uint32][]*model.Lesson, 1000)
	allFetchedLessons := make([]*model.Lesson, 0, 5000)
	wg := &sync.WaitGroup{}
	for _, teacherID := range teacherIDs {
		wg.Add(1)
		go func(teacherID uint32) {
			defer wg.Done()
			teacher, fetchedLessons, newAvailableLessons, err := n.fetchAndExtractNewAvailableLessons(teacherID)
			if err != nil {
				switch err.(type) {
				case *errors.NotFound:
					// TODO: update teacher table flag
					// TODO: Not need to log
					logger.App.Warn("Cannot find teacher", zap.Uint("teacherID", uint(teacherID)))
				default:
					logger.App.Error("Cannot fetch teacher", zap.Uint("teacherID", uint(teacherID)), zap.Error(err))
				}
				return
			}

			allFetchedLessons = append(allFetchedLessons, fetchedLessons...) // TODO: delete
			// TODO: lock
			n.teachers[teacherID] = teacher
			if _, ok := n.fetchedLessons[teacherID]; ok {
				n.fetchedLessons[teacherID] = append(n.fetchedLessons[teacherID], fetchedLessons...)
			} else {
				n.fetchedLessons[teacherID] = make([]*model.Lesson, 0, 5000)
			}
			if len(newAvailableLessons) > 0 {
				availableLessonsPerTeacher[teacherID] = newAvailableLessons
			}
		}(teacherID)

		if err != nil {
			return err
		}
	}
	wg.Wait()

	if err := n.sendNotificationToUser(user, availableLessonsPerTeacher); err != nil {
		return err
	}

	//if !n.dryRun {
	//	n.lessonService.UpdateLessons(allFetchedLessons)
	//}

	time.Sleep(500 * time.Millisecond)

	return nil
}

// Returns teacher, fetchedLessons, newAvailableLessons, error
func (n *Notifier) fetchAndExtractNewAvailableLessons(teacherID uint32) (
	*model.Teacher, []*model.Lesson, []*model.Lesson, error,
) {
	teacher, fetchedLessons, err := n.fetcher.Fetch(teacherID)
	if err != nil {
		logger.App.Error(
			"TeacherLessonFetcher.Fetch",
			zap.Uint("teacherID", uint(teacherID)), zap.Error(err),
		)
		return nil, nil, nil, err
	}
	logger.App.Info(
		"TeacherLessonFetcher.Fetch",
		zap.Uint("teacherID", uint(teacher.ID)),
		zap.String("teacherName", teacher.Name),
		zap.Int("fetchedLessons", len(fetchedLessons)),
	)

	//fmt.Printf("fetchedLessons ---\n")
	//for _, l := range fetchedLessons {
	//	fmt.Printf("teacherID=%v, datetime=%v, status=%v\n", l.TeacherId, l.Datetime, l.Status)
	//}

	now := time.Now()
	fromDate := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, config.LocalTimezone())
	toDate := fromDate.Add(24 * 6 * time.Hour)
	lastFetchedLessons, err := n.lessonService.FindLessons(teacher.ID, fromDate, toDate)
	if err != nil {
		return nil, nil, nil, err
	}
	//fmt.Printf("lastFetchedLessons ---\n")
	//for _, l := range lastFetchedLessons {
	//	fmt.Printf("teacherID=%v, datetime=%v, status=%v\n", l.TeacherId, l.Datetime, l.Status)
	//}

	newAvailableLessons := n.lessonService.GetNewAvailableLessons(lastFetchedLessons, fetchedLessons)
	//fmt.Printf("newAvailableLessons ---\n")
	//for _, l := range newAvailableLessons {
	//	fmt.Printf("teacherID=%v, datetime=%v, status=%v\n", l.TeacherId, l.Datetime, l.Status)
	//}
	return teacher, fetchedLessons, newAvailableLessons, nil
}

func (n *Notifier) sendNotificationToUser(
	user *model.User,
	lessonsPerTeacher map[uint32][]*model.Lesson,
) error {
	lessonsCount := 0
	var teacherIDs []int
	for teacherID, lessons := range lessonsPerTeacher {
		teacherIDs = append(teacherIDs, int(teacherID))
		lessonsCount += len(lessons)
	}
	if lessonsCount == 0 {
		// Don't send notification
		return nil
	}

	sort.Ints(teacherIDs)
	var teacherIDs2 []uint32
	var teacherNames []string
	for _, id := range teacherIDs {
		teacherIDs2 = append(teacherIDs2, uint32(id))
		teacherNames = append(teacherNames, n.teachers[uint32(id)].Name)
	}

	t := template.New("email")
	t = template.Must(t.Parse(getEmailTemplateJP()))
	type TemplateData struct {
		TeacherIDs        []uint32
		Teachers          map[uint32]*model.Teacher
		LessonsPerTeacher map[uint32][]*model.Lesson
		WebURL            string
	}
	data := &TemplateData{
		TeacherIDs:        teacherIDs2,
		Teachers:          n.teachers,
		LessonsPerTeacher: lessonsPerTeacher,
		WebURL:            config.WebURL(),
	}

	var body bytes.Buffer
	if err := t.Execute(&body, data); err != nil {
		return errors.InternalWrapf(err, "Failed to execute template.")
	}
	//fmt.Printf("--- mail ---\n%s", body.String())

	logger.App.Info("sendNotificationToUser", zap.String("email", user.Email.Raw()))
	//subject := "Schedule of teacher " + strings.Join(teacherNames, ", ")
	subject := strings.Join(teacherNames, ", ") + "の空きレッスンがあります"
	sender := &EmailNotificationSender{}
	return sender.Send(user, subject, body.String())
}

func getEmailTemplateJP() string {
	return strings.TrimSpace(`
{{- range $teacherID := .TeacherIDs }}
{{- $teacher := index $.Teachers $teacherID -}}
--- {{ $teacher.Name }} ---
  {{- $lessons := index $.LessonsPerTeacher $teacherID }}
  {{- range $lesson := $lessons }}
{{ $lesson.Datetime.Format "2006-01-02 15:04" }}
  {{- end }}

レッスンの予約はこちらから:
PC
Mobile

{{ end }}
空きレッスンの通知の解除はこちら
	`)
}

func getEmailTemplateEN() string {
	return strings.TrimSpace(`
{{- range $teacherID := .TeacherIDs }}
{{- $teacher := index $.Teachers $teacherID -}}
--- {{ $teacher.Name }} ---
  {{- $lessons := index $.LessonsPerTeacher $teacherID }}
  {{- range $lesson := $lessons }}
{{ $lesson.Datetime.Format "2006-01-02 15:04" }}
  {{- end }}

Reserve here:
PC
Mobile
{{ end }}
Click here if you want to stop notification of the teacher.
	`)
}

func (n *Notifier) Close() {
	defer n.fetcher.Close()
	defer func() {
		if n.dryRun {
			return
		}
		for teacherID, lessons := range n.fetchedLessons {
			if _, err := n.lessonService.UpdateLessons(lessons); err != nil {
				logger.App.Error(
					"An error ocurred in Notifier.Close",
					zap.Error(err), zap.Uint("teacherID", uint(teacherID)),
				)
			}
		}
	}()
}

type NotificationSender interface {
	Send(user *model.User, subject, body string) error
}

type EmailNotificationSender struct{}

func (s *EmailNotificationSender) Send(user *model.User, subject, body string) error {
	from := mail.NewEmail("lekcije", "lekcije@lekcije.com")
	to := mail.NewEmail(user.Name, user.Email.Raw())
	content := mail.NewContent("text/html", strings.Replace(body, "\n", "
", -1))
	m := mail.NewV3MailInit(from, subject, to, content)

	req := sendgrid.GetRequest(
		os.Getenv("SENDGRID_API_KEY"),
		"/v3/mail/send",
		"https://api.sendgrid.com",
	)
	req.Method = "POST"
	req.Body = mail.GetRequestBody(m)
	resp, err := sendgrid.API(req)
	if err != nil {
		return errors.InternalWrapf(err, "Failed to send email by sendgrid")
	}
	if resp.StatusCode >= 300 {
		message := fmt.Sprintf(
			"Failed to send email by sendgrid: statusCode=%v, body=%v",
			resp.StatusCode, strings.Replace(resp.Body, "\n", "\\n", -1),
		)
		logger.App.Error(message)
		return errors.InternalWrapf(
			err,
			"Failed to send email by sendgrid: statusCode=%v",
			resp.StatusCode,
		)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"archive/tar"
	"context"
	"encoding/hex"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/client"
	"github.com/jadr2ddude/websocket"
)

func main() {
	dcli, err := client.NewEnvClient()
	if err != nil {
		panic(err)
	}
	srv := &ContainerServer{
		DockerClient: dcli,
	}
	f, err := os.Open("langs.json")
	if err != nil {
		panic(err)
	}
	err = json.NewDecoder(f).Decode(&srv.Containers)
	if err != nil {
		panic(err)
	}
	http.HandleFunc("/term", srv.HandleTerminal)
	http.HandleFunc("/run", srv.HandleRun)
	panic(http.ListenAndServe(":80", nil))
}

// ContainerConfig is a container configuration.
type ContainerConfig struct {
	Image   string   `json:"image"`
	Command []string `json:"cmd"`
}

// pullImg pulls the docker image used by the ContainerConfig.
func (cc ContainerConfig) pullImg(ctx context.Context, cli *client.Client) (err error) {
	pr, err := cli.ImagePull(ctx, "docker.io/library/"+cc.Image, types.ImagePullOptions{})
	if err != nil {
		return err
	}
	defer func() {
		cerr := pr.Close()
		if cerr != nil && err == nil {
			err = cerr
		}
	}()
	_, err = io.Copy(os.Stdout, pr)
	if err != nil {
		return err
	}
	return nil
}

// Container is a running container.
type Container struct {
	cli *client.Client
	ID  string
	IO  io.ReadWriteCloser
}

// Close closes and removes the container.
func (c *Container) Close(ctx context.Context) error {
	// close websocket
	cerr := c.IO.Close()

	// remove container
	rerr := c.cli.ContainerRemove(ctx, c.ID, types.ContainerRemoveOptions{
		Force: true,
	})

	// handle errors
	if rerr != nil {
		log.Printf("Failed to remove container: %s", rerr.Error())
	}
	err := cerr
	if err != nil {
		err = rerr
	}
	return err
}

// Deploy deploys a container with this configuration.
func (cc ContainerConfig) Deploy(ctx context.Context, cli *client.Client, prestart func(*Container) error) (cont *Container, err error) {
	/*
		    // pull image
			err = cc.pullImg(ctx, cli)
			if err != nil {
				return nil, err
			}
	*/

	// create container
	c, err := cli.ContainerCreate(ctx, &container.Config{
		Image:           cc.Image,
		Cmd:             cc.Command,
		Tty:             true,
		OpenStdin:       true,
		NetworkDisabled: true,
	}, nil, nil, "")
	if err != nil {
		return nil, err
	}

	// cleanup container on failed startup
	defer func() {
		if err != nil {
			delctx, cancel := context.WithTimeout(context.Background(), time.Minute)
			defer cancel()
			rerr := cli.ContainerRemove(delctx, c.ID, types.ContainerRemoveOptions{
				Force: true,
			})
			if rerr != nil {
				log.Printf("Failed to remove container: %s", rerr.Error())
			}
		}
	}()

	cont = &Container{
		cli: cli,
		ID:  c.ID,
	}

	// run prestart hook
	if prestart != nil {
		err = prestart(cont)
		if err != nil {
			return nil, err
		}
	}

	// attach to container
	resp, err := cli.ContainerAttach(ctx, c.ID, types.ContainerAttachOptions{
		Stream: true,
		Stdin:  true,
		Stdout: true,
		Stderr: true,
	})
	if err != nil {
		return nil, err
	}

	// start container
	err = cli.ContainerStart(ctx, c.ID, types.ContainerStartOptions{})
	if err != nil {
		return nil, err
	}

	// convert to websocket
	cont.IO = resp.Conn

	return cont, nil
}

// Language is a configuration for a programming language.
type Language struct {
	RunContainer  ContainerConfig `json:"run"`
	TermContainer ContainerConfig `json:"term"`
}

// ContainerServer is a server that runs containers
type ContainerServer struct {
	// DockerClient is the client to Docker.
	DockerClient *client.Client

	// Containers is a map of language names to container names.
	Containers map[string]Language

	// Upgrader is a websocket Upgrader used for all websocket connections.
	Upgrader websocket.Upgrader
}

// StatusUpdate is a status message which can be sent to the client.
type StatusUpdate struct {
	Status string `json:"status"`
	Error  string `json:"err,omitempty"`
}

func copyWebSocket(dst *websocket.Conn, src *websocket.Conn, cancel context.CancelFunc, n int) {
	defer cancel()
	for {
		// read message
		t, dat, err := src.ReadMessage()
		if err != nil {
			log.Println("Stopping transfer", err.Error(), n)
			return
		}
		log.Println(t, dat)

		switch t {
		case websocket.CloseMessage:
			// shut down
			return
		case websocket.BinaryMessage:
			// forward message
			err = dst.WriteMessage(t, dat)
			if err != nil {
				return
			}
		}
	}
}

// HandleTerminal serves an interactive terminal websocket.
func (cs *ContainerServer) HandleTerminal(w http.ResponseWriter, r *http.Request) {
	log.Println("term")
	// get language
	lang, ok := cs.Containers[r.URL.Query().Get("lang")]
	if !ok {
		http.Error(w, "language not supported", http.StatusBadRequest)
		return
	}

	// upgrade websocket
	log.Println("upgrade")
	conn, err := cs.Upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()

	// send status "starting"
	log.Println("starting")
	err = conn.WriteJSON(StatusUpdate{Status: "starting"})
	if err != nil {
		return
	}

	// deploy container with 1 min timeout
	startctx, startcancel := context.WithTimeout(context.Background(), time.Minute)
	defer startcancel()
	c, err := lang.TermContainer.Deploy(startctx, cs.DockerClient, nil)
	log.Println("deploy", c, err)
	if err != nil {
		conn.WriteJSON(StatusUpdate{Status: "error", Error: err.Error()})
		err = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		if err == nil {
			donech := make(chan struct{})
			go func() {
				defer close(donech)
				// drain client messages and wait for disconnect
				var e error
				for e == nil {
					_, _, e = conn.ReadMessage()
				}
			}()
			timer := time.NewTimer(10 * time.Second)
			defer timer.Stop()
			select {
			case <-donech:
			case <-timer.C:
			}
		}
	}
	defer func() {
		stopctx, stopcancel := context.WithTimeout(context.Background(), time.Minute)
		defer stopcancel()
		cerr := c.Close(stopctx)
		if cerr != nil {
			log.Printf("Failed to stop container %q: %s", c.ID, cerr)
		}
	}()

	// update status to running
	log.Println("running")
	err = conn.WriteJSON(StatusUpdate{Status: "running"})
	if err != nil {
		return
	}

	// bridge connections
	runctx, cancel := context.WithCancel(context.Background())
	go func() { // output
		defer cancel()
		buf := make([]byte, 1024)
		for {
			n, err := c.IO.Read(buf)
			log.Println(n, err)
			if err != nil {
				return
			}

			log.Println(hex.Dump(buf[:n]))

			err = conn.WriteMessage(websocket.TextMessage, buf[:n])
			if err != nil {
				return
			}
		}
	}()
	go func() { // input
		defer cancel()
		for {
			t, r, err := conn.NextReader()
			if err != nil {
				return
			}
			switch t {
			case websocket.TextMessage:
				fallthrough
			case websocket.BinaryMessage:
				_, err = io.Copy(c.IO, r)
				if err != nil {
					return
				}
			case websocket.CloseMessage:
				return
			}
		}
	}()
	//go copyWebSocket(c.IO, conn, cancel, 2)
	log.Println("Copying")
	<-runctx.Done()
	log.Println("Done")
}

func packCodeTarball(dat []byte) io.ReadCloser {
	r, w := io.Pipe()
	go func() {
		var err error
		defer func() {
			if err == nil {
				w.Close()
			} else {
				w.CloseWithError(err)
			}
		}()
		tw := tar.NewWriter(w)
		defer tw.Close()

		err = tw.WriteHeader(&tar.Header{
			Name: "code",
			Mode: 0444,
			Size: int64(len(dat)),
		})
		if err != nil {
			return
		}

		_, err = tw.Write(dat)
		if err != nil {
			return
		}
	}()
	return r
}

// HandleRun serves an interactive terminal running user code over a websocket.
func (cs *ContainerServer) HandleRun(w http.ResponseWriter, r *http.Request) {
	// get language
	lang, ok := cs.Containers[r.URL.Query().Get("lang")]
	if !ok {
		http.Error(w, "language not supported", http.StatusBadRequest)
		return
	}

	// upgrade websocket
	conn, err := cs.Upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	defer conn.Close()

	// send status "creating"
	err = conn.WriteJSON(StatusUpdate{Status: "creating"})
	if err != nil {
		return
	}

	// deploy container with 1 min timeout
	startctx, startcancel := context.WithTimeout(context.Background(), time.Minute)
	defer startcancel()
	c, err := lang.RunContainer.Deploy(startctx, cs.DockerClient, func(c *Container) error {
		// update status to ready
		err := conn.WriteJSON(StatusUpdate{Status: "ready"})
		if err != nil {
			return err
		}

		// accept user code
		t, dat, err := conn.ReadMessage()
		if err != nil {
			return err
		}
		if t != websocket.BinaryMessage && t != websocket.TextMessage {
			log.Println("Client sent invalid message type")
			return err
		}

		// update status to uploading
		err = conn.WriteJSON(StatusUpdate{Status: "uploading"})
		if err != nil {
			return err
		}

		// send code to Docker
		tr := packCodeTarball(dat)
		err = c.cli.CopyToContainer(startctx, c.ID, "/", tr, types.CopyToContainerOptions{})
		tr.Close()
		if err != nil {
			conn.WriteJSON(StatusUpdate{Status: "error", Error: err.Error()})
			return err
		}

		// update status to starting
		err = conn.WriteJSON(StatusUpdate{Status: "starting"})
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		conn.WriteJSON(StatusUpdate{Status: "error", Error: err.Error()})
		err = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		if err == nil {
			donech := make(chan struct{})
			go func() {
				defer close(donech)
				// drain client messages and wait for disconnect
				var e error
				for e == nil {
					_, _, e = conn.ReadMessage()
				}
			}()
			timer := time.NewTimer(10 * time.Second)
			defer timer.Stop()
			select {
			case <-donech:
			case <-timer.C:
			}
		}
	}
	defer func() {
		stopctx, stopcancel := context.WithTimeout(context.Background(), time.Minute)
		defer stopcancel()
		cerr := c.Close(stopctx)
		if cerr != nil {
			log.Printf("Failed to stop container %q: %s", c.ID, cerr)
		}
	}()

	// update status to running
	err = conn.WriteJSON(StatusUpdate{Status: "running"})
	if err != nil {
		return
	}

	// bridge connections
	runctx, cancel := context.WithCancel(context.Background())
	go func() { // output
		defer cancel()
		buf := make([]byte, 1024)
		for {
			n, err := c.IO.Read(buf)
			log.Println(n, err)
			if err != nil {
				return
			}

			err = conn.WriteMessage(websocket.TextMessage, buf[:n])
			if err != nil {
				return
			}
		}
	}()
	go func() { // input
		defer cancel()
		for {
			t, r, err := conn.NextReader()
			if err != nil {
				return
			}
			switch t {
			case websocket.TextMessage:
				fallthrough
			case websocket.BinaryMessage:
				_, err = io.Copy(c.IO, r)
				if err != nil {
					return
				}
			case websocket.CloseMessage:
				return
			}
		}
	}()
	log.Println("Copying")
	<-runctx.Done()
	log.Println("Done")
}
<commit_msg>do proper websocket shutdown on terminal requests<commit_after>package main

import (
	"archive/tar"
	"context"
	"encoding/hex"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/client"
	"github.com/jadr2ddude/websocket"
)

func main() {
	dcli, err := client.NewEnvClient()
	if err != nil {
		panic(err)
	}
	srv := &ContainerServer{
		DockerClient: dcli,
	}
	f, err := os.Open("langs.json")
	if err != nil {
		panic(err)
	}
	err = json.NewDecoder(f).Decode(&srv.Containers)
	if err != nil {
		panic(err)
	}
	http.HandleFunc("/term", srv.HandleTerminal)
	http.HandleFunc("/run", srv.HandleRun)
	panic(http.ListenAndServe(":80", nil))
}

// ContainerConfig is a container configuration.
type ContainerConfig struct {
	Image   string   `json:"image"`
	Command []string `json:"cmd"`
}

// pullImg pulls the docker image used by the ContainerConfig.
func (cc ContainerConfig) pullImg(ctx context.Context, cli *client.Client) (err error) {
	pr, err := cli.ImagePull(ctx, "docker.io/library/"+cc.Image, types.ImagePullOptions{})
	if err != nil {
		return err
	}
	defer func() {
		cerr := pr.Close()
		if cerr != nil && err == nil {
			err = cerr
		}
	}()
	_, err = io.Copy(os.Stdout, pr)
	if err != nil {
		return err
	}
	return nil
}

// Container is a running container.
type Container struct {
	cli *client.Client
	ID  string
	IO  io.ReadWriteCloser
}

// Close closes and removes the container.
func (c *Container) Close(ctx context.Context) error {
	// close websocket
	cerr := c.IO.Close()

	// remove container
	rerr := c.cli.ContainerRemove(ctx, c.ID, types.ContainerRemoveOptions{
		Force: true,
	})

	// handle errors
	if rerr != nil {
		log.Printf("Failed to remove container: %s", rerr.Error())
	}
	err := cerr
	if err != nil {
		err = rerr
	}
	return err
}

// Deploy deploys a container with this configuration.
func (cc ContainerConfig) Deploy(ctx context.Context, cli *client.Client, prestart func(*Container) error) (cont *Container, err error) {
	/*
		    // pull image
			err = cc.pullImg(ctx, cli)
			if err != nil {
				return nil, err
			}
	*/

	// create container
	c, err := cli.ContainerCreate(ctx, &container.Config{
		Image:           cc.Image,
		Cmd:             cc.Command,
		Tty:             true,
		OpenStdin:       true,
		NetworkDisabled: true,
	}, nil, nil, "")
	if err != nil {
		return nil, err
	}

	// cleanup container on failed startup
	defer func() {
		if err != nil {
			delctx, cancel := context.WithTimeout(context.Background(), time.Minute)
			defer cancel()
			rerr := cli.ContainerRemove(delctx, c.ID, types.ContainerRemoveOptions{
				Force: true,
			})
			if rerr != nil {
				log.Printf("Failed to remove container: %s", rerr.Error())
			}
		}
	}()

	cont = &Container{
		cli: cli,
		ID:  c.ID,
	}

	// run prestart hook
	if prestart != nil {
		err = prestart(cont)
		if err != nil {
			return nil, err
		}
	}

	// attach to container
	resp, err := cli.ContainerAttach(ctx, c.ID, types.ContainerAttachOptions{
		Stream: true,
		Stdin:  true,
		Stdout: true,
		Stderr: true,
	})
	if err != nil {
		return nil, err
	}

	// start container
	err = cli.ContainerStart(ctx, c.ID, types.ContainerStartOptions{})
	if err != nil {
		return nil, err
	}

	// convert to websocket
	cont.IO = resp.Conn

	return cont, nil
}

// Language is a configuration for a programming language.
type Language struct {
	RunContainer  ContainerConfig `json:"run"`
	TermContainer ContainerConfig `json:"term"`
}

// ContainerServer is a server that runs containers
type ContainerServer struct {
	// DockerClient is the client to Docker.
	DockerClient *client.Client

	// Containers is a map of language names to container names.
	Containers map[string]Language

	// Upgrader is a websocket Upgrader used for all websocket connections.
	Upgrader websocket.Upgrader
}

// StatusUpdate is a status message which can be sent to the client.
type StatusUpdate struct {
	Status string `json:"status"`
	Error  string `json:"err,omitempty"`
}

func copyWebSocket(dst *websocket.Conn, src *websocket.Conn, cancel context.CancelFunc, n int) {
	defer cancel()
	for {
		// read message
		t, dat, err := src.ReadMessage()
		if err != nil {
			log.Println("Stopping transfer", err.Error(), n)
			return
		}
		log.Println(t, dat)

		switch t {
		case websocket.CloseMessage:
			// shut down
			return
		case websocket.BinaryMessage:
			// forward message
			err = dst.WriteMessage(t, dat)
			if err != nil {
				return
			}
		}
	}
}

// HandleTerminal serves an interactive terminal websocket.
func (cs *ContainerServer) HandleTerminal(w http.ResponseWriter, r *http.Request) {
	log.Println("term")
	// get language
	lang, ok := cs.Containers[r.URL.Query().Get("lang")]
	if !ok {
		http.Error(w, "language not supported", http.StatusBadRequest)
		return
	}

	// upgrade websocket
	log.Println("upgrade")
	conn, err := cs.Upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}
	defer conn.Close()
	defer func() {
		cerr := conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		if cerr == nil {
			donech := make(chan struct{})
			go func() {
				defer close(donech)
				// drain client messages and wait for disconnect
				var e error
				for e == nil {
					_, _, e = conn.ReadMessage()
				}
			}()
			timer := time.NewTimer(10 * time.Second)
			defer timer.Stop()
			select {
			case <-donech:
			case <-timer.C:
			}
		}
	}()

	// send status "starting"
	log.Println("starting")
	err = conn.WriteJSON(StatusUpdate{Status: "starting"})
	if err != nil {
		return
	}

	// deploy container with 1 min timeout
	startctx, startcancel := context.WithTimeout(context.Background(), time.Minute)
	defer startcancel()
	c, err := lang.TermContainer.Deploy(startctx, cs.DockerClient, nil)
	log.Println("deploy", c, err)
	if err != nil {
		conn.WriteJSON(StatusUpdate{Status: "error", Error: err.Error()})
		return
	}
	defer func() {
		stopctx, stopcancel := context.WithTimeout(context.Background(), time.Minute)
		defer stopcancel()
		cerr := c.Close(stopctx)
		if cerr != nil {
			log.Printf("Failed to stop container %q: %s", c.ID, cerr)
		}
	}()

	// update status to running
	log.Println("running")
	err = conn.WriteJSON(StatusUpdate{Status: "running"})
	if err != nil {
		return
	}

	// bridge connections
	runctx, cancel := context.WithCancel(context.Background())
	go func() { // output
		defer cancel()
		buf := make([]byte, 1024)
		for {
			n, err := c.IO.Read(buf)
			log.Println(n, err)
			if err != nil {
				return
			}

			log.Println(hex.Dump(buf[:n]))

			err = conn.WriteMessage(websocket.TextMessage, buf[:n])
			if err != nil {
				return
			}
		}
	}()
	go func() { // input
		defer cancel()
		for {
			t, r, err := conn.NextReader()
			if err != nil {
				return
			}
			switch t {
			case websocket.TextMessage:
				fallthrough
			case websocket.BinaryMessage:
				_, err = io.Copy(c.IO, r)
				if err != nil {
					return
				}
			case websocket.CloseMessage:
				return
			}
		}
	}()
	//go copyWebSocket(c.IO, conn, cancel, 2)
	log.Println("Copying")
	<-runctx.Done()
	log.Println("Done")
}

func packCodeTarball(dat []byte) io.ReadCloser {
	r, w := io.Pipe()
	go func() {
		var err error
		defer func() {
			if err == nil {
				w.Close()
			} else {
				w.CloseWithError(err)
			}
		}()
		tw := tar.NewWriter(w)
		defer tw.Close()

		err = tw.WriteHeader(&tar.Header{
			Name: "code",
			Mode: 0444,
			Size: int64(len(dat)),
		})
		if err != nil {
			return
		}

		_, err = tw.Write(dat)
		if err != nil {
			return
		}
	}()
	return r
}

// HandleRun serves an interactive terminal running user code over a websocket.
func (cs *ContainerServer) HandleRun(w http.ResponseWriter, r *http.Request) {
	// get language
	lang, ok := cs.Containers[r.URL.Query().Get("lang")]
	if !ok {
		http.Error(w, "language not supported", http.StatusBadRequest)
		return
	}

	// upgrade websocket
	conn, err := cs.Upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	defer conn.Close()

	// send status "creating"
	err = conn.WriteJSON(StatusUpdate{Status: "creating"})
	if err != nil {
		return
	}

	// deploy container with 1 min timeout
	startctx, startcancel := context.WithTimeout(context.Background(), time.Minute)
	defer startcancel()
	c, err := lang.RunContainer.Deploy(startctx, cs.DockerClient, func(c *Container) error {
		// update status to ready
		err := conn.WriteJSON(StatusUpdate{Status: "ready"})
		if err != nil {
			return err
		}

		// accept user code
		t, dat, err := conn.ReadMessage()
		if err != nil {
			return err
		}
		if t != websocket.BinaryMessage && t != websocket.TextMessage {
			log.Println("Client sent invalid message type")
			return err
		}

		// update status to uploading
		err = conn.WriteJSON(StatusUpdate{Status: "uploading"})
		if err != nil {
			return err
		}

		// send code to Docker
		tr := packCodeTarball(dat)
		err = c.cli.CopyToContainer(startctx, c.ID, "/", tr, types.CopyToContainerOptions{})
		tr.Close()
		if err != nil {
			conn.WriteJSON(StatusUpdate{Status: "error", Error: err.Error()})
			return err
		}

		// update status to starting
		err = conn.WriteJSON(StatusUpdate{Status: "starting"})
		if err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		conn.WriteJSON(StatusUpdate{Status: "error", Error: err.Error()})
		err = conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		if err == nil {
			donech := make(chan struct{})
			go func() {
				defer close(donech)
				// drain client messages and wait for disconnect
				var e error
				for e == nil {
					_, _, e = conn.ReadMessage()
				}
			}()
			timer := time.NewTimer(10 * time.Second)
			defer timer.Stop()
			select {
			case <-donech:
			case <-timer.C:
			}
		}
	}
	defer func() {
		stopctx, stopcancel := context.WithTimeout(context.Background(), time.Minute)
		defer stopcancel()
		cerr := c.Close(stopctx)
		if cerr != nil {
			log.Printf("Failed to stop container %q: %s", c.ID, cerr)
		}
	}()

	// update status to running
	err = conn.WriteJSON(StatusUpdate{Status: "running"})
	if err != nil {
		return
	}

	// bridge connections
	runctx, cancel := context.WithCancel(context.Background())
	go func() { // output
		defer cancel()
		buf := make([]byte, 1024)
		for {
			n, err := c.IO.Read(buf)
			log.Println(n, err)
			if err != nil {
				return
			}

			err = conn.WriteMessage(websocket.TextMessage, buf[:n])
			if err != nil {
				return
			}
		}
	}()
	go func() { // input
		defer cancel()
		for {
			t, r, err := conn.NextReader()
			if err != nil {
				return
			}
			switch t {
			case websocket.TextMessage:
				fallthrough
			case websocket.BinaryMessage:
				_, err = io.Copy(c.IO, r)
				if err != nil {
					return
				}
			case websocket.CloseMessage:
				return
			}
		}
	}()
	log.Println("Copying")
	<-runctx.Done()
	log.Println("Done")
}
<|endoftext|>


<commit_before>package nftfx

import (
	"encoding/json"
	"errors"

	"github.com/ava-labs/avalanchego/vms/types"

	"github.com/ava-labs/avalanchego/utils/units"
	"github.com/ava-labs/avalanchego/vms/components/verify"
	"github.com/ava-labs/avalanchego/vms/secp256k1fx"
)

const (
	// MaxPayloadSize is the maximum size that can be placed into a payload
	MaxPayloadSize = units.KiB
)

var (
	errNilTransferOutput              = errors.New("nil transfer output")
	errPayloadTooLarge                = errors.New("payload too large")
	_                    verify.State = &TransferOutput{}
)

type TransferOutput struct {
	GroupID                  uint32              `serialize:"true" json:"groupID"`
	Payload                  types.JSONByteSlice `serialize:"true" json:"payload"`
	secp256k1fx.OutputOwners `serialize:"true"`
}

// MarshalJSON marshals Amt and the embedded OutputOwners struct
// into a JSON readable format
// If OutputOwners cannot be serialised then this will return error
func (out *TransferOutput) MarshalJSON() ([]byte, error) {
	result, err := out.OutputOwners.Fields()
	if err != nil {
		return nil, err
	}

	result["groupID"] = out.GroupID
	result["Payload"] = out.Payload
	return json.Marshal(result)
}

func (out *TransferOutput) Verify() error {
	switch {
	case out == nil:
		return errNilTransferOutput
	case len(out.Payload) > MaxPayloadSize:
		return errPayloadTooLarge
	default:
		return out.OutputOwners.Verify()
	}
}

func (out *TransferOutput) VerifyState() error { return out.Verify() }
<commit_msg>Payload ➡ payload<commit_after>package nftfx

import (
	"encoding/json"
	"errors"

	"github.com/ava-labs/avalanchego/vms/types"

	"github.com/ava-labs/avalanchego/utils/units"
	"github.com/ava-labs/avalanchego/vms/components/verify"
	"github.com/ava-labs/avalanchego/vms/secp256k1fx"
)

const (
	// MaxPayloadSize is the maximum size that can be placed into a payload
	MaxPayloadSize = units.KiB
)

var (
	errNilTransferOutput              = errors.New("nil transfer output")
	errPayloadTooLarge                = errors.New("payload too large")
	_                    verify.State = &TransferOutput{}
)

type TransferOutput struct {
	GroupID                  uint32              `serialize:"true" json:"groupID"`
	Payload                  types.JSONByteSlice `serialize:"true" json:"payload"`
	secp256k1fx.OutputOwners `serialize:"true"`
}

// MarshalJSON marshals Amt and the embedded OutputOwners struct
// into a JSON readable format
// If OutputOwners cannot be serialised then this will return error
func (out *TransferOutput) MarshalJSON() ([]byte, error) {
	result, err := out.OutputOwners.Fields()
	if err != nil {
		return nil, err
	}

	result["groupID"] = out.GroupID
	result["payload"] = out.Payload
	return json.Marshal(result)
}

func (out *TransferOutput) Verify() error {
	switch {
	case out == nil:
		return errNilTransferOutput
	case len(out.Payload) > MaxPayloadSize:
		return errPayloadTooLarge
	default:
		return out.OutputOwners.Verify()
	}
}

func (out *TransferOutput) VerifyState() error { return out.Verify() }
<|endoftext|>


<commit_before>package ztls

import (
	"encoding/json"
	"reflect"
	"testing"
)

type ZTLSHandshakeSuite struct{}

func marshalAndUnmarshal(original interface{}, target interface{}) error {
	var b []byte
	var err error
	if b, err = json.Marshal(original); err != nil {
		return err
	}
	if err = json.Unmarshal(b, target); err != nil {
		return err
	}
	return nil
}

func marshalAndUnmarshalAndCheckEquality(original interface{}, target interface{}, t *testing.T) {
	if err := marshalAndUnmarshal(original, target); err != nil {
		t.Fatalf("unable to marshalAndUnmarshal: %s", err.Error())
	}
	if eq := reflect.DeepEqual(original, target); eq != true {
		t.Errorf("expected %+v to equal %+v", original, target)
	}
}

func TestTLSVersionEncodeDecode(t *testing.T) {
	v := TLSVersion(VersionTLS12)
	var dec TLSVersion
	marshalAndUnmarshalAndCheckEquality(&v, &dec, t)
}

func TestCipherSuiteEncodeDecode(t *testing.T) {
	v := CipherSuite(TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256)
	var dec CipherSuite
	marshalAndUnmarshalAndCheckEquality(&v, &dec, t)
}
<commit_msg>Check name in CipherSuite tests<commit_after>package ztls

import (
	"encoding/json"
	"reflect"
	"testing"
)

type ZTLSHandshakeSuite struct{}

func marshalAndUnmarshal(original interface{}, target interface{}) error {
	var b []byte
	var err error
	if b, err = json.Marshal(original); err != nil {
		return err
	}
	if err = json.Unmarshal(b, target); err != nil {
		return err
	}
	return nil
}

func marshalAndUnmarshalAndCheckEquality(original interface{}, target interface{}, t *testing.T) {
	if err := marshalAndUnmarshal(original, target); err != nil {
		t.Fatalf("unable to marshalAndUnmarshal: %s", err.Error())
	}
	if eq := reflect.DeepEqual(original, target); eq != true {
		t.Errorf("expected %+v to equal %+v", original, target)
	}
}

func TestTLSVersionEncodeDecode(t *testing.T) {
	v := TLSVersion(VersionTLS12)
	var dec TLSVersion
	marshalAndUnmarshalAndCheckEquality(&v, &dec, t)
}

func TestCipherSuiteEncodeDecode(t *testing.T) {
	v := CipherSuite(TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256)
	var dec CipherSuite
	marshalAndUnmarshalAndCheckEquality(&v, &dec, t)
	expectedName := nameForSuite(TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256)
	if decodedName := dec.String(); decodedName != expectedName {
		t.Errorf("decoded wrong name, got %s, expected %s", decodedName, expectedName)
	}
}
<|endoftext|>


<commit_before>// +build js

package bytes_test

import (
	"testing"
)

func TestEqualNearPageBoundary(t *testing.T) {
	t.Skip()
}
<commit_msg>compiler/natives/src/bytes: Skip all tests that rely on syscall.Getpagesize.<commit_after>// +build js

package bytes_test

import (
	"testing"
)

func dangerousSlice(t *testing.T) []byte {
	t.Skip("dangerousSlice relies on syscall.Getpagesize, which GopherJS doesn't implement")

	panic("unreachable")
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package simple

import (
	"fmt"
	"strings"

	"k8s.io/perf-tests/clusterloader2/pkg/measurement"
	measurementutil "k8s.io/perf-tests/clusterloader2/pkg/measurement/util"
	"k8s.io/perf-tests/clusterloader2/pkg/util"
)

const (
	memoryProfileName = "MemoryProfile"
	cpuProfileName    = "CPUProfile"
)

func init() {
	measurement.Register(memoryProfileName, createMemoryProfileMeasurement)
	measurement.Register(cpuProfileName, createCPUProfileMeasurement)
}

func createMemoryProfileMeasurement() measurement.Measurement {
	return &memoryProfileMeasurement{}
}

type memoryProfileMeasurement struct{}

// Execute gathers memory profile of a given component.
func (*memoryProfileMeasurement) Execute(config *measurement.MeasurementConfig) ([]measurement.Summary, error) {
	return createMeasurement(config, "heap")
}

// Dispose cleans up after the measurement.
func (*memoryProfileMeasurement) Dispose() {}

// String returns string representation of this measurement.
func (*memoryProfileMeasurement) String() string {
	return memoryProfileName
}

func createCPUProfileMeasurement() measurement.Measurement {
	return &cpuProfileMeasurement{}
}

type cpuProfileMeasurement struct{}

// Execute gathers cpu profile of a given component.
func (*cpuProfileMeasurement) Execute(config *measurement.MeasurementConfig) ([]measurement.Summary, error) {
	return createMeasurement(config, "profile")
}

// Dispose cleans up after the measurement.
func (*cpuProfileMeasurement) Dispose() {}

// String returns string representation of this measurement.
func (*cpuProfileMeasurement) String() string {
	return cpuProfileName
}

func createMeasurement(config *measurement.MeasurementConfig, profileKind string) ([]measurement.Summary, error) {
	var summaries []measurement.Summary
	componentName, err := util.GetString(config.Params, "componentName")
	if err != nil {
		return summaries, err
	}
	provider, err := util.GetStringOrDefault(config.Params, "provider", config.ClusterConfig.Provider)
	if err != nil {
		return summaries, err
	}
	host, err := util.GetStringOrDefault(config.Params, "host", config.ClusterConfig.MasterIP)
	if err != nil {
		return summaries, err
	}

	return gatherProfile(componentName, profileKind, host, provider)
}

func gatherProfile(componentName, profileKind, host, provider string) ([]measurement.Summary, error) {
	var summaries []measurement.Summary
	profilePort, err := getPortForComponent(componentName)
	if err != nil {
		return summaries, fmt.Errorf("profile gathering failed finding component port: %v", err)
	}

	// Get the profile data over SSH.
	getCommand := fmt.Sprintf("curl -s localhost:%v/debug/pprof/%s", profilePort, profileKind)
	sshResult, err := measurementutil.SSH(getCommand, host+":22", provider)
	if err != nil {
		return summaries, fmt.Errorf("failed to execute curl command on master through SSH: %v", err)
	}

	profilePrefix := componentName
	switch {
	case profileKind == "heap":
		profilePrefix += "_MemoryProfile"
	case strings.HasPrefix(profileKind, "profile"):
		profilePrefix += "_CPUProfile"
	default:
		return summaries, fmt.Errorf("unknown profile kind provided: %s", profileKind)
	}

	rawprofile := &profileSummary{
		name:    profilePrefix,
		content: sshResult.Stdout,
	}
	summaries = append(summaries, rawprofile)
	return summaries, nil
}

func getPortForComponent(componentName string) (int, error) {
	switch componentName {
	case "kube-apiserver":
		return 8080, nil
	case "kube-scheduler":
		return 10251, nil
	case "kube-controller-manager":
		return 10252, nil
	}
	return -1, fmt.Errorf("port for component %v unknown", componentName)
}

type profileSummary struct {
	name    string
	content string
}

// SummaryName returns name of the summary.
func (p *profileSummary) SummaryName() string {
	return p.name
}

// PrintSummary returns summary as a string.
func (p *profileSummary) PrintSummary() (string, error) {
	return p.content, nil
}
<commit_msg>logging profile error for gke<commit_after>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package simple

import (
	"fmt"
	"strings"

	"github.com/golang/glog"
	"k8s.io/perf-tests/clusterloader2/pkg/measurement"
	measurementutil "k8s.io/perf-tests/clusterloader2/pkg/measurement/util"
	"k8s.io/perf-tests/clusterloader2/pkg/util"
)

const (
	memoryProfileName = "MemoryProfile"
	cpuProfileName    = "CPUProfile"
)

func init() {
	measurement.Register(memoryProfileName, createMemoryProfileMeasurement)
	measurement.Register(cpuProfileName, createCPUProfileMeasurement)
}

func createMemoryProfileMeasurement() measurement.Measurement {
	return &memoryProfileMeasurement{}
}

type memoryProfileMeasurement struct{}

// Execute gathers memory profile of a given component.
func (c *memoryProfileMeasurement) Execute(config *measurement.MeasurementConfig) ([]measurement.Summary, error) {
	return createMeasurement(c, config, "heap")
}

// Dispose cleans up after the measurement.
func (*memoryProfileMeasurement) Dispose() {}

// String returns string representation of this measurement.
func (*memoryProfileMeasurement) String() string {
	return memoryProfileName
}

func createCPUProfileMeasurement() measurement.Measurement {
	return &cpuProfileMeasurement{}
}

type cpuProfileMeasurement struct{}

// Execute gathers cpu profile of a given component.
func (c *cpuProfileMeasurement) Execute(config *measurement.MeasurementConfig) ([]measurement.Summary, error) {
	return createMeasurement(c, config, "profile")
}

// Dispose cleans up after the measurement.
func (*cpuProfileMeasurement) Dispose() {}

// String returns string representation of this measurement.
func (*cpuProfileMeasurement) String() string {
	return cpuProfileName
}

func createMeasurement(caller measurement.Measurement, config *measurement.MeasurementConfig, profileKind string) ([]measurement.Summary, error) {
	var summaries []measurement.Summary
	componentName, err := util.GetString(config.Params, "componentName")
	if err != nil {
		return summaries, err
	}
	provider, err := util.GetStringOrDefault(config.Params, "provider", config.ClusterConfig.Provider)
	if err != nil {
		return summaries, err
	}
	host, err := util.GetStringOrDefault(config.Params, "host", config.ClusterConfig.MasterIP)
	if err != nil {
		return summaries, err
	}

	return gatherProfile(caller, componentName, profileKind, host, provider)
}

func gatherProfile(caller measurement.Measurement, componentName, profileKind, host, provider string) ([]measurement.Summary, error) {
	var summaries []measurement.Summary
	profilePort, err := getPortForComponent(componentName)
	if err != nil {
		return summaries, fmt.Errorf("profile gathering failed finding component port: %v", err)
	}

	// Get the profile data over SSH.
	getCommand := fmt.Sprintf("curl -s localhost:%v/debug/pprof/%s", profilePort, profileKind)
	sshResult, err := measurementutil.SSH(getCommand, host+":22", provider)
	if err != nil {
		if provider == "gke" {
			// Only logging error for gke. SSHing to gke master is not supported.
			glog.Errorf("%s: failed to execute curl command on master through SSH: %v", caller, err)
			return summaries, nil
		}
		return summaries, fmt.Errorf("failed to execute curl command on master through SSH: %v", err)
	}

	profilePrefix := componentName
	switch {
	case profileKind == "heap":
		profilePrefix += "_MemoryProfile"
	case strings.HasPrefix(profileKind, "profile"):
		profilePrefix += "_CPUProfile"
	default:
		return summaries, fmt.Errorf("unknown profile kind provided: %s", profileKind)
	}

	rawprofile := &profileSummary{
		name:    profilePrefix,
		content: sshResult.Stdout,
	}
	summaries = append(summaries, rawprofile)
	return summaries, nil
}

func getPortForComponent(componentName string) (int, error) {
	switch componentName {
	case "kube-apiserver":
		return 8080, nil
	case "kube-scheduler":
		return 10251, nil
	case "kube-controller-manager":
		return 10252, nil
	}
	return -1, fmt.Errorf("port for component %v unknown", componentName)
}

type profileSummary struct {
	name    string
	content string
}

// SummaryName returns name of the summary.
func (p *profileSummary) SummaryName() string {
	return p.name
}

// PrintSummary returns summary as a string.
func (p *profileSummary) PrintSummary() (string, error) {
	return p.content, nil
}
<|endoftext|>


<commit_before>// +build linux darwin

package logging

import (
	log "gopkg.in/inconshreveable/log15.v2"
)

// getSystemHandler on Linux writes messages to syslog.
func getSystemHandler(syslog string, debug bool, format log.Format) log.Handler {
	// SyslogHandler
	if syslog != "" {
		if !debug {
			return log.LvlFilterHandler(
				log.LvlInfo,
				log.Must.SyslogHandler(syslog, format),
			)
		}

		return log.Must.SyslogHandler(syslog, format)
	}

	return nil
}
<commit_msg>Temporary workaround for log15 API breakage<commit_after>// +build linux darwin

package logging

import (
	slog "log/syslog"

	log "gopkg.in/inconshreveable/log15.v2"
)

// getSystemHandler on Linux writes messages to syslog.
func getSystemHandler(syslog string, debug bool, format log.Format) log.Handler {
	// SyslogHandler
	if syslog != "" {
		if !debug {
			return log.LvlFilterHandler(
				log.LvlInfo,
				log.Must.SyslogHandler(slog.LOG_INFO, syslog, format),
			)
		}

		return log.Must.SyslogHandler(slog.LOG_INFO, syslog, format)
	}

	return nil
}
<|endoftext|>


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

// Package cli contains methods useful for implementing administrative command
// line utilities.
package cli

import (
	"context"
	"fmt"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
	"time"

	log "github.com/golang/glog"
	"github.com/golang/protobuf/ptypes"
	"google.golang.org/grpc"

	"github.com/google/fleetspeak/fleetspeak/src/admin/history"
	"github.com/google/fleetspeak/fleetspeak/src/common"

	sgrpc "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
	spb "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
)

// dateFmt is a fairly dense, 23 character date format string, suitable for
// tabular date information.
const dateFmt = "15:04:05.000 2006.01.02"

// Usage prints usage information describing the command line flags and behavior
// of programs based on Execute.
func Usage() {
	n := path.Base(os.Args[0])
	fmt.Fprintf(os.Stderr,
		"Usage:\n"+
			"    %s listclients\n"+
			"    %s listcontacts <client_id> [limit]\n"+
			"    %s analysehistory <client_id>\n"+
			"    %s blacklistclient <client_id>\n"+
			"\n", n, n, n, n)
}

// Execute examines command line flags and executes one of the standard command line
// actions, as summarized by Usage. It requires a grpc connection to an admin server
// and the command line parameters to interpret.
func Execute(conn *grpc.ClientConn, args ...string) {
	admin := sgrpc.NewAdminClient(conn)

	if len(args) == 0 {
		fmt.Fprint(os.Stderr, "A command is required.\n")
		Usage()
		os.Exit(1)
	}

	switch args[0] {
	case "listclients":
		ListClients(admin, args[1:]...)
	case "listcontacts":
		ListContacts(admin, args[1:]...)
	case "analysehistory":
		AnalyseHistory(admin, args[1:]...)
	case "blacklistclient":
		BlacklistClient(admin, args[1:]...)
	default:
		fmt.Fprintf(os.Stderr, "Unknown command: %v\n", args[0])
		Usage()
		os.Exit(1)
	}
}

// ListClients prints a list of all clients in the fleetspeak system.
func ListClients(c sgrpc.AdminClient, args ...string) {
	if len(args) > 0 {
		Usage()
		os.Exit(1)
	}
	ctx := context.Background()
	res, err := c.ListClients(ctx, &spb.ListClientsRequest{}, grpc.MaxCallRecvMsgSize(1024*1024*1024))
	if err != nil {
		log.Exitf("ListClients RPC failed: %v", err)
	}
	if len(res.Clients) == 0 {
		fmt.Println("No clients found.")
		return
	}
	sort.Sort(byContactTime(res.Clients))
	fmt.Printf("%-16s %-23s %s\n", "Client ID:", "Last Seen:", "Labels:")
	for _, cl := range res.Clients {
		id, err := common.BytesToClientID(cl.ClientId)
		if err != nil {
			log.Errorf("Ignoring invalid client id [%v], %v", cl.ClientId, err)
			continue
		}
		var ls []string
		for _, l := range cl.Labels {
			ls = append(ls, l.ServiceName+":"+l.Label)
		}
		ts, err := ptypes.Timestamp(cl.LastContactTime)
		if err != nil {
			log.Errorf("Unable to parse last contact time for %v: %v", id, err)
		}
		tag := ""
		if cl.Blacklisted {
			tag = " *blacklisted*"
		}
		fmt.Printf("%v %v [%v]%s\n", id, ts.Format(dateFmt), strings.Join(ls, ","), tag)
	}
}

// byContactTime adapts []*spb.Client for use by sort.Sort.
type byContactTime []*spb.Client

func (b byContactTime) Len() int           { return len(b) }
func (b byContactTime) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
func (b byContactTime) Less(i, j int) bool { return contactTime(b[i]).Before(contactTime(b[j])) }

func contactTime(c *spb.Client) time.Time {
	return time.Unix(c.LastContactTime.Seconds, int64(c.LastContactTime.Nanos))
}

// ListContacts prints a list contacts that the system has recorded for a
// client. args[0] must be a client id. If present, args[1] limits to the most
// recent args[1] contacts.
func ListContacts(c sgrpc.AdminClient, args ...string) {
	if len(args) == 0 || len(args) > 2 {
		Usage()
		os.Exit(1)
	}
	id, err := common.StringToClientID(args[0])
	if err != nil {
		log.Exitf("Unable to parse %s as client id: %v", args[0], err)
	}
	var lim int
	if len(args) == 2 {
		lim, err = strconv.Atoi(args[1])
		if err != nil {
			log.Exitf("Unable to parse %s as a limit: %v", args[1], err)
		}
	}

	ctx := context.Background()
	res, err := c.ListClientContacts(ctx, &spb.ListClientContactsRequest{ClientId: id.Bytes()}, grpc.MaxCallRecvMsgSize(1024*1024*1024))
	if err != nil {
		log.Exitf("ListClientContacts RPC failed: %v", err)
	}
	if len(res.Contacts) == 0 {
		fmt.Println("No contacts found.")
		return
	}

	fmt.Printf("Found %d contacts.\n", len(res.Contacts))

	sort.Sort(byTimestamp(res.Contacts))
	fmt.Printf("%-23s %s", "Timestamp:", "Observed IP:\n")
	for i, con := range res.Contacts {
		if lim > 0 && i > lim {
			break
		}
		ts, err := ptypes.Timestamp(con.Timestamp)
		if err != nil {
			log.Errorf("Unable to parse timestamp for contact: %v", err)
			continue
		}
		fmt.Printf("%s %s\n", ts.Format(dateFmt), con.ObservedAddress)
	}
}

// byTimestamp adapts []*spb.ClientContact for use by sort.Sort.
type byTimestamp []*spb.ClientContact

func (b byTimestamp) Len() int           { return len(b) }
func (b byTimestamp) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
func (b byTimestamp) Less(i, j int) bool { return timestamp(b[i]).Before(timestamp(b[j])) }

func timestamp(c *spb.ClientContact) time.Time {
	return time.Unix(c.Timestamp.Seconds, int64(c.Timestamp.Nanos))
}

// AnalyseHistory prints a summary analysis of a client's history. args[0] must
// be a client id.
func AnalyseHistory(c sgrpc.AdminClient, args ...string) {
	if len(args) != 1 {
		Usage()
		os.Exit(1)
	}
	id, err := common.StringToClientID(args[0])
	if err != nil {
		log.Exitf("Unable to parse %s as client id: %v", args[0], err)
	}
	ctx := context.Background()
	res, err := c.ListClientContacts(ctx, &spb.ListClientContactsRequest{ClientId: id.Bytes()}, grpc.MaxCallRecvMsgSize(1024*1024*1024))
	if err != nil {
		log.Exitf("ListClientContacts RPC failed: %v", err)
	}
	if len(res.Contacts) == 0 {
		fmt.Println("No contacts found.")
		return
	}
	s, err := history.Summarize(res.Contacts)
	if err != nil {
		log.Exitf("Error creating summary: %v", err)
	}
	fmt.Printf(`Raw Summary:
  First Recorded Contact: %v
  Last Recorded Contact: %v
  Contact Count: %d
  Observed IP Count: %d
  Split Points: %d
  Splits: %d
  Skips: %d
`, s.Start, s.End, s.Count, s.IPCount, s.SplitPoints, s.Splits, s.Skips)
	if s.Splits > 0 {
		fmt.Printf("This client appears to have be restored %d times from %d different backup images.\n", s.Splits, s.SplitPoints)
	}
	if s.Skips > s.Splits {
		fmt.Printf("Observed %d Skips, but only %d splits. The machine may have been cloned.\n", s.Skips, s.Splits)
	}
}

// BlacklistClient blacklists a client id, forcing any client(s) using it to
// rekey. args[0] must be a client id.
func BlacklistClient(c sgrpc.AdminClient, args ...string) {
	if len(args) != 1 {
		Usage()
		os.Exit(1)
	}
	id, err := common.StringToClientID(args[0])
	if err != nil {
		log.Exitf("Unable to parse %s as client id: %v", args[0], err)
	}
	ctx := context.Background()
	if _, err := c.BlacklistClient(ctx, &spb.BlacklistClientRequest{ClientId: id.Bytes()}); err != nil {
		log.Exitf("BlacklistClient RPC failed: %v", err)
	}
}
<commit_msg>Reverse contact order.<commit_after>// Copyright 2017 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package cli contains methods useful for implementing administrative command
// line utilities.
package cli

import (
	"context"
	"fmt"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
	"time"

	log "github.com/golang/glog"
	"github.com/golang/protobuf/ptypes"
	"google.golang.org/grpc"

	"github.com/google/fleetspeak/fleetspeak/src/admin/history"
	"github.com/google/fleetspeak/fleetspeak/src/common"

	sgrpc "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
	spb "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
)

// dateFmt is a fairly dense, 23 character date format string, suitable for
// tabular date information.
const dateFmt = "15:04:05.000 2006.01.02"

// Usage prints usage information describing the command line flags and behavior
// of programs based on Execute.
func Usage() {
	n := path.Base(os.Args[0])
	fmt.Fprintf(os.Stderr,
		"Usage:\n"+
			"    %s listclients\n"+
			"    %s listcontacts <client_id> [limit]\n"+
			"    %s analysehistory <client_id>\n"+
			"    %s blacklistclient <client_id>\n"+
			"\n", n, n, n, n)
}

// Execute examines command line flags and executes one of the standard command line
// actions, as summarized by Usage. It requires a grpc connection to an admin server
// and the command line parameters to interpret.
func Execute(conn *grpc.ClientConn, args ...string) {
	admin := sgrpc.NewAdminClient(conn)

	if len(args) == 0 {
		fmt.Fprint(os.Stderr, "A command is required.\n")
		Usage()
		os.Exit(1)
	}

	switch args[0] {
	case "listclients":
		ListClients(admin, args[1:]...)
	case "listcontacts":
		ListContacts(admin, args[1:]...)
	case "analysehistory":
		AnalyseHistory(admin, args[1:]...)
	case "blacklistclient":
		BlacklistClient(admin, args[1:]...)
	default:
		fmt.Fprintf(os.Stderr, "Unknown command: %v\n", args[0])
		Usage()
		os.Exit(1)
	}
}

// ListClients prints a list of all clients in the fleetspeak system.
func ListClients(c sgrpc.AdminClient, args ...string) {
	if len(args) > 0 {
		Usage()
		os.Exit(1)
	}
	ctx := context.Background()
	res, err := c.ListClients(ctx, &spb.ListClientsRequest{}, grpc.MaxCallRecvMsgSize(1024*1024*1024))
	if err != nil {
		log.Exitf("ListClients RPC failed: %v", err)
	}
	if len(res.Clients) == 0 {
		fmt.Println("No clients found.")
		return
	}
	sort.Sort(byContactTime(res.Clients))
	fmt.Printf("%-16s %-23s %s\n", "Client ID:", "Last Seen:", "Labels:")
	for _, cl := range res.Clients {
		id, err := common.BytesToClientID(cl.ClientId)
		if err != nil {
			log.Errorf("Ignoring invalid client id [%v], %v", cl.ClientId, err)
			continue
		}
		var ls []string
		for _, l := range cl.Labels {
			ls = append(ls, l.ServiceName+":"+l.Label)
		}
		ts, err := ptypes.Timestamp(cl.LastContactTime)
		if err != nil {
			log.Errorf("Unable to parse last contact time for %v: %v", id, err)
		}
		tag := ""
		if cl.Blacklisted {
			tag = " *blacklisted*"
		}
		fmt.Printf("%v %v [%v]%s\n", id, ts.Format(dateFmt), strings.Join(ls, ","), tag)
	}
}

// byContactTime adapts []*spb.Client for use by sort.Sort. Places most recent
// contacts first.
type byContactTime []*spb.Client

func (b byContactTime) Len() int           { return len(b) }
func (b byContactTime) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
func (b byContactTime) Less(i, j int) bool { return contactTime(b[i]).Before(contactTime(b[j])) }

func contactTime(c *spb.Client) time.Time {
	return time.Unix(c.LastContactTime.Seconds, int64(c.LastContactTime.Nanos))
}

// ListContacts prints a list contacts that the system has recorded for a
// client. args[0] must be a client id. If present, args[1] limits to the most
// recent args[1] contacts.
func ListContacts(c sgrpc.AdminClient, args ...string) {
	if len(args) == 0 || len(args) > 2 {
		Usage()
		os.Exit(1)
	}
	id, err := common.StringToClientID(args[0])
	if err != nil {
		log.Exitf("Unable to parse %s as client id: %v", args[0], err)
	}
	var lim int
	if len(args) == 2 {
		lim, err = strconv.Atoi(args[1])
		if err != nil {
			log.Exitf("Unable to parse %s as a limit: %v", args[1], err)
		}
	}

	ctx := context.Background()
	res, err := c.ListClientContacts(ctx, &spb.ListClientContactsRequest{ClientId: id.Bytes()}, grpc.MaxCallRecvMsgSize(1024*1024*1024))
	if err != nil {
		log.Exitf("ListClientContacts RPC failed: %v", err)
	}
	if len(res.Contacts) == 0 {
		fmt.Println("No contacts found.")
		return
	}

	fmt.Printf("Found %d contacts.\n", len(res.Contacts))

	sort.Sort(byTimestamp(res.Contacts))
	fmt.Printf("%-23s %s", "Timestamp:", "Observed IP:\n")
	for i, con := range res.Contacts {
		if lim > 0 && i > lim {
			break
		}
		ts, err := ptypes.Timestamp(con.Timestamp)
		if err != nil {
			log.Errorf("Unable to parse timestamp for contact: %v", err)
			continue
		}
		fmt.Printf("%s %s\n", ts.Format(dateFmt), con.ObservedAddress)
	}
}

// byTimestamp adapts []*spb.ClientContact for use by sort.Sort.
type byTimestamp []*spb.ClientContact

func (b byTimestamp) Len() int           { return len(b) }
func (b byTimestamp) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
func (b byTimestamp) Less(i, j int) bool { return timestamp(b[i]).After(timestamp(b[j])) }

func timestamp(c *spb.ClientContact) time.Time {
	return time.Unix(c.Timestamp.Seconds, int64(c.Timestamp.Nanos))
}

// AnalyseHistory prints a summary analysis of a client's history. args[0] must
// be a client id.
func AnalyseHistory(c sgrpc.AdminClient, args ...string) {
	if len(args) != 1 {
		Usage()
		os.Exit(1)
	}
	id, err := common.StringToClientID(args[0])
	if err != nil {
		log.Exitf("Unable to parse %s as client id: %v", args[0], err)
	}
	ctx := context.Background()
	res, err := c.ListClientContacts(ctx, &spb.ListClientContactsRequest{ClientId: id.Bytes()}, grpc.MaxCallRecvMsgSize(1024*1024*1024))
	if err != nil {
		log.Exitf("ListClientContacts RPC failed: %v", err)
	}
	if len(res.Contacts) == 0 {
		fmt.Println("No contacts found.")
		return
	}
	s, err := history.Summarize(res.Contacts)
	if err != nil {
		log.Exitf("Error creating summary: %v", err)
	}
	fmt.Printf(`Raw Summary:
  First Recorded Contact: %v
  Last Recorded Contact: %v
  Contact Count: %d
  Observed IP Count: %d
  Split Points: %d
  Splits: %d
  Skips: %d
`, s.Start, s.End, s.Count, s.IPCount, s.SplitPoints, s.Splits, s.Skips)
	if s.Splits > 0 {
		fmt.Printf("This client appears to have be restored %d times from %d different backup images.\n", s.Splits, s.SplitPoints)
	}
	if s.Skips > s.Splits {
		fmt.Printf("Observed %d Skips, but only %d splits. The machine may have been cloned.\n", s.Skips, s.Splits)
	}
}

// BlacklistClient blacklists a client id, forcing any client(s) using it to
// rekey. args[0] must be a client id.
func BlacklistClient(c sgrpc.AdminClient, args ...string) {
	if len(args) != 1 {
		Usage()
		os.Exit(1)
	}
	id, err := common.StringToClientID(args[0])
	if err != nil {
		log.Exitf("Unable to parse %s as client id: %v", args[0], err)
	}
	ctx := context.Background()
	if _, err := c.BlacklistClient(ctx, &spb.BlacklistClientRequest{ClientId: id.Bytes()}); err != nil {
		log.Exitf("BlacklistClient RPC failed: %v", err)
	}
}
<|endoftext|>


<commit_before>package ifelse
import "github.com/qlova/ilang/src"

func init() {
	ilang.RegisterToken([]string{
		"if",
	}, ScanIf)

	ilang.RegisterToken([]string{
		"else",
	}, ScanElse)

	ilang.RegisterToken([]string{
		"elseif",
	}, ScanElseIf)

	ilang.RegisterListener(If, IfEnd)
	ilang.RegisterListener(Else, ElseIfEnd)
	ilang.RegisterListener(ElseIf, ElseIfEnd)
}

var If = ilang.NewFlag()
var Else = ilang.NewFlag()
var ElseIf = ilang.NewFlag()

func ElseIfEnd(ic *ilang.Compiler) {
	nesting, ok := ic.Scope[len(ic.Scope)-1]["flag_nesting"]
	if ok {
		for i:=0; i < nesting.Int; i++ {
			ic.Assembly("END")
		}
	}
	ic.LoseScope()
	ic.Assembly("END")
}

func IfEnd(ic *ilang.Compiler) {
	ic.Assembly("END")
}

func ScanIf(ic *ilang.Compiler) {
	var expression = ic.ScanExpression()
	if ic.ExpressionType != ilang.Number {
		ic.RaiseError("if statements must have numeric conditions!")
	}
	ic.Assembly("IF ", expression)
	ic.GainScope()
	ic.SetFlag(If)
}

func ScanElse(ic *ilang.Compiler) {
	nesting, ok := ic.Scope[len(ic.Scope)-1]["flag_nesting"]
	if !ok {
		nesting.Int = 0
	}

	if ic.GetScopedFlag(If) {
		ic.UnsetFlag(If)
	} else if ic.GetScopedFlag(ElseIf) {
		ic.UnsetFlag(ElseIf)
		ic.LoseScope()

		nesting, ok = ic.Scope[len(ic.Scope)-1]["flag_nesting"]
		if !ok {
			nesting.Int = 0
		}

	} else {
		ic.RaiseError("You cannot have an else without an if!")
	}

	ic.LoseScope()
	ic.Assembly("ELSE")
	ic.GainScope()
	ic.SetVariable("flag_nesting", ilang.Type{Int:nesting.Int})
	ic.GainScope()
	ic.SetFlag(Else)
}

func ScanElseIf(ic *ilang.Compiler) {
	nesting, ok := ic.Scope[len(ic.Scope)-1]["flag_nesting"]
	if !ok {
		nesting.Int = 0
	}

	if ic.GetScopedFlag(If) {
		ic.UnsetFlag(If)
	} else if ic.GetScopedFlag(ElseIf) {
		ic.UnsetFlag(ElseIf)
		ic.LoseScope()

		nesting, ok = ic.Scope[len(ic.Scope)-1]["flag_nesting"]
		if !ok {
			nesting.Int = 0
		}

	} else if ic.GetScopedFlag(Else) {
		ic.RaiseError("Cannot use ifelse after an else...")
	} else {
		ic.RaiseError("Cannot use ifelse without an if!")
	}

	ic.LoseScope()
	ic.Assembly("ELSE")
	var expression = ic.ScanExpression()
	ic.Assembly("IF ", expression)
	ic.GainScope()
	ic.SetVariable("flag_nesting", ilang.Type{Int:nesting.Int+1})
	ic.GainScope()
	ic.SetFlag(ElseIf)
}
<commit_msg>Allow if statements to take an array or list as condition.<commit_after>package ifelse
import "github.com/qlova/ilang/src"

func init() {
	ilang.RegisterToken([]string{
		"if",
	}, ScanIf)

	ilang.RegisterToken([]string{
		"else",
	}, ScanElse)

	ilang.RegisterToken([]string{
		"elseif",
	}, ScanElseIf)

	ilang.RegisterListener(If, IfEnd)
	ilang.RegisterListener(Else, ElseIfEnd)
	ilang.RegisterListener(ElseIf, ElseIfEnd)
}

var If = ilang.NewFlag()
var Else = ilang.NewFlag()
var ElseIf = ilang.NewFlag()

func ElseIfEnd(ic *ilang.Compiler) {
	nesting, ok := ic.Scope[len(ic.Scope)-1]["flag_nesting"]
	if ok {
		for i:=0; i < nesting.Int; i++ {
			ic.Assembly("END")
		}
	}
	ic.LoseScope()
	ic.Assembly("END")
}

func IfEnd(ic *ilang.Compiler) {
	ic.Assembly("END")
}

func ScanIf(ic *ilang.Compiler) {
	var expression = ic.ScanExpression()
	if !ic.ExpressionType.Equals(ilang.Number) && !ic.ExpressionType.Equals(ilang.Text) {
		ic.RaiseError("if statements must have numeric conditions!")
	}

	var condition = expression
	if ic.ExpressionType.Equals(ilang.Text) {
		condition = "#"+condition
	}

	ic.Assembly("IF ", condition)
	ic.GainScope()
	ic.SetFlag(If)
}

func ScanElse(ic *ilang.Compiler) {
	nesting, ok := ic.Scope[len(ic.Scope)-1]["flag_nesting"]
	if !ok {
		nesting.Int = 0
	}

	if ic.GetScopedFlag(If) {
		ic.UnsetFlag(If)
	} else if ic.GetScopedFlag(ElseIf) {
		ic.UnsetFlag(ElseIf)
		ic.LoseScope()

		nesting, ok = ic.Scope[len(ic.Scope)-1]["flag_nesting"]
		if !ok {
			nesting.Int = 0
		}

	} else {
		ic.RaiseError("You cannot have an else without an if!")
	}

	ic.LoseScope()
	ic.Assembly("ELSE")
	ic.GainScope()
	ic.SetVariable("flag_nesting", ilang.Type{Int:nesting.Int})
	ic.GainScope()
	ic.SetFlag(Else)
}

func ScanElseIf(ic *ilang.Compiler) {
	nesting, ok := ic.Scope[len(ic.Scope)-1]["flag_nesting"]
	if !ok {
		nesting.Int = 0
	}

	if ic.GetScopedFlag(If) {
		ic.UnsetFlag(If)
	} else if ic.GetScopedFlag(ElseIf) {
		ic.UnsetFlag(ElseIf)
		ic.LoseScope()

		nesting, ok = ic.Scope[len(ic.Scope)-1]["flag_nesting"]
		if !ok {
			nesting.Int = 0
		}

	} else if ic.GetScopedFlag(Else) {
		ic.RaiseError("Cannot use ifelse after an else...")
	} else {
		ic.RaiseError("Cannot use ifelse without an if!")
	}

	ic.LoseScope()
	ic.Assembly("ELSE")
	var expression = ic.ScanExpression()
	ic.Assembly("IF ", expression)
	ic.GainScope()
	ic.SetVariable("flag_nesting", ilang.Type{Int:nesting.Int+1})
	ic.GainScope()
	ic.SetFlag(ElseIf)
}
<|endoftext|>


<commit_before>package slackbot

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"

	cloudbuild "google.golang.org/api/cloudbuild/v1"
)

// Notify posts a notification to Slack that the build is complete.
func Notify(b *cloudbuild.Build, webhook string, project string) {
	url := fmt.Sprintf("https://console.cloud.google.com/cloud-build/builds/%s", b.Id)
	var i string
	switch b.Status {
	case "WORKING":
		i = ":hammer:"
	case "SUCCESS":
		i = ":white_check_mark:"
	case "FAILURE":
		i = ":x:"
	case "CANCELLED":
		i = ":wastebasket:"
	case "TIMEOUT":
		i = ":hourglass:"
	case "STATUS_UNKNOWN", "INTERNAL_ERROR":
		i = ":interrobang:"
	default:
		i = ":question:"
	}

	startTime, err := time.Parse(time.RFC3339, b.StartTime)
	if err != nil {
		log.Fatalf("Failed to parse Build.StartTime: %v", err)
	}
	finishTime, err := time.Parse(time.RFC3339, b.FinishTime)
	if err != nil {
		log.Fatalf("Failed to parse Build.FinishTime: %v", err)
	}
	buildDuration := finishTime.Sub(startTime).Truncate(time.Second)

	var msg string
	if b.Status == "WORKING" {
		msgFmt := `{
			"text": "%s *%s* build started",
			"attachments": [{
				"fallback": "Open build details at %s",
				"actions": [{
					"type": "button",
					"text": "Open details",
					"url": "%s"
				}]
			}]
		}`
		msg = fmt.Sprintf(msgFmt, i, project, url, url)
	} else {
		msgFmt := `{
			"text": "%s *%s* build _%s_ after _%s_",
			"attachments": [{
				"fallback": "Open build details at %s",
				"actions": [{
					"type": "button",
					"text": "Open details",
					"url": "%s"
				}]
			}]
		}`
		msg = fmt.Sprintf(msgFmt, i, project, b.Status, buildDuration, url, url)
	}

	r := strings.NewReader(msg)
	resp, err := http.Post(webhook, "application/json", r)
	if err != nil {
		log.Fatalf("Failed to post to Slack: %v", err)
	}
	defer resp.Body.Close()
	body, _ := ioutil.ReadAll(resp.Body)
	log.Printf("Posted message to Slack: [%v], got response [%s]", msg, body)
}
<commit_msg>Do not try to use finish time when WORKING is used<commit_after>package slackbot

import (
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"time"

	cloudbuild "google.golang.org/api/cloudbuild/v1"
)

// Notify posts a notification to Slack that the build is complete.
func Notify(b *cloudbuild.Build, webhook string, project string) {
	url := fmt.Sprintf("https://console.cloud.google.com/cloud-build/builds/%s", b.Id)
	var i string
	switch b.Status {
	case "WORKING":
		i = ":hammer:"
	case "SUCCESS":
		i = ":white_check_mark:"
	case "FAILURE":
		i = ":x:"
	case "CANCELLED":
		i = ":wastebasket:"
	case "TIMEOUT":
		i = ":hourglass:"
	case "STATUS_UNKNOWN", "INTERNAL_ERROR":
		i = ":interrobang:"
	default:
		i = ":question:"
	}

	var msg string
	if b.Status == "WORKING" {
		msgFmt := `{
			"text": "%s *%s* build started",
			"attachments": [{
				"fallback": "Open build details at %s",
				"actions": [{
					"type": "button",
					"text": "Open details",
					"url": "%s"
				}]
			}]
		}`
		msg = fmt.Sprintf(msgFmt, i, project, url, url)
	} else {
		startTime, err := time.Parse(time.RFC3339, b.StartTime)
		if err != nil {
			log.Fatalf("Failed to parse Build.StartTime: %v", err)
		}
		finishTime, err := time.Parse(time.RFC3339, b.FinishTime)
		if err != nil {
			log.Fatalf("Failed to parse Build.FinishTime: %v", err)
		}
		buildDuration := finishTime.Sub(startTime).Truncate(time.Second)

		msgFmt := `{
			"text": "%s *%s* build _%s_ after _%s_",
			"attachments": [{
				"fallback": "Open build details at %s",
				"actions": [{
					"type": "button",
					"text": "Open details",
					"url": "%s"
				}]
			}]
		}`
		msg = fmt.Sprintf(msgFmt, i, project, b.Status, buildDuration, url, url)
	}

	r := strings.NewReader(msg)
	resp, err := http.Post(webhook, "application/json", r)
	if err != nil {
		log.Fatalf("Failed to post to Slack: %v", err)
	}
	defer resp.Body.Close()
	body, _ := ioutil.ReadAll(resp.Body)
	log.Printf("Posted message to Slack: [%v], got response [%s]", msg, body)
}
<|endoftext|>