<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/BurntSushi/toml"
	_ "github.com/go-sql-driver/mysql"
	"os"
)

type Config struct {
	Access   Access
	Database []Database
	User     []User
}

const selectAllUsers = "SELECT User, Host FROM mysql.user"
const showAllDatabases = "SHOW DATABASES"

func main() {
	var execute bool
	var configFile string

	flag.BoolVar(&execute, "e", false, "Execute. If specified - changes will be applied")
	flag.StringVar(&configFile, "c", "/etc/mmdu/mmdu.toml", "Path to config file.")
	flag.Parse()

	var conf Config
	if _, err := toml.DecodeFile(configFile, &conf); err != nil {
		fmt.Println("Failed to parse config file", err.Error())
		os.Exit(1)
	}

	defaultDatabases := []Database{Database{"information_schema"}, Database{"mysql"}, Database{"performance_schema"}}
	validatedUsers, err := validateUsers(conf.User)
	if err != nil {
		fmt.Println("Error during validation of user list:", err.Error())
		os.Exit(1)
	}

	db := conf.Access.connectAndCheck()
	defer db.Close()

	usersFromDB, err := getAllUsersFromDB(db)
	if err != nil {
		fmt.Println("Failed to get users from DB" , err.Error())
		os.Exit(2)
	}

	databasesFromDB, err := getDatabasesFromDB(db)
	if err != nil {
		fmt.Println("Failed to get databases from DB", err.Error())
		os.Exit(2)
	}

	// Merge from 3 sources: default, from DbConfig and from UserConfg
	databasesFromConf := removeDuplicateDatabases(
		append(append(defaultDatabases, conf.Database...), getDatabasesFromUsers(validatedUsers)...))

	usersToRemove := getUsersToRemove(validatedUsers, usersFromDB)
	usersToAdd := getUsersToAdd(validatedUsers, usersFromDB)

	databasesToRemove := getDatabasesToRemove(databasesFromConf, databasesFromDB)
	databasesToAdd := getDatabasesToAdd(databasesFromConf, databasesFromDB)

	if len(usersToRemove) == 0 && len(usersToAdd) == 0 && len(databasesToRemove) == 0 && len(databasesToAdd) == 0 {
		fmt.Println("Nothing to do")
	} else {
		tx, err := db.Begin()
		if err != nil {
			fmt.Println("Failed to start transaction", err.Error())
			os.Exit(2)
		}

		for _, database := range databasesToRemove {
			if !database.dropDatabase(tx, execute) {
				tx.Rollback()
			}
		}

		for _, database := range databasesToAdd {
			if !database.addDatabase(tx, execute) {
				tx.Rollback()
			}
		}

		for _, user := range usersToRemove {
			if !user.dropUser(tx, execute) {
				tx.Rollback()
			}
		}

		for _, user := range usersToAdd {
			if !user.addUser(tx, execute) {
				tx.Rollback()
			}
		}
		tx.Commit()
	}
}
<commit_msg>Fixed comment<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/BurntSushi/toml"
	_ "github.com/go-sql-driver/mysql"
	"os"
)

type Config struct {
	Access   Access
	Database []Database
	User     []User
}

const selectAllUsers = "SELECT User, Host FROM mysql.user"
const showAllDatabases = "SHOW DATABASES"

func main() {
	var execute bool
	var configFile string

	flag.BoolVar(&execute, "e", false, "Execute. If specified - changes will be applied")
	flag.StringVar(&configFile, "c", "/etc/mmdu/mmdu.toml", "Path to config file.")
	flag.Parse()

	var conf Config
	if _, err := toml.DecodeFile(configFile, &conf); err != nil {
		fmt.Println("Failed to parse config file", err.Error())
		os.Exit(1)
	}

	defaultDatabases := []Database{Database{"information_schema"}, Database{"mysql"}, Database{"performance_schema"}}
	validatedUsers, err := validateUsers(conf.User)
	if err != nil {
		fmt.Println("Error during validation of user list:", err.Error())
		os.Exit(1)
	}

	db := conf.Access.connectAndCheck()
	defer db.Close()

	usersFromDB, err := getAllUsersFromDB(db)
	if err != nil {
		fmt.Println("Failed to get users from DB" , err.Error())
		os.Exit(2)
	}

	databasesFromDB, err := getDatabasesFromDB(db)
	if err != nil {
		fmt.Println("Failed to get databases from DB", err.Error())
		os.Exit(2)
	}

	// Merge from 3 sources: default, from DbConfig and from UserConfig
	databasesFromConf := removeDuplicateDatabases(
		append(append(defaultDatabases, conf.Database...), getDatabasesFromUsers(validatedUsers)...))

	usersToRemove := getUsersToRemove(validatedUsers, usersFromDB)
	usersToAdd := getUsersToAdd(validatedUsers, usersFromDB)

	databasesToRemove := getDatabasesToRemove(databasesFromConf, databasesFromDB)
	databasesToAdd := getDatabasesToAdd(databasesFromConf, databasesFromDB)

	if len(usersToRemove) == 0 && len(usersToAdd) == 0 && len(databasesToRemove) == 0 && len(databasesToAdd) == 0 {
		fmt.Println("Nothing to do")
	} else {
		tx, err := db.Begin()
		if err != nil {
			fmt.Println("Failed to start transaction", err.Error())
			os.Exit(2)
		}

		for _, database := range databasesToRemove {
			if !database.dropDatabase(tx, execute) {
				tx.Rollback()
			}
		}

		for _, database := range databasesToAdd {
			if !database.addDatabase(tx, execute) {
				tx.Rollback()
			}
		}

		for _, user := range usersToRemove {
			if !user.dropUser(tx, execute) {
				tx.Rollback()
			}
		}

		for _, user := range usersToAdd {
			if !user.addUser(tx, execute) {
				tx.Rollback()
			}
		}
		_, err = tx.Exec("select * from mysql.user ; select * from mysql.user;")
		tx.Commit()
	}
}
<|endoftext|>


<commit_before>package mixpanel

import (
	"errors"
	"fmt"
	"time"
)

// Mocked version of Mixpanel which can be used in unit tests.
type Mock struct {
	// All People identified, mapped by distinctId
	People map[string]*MockPeople
}

func NewMock() *Mock {
	return &Mock{
		People: map[string]*MockPeople{},
	}
}

func (m *Mock) String() string {
	str := ""
	for id, p := range m.People {
		str += id + ":\n" + p.String()
	}
	return str
}

// Identifies a user. The user will be added to the People map.
func (m *Mock) people(distinctId string) *MockPeople {
	p := m.People[distinctId]
	if p == nil {
		p = &MockPeople{
			Properties: map[string]interface{}{},
		}
		m.People[distinctId] = p
	}

	return p
}

func (m *Mock) Track(distinctId, eventName string, e *Event) error {
	p := m.people(distinctId)
	p.Events = append(p.Events, MockEvent{
		Event: *e,
		Name:  eventName,
	})
	return nil
}

func (m *Mock) Import(distinctId, eventName string, e *Event) error {
	p := m.people(distinctId)
	p.Events = append(p.Events, MockEvent{
		Event: *e,
		Name:  eventName,
	})
	return nil
}

type MockPeople struct {
	Properties map[string]interface{}
	Time       *time.Time
	IP         string
	Events     []MockEvent
}

func (mp *MockPeople) String() string {
	timeStr := ""
	if mp.Time != nil {
		timeStr = mp.Time.Format(time.RFC3339)
	}

	str := fmt.Sprintf("  ip: %s\n  time: %s\n", mp.IP, timeStr)
	str += "  properties:\n"
	for key, val := range mp.Properties {
		str += fmt.Sprintf("    %s: %v\n", key, val)
	}
	str += "  events:\n"
	for _, event := range mp.Events {
		str += "    " + event.Name + ":\n"
		str += fmt.Sprintf("      IP: %s\n", event.IP)
		if event.Timestamp != nil {
			str += fmt.Sprintf(
				"      Timestamp: %s\n", event.Timestamp.Format(time.RFC3339),
			)
		} else {
			str += "      Timestamp:\n"
		}
		for key, val := range event.Properties {
			str += fmt.Sprintf("      %s: %v\n", key, val)
		}
	}
	return str
}

func (m *Mock) Update(distinctId string, u *Update) error {
	return m.UpdateUser(distinctId, u)
}

func (m *Mock) UpdateUser(distinctId string, u *Update) error {
	p := m.people(distinctId)

	if u.IP != "" {
		p.IP = u.IP
	}
	if u.Timestamp != nil && u.Timestamp != IgnoreTime {
		p.Time = u.Timestamp
	}

	switch u.Operation {
	case "$set", "$set_once":
		for key, val := range u.Properties {
			p.Properties[key] = val
		}
	default:
		return errors.New("mixpanel.Mock only supports the $set and $set_once operations")
	}

	return nil
}

func (m *Mock) UnionUser(userID string, u *Update) error {
	return nil
}

func (m *Mock) UpdateGroup(groupKey, groupUser string, u *Update) error {
	return nil
}

func (m *Mock) UnionGroup(groupKey, groupUser string, u *Update) error {
	return nil
}

func (m *Mock) Alias(distinctId, newId string) error {
	return nil
}

type MockEvent struct {
	Event
	Name string
}
<commit_msg>use thread safe map in the mock<commit_after>package mixpanel

import (
	"errors"
	"fmt"
	"sync"
	"time"
)

// Mock Mixpanel client which can be used in unit tests.
type Mock struct {
	// All people identified, mapped by distinctId
	people sync.Map
}

func NewMock() *Mock {
	return &Mock{
		people: sync.Map{},
	}
}

func (m *Mock) String() string {
	str := ""
	m.people.Range(func(id, p interface{}) bool {
		str += id.(string) + ":\n" + p.(*MockPeople).String()
		return true
	})
	return str
}

func (m *Mock) Reset() {
	m.people = sync.Map{}
}

// People identifies a user. The user will be added to the people map.
func (m *Mock) People(distinctId string) *MockPeople {
	p, _ := m.people.LoadOrStore(distinctId, &MockPeople{
		Properties: map[string]interface{}{},
	})

	return p.(*MockPeople)
}

func (m *Mock) Track(distinctId, eventName string, e *Event) error {
	p := m.People(distinctId)
	p.Events = append(p.Events, MockEvent{
		Event: *e,
		Name:  eventName,
	})
	return nil
}

func (m *Mock) Import(distinctId, eventName string, e *Event) error {
	p := m.People(distinctId)
	p.Events = append(p.Events, MockEvent{
		Event: *e,
		Name:  eventName,
	})
	return nil
}

type MockPeople struct {
	Properties map[string]interface{}
	Time       *time.Time
	IP         string
	Events     []MockEvent
}

func (mp *MockPeople) String() string {
	timeStr := ""
	if mp.Time != nil {
		timeStr = mp.Time.Format(time.RFC3339)
	}

	str := fmt.Sprintf("  ip: %s\n  time: %s\n", mp.IP, timeStr)
	str += "  properties:\n"
	for key, val := range mp.Properties {
		str += fmt.Sprintf("    %s: %v\n", key, val)
	}
	str += "  events:\n"
	for _, event := range mp.Events {
		str += "    " + event.Name + ":\n"
		str += fmt.Sprintf("      IP: %s\n", event.IP)
		if event.Timestamp != nil {
			str += fmt.Sprintf(
				"      Timestamp: %s\n", event.Timestamp.Format(time.RFC3339),
			)
		} else {
			str += "      Timestamp:\n"
		}
		for key, val := range event.Properties {
			str += fmt.Sprintf("      %s: %v\n", key, val)
		}
	}
	return str
}

func (m *Mock) Update(distinctId string, u *Update) error {
	return m.UpdateUser(distinctId, u)
}

func (m *Mock) UpdateUser(distinctId string, u *Update) error {
	p := m.People(distinctId)

	if u.IP != "" {
		p.IP = u.IP
	}
	if u.Timestamp != nil && u.Timestamp != IgnoreTime {
		p.Time = u.Timestamp
	}

	switch u.Operation {
	case "$set", "$set_once":
		for key, val := range u.Properties {
			p.Properties[key] = val
		}
	default:
		return errors.New("mixpanel.Mock only supports the $set and $set_once operations")
	}

	return nil
}

func (m *Mock) UnionUser(userID string, u *Update) error {
	return nil
}

func (m *Mock) UpdateGroup(groupKey, groupUser string, u *Update) error {
	return nil
}

func (m *Mock) UnionGroup(groupKey, groupUser string, u *Update) error {
	return nil
}

func (m *Mock) Alias(distinctId, newId string) error {
	return nil
}

type MockEvent struct {
	Event
	Name string
}
<|endoftext|>


<commit_before>package minodedb

import (
	"fmt"
	"net/http"

	"github.com/nyaxt/otaru/inodedb"
	"github.com/nyaxt/otaru/mgmt"
)

func Install(srv *mgmt.Server, h inodedb.DBHandler) {
	rtr := srv.APIRouter().PathPrefix("/inodedb").Subrouter()

	rtr.HandleFunc("/stats", mgmt.JSONHandler(func(req *http.Request) interface{} {
		prov, ok := h.(inodedb.DBServiceStatsProvider)
		if !ok {
			return fmt.Errorf("Active inodedb doesn't support /stats")
		}
		return prov.GetStats()
	}))
	rtr.HandleFunc("/recenttxs", mgmt.JSONHandler(func(req *http.Request) interface{} {
		prov, ok := h.(inodedb.QueryRecentTransactionsProvider)
		if !ok {
			return fmt.Errorf("Active inodedb doesn't support /recenttxs")
		}
		txs, err := prov.QueryRecentTransactions()
		if err != nil {
			return fmt.Errorf("QueryRecentTransactions failed: %v", err)
		}
		for _, tx := range txs {
			if err := inodedb.SetOpMetas(tx.Ops); err != nil {
				return fmt.Errorf("SetOpMetas failed: %v", err)
			}
		}

		return txs
	}))
}
<commit_msg>query inode from mgmt<commit_after>package minodedb

import (
	"fmt"
	"strconv"

	"net/http"

	"github.com/gorilla/mux"

	"github.com/nyaxt/otaru/inodedb"
	"github.com/nyaxt/otaru/mgmt"
)

func Install(srv *mgmt.Server, h inodedb.DBHandler) {
	rtr := srv.APIRouter().PathPrefix("/inodedb").Subrouter()

	rtr.HandleFunc("/stats", mgmt.JSONHandler(func(req *http.Request) interface{} {
		prov, ok := h.(inodedb.DBServiceStatsProvider)
		if !ok {
			return fmt.Errorf("Active inodedb doesn't support /stats")
		}
		return prov.GetStats()
	}))
	rtr.HandleFunc("/recenttxs", mgmt.JSONHandler(func(req *http.Request) interface{} {
		prov, ok := h.(inodedb.QueryRecentTransactionsProvider)
		if !ok {
			return fmt.Errorf("Active inodedb doesn't support /recenttxs")
		}
		txs, err := prov.QueryRecentTransactions()
		if err != nil {
			return fmt.Errorf("QueryRecentTransactions failed: %v", err)
		}
		for _, tx := range txs {
			if err := inodedb.SetOpMetas(tx.Ops); err != nil {
				return fmt.Errorf("SetOpMetas failed: %v", err)
			}
		}

		return txs
	}))
	rtr.HandleFunc("/inode/{id:[0-9]+}", mgmt.JSONHandler(func(req *http.Request) interface{} {
		vars := mux.Vars(req)
		nid, err := strconv.ParseUint(vars["id"], 10, 32)
		if err != nil {
			return err
		}
		id := inodedb.ID(nid)
		v, _, err := h.QueryNode(id, false)
		if err != nil {
			return err
		}

		return v
	}))
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"time"
	"unsafe"

	"github.com/docker/docker/pkg/system"
	"github.com/sirupsen/logrus"
	"github.com/spf13/pflag"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/debug"
	"golang.org/x/sys/windows/svc/eventlog"
	"golang.org/x/sys/windows/svc/mgr"
)

var (
	flServiceName       *string
	flRegisterService   *bool
	flUnregisterService *bool
	flRunService        *bool

	setStdHandle = windows.NewLazySystemDLL("kernel32.dll").NewProc("SetStdHandle")
	oldStderr    windows.Handle
	panicFile    *os.File

	service *handler
)

const (
	// These should match the values in event_messages.mc.
	eventInfo  = 1
	eventWarn  = 1
	eventError = 1
	eventDebug = 2
	eventPanic = 3
	eventFatal = 4

	eventExtraOffset = 10 // Add this to any event to get a string that supports extended data
)

func installServiceFlags(flags *pflag.FlagSet) {
	flServiceName = flags.String("service-name", "docker", "Set the Windows service name")
	flRegisterService = flags.Bool("register-service", false, "Register the service and exit")
	flUnregisterService = flags.Bool("unregister-service", false, "Unregister the service and exit")
	flRunService = flags.Bool("run-service", false, "")
	flags.MarkHidden("run-service")
}

type handler struct {
	tosvc     chan bool
	fromsvc   chan error
	daemonCli *DaemonCli
}

type etwHook struct {
	log *eventlog.Log
}

func (h *etwHook) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.PanicLevel,
		logrus.FatalLevel,
		logrus.ErrorLevel,
		logrus.WarnLevel,
		logrus.InfoLevel,
		logrus.DebugLevel,
	}
}

func (h *etwHook) Fire(e *logrus.Entry) error {
	var (
		etype uint16
		eid   uint32
	)

	switch e.Level {
	case logrus.PanicLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventPanic
	case logrus.FatalLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventFatal
	case logrus.ErrorLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventError
	case logrus.WarnLevel:
		etype = windows.EVENTLOG_WARNING_TYPE
		eid = eventWarn
	case logrus.InfoLevel:
		etype = windows.EVENTLOG_INFORMATION_TYPE
		eid = eventInfo
	case logrus.DebugLevel:
		etype = windows.EVENTLOG_INFORMATION_TYPE
		eid = eventDebug
	default:
		return errors.New("unknown level")
	}

	// If there is additional data, include it as a second string.
	exts := ""
	if len(e.Data) > 0 {
		fs := bytes.Buffer{}
		for k, v := range e.Data {
			fs.WriteString(k)
			fs.WriteByte('=')
			fmt.Fprint(&fs, v)
			fs.WriteByte(' ')
		}

		exts = fs.String()[:fs.Len()-1]
		eid += eventExtraOffset
	}

	if h.log == nil {
		fmt.Fprintf(os.Stderr, "%s [%s]\n", e.Message, exts)
		return nil
	}

	var (
		ss  [2]*uint16
		err error
	)

	ss[0], err = windows.UTF16PtrFromString(e.Message)
	if err != nil {
		return err
	}

	count := uint16(1)
	if exts != "" {
		ss[1], err = windows.UTF16PtrFromString(exts)
		if err != nil {
			return err
		}

		count++
	}

	return windows.ReportEvent(h.log.Handle, etype, 0, eid, 0, count, 0, &ss[0], nil)
}

func getServicePath() (string, error) {
	p, err := exec.LookPath(os.Args[0])
	if err != nil {
		return "", err
	}
	return filepath.Abs(p)
}

func registerService() error {
	p, err := getServicePath()
	if err != nil {
		return err
	}
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	depends := []string{}

	// This dependency is required on build 14393 (RS1)
	// it is added to the platform in newer builds
	if system.GetOSVersion().Build == 14393 {
		depends = append(depends, "ConDrv")
	}

	c := mgr.Config{
		ServiceType:  windows.SERVICE_WIN32_OWN_PROCESS,
		StartType:    mgr.StartAutomatic,
		ErrorControl: mgr.ErrorNormal,
		Dependencies: depends,
		DisplayName:  "Docker Engine",
	}

	// Configure the service to launch with the arguments that were just passed.
	args := []string{"--run-service"}
	for _, a := range os.Args[1:] {
		if a != "--register-service" && a != "--unregister-service" {
			args = append(args, a)
		}
	}

	s, err := m.CreateService(*flServiceName, p, c, args...)
	if err != nil {
		return err
	}
	defer s.Close()

	// See http://stackoverflow.com/questions/35151052/how-do-i-configure-failure-actions-of-a-windows-service-written-in-go
	const (
		scActionNone       = 0
		scActionRestart    = 1
		scActionReboot     = 2
		scActionRunCommand = 3

		serviceConfigFailureActions = 2
	)

	type serviceFailureActions struct {
		ResetPeriod  uint32
		RebootMsg    *uint16
		Command      *uint16
		ActionsCount uint32
		Actions      uintptr
	}

	type scAction struct {
		Type  uint32
		Delay uint32
	}
	t := []scAction{
		{Type: scActionRestart, Delay: uint32(60 * time.Second / time.Millisecond)},
		{Type: scActionRestart, Delay: uint32(60 * time.Second / time.Millisecond)},
		{Type: scActionNone},
	}
	lpInfo := serviceFailureActions{ResetPeriod: uint32(24 * time.Hour / time.Second), ActionsCount: uint32(3), Actions: uintptr(unsafe.Pointer(&t[0]))}
	err = windows.ChangeServiceConfig2(s.Handle, serviceConfigFailureActions, (*byte)(unsafe.Pointer(&lpInfo)))
	if err != nil {
		return err
	}

	return eventlog.Install(*flServiceName, p, false, eventlog.Info|eventlog.Warning|eventlog.Error)
}

func unregisterService() error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	s, err := m.OpenService(*flServiceName)
	if err != nil {
		return err
	}
	defer s.Close()

	eventlog.Remove(*flServiceName)
	err = s.Delete()
	if err != nil {
		return err
	}
	return nil
}

// initService is the entry point for running the daemon as a Windows
// service. It returns an indication to stop (if registering/un-registering);
// an indication of whether it is running as a service; and an error.
func initService(daemonCli *DaemonCli) (bool, bool, error) {
	if *flUnregisterService {
		if *flRegisterService {
			return true, false, errors.New("--register-service and --unregister-service cannot be used together")
		}
		return true, false, unregisterService()
	}

	if *flRegisterService {
		return true, false, registerService()
	}

	if !*flRunService {
		return false, false, nil
	}

	interactive, err := svc.IsAnInteractiveSession()
	if err != nil {
		return false, false, err
	}

	h := &handler{
		tosvc:     make(chan bool),
		fromsvc:   make(chan error),
		daemonCli: daemonCli,
	}

	var log *eventlog.Log
	if !interactive {
		log, err = eventlog.Open(*flServiceName)
		if err != nil {
			return false, false, err
		}
	}

	logrus.AddHook(&etwHook{log})
	logrus.SetOutput(ioutil.Discard)

	service = h
	go func() {
		if interactive {
			err = debug.Run(*flServiceName, h)
		} else {
			err = svc.Run(*flServiceName, h)
		}

		h.fromsvc <- err
	}()

	// Wait for the first signal from the service handler.
	err = <-h.fromsvc
	if err != nil {
		return false, false, err
	}
	return false, true, nil
}

func (h *handler) started() error {
	// This must be delayed until daemonCli initializes Config.Root
	err := initPanicFile(filepath.Join(h.daemonCli.Config.Root, "panic.log"))
	if err != nil {
		return err
	}

	h.tosvc <- false
	return nil
}

func (h *handler) stopped(err error) {
	logrus.Debugf("Stopping service: %v", err)
	h.tosvc <- err != nil
	<-h.fromsvc
}

func (h *handler) Execute(_ []string, r <-chan svc.ChangeRequest, s chan<- svc.Status) (bool, uint32) {
	s <- svc.Status{State: svc.StartPending, Accepts: 0}
	// Unblock initService()
	h.fromsvc <- nil

	// Wait for initialization to complete.
	failed := <-h.tosvc
	if failed {
		logrus.Debug("Aborting service start due to failure during initialization")
		return true, 1
	}

	s <- svc.Status{State: svc.Running, Accepts: svc.AcceptStop | svc.AcceptShutdown | svc.Accepted(windows.SERVICE_ACCEPT_PARAMCHANGE)}
	logrus.Debug("Service running")
Loop:
	for {
		select {
		case failed = <-h.tosvc:
			break Loop
		case c := <-r:
			switch c.Cmd {
			case svc.Cmd(windows.SERVICE_CONTROL_PARAMCHANGE):
				h.daemonCli.reloadConfig()
			case svc.Interrogate:
				s <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s <- svc.Status{State: svc.StopPending, Accepts: 0}
				h.daemonCli.stop()
			}
		}
	}

	removePanicFile()
	if failed {
		return true, 1
	}
	return false, 0
}

func initPanicFile(path string) error {
	var err error
	panicFile, err = os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0)
	if err != nil {
		return err
	}

	st, err := panicFile.Stat()
	if err != nil {
		return err
	}

	// If there are contents in the file already, move the file out of the way
	// and replace it.
	if st.Size() > 0 {
		panicFile.Close()
		os.Rename(path, path+".old")
		panicFile, err = os.Create(path)
		if err != nil {
			return err
		}
	}

	// Update STD_ERROR_HANDLE to point to the panic file so that Go writes to
	// it when it panics. Remember the old stderr to restore it before removing
	// the panic file.
	sh := windows.STD_ERROR_HANDLE
	h, err := windows.GetStdHandle(uint32(sh))
	if err != nil {
		return err
	}

	oldStderr = h

	r, _, err := setStdHandle.Call(uintptr(sh), uintptr(panicFile.Fd()))
	if r == 0 && err != nil {
		return err
	}

	// Reset os.Stderr to the panic file (so fmt.Fprintf(os.Stderr,...) actually gets redirected)
	os.Stderr = os.NewFile(uintptr(panicFile.Fd()), "/dev/stderr")

	// Force threads that panic to write to stderr (the panicFile handle now), otherwise it will go into the ether
	log.SetOutput(os.Stderr)

	return nil
}

func removePanicFile() {
	if st, err := panicFile.Stat(); err == nil {
		if st.Size() == 0 {
			sh := windows.STD_ERROR_HANDLE
			setStdHandle.Call(uintptr(sh), uintptr(oldStderr))
			panicFile.Close()
			os.Remove(panicFile.Name())
		}
	}
}
<commit_msg>Properly type handles to prevent overflows<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"time"
	"unsafe"

	"github.com/docker/docker/pkg/system"
	"github.com/sirupsen/logrus"
	"github.com/spf13/pflag"
	"golang.org/x/sys/windows"
	"golang.org/x/sys/windows/svc"
	"golang.org/x/sys/windows/svc/debug"
	"golang.org/x/sys/windows/svc/eventlog"
	"golang.org/x/sys/windows/svc/mgr"
)

var (
	flServiceName       *string
	flRegisterService   *bool
	flUnregisterService *bool
	flRunService        *bool

	setStdHandle = windows.NewLazySystemDLL("kernel32.dll").NewProc("SetStdHandle")
	oldStderr    windows.Handle
	panicFile    *os.File

	service *handler
)

const (
	// These should match the values in event_messages.mc.
	eventInfo  = 1
	eventWarn  = 1
	eventError = 1
	eventDebug = 2
	eventPanic = 3
	eventFatal = 4

	eventExtraOffset = 10 // Add this to any event to get a string that supports extended data
)

func installServiceFlags(flags *pflag.FlagSet) {
	flServiceName = flags.String("service-name", "docker", "Set the Windows service name")
	flRegisterService = flags.Bool("register-service", false, "Register the service and exit")
	flUnregisterService = flags.Bool("unregister-service", false, "Unregister the service and exit")
	flRunService = flags.Bool("run-service", false, "")
	flags.MarkHidden("run-service")
}

type handler struct {
	tosvc     chan bool
	fromsvc   chan error
	daemonCli *DaemonCli
}

type etwHook struct {
	log *eventlog.Log
}

func (h *etwHook) Levels() []logrus.Level {
	return []logrus.Level{
		logrus.PanicLevel,
		logrus.FatalLevel,
		logrus.ErrorLevel,
		logrus.WarnLevel,
		logrus.InfoLevel,
		logrus.DebugLevel,
	}
}

func (h *etwHook) Fire(e *logrus.Entry) error {
	var (
		etype uint16
		eid   uint32
	)

	switch e.Level {
	case logrus.PanicLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventPanic
	case logrus.FatalLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventFatal
	case logrus.ErrorLevel:
		etype = windows.EVENTLOG_ERROR_TYPE
		eid = eventError
	case logrus.WarnLevel:
		etype = windows.EVENTLOG_WARNING_TYPE
		eid = eventWarn
	case logrus.InfoLevel:
		etype = windows.EVENTLOG_INFORMATION_TYPE
		eid = eventInfo
	case logrus.DebugLevel:
		etype = windows.EVENTLOG_INFORMATION_TYPE
		eid = eventDebug
	default:
		return errors.New("unknown level")
	}

	// If there is additional data, include it as a second string.
	exts := ""
	if len(e.Data) > 0 {
		fs := bytes.Buffer{}
		for k, v := range e.Data {
			fs.WriteString(k)
			fs.WriteByte('=')
			fmt.Fprint(&fs, v)
			fs.WriteByte(' ')
		}

		exts = fs.String()[:fs.Len()-1]
		eid += eventExtraOffset
	}

	if h.log == nil {
		fmt.Fprintf(os.Stderr, "%s [%s]\n", e.Message, exts)
		return nil
	}

	var (
		ss  [2]*uint16
		err error
	)

	ss[0], err = windows.UTF16PtrFromString(e.Message)
	if err != nil {
		return err
	}

	count := uint16(1)
	if exts != "" {
		ss[1], err = windows.UTF16PtrFromString(exts)
		if err != nil {
			return err
		}

		count++
	}

	return windows.ReportEvent(h.log.Handle, etype, 0, eid, 0, count, 0, &ss[0], nil)
}

func getServicePath() (string, error) {
	p, err := exec.LookPath(os.Args[0])
	if err != nil {
		return "", err
	}
	return filepath.Abs(p)
}

func registerService() error {
	p, err := getServicePath()
	if err != nil {
		return err
	}
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	depends := []string{}

	// This dependency is required on build 14393 (RS1)
	// it is added to the platform in newer builds
	if system.GetOSVersion().Build == 14393 {
		depends = append(depends, "ConDrv")
	}

	c := mgr.Config{
		ServiceType:  windows.SERVICE_WIN32_OWN_PROCESS,
		StartType:    mgr.StartAutomatic,
		ErrorControl: mgr.ErrorNormal,
		Dependencies: depends,
		DisplayName:  "Docker Engine",
	}

	// Configure the service to launch with the arguments that were just passed.
	args := []string{"--run-service"}
	for _, a := range os.Args[1:] {
		if a != "--register-service" && a != "--unregister-service" {
			args = append(args, a)
		}
	}

	s, err := m.CreateService(*flServiceName, p, c, args...)
	if err != nil {
		return err
	}
	defer s.Close()

	// See http://stackoverflow.com/questions/35151052/how-do-i-configure-failure-actions-of-a-windows-service-written-in-go
	const (
		scActionNone       = 0
		scActionRestart    = 1
		scActionReboot     = 2
		scActionRunCommand = 3

		serviceConfigFailureActions = 2
	)

	type serviceFailureActions struct {
		ResetPeriod  uint32
		RebootMsg    *uint16
		Command      *uint16
		ActionsCount uint32
		Actions      uintptr
	}

	type scAction struct {
		Type  uint32
		Delay uint32
	}
	t := []scAction{
		{Type: scActionRestart, Delay: uint32(60 * time.Second / time.Millisecond)},
		{Type: scActionRestart, Delay: uint32(60 * time.Second / time.Millisecond)},
		{Type: scActionNone},
	}
	lpInfo := serviceFailureActions{ResetPeriod: uint32(24 * time.Hour / time.Second), ActionsCount: uint32(3), Actions: uintptr(unsafe.Pointer(&t[0]))}
	err = windows.ChangeServiceConfig2(s.Handle, serviceConfigFailureActions, (*byte)(unsafe.Pointer(&lpInfo)))
	if err != nil {
		return err
	}

	return eventlog.Install(*flServiceName, p, false, eventlog.Info|eventlog.Warning|eventlog.Error)
}

func unregisterService() error {
	m, err := mgr.Connect()
	if err != nil {
		return err
	}
	defer m.Disconnect()

	s, err := m.OpenService(*flServiceName)
	if err != nil {
		return err
	}
	defer s.Close()

	eventlog.Remove(*flServiceName)
	err = s.Delete()
	if err != nil {
		return err
	}
	return nil
}

// initService is the entry point for running the daemon as a Windows
// service. It returns an indication to stop (if registering/un-registering);
// an indication of whether it is running as a service; and an error.
func initService(daemonCli *DaemonCli) (bool, bool, error) {
	if *flUnregisterService {
		if *flRegisterService {
			return true, false, errors.New("--register-service and --unregister-service cannot be used together")
		}
		return true, false, unregisterService()
	}

	if *flRegisterService {
		return true, false, registerService()
	}

	if !*flRunService {
		return false, false, nil
	}

	interactive, err := svc.IsAnInteractiveSession()
	if err != nil {
		return false, false, err
	}

	h := &handler{
		tosvc:     make(chan bool),
		fromsvc:   make(chan error),
		daemonCli: daemonCli,
	}

	var log *eventlog.Log
	if !interactive {
		log, err = eventlog.Open(*flServiceName)
		if err != nil {
			return false, false, err
		}
	}

	logrus.AddHook(&etwHook{log})
	logrus.SetOutput(ioutil.Discard)

	service = h
	go func() {
		if interactive {
			err = debug.Run(*flServiceName, h)
		} else {
			err = svc.Run(*flServiceName, h)
		}

		h.fromsvc <- err
	}()

	// Wait for the first signal from the service handler.
	err = <-h.fromsvc
	if err != nil {
		return false, false, err
	}
	return false, true, nil
}

func (h *handler) started() error {
	// This must be delayed until daemonCli initializes Config.Root
	err := initPanicFile(filepath.Join(h.daemonCli.Config.Root, "panic.log"))
	if err != nil {
		return err
	}

	h.tosvc <- false
	return nil
}

func (h *handler) stopped(err error) {
	logrus.Debugf("Stopping service: %v", err)
	h.tosvc <- err != nil
	<-h.fromsvc
}

func (h *handler) Execute(_ []string, r <-chan svc.ChangeRequest, s chan<- svc.Status) (bool, uint32) {
	s <- svc.Status{State: svc.StartPending, Accepts: 0}
	// Unblock initService()
	h.fromsvc <- nil

	// Wait for initialization to complete.
	failed := <-h.tosvc
	if failed {
		logrus.Debug("Aborting service start due to failure during initialization")
		return true, 1
	}

	s <- svc.Status{State: svc.Running, Accepts: svc.AcceptStop | svc.AcceptShutdown | svc.Accepted(windows.SERVICE_ACCEPT_PARAMCHANGE)}
	logrus.Debug("Service running")
Loop:
	for {
		select {
		case failed = <-h.tosvc:
			break Loop
		case c := <-r:
			switch c.Cmd {
			case svc.Cmd(windows.SERVICE_CONTROL_PARAMCHANGE):
				h.daemonCli.reloadConfig()
			case svc.Interrogate:
				s <- c.CurrentStatus
			case svc.Stop, svc.Shutdown:
				s <- svc.Status{State: svc.StopPending, Accepts: 0}
				h.daemonCli.stop()
			}
		}
	}

	removePanicFile()
	if failed {
		return true, 1
	}
	return false, 0
}

func initPanicFile(path string) error {
	var err error
	panicFile, err = os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0)
	if err != nil {
		return err
	}

	st, err := panicFile.Stat()
	if err != nil {
		return err
	}

	// If there are contents in the file already, move the file out of the way
	// and replace it.
	if st.Size() > 0 {
		panicFile.Close()
		os.Rename(path, path+".old")
		panicFile, err = os.Create(path)
		if err != nil {
			return err
		}
	}

	// Update STD_ERROR_HANDLE to point to the panic file so that Go writes to
	// it when it panics. Remember the old stderr to restore it before removing
	// the panic file.
	sh := uint32(windows.STD_ERROR_HANDLE)
	h, err := windows.GetStdHandle(sh)
	if err != nil {
		return err
	}

	oldStderr = h

	r, _, err := setStdHandle.Call(uintptr(sh), uintptr(panicFile.Fd()))
	if r == 0 && err != nil {
		return err
	}

	// Reset os.Stderr to the panic file (so fmt.Fprintf(os.Stderr,...) actually gets redirected)
	os.Stderr = os.NewFile(uintptr(panicFile.Fd()), "/dev/stderr")

	// Force threads that panic to write to stderr (the panicFile handle now), otherwise it will go into the ether
	log.SetOutput(os.Stderr)

	return nil
}

func removePanicFile() {
	if st, err := panicFile.Stat(); err == nil {
		if st.Size() == 0 {
			sh := uint32(windows.STD_ERROR_HANDLE)
			setStdHandle.Call(uintptr(sh), uintptr(oldStderr))
			panicFile.Close()
			os.Remove(panicFile.Name())
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// This is necessary since it must test a recursive unexported function,
// i.e., the function cannot be exported via a var
package action

import (
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
)

type CommonSuite struct{}

var _ = gc.Suite(&CommonSuite{})

func (s *CommonSuite) TestConform(c *gc.C) {
	var goodInterfaceTests = []struct {
		description       string
		inputInterface    interface{}
		expectedInterface map[string]interface{}
		expectedError     string
	}{{
		description: "An interface requiring no changes.",
		inputInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[string]interface{}{
				"foo1": "val1",
				"foo2": "val2"}},
		expectedInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[string]interface{}{
				"foo1": "val1",
				"foo2": "val2"}},
	}, {
		description: "Substitute a single inner map[i]i.",
		inputInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[interface{}]interface{}{
				"foo1": "val1",
				"foo2": "val2"}},
		expectedInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[string]interface{}{
				"foo1": "val1",
				"foo2": "val2"}},
	}, {
		description: "Substitute nested inner map[i]i.",
		inputInterface: map[string]interface{}{
			"key1a": "val1a",
			"key2a": "val2a",
			"key3a": map[interface{}]interface{}{
				"key1b": "val1b",
				"key2b": map[interface{}]interface{}{
					"key1c": "val1c"}}},
		expectedInterface: map[string]interface{}{
			"key1a": "val1a",
			"key2a": "val2a",
			"key3a": map[string]interface{}{
				"key1b": "val1b",
				"key2b": map[string]interface{}{
					"key1c": "val1c"}}},
	}, {
		description: "Substitute nested map[i]i within []i.",
		inputInterface: map[string]interface{}{
			"key1a": "val1a",
			"key2a": []interface{}{5, "foo", map[string]interface{}{
				"key1b": "val1b",
				"key2b": map[interface{}]interface{}{
					"key1c": "val1c"}}}},
		expectedInterface: map[string]interface{}{
			"key1a": "val1a",
			"key2a": []interface{}{5, "foo", map[string]interface{}{
				"key1b": "val1b",
				"key2b": map[string]interface{}{
					"key1c": "val1c"}}}},
	}, {
		description: "An inner map[interface{}]interface{} with an int key.",
		inputInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[interface{}]interface{}{
				"foo1": "val1",
				5:      "val2"}},
		expectedError: "map keyed with non-string value",
	}, {
		description: "An inner []interface{} containing a map[i]i with an int key.",
		inputInterface: map[string]interface{}{
			"key1a": "val1b",
			"key2a": "val2b",
			"key3a": []interface{}{"foo1", 5, map[interface{}]interface{}{
				"key1b": "val1b",
				"key2b": map[interface{}]interface{}{
					"key1c": "val1c",
					5:       "val2c"}}}},
		expectedError: "map keyed with non-string value",
	}}

	for i, test := range goodInterfaceTests {
		c.Logf("test %d: %s", i, test.description)
		input := test.inputInterface
		cleansedInterfaceMap, err := conform(input)
		if test.expectedError == "" {
			if !c.Check(err, jc.ErrorIsNil) {
				continue
			}
			c.Check(cleansedInterfaceMap, gc.DeepEquals, test.expectedInterface)
		} else {
			c.Check(err, gc.ErrorMatches, test.expectedError)
		}
	}
}

type insertSliceValue struct {
	atSlice []string
	value   interface{}
}

func (s *CommonSuite) TestAddValueToMap(c *gc.C) {
	for i, t := range []struct {
		should       string
		startingMap  map[string]interface{}
		insertSlices []insertSliceValue
		expectedMap  map[string]interface{}
	}{{
		should: "insert a couple of values",
		startingMap: map[string]interface{}{
			"foo": "bar",
			"bar": map[string]interface{}{
				"baz": "bo",
				"bur": "bor",
			},
		},
		insertSlices: []insertSliceValue{
			{
				atSlice: []string{"well", "now"},
				value:   5,
			},
			{
				atSlice: []string{"foo"},
				value:   "kek",
			},
		},
		expectedMap: map[string]interface{}{
			"foo": "kek",
			"bar": map[string]interface{}{
				"baz": "bo",
				"bur": "bor",
			},
			"well": map[string]interface{}{
				"now": 5,
			},
		},
	}} {
		c.Logf("test %d: should %s", i, t.should)
		for _, sVal := range t.insertSlices {
			addValueToMap(sVal.atSlice, sVal.value, t.startingMap)
		}
		// note addValueToMap mutates target.
		c.Check(t.startingMap, jc.DeepEquals, t.expectedMap)
	}
}
<commit_msg>Rename action CLI test variable per TheMue<commit_after>// Copyright 2014-2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

// This is necessary since it must test a recursive unexported function,
// i.e., the function cannot be exported via a var
package action

import (
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
)

type CommonSuite struct{}

var _ = gc.Suite(&CommonSuite{})

func (s *CommonSuite) TestConform(c *gc.C) {
	var goodInterfaceTests = []struct {
		description       string
		inputInterface    interface{}
		expectedInterface map[string]interface{}
		expectedError     string
	}{{
		description: "An interface requiring no changes.",
		inputInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[string]interface{}{
				"foo1": "val1",
				"foo2": "val2"}},
		expectedInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[string]interface{}{
				"foo1": "val1",
				"foo2": "val2"}},
	}, {
		description: "Substitute a single inner map[i]i.",
		inputInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[interface{}]interface{}{
				"foo1": "val1",
				"foo2": "val2"}},
		expectedInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[string]interface{}{
				"foo1": "val1",
				"foo2": "val2"}},
	}, {
		description: "Substitute nested inner map[i]i.",
		inputInterface: map[string]interface{}{
			"key1a": "val1a",
			"key2a": "val2a",
			"key3a": map[interface{}]interface{}{
				"key1b": "val1b",
				"key2b": map[interface{}]interface{}{
					"key1c": "val1c"}}},
		expectedInterface: map[string]interface{}{
			"key1a": "val1a",
			"key2a": "val2a",
			"key3a": map[string]interface{}{
				"key1b": "val1b",
				"key2b": map[string]interface{}{
					"key1c": "val1c"}}},
	}, {
		description: "Substitute nested map[i]i within []i.",
		inputInterface: map[string]interface{}{
			"key1a": "val1a",
			"key2a": []interface{}{5, "foo", map[string]interface{}{
				"key1b": "val1b",
				"key2b": map[interface{}]interface{}{
					"key1c": "val1c"}}}},
		expectedInterface: map[string]interface{}{
			"key1a": "val1a",
			"key2a": []interface{}{5, "foo", map[string]interface{}{
				"key1b": "val1b",
				"key2b": map[string]interface{}{
					"key1c": "val1c"}}}},
	}, {
		description: "An inner map[interface{}]interface{} with an int key.",
		inputInterface: map[string]interface{}{
			"key1": "value1",
			"key2": "value2",
			"key3": map[interface{}]interface{}{
				"foo1": "val1",
				5:      "val2"}},
		expectedError: "map keyed with non-string value",
	}, {
		description: "An inner []interface{} containing a map[i]i with an int key.",
		inputInterface: map[string]interface{}{
			"key1a": "val1b",
			"key2a": "val2b",
			"key3a": []interface{}{"foo1", 5, map[interface{}]interface{}{
				"key1b": "val1b",
				"key2b": map[interface{}]interface{}{
					"key1c": "val1c",
					5:       "val2c"}}}},
		expectedError: "map keyed with non-string value",
	}}

	for i, test := range goodInterfaceTests {
		c.Logf("test %d: %s", i, test.description)
		input := test.inputInterface
		cleansedInterfaceMap, err := conform(input)
		if test.expectedError == "" {
			if !c.Check(err, jc.ErrorIsNil) {
				continue
			}
			c.Check(cleansedInterfaceMap, gc.DeepEquals, test.expectedInterface)
		} else {
			c.Check(err, gc.ErrorMatches, test.expectedError)
		}
	}
}

type insertSliceValue struct {
	valuePath []string
	value     interface{}
}

func (s *CommonSuite) TestAddValueToMap(c *gc.C) {
	for i, t := range []struct {
		should       string
		startingMap  map[string]interface{}
		insertSlices []insertSliceValue
		expectedMap  map[string]interface{}
	}{{
		should: "insert a couple of values",
		startingMap: map[string]interface{}{
			"foo": "bar",
			"bar": map[string]interface{}{
				"baz": "bo",
				"bur": "bor",
			},
		},
		insertSlices: []insertSliceValue{
			{
				valuePath: []string{"well", "now"},
				value:     5,
			},
			{
				valuePath: []string{"foo"},
				value:     "kek",
			},
		},
		expectedMap: map[string]interface{}{
			"foo": "kek",
			"bar": map[string]interface{}{
				"baz": "bo",
				"bur": "bor",
			},
			"well": map[string]interface{}{
				"now": 5,
			},
		},
	}} {
		c.Logf("test %d: should %s", i, t.should)
		for _, sVal := range t.insertSlices {
			addValueToMap(sVal.valuePath, sVal.value, t.startingMap)
		}
		// note addValueToMap mutates target.
		c.Check(t.startingMap, jc.DeepEquals, t.expectedMap)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package commands

import (
	"bytes"
	"io"

	"github.com/juju/cmd"
	"github.com/juju/loggo"
	"github.com/juju/version"
	"launchpad.net/gnuflag"

	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
	"github.com/juju/juju/environs/filestorage"
	"github.com/juju/juju/environs/sync"
	envtools "github.com/juju/juju/environs/tools"
	coretools "github.com/juju/juju/tools"
)

var syncTools = sync.SyncTools

func newSyncToolsCommand() cmd.Command {
	return modelcmd.Wrap(&syncToolsCommand{})
}

// syncToolsCommand copies all the tools from the us-east-1 bucket to the local
// bucket.
type syncToolsCommand struct {
	modelcmd.ModelCommandBase
	allVersions  bool
	versionStr   string
	majorVersion int
	minorVersion int
	dryRun       bool
	dev          bool
	public       bool
	source       string
	stream       string
	localDir     string
	destination  string
}

var _ cmd.Command = (*syncToolsCommand)(nil)

func (c *syncToolsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "sync-tools",
		Purpose: "copy tools from the official tool store into a local model",
		Doc: `
This copies the Juju tools tarball from the official tools store (located
at https://streams.canonical.com/juju) into your model.
This is generally done when you want Juju to be able to run without having to
access the Internet. Alternatively you can specify a local directory as source.

Sometimes this is because the model does not have public access,
and sometimes you just want to avoid having to access data outside of
the local cloud.
`,
	}
}

func (c *syncToolsCommand) SetFlags(f *gnuflag.FlagSet) {
	f.BoolVar(&c.allVersions, "all", false, "copy all versions, not just the latest")
	f.StringVar(&c.versionStr, "version", "", "copy a specific major[.minor] version")
	f.BoolVar(&c.dryRun, "dry-run", false, "don't copy, just print what would be copied")
	f.BoolVar(&c.dev, "dev", false, "consider development versions as well as released ones\n    DEPRECATED: use --stream instead")
	f.BoolVar(&c.public, "public", false, "tools are for a public cloud, so generate mirrors information")
	f.StringVar(&c.source, "source", "", "local source directory")
	f.StringVar(&c.stream, "stream", "", "simplestreams stream for which to sync metadata")
	f.StringVar(&c.localDir, "local-dir", "", "local destination directory")
	f.StringVar(&c.destination, "destination", "", "local destination directory")
}

func (c *syncToolsCommand) Init(args []string) error {
	if c.destination != "" {
		// Override localDir with destination as localDir now replaces destination
		c.localDir = c.destination
		logger.Infof("Use of the --destination flag is deprecated in 1.18. Please use --local-dir instead.")
	}
	if c.versionStr != "" {
		var err error
		if c.majorVersion, c.minorVersion, err = version.ParseMajorMinor(c.versionStr); err != nil {
			return err
		}
	}
	if c.dev {
		c.stream = envtools.TestingStream
	}
	return cmd.CheckEmpty(args)
}

// syncToolsAPI provides an interface with a subset of the
// api.Client API. This exists to enable mocking.
type syncToolsAPI interface {
	FindTools(majorVersion, minorVersion int, series, arch string) (params.FindToolsResult, error)
	UploadTools(r io.ReadSeeker, v version.Binary, series ...string) (coretools.List, error)
	Close() error
}

var getSyncToolsAPI = func(c *syncToolsCommand) (syncToolsAPI, error) {
	return c.NewAPIClient()
}

func (c *syncToolsCommand) Run(ctx *cmd.Context) (resultErr error) {
	// Register writer for output on screen.
	loggo.RegisterWriter("synctools", cmd.NewCommandLogWriter("juju.environs.sync", ctx.Stdout, ctx.Stderr), loggo.INFO)
	defer loggo.RemoveWriter("synctools")

	sctx := &sync.SyncContext{
		AllVersions:  c.allVersions,
		MajorVersion: c.majorVersion,
		MinorVersion: c.minorVersion,
		DryRun:       c.dryRun,
		Stream:       c.stream,
		Source:       c.source,
	}

	if c.localDir != "" {
		stor, err := filestorage.NewFileStorageWriter(c.localDir)
		if err != nil {
			return err
		}
		writeMirrors := envtools.DoNotWriteMirrors
		if c.public {
			writeMirrors = envtools.WriteMirrors
		}
		sctx.TargetToolsFinder = sync.StorageToolsFinder{Storage: stor}
		sctx.TargetToolsUploader = sync.StorageToolsUploader{
			Storage:       stor,
			WriteMetadata: true,
			WriteMirrors:  writeMirrors,
		}
	} else {
		if c.public {
			logger.Infof("--public is ignored unless --local-dir is specified")
		}
		api, err := getSyncToolsAPI(c)
		if err != nil {
			return err
		}
		defer api.Close()
		adapter := syncToolsAPIAdapter{api}
		sctx.TargetToolsFinder = adapter
		sctx.TargetToolsUploader = adapter
	}
	return block.ProcessBlockedError(syncTools(sctx), block.BlockChange)
}

// syncToolsAPIAdapter implements sync.ToolsFinder and
// sync.ToolsUploader, adapting a syncToolsAPI. This
// enables the use of sync.SyncTools with the client
// API.
type syncToolsAPIAdapter struct {
	syncToolsAPI
}

func (s syncToolsAPIAdapter) FindTools(majorVersion int, stream string) (coretools.List, error) {
	result, err := s.syncToolsAPI.FindTools(majorVersion, -1, "", "")
	if err != nil {
		return nil, err
	}
	if result.Error != nil {
		if params.IsCodeNotFound(result.Error) {
			return nil, coretools.ErrNoMatches
		}
		return nil, result.Error
	}
	return result.List, nil
}

func (s syncToolsAPIAdapter) UploadTools(toolsDir, stream string, tools *coretools.Tools, data []byte) error {
	_, err := s.syncToolsAPI.UploadTools(bytes.NewReader(data), tools.Version)
	return err
}
<commit_msg>Improve sync-tools command documentation<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package commands

import (
	"bytes"
	"io"

	"github.com/juju/cmd"
	"github.com/juju/loggo"
	"github.com/juju/version"
	"launchpad.net/gnuflag"

	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/cmd/juju/block"
	"github.com/juju/juju/cmd/modelcmd"
	"github.com/juju/juju/environs/filestorage"
	"github.com/juju/juju/environs/sync"
	envtools "github.com/juju/juju/environs/tools"
	coretools "github.com/juju/juju/tools"
)

var syncTools = sync.SyncTools

func newSyncToolsCommand() cmd.Command {
	return modelcmd.Wrap(&syncToolsCommand{})
}

// syncToolsCommand copies all the tools from the us-east-1 bucket to the local
// bucket.
type syncToolsCommand struct {
	modelcmd.ModelCommandBase
	allVersions  bool
	versionStr   string
	majorVersion int
	minorVersion int
	dryRun       bool
	dev          bool
	public       bool
	source       string
	stream       string
	localDir     string
	destination  string
}

var _ cmd.Command = (*syncToolsCommand)(nil)

const synctoolsDoc = `
This copies the Juju agent software from the official tools store (located
at https://streams.canonical.com/juju) into a model. It is generally done
when the model is without Internet access.

Instead of the above site, a local directory can be specified as source.
The online store will, of course, need to be contacted at some point to get
the software.

Examples:

Download the software (version auto-selected) to the model:

    juju sync-tools --debug

Download a specific version of the software locally:

    juju sync-tools --debug --version 2.0 --local-dir=/home/ubuntu/sync-tools

Get locally available software to the model:

    juju sync-tools --debug --source=/home/ubuntu/sync-tools

See also: upgrade-juju

`

func (c *syncToolsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "sync-tools",
		Purpose: "Copy tools from the official tool store into a local model.",
		Doc:     synctoolsDoc,
	}
}

func (c *syncToolsCommand) SetFlags(f *gnuflag.FlagSet) {
	f.BoolVar(&c.allVersions, "all", false, "copy all versions, not just the latest")
	f.StringVar(&c.versionStr, "version", "", "copy a specific major[.minor] version")
	f.BoolVar(&c.dryRun, "dry-run", false, "don't copy, just print what would be copied")
	f.BoolVar(&c.dev, "dev", false, "consider development versions as well as released ones\n    DEPRECATED: use --stream instead")
	f.BoolVar(&c.public, "public", false, "tools are for a public cloud, so generate mirrors information")
	f.StringVar(&c.source, "source", "", "local source directory")
	f.StringVar(&c.stream, "stream", "", "simplestreams stream for which to sync metadata")
	f.StringVar(&c.localDir, "local-dir", "", "local destination directory")
	f.StringVar(&c.destination, "destination", "", "local destination directory\n    DEPRECATED: use --local-dir instead")
}

func (c *syncToolsCommand) Init(args []string) error {
	if c.destination != "" {
		// Override localDir with destination as localDir now replaces destination
		c.localDir = c.destination
		logger.Infof("Use of the --destination flag is deprecated in 1.18. Please use --local-dir instead.")
	}
	if c.versionStr != "" {
		var err error
		if c.majorVersion, c.minorVersion, err = version.ParseMajorMinor(c.versionStr); err != nil {
			return err
		}
	}
	if c.dev {
		c.stream = envtools.TestingStream
	}
	return cmd.CheckEmpty(args)
}

// syncToolsAPI provides an interface with a subset of the
// api.Client API. This exists to enable mocking.
type syncToolsAPI interface {
	FindTools(majorVersion, minorVersion int, series, arch string) (params.FindToolsResult, error)
	UploadTools(r io.ReadSeeker, v version.Binary, series ...string) (coretools.List, error)
	Close() error
}

var getSyncToolsAPI = func(c *syncToolsCommand) (syncToolsAPI, error) {
	return c.NewAPIClient()
}

func (c *syncToolsCommand) Run(ctx *cmd.Context) (resultErr error) {
	// Register writer for output on screen.
	loggo.RegisterWriter("synctools", cmd.NewCommandLogWriter("juju.environs.sync", ctx.Stdout, ctx.Stderr), loggo.INFO)
	defer loggo.RemoveWriter("synctools")

	sctx := &sync.SyncContext{
		AllVersions:  c.allVersions,
		MajorVersion: c.majorVersion,
		MinorVersion: c.minorVersion,
		DryRun:       c.dryRun,
		Stream:       c.stream,
		Source:       c.source,
	}

	if c.localDir != "" {
		stor, err := filestorage.NewFileStorageWriter(c.localDir)
		if err != nil {
			return err
		}
		writeMirrors := envtools.DoNotWriteMirrors
		if c.public {
			writeMirrors = envtools.WriteMirrors
		}
		sctx.TargetToolsFinder = sync.StorageToolsFinder{Storage: stor}
		sctx.TargetToolsUploader = sync.StorageToolsUploader{
			Storage:       stor,
			WriteMetadata: true,
			WriteMirrors:  writeMirrors,
		}
	} else {
		if c.public {
			logger.Infof("--public is ignored unless --local-dir is specified")
		}
		api, err := getSyncToolsAPI(c)
		if err != nil {
			return err
		}
		defer api.Close()
		adapter := syncToolsAPIAdapter{api}
		sctx.TargetToolsFinder = adapter
		sctx.TargetToolsUploader = adapter
	}
	return block.ProcessBlockedError(syncTools(sctx), block.BlockChange)
}

// syncToolsAPIAdapter implements sync.ToolsFinder and
// sync.ToolsUploader, adapting a syncToolsAPI. This
// enables the use of sync.SyncTools with the client
// API.
type syncToolsAPIAdapter struct {
	syncToolsAPI
}

func (s syncToolsAPIAdapter) FindTools(majorVersion int, stream string) (coretools.List, error) {
	result, err := s.syncToolsAPI.FindTools(majorVersion, -1, "", "")
	if err != nil {
		return nil, err
	}
	if result.Error != nil {
		if params.IsCodeNotFound(result.Error) {
			return nil, coretools.ErrNoMatches
		}
		return nil, result.Error
	}
	return result.List, nil
}

func (s syncToolsAPIAdapter) UploadTools(toolsDir, stream string, tools *coretools.Tools, data []byte) error {
	_, err := s.syncToolsAPI.UploadTools(bytes.NewReader(data), tools.Version)
	return err
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 The DevMine authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package repotool-db is able to fetch information from a source code repository.
// Typically, it can get all commits, their authors and commiters and so on
// and is able to populate the information into a PostgreSQL database.
// Currently, on the Git VCS is supported.
package main

import (
	"bytes"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strings"
	"sync"

	"github.com/golang/glog"
	"github.com/lib/pq"

	"github.com/DevMine/repotool/config"
	"github.com/DevMine/repotool/model"
	"github.com/DevMine/repotool/repo"
)

const version = "0.1.0"

// database fields per tables
var (
	diffDeltaFields = []string{
		"commit_id",
		"file_status",
		"is_file_binary",
		"similarity",
		"old_file_path",
		"new_file_path"}

	commitFields = []string{
		"repository_id",
		"author_id",
		"committer_id",
		"vcs_id",
		"message",
		"author_date",
		"commit_date",
		"file_changed_count",
		"insertions_count",
		"deletions_count"}
)

// program flags
var (
	configPath    = flag.String("c", "", "configuration file")
	vflag         = flag.Bool("V", false, "print version.")
	cpuprofile    = flag.String("cpuprofile", "", "write cpu profile to file")
	depthflag     = flag.Uint("d", 0, "depth level where to find repositories")
	numGoroutines = flag.Uint("g", uint(runtime.NumCPU()), "max number of goroutines to spawn")
)

// globals
var (
	commitsCount uint
	userIDs      = map[string]uint64{}
	repoIDs      = map[string]uint64{}
)

type commit struct {
	repoID     uint64
	authorID   uint64
	commiterID uint64
	model.Commit
}

func main() {
	var err error

	flag.Usage = func() {
		fmt.Printf("usage: %s [OPTION(S)] [REPOSITORIES ROOT FOLDER]\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()

	if *vflag {
		fmt.Printf("%s - %s\n", filepath.Base(os.Args[0]), version)
		os.Exit(0)
	}

	if *cpuprofile != "" {
		var f *os.File
		f, err = os.Create(*cpuprofile)
		if err != nil {
			glog.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if len(flag.Args()) != 1 {
		fmt.Fprintln(os.Stderr, "invalid # of arguments")
		flag.Usage()
	}

	if len(*configPath) == 0 {
		glog.Fatal(errors.New("a configuration file must be specified"))
	}

	var cfg *config.Config
	cfg, err = config.ReadConfig(*configPath)
	if err != nil {
		glog.Fatal(err)
	}
	commitsCount = cfg.Database.CommitsPerTransaction

	// Make sure we finish writing logs before exiting.
	defer glog.Flush()

	var db *sql.DB
	db, err = openDBSession(*cfg.Database)
	if err != nil {
		return
	}
	defer func() {
		db.Close()
		if err != nil {
			glog.Fatal(err)
		}
	}()

	var empty bool
	if empty, err = isTableEmpty(db, "commits"); !empty || (err != nil) {
		if err == nil {
			err = errors.New("commits table is not empty")
		}
		return
	}

	if err = fetchAllUsers(db); err != nil {
		return
	}
	if err = fetchAllRepos(db); err != nil {
		return
	}

	var w sync.WaitGroup
	var commitsChan chan commit
	if !cfg.Data.CommitDeltas {
		w.Add(1)
		*numGoroutines--
		// we can use pq.CopyIn() to improve db imports
		commitsChan = make(chan commit, commitsCount)
		go func() {
			dbCopyRoutine(db, commitsChan)
			w.Done()
		}()
	}

	reposPathChan := make(chan string)
	var wg sync.WaitGroup
	for w := uint(0); w < *numGoroutines; w++ {
		wg.Add(1)
		go func() {
			repoRoutine(db, cfg.Data, commitsChan, reposPathChan)
			wg.Done()
		}()
	}

	reposDir := flag.Arg(0)
	iterateRepos(reposPathChan, reposDir, *depthflag)

	close(reposPathChan)
	wg.Wait()

	if !cfg.Data.CommitDeltas {
		close(commitsChan)
		w.Wait()
	}
}

func iterateRepos(reposPathChan chan string, path string, depth uint) {
	fis, err := ioutil.ReadDir(path)
	if err != nil {
		glog.Fatal(err)
	}

	if depth == 0 {
		for _, fi := range fis {
			if !fi.IsDir() {
				if filepath.Ext(fi.Name()) != ".tar" {
					continue
				}
			}

			repoPath := filepath.Join(path, fi.Name())
			glog.Info("adding repository:", repoPath, "to the pool")
			reposPathChan <- repoPath
		}
		return
	}

	for _, fi := range fis {
		if !fi.IsDir() {
			continue
		}

		iterateRepos(reposPathChan, filepath.Join(path, fi.Name()), depth-1)
	}
}

func dbCopyRoutine(db *sql.DB, commitsChan chan commit) {
	var err error
	defer func() {
		if err != nil {
			glog.Fatal(err)
		}
	}()

	initTx := func() (*sql.Tx, *sql.Stmt, error) {
		tx, err := db.Begin()
		if err != nil {
			return nil, nil, err
		}
		stmt, err := tx.Prepare(pq.CopyIn("commits", commitFields...))
		if err != nil {
			return nil, nil, err
		}
		return tx, stmt, nil
	}

	commitTx := func(tx *sql.Tx, stmt *sql.Stmt) error {
		defer tx.Rollback()
		if err := stmt.Close(); err != nil {
			return err
		}

		if err := tx.Commit(); err != nil {
			return err
		}
		return nil
	}

	dbExec := func(query string) {
		if err == nil {
			_, err = db.Exec(query)
		}
	}
	// disable constraints and indexes
	dbExec("ALTER TABLE ONLY commit_diff_deltas DROP CONSTRAINT commit_diff_deltas_fk_commits")
	dbExec("ALTER TABLE ONLY commits DROP CONSTRAINT commits_pk")
	dbExec("ALTER TABLE ONLY commits DROP CONSTRAINT commits_fk_repositories")
	dbExec("DROP INDEX fki_commit_diff_deltas_fk_commits")
	dbExec("DROP INDEX fki_commits_fk_repositories")
	defer func() {
		dbExec("ALTER TABLE ONLY commits ADD CONSTRAINT commits_pk PRIMARY KEY (id)")
		dbExec("ALTER TABLE ONLY commit_diff_deltas ADD CONSTRAINT commit_diff_deltas_fk_commits FOREIGN KEY (commit_id) REFERENCES commits(id)")
		dbExec("ALTER TABLE ONLY commits ADD CONSTRAINT commits_fk_repositories FOREIGN KEY (repository_id) REFERENCES repositories(id)")
		dbExec("CREATE INDEX fki_commit_diff_deltas_fk_commits ON commit_diff_deltas USING btree (commit_id)")
		dbExec("CREATE INDEX fki_commits_fk_repositories ON commits USING btree (repository_id)")
	}()

	var tx *sql.Tx
	var stmt *sql.Stmt
	tx, stmt, err = initTx()
	if err != nil {
		return
	}

	var i uint
	for c := range commitsChan {
		_, err = stmt.Exec(
			c.repoID,
			c.authorID,
			c.commiterID,
			c.VCSID,
			c.Message,
			c.AuthorDate,
			c.CommitDate,
			c.FileChangedCount,
			c.InsertionsCount,
			c.DeletionsCount)
		if err != nil {
			tx.Rollback()
			return
		}

		i++

		if i == commitsCount {
			glog.Info("committing", i, "repository commits...")
			if err = commitTx(tx, stmt); err != nil {
				return
			}

			i = 0
			if tx, stmt, err = initTx(); err != nil {
				return
			}
		}
	}

	if i > 0 {
		glog.Info("committing", i, "repository commits...")
		err = commitTx(tx, stmt)
	}
}
func repoRoutine(db *sql.DB, cfg config.DataConfig, commitsChan chan commit, reposPathChan chan string) {
	for path := range reposPathChan {
		work := func() error {
			repository, err := repo.New(cfg, path)
			if err != nil {
				return err
			}
			defer repository.Cleanup()

			if err = repository.FetchCommits(); err != nil {
				return err
			}

			if cfg.CommitDeltas {
				if err = insertRepoData(db, repository); err != nil {
					return err
				}
			} else {
				repoID, ok := repoIDs[repository.GetCloneURL()]
				if !ok {
					return errors.New("cannot find corresponding repository in database")
				}
				for _, c := range repository.GetCommits() {
					commitsChan <- commit{repoID, userIDs[c.Author.Email], userIDs[c.Committer.Email], c}
				}
			}
			return nil
		}
		if err := work(); err != nil {
			glog.Error(err)
		}
	}
}

// openDBSession creates a session to the database.
func openDBSession(cfg config.DatabaseConfig) (*sql.DB, error) {
	dbURL := fmt.Sprintf(
		"user='%s' password='%s' host='%s' port=%d dbname='%s' sslmode='%s'",
		cfg.UserName, cfg.Password, cfg.HostName, cfg.Port, cfg.DBName, cfg.SSLMode)

	return sql.Open("postgres", dbURL)
}

// insertRepoData inserts repository data into the database, or updates it
// if it is already there.
func insertRepoData(db *sql.DB, r repo.Repo) error {
	if db == nil {
		return errors.New("nil database given")
	}

	repoID, ok := repoIDs[r.GetCloneURL()]
	if !ok {
		return errors.New("cannot find corresponding repository in database")
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	commitStmt, err := tx.Prepare(genInsQuery("commits", commitFields...) + " RETURNING id")
	if err != nil {
		return err
	}

	deltaStmt, err := tx.Prepare(genInsQuery("commit_diff_deltas", diffDeltaFields...))
	if err != nil {
		return err
	}

	for _, c := range r.GetCommits() {
		if err := insertCommit(repoID, c, tx, commitStmt, deltaStmt); err != nil {
			return err
		}
	}

	if err := commitStmt.Close(); err != nil {
		return err
	}
	if err := deltaStmt.Close(); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

// insertCommit inserts a commit into the database
func insertCommit(repoID uint64, c model.Commit, tx *sql.Tx, commitStmt, deltaStmt *sql.Stmt) error {
	authorID := userIDs[c.Author.Email]
	committerID := userIDs[c.Committer.Email]

	var commitID uint64
	err := commitStmt.QueryRow(
		repoID, authorID, committerID,
		c.VCSID, c.Message, c.AuthorDate, c.CommitDate,
		c.FileChangedCount, c.InsertionsCount, c.DeletionsCount).Scan(&commitID)
	if err != nil {
		return err
	}

	for _, d := range c.DiffDelta {
		if err := insertDiffDelta(commitID, d, deltaStmt); err != nil {
			return err
		}
	}

	return nil
}

// insertDiffDelta inserts a commit diff delta into the database.
func insertDiffDelta(commitID uint64, d model.DiffDelta, stmt *sql.Stmt) error {
	_, err := stmt.Exec(commitID, d.Status, d.Binary, d.Similarity, d.OldFilePath, d.NewFilePath)
	if err != nil {
		return err
	}
	return nil
}

// fetchAllUsers fetch users IDs and put them into the userIDs global hashmap
// with their email address as keys.
func fetchAllUsers(db *sql.DB) error {
	rows, err := db.Query("SELECT id, email FROM users WHERE email IS NOT NULL AND email != ''")
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var email string
		var id uint64
		if err := rows.Scan(&id, &email); err != nil {
			return err
		}
		userIDs[email] = id
	}

	return nil
}

func fetchAllRepos(db *sql.DB) error {
	rows, err := db.Query("SELECT id, clone_url FROM repositories")
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var id uint64
		var cloneURL string
		if err := rows.Scan(&id, &cloneURL); err != nil {
			return err
		}
		repoIDs[cloneURL] = id
	}

	return nil
}

// genInsQuery generates a query string for an insertion in the database.
func genInsQuery(tableName string, fields ...string) string {
	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("INSERT INTO %s(%s)\n",
		tableName, strings.Join(fields, ",")))
	buf.WriteString("VALUES(")

	for ind := range fields {
		if ind > 0 {
			buf.WriteString(",")
		}

		buf.WriteString(fmt.Sprintf("$%d", ind+1))
	}

	buf.WriteString(")\n")

	return buf.String()
}

// isTableEmpty returns true of the table tableName is empty, false otherwise.
func isTableEmpty(db *sql.DB, tableName string) (bool, error) {
	var state bool
	err := db.QueryRow("SELECT EXISTS(SELECT 1 FROM " + tableName + " LIMIT 1)").Scan(&state)
	return !state, err
}
<commit_msg>repotool-db: update version number and flag<commit_after>// Copyright 2014-2015 The DevMine authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package repotool-db is able to fetch information from a source code repository.
// Typically, it can get all commits, their authors and commiters and so on
// and is able to populate the information into a PostgreSQL database.
// Currently, on the Git VCS is supported.
package main

import (
	"bytes"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strings"
	"sync"

	"github.com/golang/glog"
	"github.com/lib/pq"

	"github.com/DevMine/repotool/config"
	"github.com/DevMine/repotool/model"
	"github.com/DevMine/repotool/repo"
)

const version = "1.0.0"

// database fields per tables
var (
	diffDeltaFields = []string{
		"commit_id",
		"file_status",
		"is_file_binary",
		"similarity",
		"old_file_path",
		"new_file_path"}

	commitFields = []string{
		"repository_id",
		"author_id",
		"committer_id",
		"vcs_id",
		"message",
		"author_date",
		"commit_date",
		"file_changed_count",
		"insertions_count",
		"deletions_count"}
)

// program flags
var (
	configPath    = flag.String("c", "", "configuration file")
	versionflag   = flag.Bool("version", false, "print version.")
	cpuprofile    = flag.String("cpuprofile", "", "write cpu profile to file")
	depthflag     = flag.Uint("d", 0, "depth level where to find repositories")
	numGoroutines = flag.Uint("g", uint(runtime.NumCPU()), "max number of goroutines to spawn")
)

// globals
var (
	commitsCount uint
	userIDs      = map[string]uint64{}
	repoIDs      = map[string]uint64{}
)

type commit struct {
	repoID     uint64
	authorID   uint64
	commiterID uint64
	model.Commit
}

func main() {
	var err error

	flag.Usage = func() {
		fmt.Printf("usage: %s [OPTION(S)] [REPOSITORIES ROOT FOLDER]\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()

	if *versionflag {
		fmt.Printf("%s - %s\n", filepath.Base(os.Args[0]), version)
		os.Exit(0)
	}

	if *cpuprofile != "" {
		var f *os.File
		f, err = os.Create(*cpuprofile)
		if err != nil {
			glog.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	if len(flag.Args()) != 1 {
		fmt.Fprintln(os.Stderr, "invalid # of arguments")
		flag.Usage()
	}

	if len(*configPath) == 0 {
		glog.Fatal(errors.New("a configuration file must be specified"))
	}

	var cfg *config.Config
	cfg, err = config.ReadConfig(*configPath)
	if err != nil {
		glog.Fatal(err)
	}
	commitsCount = cfg.Database.CommitsPerTransaction

	// Make sure we finish writing logs before exiting.
	defer glog.Flush()

	var db *sql.DB
	db, err = openDBSession(*cfg.Database)
	if err != nil {
		return
	}
	defer func() {
		db.Close()
		if err != nil {
			glog.Fatal(err)
		}
	}()

	var empty bool
	if empty, err = isTableEmpty(db, "commits"); !empty || (err != nil) {
		if err == nil {
			err = errors.New("commits table is not empty")
		}
		return
	}

	if err = fetchAllUsers(db); err != nil {
		return
	}
	if err = fetchAllRepos(db); err != nil {
		return
	}

	var w sync.WaitGroup
	var commitsChan chan commit
	if !cfg.Data.CommitDeltas {
		w.Add(1)
		*numGoroutines--
		// we can use pq.CopyIn() to improve db imports
		commitsChan = make(chan commit, commitsCount)
		go func() {
			dbCopyRoutine(db, commitsChan)
			w.Done()
		}()
	}

	reposPathChan := make(chan string)
	var wg sync.WaitGroup
	for w := uint(0); w < *numGoroutines; w++ {
		wg.Add(1)
		go func() {
			repoRoutine(db, cfg.Data, commitsChan, reposPathChan)
			wg.Done()
		}()
	}

	reposDir := flag.Arg(0)
	iterateRepos(reposPathChan, reposDir, *depthflag)

	close(reposPathChan)
	wg.Wait()

	if !cfg.Data.CommitDeltas {
		close(commitsChan)
		w.Wait()
	}
}

func iterateRepos(reposPathChan chan string, path string, depth uint) {
	fis, err := ioutil.ReadDir(path)
	if err != nil {
		glog.Fatal(err)
	}

	if depth == 0 {
		for _, fi := range fis {
			if !fi.IsDir() {
				if filepath.Ext(fi.Name()) != ".tar" {
					continue
				}
			}

			repoPath := filepath.Join(path, fi.Name())
			glog.Info("adding repository:", repoPath, "to the pool")
			reposPathChan <- repoPath
		}
		return
	}

	for _, fi := range fis {
		if !fi.IsDir() {
			continue
		}

		iterateRepos(reposPathChan, filepath.Join(path, fi.Name()), depth-1)
	}
}

func dbCopyRoutine(db *sql.DB, commitsChan chan commit) {
	var err error
	defer func() {
		if err != nil {
			glog.Fatal(err)
		}
	}()

	initTx := func() (*sql.Tx, *sql.Stmt, error) {
		tx, err := db.Begin()
		if err != nil {
			return nil, nil, err
		}
		stmt, err := tx.Prepare(pq.CopyIn("commits", commitFields...))
		if err != nil {
			return nil, nil, err
		}
		return tx, stmt, nil
	}

	commitTx := func(tx *sql.Tx, stmt *sql.Stmt) error {
		defer tx.Rollback()
		if err := stmt.Close(); err != nil {
			return err
		}

		if err := tx.Commit(); err != nil {
			return err
		}
		return nil
	}

	dbExec := func(query string) {
		if err == nil {
			_, err = db.Exec(query)
		}
	}
	// disable constraints and indexes
	dbExec("ALTER TABLE ONLY commit_diff_deltas DROP CONSTRAINT commit_diff_deltas_fk_commits")
	dbExec("ALTER TABLE ONLY commits DROP CONSTRAINT commits_pk")
	dbExec("ALTER TABLE ONLY commits DROP CONSTRAINT commits_fk_repositories")
	dbExec("DROP INDEX fki_commit_diff_deltas_fk_commits")
	dbExec("DROP INDEX fki_commits_fk_repositories")
	defer func() {
		dbExec("ALTER TABLE ONLY commits ADD CONSTRAINT commits_pk PRIMARY KEY (id)")
		dbExec("ALTER TABLE ONLY commit_diff_deltas ADD CONSTRAINT commit_diff_deltas_fk_commits FOREIGN KEY (commit_id) REFERENCES commits(id)")
		dbExec("ALTER TABLE ONLY commits ADD CONSTRAINT commits_fk_repositories FOREIGN KEY (repository_id) REFERENCES repositories(id)")
		dbExec("CREATE INDEX fki_commit_diff_deltas_fk_commits ON commit_diff_deltas USING btree (commit_id)")
		dbExec("CREATE INDEX fki_commits_fk_repositories ON commits USING btree (repository_id)")
	}()

	var tx *sql.Tx
	var stmt *sql.Stmt
	tx, stmt, err = initTx()
	if err != nil {
		return
	}

	var i uint
	for c := range commitsChan {
		_, err = stmt.Exec(
			c.repoID,
			c.authorID,
			c.commiterID,
			c.VCSID,
			c.Message,
			c.AuthorDate,
			c.CommitDate,
			c.FileChangedCount,
			c.InsertionsCount,
			c.DeletionsCount)
		if err != nil {
			tx.Rollback()
			return
		}

		i++

		if i == commitsCount {
			glog.Info("committing", i, "repository commits...")
			if err = commitTx(tx, stmt); err != nil {
				return
			}

			i = 0
			if tx, stmt, err = initTx(); err != nil {
				return
			}
		}
	}

	if i > 0 {
		glog.Info("committing", i, "repository commits...")
		err = commitTx(tx, stmt)
	}
}
func repoRoutine(db *sql.DB, cfg config.DataConfig, commitsChan chan commit, reposPathChan chan string) {
	for path := range reposPathChan {
		work := func() error {
			repository, err := repo.New(cfg, path)
			if err != nil {
				return err
			}
			defer repository.Cleanup()

			if err = repository.FetchCommits(); err != nil {
				return err
			}

			if cfg.CommitDeltas {
				if err = insertRepoData(db, repository); err != nil {
					return err
				}
			} else {
				repoID, ok := repoIDs[repository.GetCloneURL()]
				if !ok {
					return errors.New("cannot find corresponding repository in database")
				}
				for _, c := range repository.GetCommits() {
					commitsChan <- commit{repoID, userIDs[c.Author.Email], userIDs[c.Committer.Email], c}
				}
			}
			return nil
		}
		if err := work(); err != nil {
			glog.Error(err)
		}
	}
}

// openDBSession creates a session to the database.
func openDBSession(cfg config.DatabaseConfig) (*sql.DB, error) {
	dbURL := fmt.Sprintf(
		"user='%s' password='%s' host='%s' port=%d dbname='%s' sslmode='%s'",
		cfg.UserName, cfg.Password, cfg.HostName, cfg.Port, cfg.DBName, cfg.SSLMode)

	return sql.Open("postgres", dbURL)
}

// insertRepoData inserts repository data into the database, or updates it
// if it is already there.
func insertRepoData(db *sql.DB, r repo.Repo) error {
	if db == nil {
		return errors.New("nil database given")
	}

	repoID, ok := repoIDs[r.GetCloneURL()]
	if !ok {
		return errors.New("cannot find corresponding repository in database")
	}

	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	commitStmt, err := tx.Prepare(genInsQuery("commits", commitFields...) + " RETURNING id")
	if err != nil {
		return err
	}

	deltaStmt, err := tx.Prepare(genInsQuery("commit_diff_deltas", diffDeltaFields...))
	if err != nil {
		return err
	}

	for _, c := range r.GetCommits() {
		if err := insertCommit(repoID, c, tx, commitStmt, deltaStmt); err != nil {
			return err
		}
	}

	if err := commitStmt.Close(); err != nil {
		return err
	}
	if err := deltaStmt.Close(); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

// insertCommit inserts a commit into the database
func insertCommit(repoID uint64, c model.Commit, tx *sql.Tx, commitStmt, deltaStmt *sql.Stmt) error {
	authorID := userIDs[c.Author.Email]
	committerID := userIDs[c.Committer.Email]

	var commitID uint64
	err := commitStmt.QueryRow(
		repoID, authorID, committerID,
		c.VCSID, c.Message, c.AuthorDate, c.CommitDate,
		c.FileChangedCount, c.InsertionsCount, c.DeletionsCount).Scan(&commitID)
	if err != nil {
		return err
	}

	for _, d := range c.DiffDelta {
		if err := insertDiffDelta(commitID, d, deltaStmt); err != nil {
			return err
		}
	}

	return nil
}

// insertDiffDelta inserts a commit diff delta into the database.
func insertDiffDelta(commitID uint64, d model.DiffDelta, stmt *sql.Stmt) error {
	_, err := stmt.Exec(commitID, d.Status, d.Binary, d.Similarity, d.OldFilePath, d.NewFilePath)
	if err != nil {
		return err
	}
	return nil
}

// fetchAllUsers fetch users IDs and put them into the userIDs global hashmap
// with their email address as keys.
func fetchAllUsers(db *sql.DB) error {
	rows, err := db.Query("SELECT id, email FROM users WHERE email IS NOT NULL AND email != ''")
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var email string
		var id uint64
		if err := rows.Scan(&id, &email); err != nil {
			return err
		}
		userIDs[email] = id
	}

	return nil
}

func fetchAllRepos(db *sql.DB) error {
	rows, err := db.Query("SELECT id, clone_url FROM repositories")
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var id uint64
		var cloneURL string
		if err := rows.Scan(&id, &cloneURL); err != nil {
			return err
		}
		repoIDs[cloneURL] = id
	}

	return nil
}

// genInsQuery generates a query string for an insertion in the database.
func genInsQuery(tableName string, fields ...string) string {
	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("INSERT INTO %s(%s)\n",
		tableName, strings.Join(fields, ",")))
	buf.WriteString("VALUES(")

	for ind := range fields {
		if ind > 0 {
			buf.WriteString(",")
		}

		buf.WriteString(fmt.Sprintf("$%d", ind+1))
	}

	buf.WriteString(")\n")

	return buf.String()
}

// isTableEmpty returns true of the table tableName is empty, false otherwise.
func isTableEmpty(db *sql.DB, tableName string) (bool, error) {
	var state bool
	err := db.QueryRow("SELECT EXISTS(SELECT 1 FROM " + tableName + " LIMIT 1)").Scan(&state)
	return !state, err
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"

	"github.com/Impyy/scytale"
	"github.com/Impyy/scytale/auth"
)

var (
	assetMap = GetAssets()
)

const (
	imgDir             = "./img/"
	uploadReqMaxSize   = 5000000 //in bytes
	extensionMaxLength = 10      //in chars
)

type Settings struct {
	Port int
	Keys auth.KeyList
}

type Server struct {
	settings *Settings
}

func New(settings *Settings) *Server {
	return &Server{settings: settings}
}

func (s *Server) Serve() error {
	//create the img directory if it doesn't exist
	if _, err := os.Stat(imgDir); os.IsNotExist(err) {
		err = os.Mkdir(imgDir, 0777)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	http.HandleFunc("/", s.handleHTTPRequest)
	http.HandleFunc("/ul", s.handleUploadRequest)
	http.HandleFunc("/dl", s.handleDownloadRequest)
	return http.ListenAndServe(fmt.Sprintf(":%d", s.settings.Port), nil)
}

func (s *Server) handleHTTPRequest(w http.ResponseWriter, r *http.Request) {
	urlPath := r.URL.Path[1:]
	switch urlPath {
	case "":
		http.Error(w, http.StatusText(403), 403)
		return
	case "view":
		urlPath = "view.html"
	}

	data, exists := assetMap[urlPath]
	if !exists {
		http.Error(w, http.StatusText(404), 404)
	} else {
		w.Header().Set("Content-Type", mimeTypeByExtension(urlPath))
		w.Write(data)
	}
}

//TODO: clean and split up this handler
func (s *Server) handleUploadRequest(w http.ResponseWriter, r *http.Request) {
	var key auth.Key
	var filenameString string
	var file *os.File
	var err error
	res := scytale.UploadResponse{ErrorCode: scytale.ErrorCodeOK}
	req := scytale.UploadRequest{}
	ext := ""

	if r.Method != "POST" {
		res.ErrorCode = scytale.ErrorCodeFormat
		goto sendRes
	}

	if r.ContentLength > uploadReqMaxSize {
		res.ErrorCode = scytale.ErrorCodeSize
		goto sendRes
	}

	key, err = auth.ParseKey([]byte(r.Header.Get("X-Key")))
	if err != nil {
		res.ErrorCode = scytale.ErrorCodeFormat
		goto sendRes
	}

	if !s.settings.Keys.Contains(key) {
		res.ErrorCode = scytale.ErrorCodePermissionDenied
		goto sendRes
	}

	err = json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		res.ErrorCode = scytale.ErrorCodeFormat
		goto sendRes
	}

	if !req.IsEncrypted {
		if len(req.Extension) > extensionMaxLength {
			res.ErrorCode = scytale.ErrorCodeExtensionTooLong
			goto sendRes
		}
		ext = req.Extension
	}

	filenameString, err = generateFilename(ext)
	if err != nil {
		res.ErrorCode = scytale.ErrorCodeInternal
		goto sendRes
	}

	file, err = os.Create(path.Join(imgDir, filenameString))
	if err != nil {
		fmt.Printf("file create err: %s\n", err.Error())
		res.ErrorCode = scytale.ErrorCodeInternal
		goto sendRes
	}
	defer file.Close()

	_, err = io.Copy(file, base64.NewDecoder(base64.StdEncoding, bytes.NewReader([]byte(req.Data))))
	if err != nil {
		fmt.Printf("io copy err: %s\n", err.Error())
		res.ErrorCode = scytale.ErrorCodeInternal
		goto sendRes
	}

	if req.IsEncrypted {
		res.Location = fmt.Sprintf("/view?l=%s", filenameString)
	} else {
		res.Location = fmt.Sprintf("/dl?l=%s", filenameString)
	}

sendRes:
	resBytes, err := json.Marshal(res)
	if err != nil {
		http.Error(w, "", 500)
		return
	}
	w.Write(resBytes)
}

func (s *Server) handleDownloadRequest(w http.ResponseWriter, r *http.Request) {
	loc := r.URL.Query().Get("l")
	if loc == "" {
		fmt.Printf("error: empty loc\n")
	}

	p := path.Join(imgDir, loc)
	data, err := ioutil.ReadFile(p)

	if err != nil {
		http.Error(w, http.StatusText(404), 404)
	} else {
		w.Header().Set("Content-Type", mimeTypeByExtension(p))
		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(data)))
		w.Write(data)
	}
}
<commit_msg>Close path traversal hole<commit_after>package server

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"

	"github.com/Impyy/scytale"
	"github.com/Impyy/scytale/auth"
)

var (
	assetMap = GetAssets()
)

const (
	imgDir             = "./img/"
	uploadReqMaxSize   = 5000000 //in bytes
	extensionMaxLength = 10      //in chars
)

type Settings struct {
	Port int
	Keys auth.KeyList
}

type Server struct {
	settings *Settings
}

func New(settings *Settings) *Server {
	return &Server{settings: settings}
}

func (s *Server) Serve() error {
	//create the img directory if it doesn't exist
	if _, err := os.Stat(imgDir); os.IsNotExist(err) {
		err = os.Mkdir(imgDir, 0777)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	http.HandleFunc("/", s.handleHTTPRequest)
	http.HandleFunc("/ul", s.handleUploadRequest)
	http.HandleFunc("/dl", s.handleDownloadRequest)
	return http.ListenAndServe(fmt.Sprintf(":%d", s.settings.Port), nil)
}

func (s *Server) handleHTTPRequest(w http.ResponseWriter, r *http.Request) {
	urlPath := r.URL.Path[1:]
	switch urlPath {
	case "":
		http.Error(w, http.StatusText(403), 403)
		return
	case "view":
		urlPath = "view.html"
	}

	data, exists := assetMap[urlPath]
	if !exists {
		http.Error(w, http.StatusText(404), 404)
	} else {
		w.Header().Set("Content-Type", mimeTypeByExtension(urlPath))
		w.Write(data)
	}
}

//TODO: clean and split up this handler
func (s *Server) handleUploadRequest(w http.ResponseWriter, r *http.Request) {
	var key auth.Key
	var filenameString string
	var file *os.File
	var err error
	res := scytale.UploadResponse{ErrorCode: scytale.ErrorCodeOK}
	req := scytale.UploadRequest{}
	ext := ""

	if r.Method != "POST" {
		res.ErrorCode = scytale.ErrorCodeFormat
		goto sendRes
	}

	if r.ContentLength > uploadReqMaxSize {
		res.ErrorCode = scytale.ErrorCodeSize
		goto sendRes
	}

	key, err = auth.ParseKey([]byte(r.Header.Get("X-Key")))
	if err != nil {
		res.ErrorCode = scytale.ErrorCodeFormat
		goto sendRes
	}

	if !s.settings.Keys.Contains(key) {
		res.ErrorCode = scytale.ErrorCodePermissionDenied
		goto sendRes
	}

	err = json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		res.ErrorCode = scytale.ErrorCodeFormat
		goto sendRes
	}

	if !req.IsEncrypted {
		if len(req.Extension) > extensionMaxLength {
			res.ErrorCode = scytale.ErrorCodeExtensionTooLong
			goto sendRes
		}
		ext = req.Extension
	}

	filenameString, err = generateFilename(ext)
	if err != nil {
		res.ErrorCode = scytale.ErrorCodeInternal
		goto sendRes
	}

	file, err = os.Create(path.Join(imgDir, filenameString))
	if err != nil {
		fmt.Printf("file create err: %s\n", err.Error())
		res.ErrorCode = scytale.ErrorCodeInternal
		goto sendRes
	}
	defer file.Close()

	_, err = io.Copy(file, base64.NewDecoder(base64.StdEncoding, bytes.NewReader([]byte(req.Data))))
	if err != nil {
		fmt.Printf("io copy err: %s\n", err.Error())
		res.ErrorCode = scytale.ErrorCodeInternal
		goto sendRes
	}

	if req.IsEncrypted {
		res.Location = fmt.Sprintf("/view?l=%s", filenameString)
	} else {
		res.Location = fmt.Sprintf("/dl?l=%s", filenameString)
	}

sendRes:
	resBytes, err := json.Marshal(res)
	if err != nil {
		http.Error(w, "", 500)
		return
	}
	w.Write(resBytes)
}

func (s *Server) handleDownloadRequest(w http.ResponseWriter, r *http.Request) {
	loc := r.URL.Query().Get("l")
	if loc == "" {
		fmt.Printf("error: empty loc\n")
	}

	p := path.Join(imgDir, filepath.Base(loc))
	data, err := ioutil.ReadFile(p)

	if err != nil {
		http.Error(w, http.StatusText(404), 404)
	} else {
		w.Header().Set("Content-Type", mimeTypeByExtension(p))
		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(data)))
		w.Write(data)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"runtime"

	sdk "github.com/operator-framework/operator-sdk/pkg/sdk"
	k8sutil "github.com/operator-framework/operator-sdk/pkg/util/k8sutil"
	sdkVersion "github.com/operator-framework/operator-sdk/version"
	stub "github.com/pravega/zookeeper-operator/pkg/stub"

	"fmt"
	"github.com/sirupsen/logrus"
	"os"
)

func printVersion() {
	logrus.Infof("Go Version: %s", runtime.Version())
	logrus.Infof("Go OS/Arch: %s/%s", runtime.GOOS, runtime.GOARCH)
	logrus.Infof("operator-sdk Version: %v", sdkVersion.Version)
}

func main() {
	printVersion()

	resource := "zookeeper.pravega.io/v1beta1"
	kind := "ZookeeperCluster"
	namespace, err := getWatchNamespaceAllowBlank()
	if err != nil {
		logrus.Fatalf("Failed to get watch namespace: %v", err)
	}
	resyncPeriod := 5
	logrus.Infof("Watching %s, %s, %s, %d", resource, kind, namespace, resyncPeriod)
	sdk.Watch(resource, kind, namespace, resyncPeriod)
	sdk.Handle(stub.NewHandler())
	sdk.Run(context.TODO())
}

// GetWatchNamespaceAllowBlank returns the namespace the operator should be watching for changes
func getWatchNamespaceAllowBlank() (string, error) {
	ns, found := os.LookupEnv(k8sutil.WatchNamespaceEnvVar)
	if !found {
		return "", fmt.Errorf("%s must be set", k8sutil.WatchNamespaceEnvVar)
	}
	return ns, nil
}
<commit_msg>Watch all namespaces if one is not specified<commit_after>package main

import (
	"context"
	"runtime"

	"github.com/operator-framework/operator-sdk/pkg/sdk"
	"github.com/operator-framework/operator-sdk/pkg/util/k8sutil"
	sdkVersion "github.com/operator-framework/operator-sdk/version"
	"github.com/pravega/zookeeper-operator/pkg/stub"

	"github.com/sirupsen/logrus"
	"os"
)

func printVersion() {
	logrus.Infof("Go Version: %s", runtime.Version())
	logrus.Infof("Go OS/Arch: %s/%s", runtime.GOOS, runtime.GOARCH)
	logrus.Infof("operator-sdk Version: %v", sdkVersion.Version)
}

func main() {
	printVersion()

	resource := "zookeeper.pravega.io/v1beta1"
	kind := "ZookeeperCluster"
	namespace := getWatchNamespaceAllowBlank()
	resyncPeriod := 5
	logrus.Infof("Watching %s, %s, %s, %d", resource, kind, namespace, resyncPeriod)
	sdk.Watch(resource, kind, namespace, resyncPeriod)
	sdk.Handle(stub.NewHandler())
	sdk.Run(context.TODO())
}

// GetWatchNamespaceAllowBlank returns the namespace the operator should be watching for changes
func getWatchNamespaceAllowBlank() string {
	ns, found := os.LookupEnv(k8sutil.WatchNamespaceEnvVar)
	if !found {
		logrus.Infof("%s is not set, watching all namespaces", k8sutil.WatchNamespaceEnvVar)
		ns = ""
	}
	return ns
}
<|endoftext|>