<commit_before>// Copyright 2013-2014 Alexandre Fiori
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Web server of freegeoip.net

package main

import (
	"database/sql"
	"encoding/binary"
	"encoding/json"
	"encoding/xml"
	"expvar"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"strings"
	"sync"
	"time"

	"github.com/fiorix/go-web/httpxtra"

	_ "github.com/mattn/go-sqlite3"
	//_ "code.google.com/p/gosqlite/sqlite3"
)

var (
	conf        = ConfigFile{}
	protoCount  = expvar.NewMap("Protocol") // HTTP or HTTPS
	outputCount = expvar.NewMap("Output")   // json, xml, csv or other
	statusCount = expvar.NewMap("Status")   // 200, 403, 404, etc
)

func main() {
	cf := flag.String("config", "freegeoip.conf", "set config file")
	prof := flag.Bool("profile", false, "run cpu and mem profiling")
	flag.Parse()

	if fd, err := os.Open(*cf); err != nil {
		log.Fatal(err)
	} else {
		if err = xml.NewDecoder(fd).Decode(&conf); err != nil {
			log.Fatal(err)
		}
		fd.Close()
	}

	if *prof {
		profile()
	}

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Printf("FreeGeoIP server starting. debug=%t", conf.Debug)

	if conf.Debug && len(conf.DebugSrv) > 0 {
		go func() {
			// server for expvar's /debug/vars only
			log.Printf("Starting DEBUG HTTP server on %s", conf.DebugSrv)
			log.Fatal(http.ListenAndServe(conf.DebugSrv, nil))
		}()
	}

	mux := http.NewServeMux()
	mux.Handle("/", http.FileServer(http.Dir(conf.DocumentRoot)))

	h := LookupHandler()
	mux.HandleFunc("/csv/", h)
	mux.HandleFunc("/xml/", h)
	mux.HandleFunc("/json/", h)

	wg := new(sync.WaitGroup)
	for _, l := range conf.Listen {
		if l.Addr == "" {
			continue
		}
		wg.Add(1)
		h := httpxtra.Handler{Handler: mux, XHeaders: l.XHeaders}
		if l.Log {
			h.Logger = logger
		}

		s := http.Server{
			Addr:         l.Addr,
			Handler:      h,
			ReadTimeout:  15 * time.Second,
			WriteTimeout: 15 * time.Second,
		}

		if l.KeyFile == "" && l.CertFile == "" {
			log.Printf("Starting HTTP server on %s "+
				"log=%t xheaders=%t",
				l.Addr, l.Log, l.XHeaders)
			go func() {
				log.Fatal(httpxtra.ListenAndServe(s))
			}()
		} else {
			log.Printf("Starting HTTPS server on %s "+
				"log=%t xheaders=%t cert=%s key=%s",
				l.Addr, l.Log, l.XHeaders,
				l.CertFile, l.KeyFile)
			go func() {
				log.Fatal(s.ListenAndServeTLS(
					l.CertFile,
					l.KeyFile,
				))
			}()
		}
	}

	wg.Wait()
}

// LookupHandler handles GET on /csv, /xml and /json.
func LookupHandler() http.HandlerFunc {
	db, err := sql.Open("sqlite3", conf.IPDB.File)
	if err != nil {
		log.Fatal(err)
	}

	_, err = db.Exec("PRAGMA cache_size=" + conf.IPDB.CacheSize)
	if err != nil {
		log.Fatal(err)
	}

	stmt, err := db.Prepare(ipdb_query)
	if err != nil {
		log.Fatal(err)
	}

	st := time.Now()
	cache := NewCache(db)
	log.Println("IPDB cached in", time.Since(st))

	//defer stmt.Close()

	var quota Quota
	if len(conf.Redis) == 0 {
		quota = new(MapQuota)
		quota.Setup()
		log.Printf("Using internal map to manage quota.")
	} else {
		quota = new(RedisQuota)
		quota.Setup(conf.Redis...)
		log.Printf("Using redis to manage quota: %s", conf.Redis)
	}

	return func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case "GET":
			w.Header().Set("Access-Control-Allow-Origin", "*")
		case "OPTIONS":
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Content-Type", "text/plain")
			w.Header().Set("Access-Control-Allow-Methods", "GET")
			w.Header().Set("Access-Control-Allow-Headers", "X-Requested-With")
			w.WriteHeader(200)
			return
		default:
			w.Header().Set("Allow", "GET, OPTIONS")
			http.Error(w, http.StatusText(405), 405)
			return
		}

		// GET continues...
		var srcIP net.IP

		// If xheaders is enabled, RemoteAddr might be a copy of
		// the X-Real-IP or X-Forwarded-For HTTP headers, which
		// can be a comma separated list of IPs. In this case,
		// only the first IP in the list is used.
		r.RemoteAddr = strings.SplitN(r.RemoteAddr, ",", 2)[0]

		if ip, _, err := net.SplitHostPort(r.RemoteAddr); err != nil {
			srcIP = net.ParseIP(r.RemoteAddr) // Use X-Real-IP
		} else {
			srcIP = net.ParseIP(ip)
		}

		if srcIP == nil {
			http.Error(w, http.StatusText(400), 400)
			return
		}

		nsrcIP, err := ip2int(srcIP)
		if err != nil {
			if conf.Debug {
				log.Println(err)
			}
			http.Error(w, http.StatusText(400), 400)
			return
		}

		// Check quota.
		if conf.Limit.MaxRequests > 0 {
			var ok bool
			if ok, err = quota.Ok(nsrcIP); err != nil {
				if conf.Debug {
					log.Println(err) // redis error
				}
				http.Error(w, http.StatusText(503), 503)
				return
			} else if !ok {
				// Over quota, soz :(
				http.Error(w, http.StatusText(403), 403)
				return
			}
		}

		var (
			queryIP  net.IP
			nqueryIP uint32
		)

		// Parse URL (e.g. /csv/ip, /xml/)
		a := strings.SplitN(r.URL.Path, "/", 3)
		if len(a) == 3 && a[2] != "" {
			addrs, err := net.LookupHost(a[2])
			if err != nil {
				// DNS lookup failed, assume host not found.
				http.Error(w, http.StatusText(404), 404)
				return
			}

			if queryIP = net.ParseIP(addrs[0]); queryIP == nil {
				http.Error(w, http.StatusText(400), 400)
				return
			}

			nqueryIP, err = ip2int(net.ParseIP(addrs[0]))
			if err != nil {
				if conf.Debug {
					log.Println(err)
				}
				http.Error(w, http.StatusText(400), 400)
				return
			}
		} else {
			queryIP = srcIP
			nqueryIP = nsrcIP
		}

		// Query the db.
		geoip, err := ipdb_lookup(stmt, cache, queryIP, nqueryIP)
		if err != nil {
			http.NotFound(w, r)
			return
		}

		switch a[1][0] {
		case 'j': // json
			enc := json.NewEncoder(w)
			if cb := r.FormValue("callback"); cb != "" {
				w.Header().Set("Content-Type", "text/javascript")
				fmt.Fprintf(w, "%s(", cb)
				enc.Encode(geoip)
				fmt.Fprintf(w, ");")
			} else {
				w.Header().Set("Content-Type", "application/json")
				enc.Encode(geoip)
			}
		case 'x': // xml
			w.Header().Set("Content-Type", "application/xml")
			fmt.Fprintf(w, xml.Header)
			enc := xml.NewEncoder(w)
			enc.Indent("", " ")
			enc.Encode(geoip)
			fmt.Fprintf(w, "\n")
		case 'c': // csv
			w.Header().Set("Content-Type", "application/csv")
			fmt.Fprintf(w, `"%s","%s","%s","%s","%s","%s",`+
				`"%s","%0.4f","%0.4f","%s","%s"`+"\r\n",
				geoip.Ip,
				geoip.CountryCode, geoip.CountryName,
				geoip.RegionCode, geoip.RegionName,
				geoip.CityName, geoip.ZipCode,
				geoip.Latitude, geoip.Longitude,
				geoip.MetroCode, geoip.AreaCode)
		}
	}
}

func ip2int(ip net.IP) (uint32, error) {
	ipv4 := ip.To4()
	if ipv4 == nil {
		return 0, fmt.Errorf("IP %s is not IPv4", ip.String())
	}

	return binary.BigEndian.Uint32(ipv4), nil
}

func profile() {
	f, err := os.Create("freegeoip.cpu.prof")
	if err != nil {
		log.Fatal(err)
	}

	pprof.StartCPUProfile(f)
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)

	go func() {
		<-sig
		pprof.StopCPUProfile()
		f.Close()
		f, err = os.Create("freegeoip.mem.prof")
		if err != nil {
			log.Fatal(err)
		}
		pprof.WriteHeapProfile(f)
		os.Exit(0)
	}()
}

func logger(r *http.Request, created time.Time, status, bytes int) {
	//fmt.Println(httpxtra.ApacheCommonLog(r, created, status, bytes))
	var (
		s, ip string
		err   error
	)
	if r.TLS == nil {
		s = "HTTP"
	} else {
		s = "HTTPS"
	}
	if ip, _, err = net.SplitHostPort(r.RemoteAddr); err != nil {
		ip = r.RemoteAddr
	}
	log.Printf("%s %d %s %q (%s) :: %d bytes in %s",
		s,
		status,
		r.Method,
		r.URL.Path,
		ip,
		bytes,
		time.Since(created),
	)
	if conf.Debug {
		go func() {
			protoCount.Add(s, 1)
			statusCount.Add(fmt.Sprintf("%d", status), 1)
			switch strings.SplitN(r.URL.Path, "/", 3)[1] {
			case "json":
				outputCount.Add("json", 1)
			case "xml":
				outputCount.Add("xml", 1)
			case "csv":
				outputCount.Add("csv", 1)
			default:
				outputCount.Add("other", 1)
			}
		}()
	}
}

type ConfigFile struct {
	XMLName      xml.Name `xml:"Server"`
	Debug        bool     `xml:"debug,attr"`
	DebugSrv     string   `xml:"debugsrv,attr"`
	DocumentRoot string

	Listen []*struct {
		Log      bool   `xml:"log,attr"`
		XHeaders bool   `xml:"xheaders,attr"`
		Addr     string `xml:"addr,attr"`
		CertFile string
		KeyFile  string
	}

	IPDB struct {
		File      string `xml:",attr"`
		CacheSize string `xml:",attr"`
	}

	Limit struct {
		MaxRequests int `xml:",attr"`
		Expire      int `xml:",attr"`
	}

	Redis []string `xml:"Redis>Addr"`
}
<commit_msg>Updated WaitGroup<commit_after>// Copyright 2013-2014 Alexandre Fiori
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Web server of freegeoip.net

package main

import (
	"database/sql"
	"encoding/binary"
	"encoding/json"
	"encoding/xml"
	"expvar"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"strings"
	"sync"
	"time"

	"github.com/fiorix/go-web/httpxtra"

	_ "github.com/mattn/go-sqlite3"
	//_ "code.google.com/p/gosqlite/sqlite3"
)

var (
	conf        = ConfigFile{}
	protoCount  = expvar.NewMap("Protocol") // HTTP or HTTPS
	outputCount = expvar.NewMap("Output")   // json, xml, csv or other
	statusCount = expvar.NewMap("Status")   // 200, 403, 404, etc
)

func main() {
	cf := flag.String("config", "freegeoip.conf", "set config file")
	prof := flag.Bool("profile", false, "run cpu and mem profiling")
	flag.Parse()

	if fd, err := os.Open(*cf); err != nil {
		log.Fatal(err)
	} else {
		if err = xml.NewDecoder(fd).Decode(&conf); err != nil {
			log.Fatal(err)
		}
		fd.Close()
	}

	if *prof {
		profile()
	}

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Printf("FreeGeoIP server starting. debug=%t", conf.Debug)

	if conf.Debug && len(conf.DebugSrv) > 0 {
		go func() {
			// server for expvar's /debug/vars only
			log.Printf("Starting DEBUG HTTP server on %s", conf.DebugSrv)
			log.Fatal(http.ListenAndServe(conf.DebugSrv, nil))
		}()
	}

	mux := http.NewServeMux()
	mux.Handle("/", http.FileServer(http.Dir(conf.DocumentRoot)))

	h := LookupHandler()
	mux.HandleFunc("/csv/", h)
	mux.HandleFunc("/xml/", h)
	mux.HandleFunc("/json/", h)

	wg := sync.WaitGroup{}
	for _, l := range conf.Listen {
		if l.Addr == "" {
			continue
		}
		wg.Add(1)
		h := httpxtra.Handler{Handler: mux, XHeaders: l.XHeaders}
		if l.Log {
			h.Logger = logger
		}

		s := http.Server{
			Addr:         l.Addr,
			Handler:      h,
			ReadTimeout:  15 * time.Second,
			WriteTimeout: 15 * time.Second,
		}

		if l.KeyFile == "" && l.CertFile == "" {
			log.Printf("Starting HTTP server on %s "+
				"log=%t xheaders=%t",
				l.Addr, l.Log, l.XHeaders)
			go func() {
				log.Fatal(httpxtra.ListenAndServe(s))
			}()
		} else {
			log.Printf("Starting HTTPS server on %s "+
				"log=%t xheaders=%t cert=%s key=%s",
				l.Addr, l.Log, l.XHeaders,
				l.CertFile, l.KeyFile)
			go func() {
				log.Fatal(s.ListenAndServeTLS(
					l.CertFile,
					l.KeyFile,
				))
			}()
		}
	}
	wg.Wait()
}

// LookupHandler handles GET on /csv, /xml and /json.
func LookupHandler() http.HandlerFunc {
	db, err := sql.Open("sqlite3", conf.IPDB.File)
	if err != nil {
		log.Fatal(err)
	}

	_, err = db.Exec("PRAGMA cache_size=" + conf.IPDB.CacheSize)
	if err != nil {
		log.Fatal(err)
	}

	stmt, err := db.Prepare(ipdb_query)
	if err != nil {
		log.Fatal(err)
	}

	st := time.Now()
	cache := NewCache(db)
	log.Println("IPDB cached in", time.Since(st))

	//defer stmt.Close()

	var quota Quota
	if len(conf.Redis) == 0 {
		quota = new(MapQuota)
		quota.Setup()
		log.Printf("Using internal map to manage quota.")
	} else {
		quota = new(RedisQuota)
		quota.Setup(conf.Redis...)
		log.Printf("Using redis to manage quota: %s", conf.Redis)
	}

	return func(w http.ResponseWriter, r *http.Request) {
		switch r.Method {
		case "GET":
			w.Header().Set("Access-Control-Allow-Origin", "*")
		case "OPTIONS":
			w.Header().Set("Access-Control-Allow-Origin", "*")
			w.Header().Set("Content-Type", "text/plain")
			w.Header().Set("Access-Control-Allow-Methods", "GET")
			w.Header().Set("Access-Control-Allow-Headers", "X-Requested-With")
			w.WriteHeader(200)
			return
		default:
			w.Header().Set("Allow", "GET, OPTIONS")
			http.Error(w, http.StatusText(405), 405)
			return
		}

		// GET continues...
		var srcIP net.IP

		// If xheaders is enabled, RemoteAddr might be a copy of
		// the X-Real-IP or X-Forwarded-For HTTP headers, which
		// can be a comma separated list of IPs. In this case,
		// only the first IP in the list is used.
		r.RemoteAddr = strings.SplitN(r.RemoteAddr, ",", 2)[0]

		if ip, _, err := net.SplitHostPort(r.RemoteAddr); err != nil {
			srcIP = net.ParseIP(r.RemoteAddr) // Use X-Real-IP
		} else {
			srcIP = net.ParseIP(ip)
		}

		if srcIP == nil {
			http.Error(w, http.StatusText(400), 400)
			return
		}

		nsrcIP, err := ip2int(srcIP)
		if err != nil {
			if conf.Debug {
				log.Println(err)
			}
			http.Error(w, http.StatusText(400), 400)
			return
		}

		// Check quota.
		if conf.Limit.MaxRequests > 0 {
			var ok bool
			if ok, err = quota.Ok(nsrcIP); err != nil {
				if conf.Debug {
					log.Println(err) // redis error
				}
				http.Error(w, http.StatusText(503), 503)
				return
			} else if !ok {
				// Over quota, soz :(
				http.Error(w, http.StatusText(403), 403)
				return
			}
		}

		var (
			queryIP  net.IP
			nqueryIP uint32
		)

		// Parse URL (e.g. /csv/ip, /xml/)
		a := strings.SplitN(r.URL.Path, "/", 3)
		if len(a) == 3 && a[2] != "" {
			addrs, err := net.LookupHost(a[2])
			if err != nil {
				// DNS lookup failed, assume host not found.
				http.Error(w, http.StatusText(404), 404)
				return
			}

			if queryIP = net.ParseIP(addrs[0]); queryIP == nil {
				http.Error(w, http.StatusText(400), 400)
				return
			}

			nqueryIP, err = ip2int(net.ParseIP(addrs[0]))
			if err != nil {
				if conf.Debug {
					log.Println(err)
				}
				http.Error(w, http.StatusText(400), 400)
				return
			}
		} else {
			queryIP = srcIP
			nqueryIP = nsrcIP
		}

		// Query the db.
		geoip, err := ipdb_lookup(stmt, cache, queryIP, nqueryIP)
		if err != nil {
			http.NotFound(w, r)
			return
		}

		switch a[1][0] {
		case 'j': // json
			enc := json.NewEncoder(w)
			if cb := r.FormValue("callback"); cb != "" {
				w.Header().Set("Content-Type", "text/javascript")
				fmt.Fprintf(w, "%s(", cb)
				enc.Encode(geoip)
				fmt.Fprintf(w, ");")
			} else {
				w.Header().Set("Content-Type", "application/json")
				enc.Encode(geoip)
			}
		case 'x': // xml
			w.Header().Set("Content-Type", "application/xml")
			fmt.Fprintf(w, xml.Header)
			enc := xml.NewEncoder(w)
			enc.Indent("", " ")
			enc.Encode(geoip)
			fmt.Fprintf(w, "\n")
		case 'c': // csv
			w.Header().Set("Content-Type", "application/csv")
			fmt.Fprintf(w, `"%s","%s","%s","%s","%s","%s",`+
				`"%s","%0.4f","%0.4f","%s","%s"`+"\r\n",
				geoip.Ip,
				geoip.CountryCode, geoip.CountryName,
				geoip.RegionCode, geoip.RegionName,
				geoip.CityName, geoip.ZipCode,
				geoip.Latitude, geoip.Longitude,
				geoip.MetroCode, geoip.AreaCode)
		}
	}
}

func ip2int(ip net.IP) (uint32, error) {
	ipv4 := ip.To4()
	if ipv4 == nil {
		return 0, fmt.Errorf("IP %s is not IPv4", ip.String())
	}

	return binary.BigEndian.Uint32(ipv4), nil
}

func profile() {
	f, err := os.Create("freegeoip.cpu.prof")
	if err != nil {
		log.Fatal(err)
	}

	pprof.StartCPUProfile(f)
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)

	go func() {
		<-sig
		pprof.StopCPUProfile()
		f.Close()
		f, err = os.Create("freegeoip.mem.prof")
		if err != nil {
			log.Fatal(err)
		}
		pprof.WriteHeapProfile(f)
		os.Exit(0)
	}()
}

func logger(r *http.Request, created time.Time, status, bytes int) {
	//fmt.Println(httpxtra.ApacheCommonLog(r, created, status, bytes))
	var (
		s, ip string
		err   error
	)
	if r.TLS == nil {
		s = "HTTP"
	} else {
		s = "HTTPS"
	}
	if ip, _, err = net.SplitHostPort(r.RemoteAddr); err != nil {
		ip = r.RemoteAddr
	}
	log.Printf("%s %d %s %q (%s) :: %d bytes in %s",
		s,
		status,
		r.Method,
		r.URL.Path,
		ip,
		bytes,
		time.Since(created),
	)
	if conf.Debug {
		go func() {
			protoCount.Add(s, 1)
			statusCount.Add(fmt.Sprintf("%d", status), 1)
			switch strings.SplitN(r.URL.Path, "/", 3)[1] {
			case "json":
				outputCount.Add("json", 1)
			case "xml":
				outputCount.Add("xml", 1)
			case "csv":
				outputCount.Add("csv", 1)
			default:
				outputCount.Add("other", 1)
			}
		}()
	}
}

type ConfigFile struct {
	XMLName      xml.Name `xml:"Server"`
	Debug        bool     `xml:"debug,attr"`
	DebugSrv     string   `xml:"debugsrv,attr"`
	DocumentRoot string

	Listen []*struct {
		Log      bool   `xml:"log,attr"`
		XHeaders bool   `xml:"xheaders,attr"`
		Addr     string `xml:"addr,attr"`
		CertFile string
		KeyFile  string
	}

	IPDB struct {
		File      string `xml:",attr"`
		CacheSize string `xml:",attr"`
	}

	Limit struct {
		MaxRequests int `xml:",attr"`
		Expire      int `xml:",attr"`
	}

	Redis []string `xml:"Redis>Addr"`
}
<|endoftext|>


<commit_before>// Read and write the config file
package fs

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"os"
	"os/user"
	"path"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/Unknwon/goconfig"
)

const (
	configFileName = ".rclone.conf"
)

// Global
var (
	// Config file
	ConfigFile *goconfig.ConfigFile
	// Config file path
	ConfigPath string
	// Global config
	Config = &ConfigInfo{}
	// Home directory
	HomeDir string
	// Flags
	verbose      = flag.Bool("verbose", false, "Print lots more stuff")
	quiet        = flag.Bool("quiet", false, "Print as little stuff as possible")
	modifyWindow = flag.Duration("modify-window", time.Nanosecond, "Max time diff to be considered the same")
	checkers     = flag.Int("checkers", 8, "Number of checkers to run in parallel.")
	transfers    = flag.Int("transfers", 4, "Number of file transfers to run in parallel.")
)

// Filesystem config options
type ConfigInfo struct {
	Verbose      bool
	Quiet        bool
	ModifyWindow time.Duration
	Checkers     int
	Transfers    int
}

// Loads the config file
func LoadConfig() {
	// Read some flags if set
	//
	// FIXME read these from the config file too
	Config.Verbose = *verbose
	Config.Quiet = *quiet
	Config.ModifyWindow = *modifyWindow
	Config.Checkers = *checkers
	Config.Transfers = *transfers

	// Find users home directory
	usr, err := user.Current()
	if err != nil {
		log.Printf("Couldn't find home directory: %v", err)
		return
	}
	HomeDir = usr.HomeDir
	ConfigPath = path.Join(HomeDir, configFileName)

	// Load configuration file.
	ConfigFile, err = goconfig.LoadConfigFile(ConfigPath)
	if err != nil {
		log.Printf("Failed to load config file %v - using defaults", ConfigPath)
	}
}

// Save configuration file.
func SaveConfig() {
	err := goconfig.SaveConfigFile(ConfigFile, ConfigPath)
	if err != nil {
		log.Fatalf("Failed to save config file: %v", err)
	}
}

// Show an overview of the config file
func ShowConfig() {
	remotes := ConfigFile.GetSectionList()
	sort.Strings(remotes)
	fmt.Printf("%-20s %s\n", "Name", "Type")
	fmt.Printf("%-20s %s\n", "====", "====")
	for _, remote := range remotes {
		fmt.Printf("%-20s %s\n", remote, ConfigFile.MustValue(remote, "type"))
	}
}

// ChooseRemote chooses a remote name
func ChooseRemote() string {
	remotes := ConfigFile.GetSectionList()
	sort.Strings(remotes)
	return Choose("remote", remotes, nil, false)
}

// Read some input
func ReadLine() string {
	buf := bufio.NewReader(os.Stdin)
	line, err := buf.ReadString('\n')
	if err != nil {
		log.Fatalf("Failed to read line: %v", err)
	}
	return strings.TrimSpace(line)
}

// Command - choose one
func Command(commands []string) int {
	opts := []string{}
	for _, text := range commands {
		fmt.Printf("%c) %s\n", text[0], text[1:])
		opts = append(opts, text[:1])
	}
	optString := strings.Join(opts, "")
	optHelp := strings.Join(opts, "/")
	for {
		fmt.Printf("%s> ", optHelp)
		result := strings.ToLower(ReadLine())
		if len(result) != 1 {
			continue
		}
		i := strings.IndexByte(optString, result[0])
		if i >= 0 {
			return i
		}
	}
}

// Choose one of the defaults or type a new string if newOk is set
func Choose(what string, defaults, help []string, newOk bool) string {
	fmt.Printf("Choose a number from below")
	if newOk {
		fmt.Printf(", or type in your own value")
	}
	fmt.Println()
	for i, text := range defaults {
		if help != nil {
			parts := strings.Split(help[i], "\n")
			for _, part := range parts {
				fmt.Printf(" * %s\n", part)
			}
		}
		fmt.Printf("%2d) %s\n", i+1, text)
	}
	for {
		fmt.Printf("%s> ", what)
		result := ReadLine()
		i, err := strconv.Atoi(result)
		if err != nil {
			if newOk {
				return result
			}
			continue
		}
		if i >= 1 && i <= len(defaults) {
			return defaults[i-1]
		}
	}
}

// Show the contents of the remote
func ShowRemote(name string) {
	fmt.Printf("--------------------\n")
	fmt.Printf("[%s]\n", name)
	for _, key := range ConfigFile.GetKeyList(name) {
		fmt.Printf("%s = %s\n", key, ConfigFile.MustValue(name, key))
	}
	fmt.Printf("--------------------\n")
}

// Print the contents of the remote and ask if it is OK
func OkRemote(name string) bool {
	ShowRemote(name)
	switch i := Command([]string{"yYes this is OK", "eEdit this remote", "dDelete this remote"}); i {
	case 0:
		return true
	case 1:
		return false
	case 2:
		ConfigFile.DeleteSection(name)
		return true
	default:
		log.Printf("Bad choice %d", i)
	}
	return false
}

// Make a new remote
func NewRemote(name string) {
	fmt.Printf("What type of source is it?\n")
	types := []string{}
	for _, item := range fsRegistry {
		types = append(types, item.Name)
	}
	newType := Choose("type", types, nil, false)
	ConfigFile.SetValue(name, "type", newType)
	fs, err := Find(newType)
	if err != nil {
		log.Fatalf("Failed to find fs: %v", err)
	}
	for _, option := range fs.Options {
		ConfigFile.SetValue(name, option.Name, option.Choose())
	}
	if OkRemote(name) {
		SaveConfig()
		return
	}
	EditRemote(name)
}

// Edit a remote
func EditRemote(name string) {
	ShowRemote(name)
	fmt.Printf("Edit remote\n")
	for {
		for _, key := range ConfigFile.GetKeyList(name) {
			value := ConfigFile.MustValue(name, key)
			fmt.Printf("Press enter to accept current value, or type in a new one\n")
			fmt.Printf("%s = %s>", key, value)
			newValue := ReadLine()
			if newValue != "" {
				ConfigFile.SetValue(name, key, newValue)
			}
		}
		if OkRemote(name) {
			break
		}
	}
	SaveConfig()
}

// Edit the config file interactively
func EditConfig() {
	for {
		fmt.Printf("Current remotes:\n\n")
		ShowConfig()
		fmt.Printf("\n")
		switch i := Command([]string{"eEdit existing remote", "nNew remote", "dDelete remote", "qQuit config"}); i {
		case 0:
			name := ChooseRemote()
			EditRemote(name)
		case 1:
			fmt.Printf("name> ")
			name := ReadLine()
			NewRemote(name)
		case 2:
			name := ChooseRemote()
			ConfigFile.DeleteSection(name)
		case 3:
			return
		}
	}
}
<commit_msg>Config fixes<commit_after>// Read and write the config file
package fs

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"os"
	"os/user"
	"path"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/Unknwon/goconfig"
)

const (
	configFileName = ".rclone.conf"
)

// Global
var (
	// Config file
	ConfigFile *goconfig.ConfigFile
	// Config file path
	ConfigPath string
	// Global config
	Config = &ConfigInfo{}
	// Home directory
	HomeDir string
	// Flags
	verbose      = flag.Bool("verbose", false, "Print lots more stuff")
	quiet        = flag.Bool("quiet", false, "Print as little stuff as possible")
	modifyWindow = flag.Duration("modify-window", time.Nanosecond, "Max time diff to be considered the same")
	checkers     = flag.Int("checkers", 8, "Number of checkers to run in parallel.")
	transfers    = flag.Int("transfers", 4, "Number of file transfers to run in parallel.")
)

// Filesystem config options
type ConfigInfo struct {
	Verbose      bool
	Quiet        bool
	ModifyWindow time.Duration
	Checkers     int
	Transfers    int
}

// Loads the config file
func LoadConfig() {
	// Read some flags if set
	//
	// FIXME read these from the config file too
	Config.Verbose = *verbose
	Config.Quiet = *quiet
	Config.ModifyWindow = *modifyWindow
	Config.Checkers = *checkers
	Config.Transfers = *transfers

	// Find users home directory
	usr, err := user.Current()
	if err != nil {
		log.Printf("Couldn't find home directory: %v", err)
		return
	}
	HomeDir = usr.HomeDir
	ConfigPath = path.Join(HomeDir, configFileName)

	// Load configuration file.
	ConfigFile, err = goconfig.LoadConfigFile(ConfigPath)
	if err != nil {
		log.Printf("Failed to load config file %v - using defaults", ConfigPath)
		ConfigFile, err = goconfig.LoadConfigFile(os.DevNull)
		if err != nil {
			log.Fatalf("Failed to read null config file: %v", err)
		}
	}
}

// Save configuration file.
func SaveConfig() {
	err := goconfig.SaveConfigFile(ConfigFile, ConfigPath)
	if err != nil {
		log.Fatalf("Failed to save config file: %v", err)
	}
}

// Show an overview of the config file
func ShowRemotes() {
	remotes := ConfigFile.GetSectionList()
	if len(remotes) == 0 {
		return
	}
	sort.Strings(remotes)
	fmt.Printf("%-20s %s\n", "Name", "Type")
	fmt.Printf("%-20s %s\n", "====", "====")
	for _, remote := range remotes {
		fmt.Printf("%-20s %s\n", remote, ConfigFile.MustValue(remote, "type"))
	}
}

// ChooseRemote chooses a remote name
func ChooseRemote() string {
	remotes := ConfigFile.GetSectionList()
	sort.Strings(remotes)
	return Choose("remote", remotes, nil, false)
}

// Read some input
func ReadLine() string {
	buf := bufio.NewReader(os.Stdin)
	line, err := buf.ReadString('\n')
	if err != nil {
		log.Fatalf("Failed to read line: %v", err)
	}
	return strings.TrimSpace(line)
}

// Command - choose one
func Command(commands []string) byte {
	opts := []string{}
	for _, text := range commands {
		fmt.Printf("%c) %s\n", text[0], text[1:])
		opts = append(opts, text[:1])
	}
	optString := strings.Join(opts, "")
	optHelp := strings.Join(opts, "/")
	for {
		fmt.Printf("%s> ", optHelp)
		result := strings.ToLower(ReadLine())
		if len(result) != 1 {
			continue
		}
		i := strings.IndexByte(optString, result[0])
		if i >= 0 {
			return result[0]
		}
	}
}

// Choose one of the defaults or type a new string if newOk is set
func Choose(what string, defaults, help []string, newOk bool) string {
	fmt.Printf("Choose a number from below")
	if newOk {
		fmt.Printf(", or type in your own value")
	}
	fmt.Println()
	for i, text := range defaults {
		if help != nil {
			parts := strings.Split(help[i], "\n")
			for _, part := range parts {
				fmt.Printf(" * %s\n", part)
			}
		}
		fmt.Printf("%2d) %s\n", i+1, text)
	}
	for {
		fmt.Printf("%s> ", what)
		result := ReadLine()
		i, err := strconv.Atoi(result)
		if err != nil {
			if newOk {
				return result
			}
			continue
		}
		if i >= 1 && i <= len(defaults) {
			return defaults[i-1]
		}
	}
}

// Show the contents of the remote
func ShowRemote(name string) {
	fmt.Printf("--------------------\n")
	fmt.Printf("[%s]\n", name)
	for _, key := range ConfigFile.GetKeyList(name) {
		fmt.Printf("%s = %s\n", key, ConfigFile.MustValue(name, key))
	}
	fmt.Printf("--------------------\n")
}

// Print the contents of the remote and ask if it is OK
func OkRemote(name string) bool {
	ShowRemote(name)
	switch i := Command([]string{"yYes this is OK", "eEdit this remote", "dDelete this remote"}); i {
	case 'y':
		return true
	case 'e':
		return false
	case 'd':
		ConfigFile.DeleteSection(name)
		return true
	default:
		log.Printf("Bad choice %d", i)
	}
	return false
}

// Make a new remote
func NewRemote(name string) {
	fmt.Printf("What type of source is it?\n")
	types := []string{}
	for _, item := range fsRegistry {
		types = append(types, item.Name)
	}
	newType := Choose("type", types, nil, false)
	ConfigFile.SetValue(name, "type", newType)
	fs, err := Find(newType)
	if err != nil {
		log.Fatalf("Failed to find fs: %v", err)
	}
	for _, option := range fs.Options {
		ConfigFile.SetValue(name, option.Name, option.Choose())
	}
	if OkRemote(name) {
		SaveConfig()
		return
	}
	EditRemote(name)
}

// Edit a remote
func EditRemote(name string) {
	ShowRemote(name)
	fmt.Printf("Edit remote\n")
	for {
		for _, key := range ConfigFile.GetKeyList(name) {
			value := ConfigFile.MustValue(name, key)
			fmt.Printf("Press enter to accept current value, or type in a new one\n")
			fmt.Printf("%s = %s>", key, value)
			newValue := ReadLine()
			if newValue != "" {
				ConfigFile.SetValue(name, key, newValue)
			}
		}
		if OkRemote(name) {
			break
		}
	}
	SaveConfig()
}

// Delete a remote
func DeleteRemote(name string) {
	ConfigFile.DeleteSection(name)
	SaveConfig()
}

// Edit the config file interactively
func EditConfig() {
	for {
		haveRemotes := len(ConfigFile.GetSectionList()) != 0
		what := []string{"eEdit existing remote", "nNew remote", "dDelete remote", "qQuit config"}
		if haveRemotes {
			fmt.Printf("Current remotes:\n\n")
			ShowRemotes()
			fmt.Printf("\n")
		} else {
			fmt.Printf("No remotes found - make a new one\n")
			what = append(what[1:2], what[3])
		}
		switch i := Command(what); i {
		case 'e':
			name := ChooseRemote()
			EditRemote(name)
		case 'n':
			fmt.Printf("name> ")
			name := ReadLine()
			NewRemote(name)
		case 'd':
			name := ChooseRemote()
			DeleteRemote(name)
		case 'q':
			return
		}
	}
}
<|endoftext|>


<commit_before>package fsmonitor

import (
	"github.com/howeyc/fsnotify"
	"path/filepath"
	"os"
)

func NewWatcher() (*Watcher, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}
	event := make(chan *fsnotify.FileEvent)
	watcherError := make(chan error)
	monitorWatcher := &Watcher{Event: event, Error: watcherError, watcher: watcher}
	go func() {
		for {
			select {
			case ev := <- watcher.Event:
				event <- ev
				if ev.IsCreate() {
					go func() {
						monitorWatcher.watchAllFolders(ev.Name)
					}()
				}
				if ev.IsDelete() {
					go func() {
						watcher.RemoveWatch(ev.Name)
					}()
				}
			case e := <- watcher.Error:
				watcherError  <- e
			}
		}
	}()
	return monitorWatcher, nil
}

type Watcher struct {
	Event chan *fsnotify.FileEvent
	Error chan error
	watcher *fsnotify.Watcher
}

func (self *Watcher) Watch(path string) error {
	err := self.watchAllFolders(path)
	if err != nil {
		return err
	}
	return nil
}

func (self *Watcher) watchAllFolders(path string) (err error) {
	err = filepath.Walk(path, func(path string, f os.FileInfo, err error) error {
		if f.IsDir() {
			err := self.addWatcher(path)
			if err != nil {
				return err
			}
		}
		return nil
	})
	return
}

func (self *Watcher) addWatcher(path string) (err error) {
	err = self.watcher.Watch(path)
	return
}
<commit_msg>Changed fsnotify package<commit_after>package fsmonitor

import (
	"code.google.com/p/go.exp/fsnotify"
	"path/filepath"
	"os"
)

func NewWatcher() (*Watcher, error) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return nil, err
	}
	event := make(chan *fsnotify.FileEvent)
	watcherError := make(chan error)
	monitorWatcher := &Watcher{Event: event, Error: watcherError, watcher: watcher}
	go func() {
		for {
			select {
			case ev := <- watcher.Event:
				event <- ev
				if ev.IsCreate() {
					go func() {
						monitorWatcher.watchAllFolders(ev.Name)
					}()
				}
				if ev.IsDelete() {
					go func() {
						watcher.RemoveWatch(ev.Name)
					}()
				}
			case e := <- watcher.Error:
				watcherError  <- e
			}
		}
	}()
	return monitorWatcher, nil
}

type Watcher struct {
	Event chan *fsnotify.FileEvent
	Error chan error
	watcher *fsnotify.Watcher
}

func (self *Watcher) Watch(path string) error {
	err := self.watchAllFolders(path)
	if err != nil {
		return err
	}
	return nil
}

func (self *Watcher) watchAllFolders(path string) (err error) {
	err = filepath.Walk(path, func(path string, f os.FileInfo, err error) error {
		if f.IsDir() {
			err := self.addWatcher(path)
			if err != nil {
				return err
			}
		}
		return nil
	})
	return
}

func (self *Watcher) addWatcher(path string) (err error) {
	err = self.watcher.Watch(path)
	return
}
<|endoftext|>


<commit_before>// An experimental FTP server framework. By providing a simple driver class that
// responds to a handful of methods you can have a complete FTP server.
//
// Some sample use cases include persisting data to an Amazon S3 bucket, a
// relational database, redis or memory.
//
// There is a sample in-memory driver available - see the documentation for the
// graval-mem package or the graval READEME for more details.
package graval

import (
	"crypto/tls"
	"net"
	"strconv"
	"strings"
)

type CryptoConfig struct {
	Implicit  bool
	Force     bool
	TlsConfig *tls.Config
}

// serverOpts contains parameters for graval.NewFTPServer()
type FTPServerOpts struct {
	// Server name will be used for welcome message
	ServerName string

	// The factory that will be used to create a new FTPDriver instance for
	// each client connection. This is a mandatory option.
	Factory FTPDriverFactory

	// The hostname that the FTP server should listen on. Optional, defaults to
	// "::", which means all hostnames on ipv4 and ipv6.
	Hostname string

	// The port that the FTP should listen on. Optional, defaults to 3000. In
	// a production environment you will probably want to change this to 21.
	Port int

	// Options for passive data connections
	PassiveOpts *PassiveOpts

	// Options for FTPS and FTPES
	CryptoConfig *CryptoConfig

	// Disable logging (useful for tests)
	Quiet bool
}

// FTPServer is the root of your FTP application. You should instantiate one
// of these and call ListenAndServe() to start accepting client connections.
//
// Always use the NewFTPServer() method to create a new FTPServer.
type FTPServer struct {
	serverName    string
	listenTo      string
	listener      net.Listener
	closed        bool
	driverFactory FTPDriverFactory
	quiet         bool
	logger        *ftpLogger
	passiveOpts   *PassiveOpts
	cryptoConfig  *CryptoConfig
}

// serverOptsWithDefaults copies an FTPServerOpts struct into a new struct,
// then adds any default values that are missing and returns the new data.
func serverOptsWithDefaults(opts *FTPServerOpts) *FTPServerOpts {
	var newOpts FTPServerOpts

	if opts == nil {
		opts = &FTPServerOpts{}
	}

	if opts.ServerName == "" {
		newOpts.ServerName = "Go FTP Server"
	} else {
		newOpts.ServerName = opts.ServerName
	}

	if opts.Hostname == "" {
		newOpts.Hostname = "::"
	} else {
		newOpts.Hostname = opts.Hostname
	}

	if opts.Port == 0 {
		newOpts.Port = 3000
	} else {
		newOpts.Port = opts.Port
	}

	newOpts.Factory = opts.Factory

	if opts.PassiveOpts == nil {
		newOpts.PassiveOpts = &PassiveOpts{}
	} else {
		newOpts.PassiveOpts = opts.PassiveOpts
	}

	newOpts.CryptoConfig = opts.CryptoConfig

	newOpts.Quiet = opts.Quiet

	return &newOpts
}

// NewFTPServer initialises a new FTP server. Configuration options are provided
// via an instance of FTPServerOpts. Calling this function in your code will
// probably look something like this:
//
//     factory := &MyDriverFactory{}
//     server  := graval.NewFTPServer(&graval.FTPServerOpts{ Factory: factory })
//
// or:
//
//     factory := &MyDriverFactory{}
//     opts    := &graval.FTPServerOpts{
//       Factory: factory,
//       Port: 2000,
//       Hostname: "127.0.0.1",
//     }
//     server  := graval.NewFTPServer(opts)
//
func NewFTPServer(opts *FTPServerOpts) *FTPServer {
	opts = serverOptsWithDefaults(opts)
	s := new(FTPServer)
	s.listenTo = buildTcpString(opts.Hostname, opts.Port)
	s.serverName = opts.ServerName
	s.driverFactory = opts.Factory
	s.quiet = opts.Quiet
	s.logger = newFtpLogger("", s.quiet)
	s.passiveOpts = opts.PassiveOpts
	s.cryptoConfig = opts.CryptoConfig
	return s
}

// ListenAndServe asks a new FTPServer to begin accepting client connections. It
// accepts no arguments - all configuration is provided via the NewFTPServer
// function.
//
// If the server fails to start for any reason, an error will be returned. Common
// errors are trying to bind to a privileged port or something else is already
// listening on the same port.
//
func (ftpServer *FTPServer) ListenAndServe() error {
	laddr, err := net.ResolveTCPAddr("tcp", ftpServer.listenTo)
	if err != nil {
		return err
	}
	listener, err := net.ListenTCP("tcp", laddr)
	if err != nil {
		return err
	}
	ftpServer.listener = listener
	for {
		tcpConn, err := listener.AcceptTCP()

		if ftpServer.closed {
			return nil
		}

		if err != nil {
			ftpServer.logger.Print("listening error")
			break
		}
		driver, err := ftpServer.driverFactory.NewDriver()
		if err != nil {
			ftpServer.logger.Print("Error creating driver, aborting client connection")
		} else {
			ftpConn := newftpConn(tcpConn, driver, ftpServer.serverName, ftpServer.passiveOpts, ftpServer.cryptoConfig, ftpServer.quiet)
			go ftpConn.Serve()
		}
	}
	return nil
}

func (ftpServer *FTPServer) Close() error {
	ftpServer.closed = true
	return ftpServer.listener.Close()
}

func buildTcpString(hostname string, port int) (result string) {
	if strings.Contains(hostname, ":") {
		// ipv6
		if port == 0 {
			result = "[" + hostname + "]"
		} else {
			result = "[" + hostname + "]:" + strconv.Itoa(port)
		}
	} else {
		// ipv4
		if port == 0 {
			result = hostname
		} else {
			result = hostname + ":" + strconv.Itoa(port)
		}
	}
	return
}
<commit_msg>Server default crypto config<commit_after>// An experimental FTP server framework. By providing a simple driver class that
// responds to a handful of methods you can have a complete FTP server.
//
// Some sample use cases include persisting data to an Amazon S3 bucket, a
// relational database, redis or memory.
//
// There is a sample in-memory driver available - see the documentation for the
// graval-mem package or the graval READEME for more details.
package graval

import (
	"crypto/tls"
	"net"
	"strconv"
	"strings"
)

type CryptoConfig struct {
	Implicit  bool
	Force     bool
	TlsConfig *tls.Config
}

// serverOpts contains parameters for graval.NewFTPServer()
type FTPServerOpts struct {
	// Server name will be used for welcome message
	ServerName string

	// The factory that will be used to create a new FTPDriver instance for
	// each client connection. This is a mandatory option.
	Factory FTPDriverFactory

	// The hostname that the FTP server should listen on. Optional, defaults to
	// "::", which means all hostnames on ipv4 and ipv6.
	Hostname string

	// The port that the FTP should listen on. Optional, defaults to 3000. In
	// a production environment you will probably want to change this to 21.
	Port int

	// Options for passive data connections
	PassiveOpts *PassiveOpts

	// Options for FTPS and FTPES
	CryptoConfig *CryptoConfig

	// Disable logging (useful for tests)
	Quiet bool
}

// FTPServer is the root of your FTP application. You should instantiate one
// of these and call ListenAndServe() to start accepting client connections.
//
// Always use the NewFTPServer() method to create a new FTPServer.
type FTPServer struct {
	serverName    string
	listenTo      string
	listener      net.Listener
	closed        bool
	driverFactory FTPDriverFactory
	quiet         bool
	logger        *ftpLogger
	passiveOpts   *PassiveOpts
	cryptoConfig  *CryptoConfig
}

// serverOptsWithDefaults copies an FTPServerOpts struct into a new struct,
// then adds any default values that are missing and returns the new data.
func serverOptsWithDefaults(opts *FTPServerOpts) *FTPServerOpts {
	var newOpts FTPServerOpts

	if opts == nil {
		opts = &FTPServerOpts{}
	}

	if opts.ServerName == "" {
		newOpts.ServerName = "Go FTP Server"
	} else {
		newOpts.ServerName = opts.ServerName
	}

	if opts.Hostname == "" {
		newOpts.Hostname = "::"
	} else {
		newOpts.Hostname = opts.Hostname
	}

	if opts.Port == 0 {
		newOpts.Port = 3000
	} else {
		newOpts.Port = opts.Port
	}

	newOpts.Factory = opts.Factory

	if opts.PassiveOpts == nil {
		newOpts.PassiveOpts = &PassiveOpts{}
	} else {
		newOpts.PassiveOpts = opts.PassiveOpts
	}

	if opts.CryptoConfig == nil {
		newOpts.CryptoConfig = &CryptoConfig{
			Implicit:  false,
			Force:     false,
			TlsConfig: nil,
		}
	} else {
		newOpts.CryptoConfig = opts.CryptoConfig
	}

	newOpts.Quiet = opts.Quiet

	return &newOpts
}

// NewFTPServer initialises a new FTP server. Configuration options are provided
// via an instance of FTPServerOpts. Calling this function in your code will
// probably look something like this:
//
//     factory := &MyDriverFactory{}
//     server  := graval.NewFTPServer(&graval.FTPServerOpts{ Factory: factory })
//
// or:
//
//     factory := &MyDriverFactory{}
//     opts    := &graval.FTPServerOpts{
//       Factory: factory,
//       Port: 2000,
//       Hostname: "127.0.0.1",
//     }
//     server  := graval.NewFTPServer(opts)
//
func NewFTPServer(opts *FTPServerOpts) *FTPServer {
	opts = serverOptsWithDefaults(opts)
	s := new(FTPServer)
	s.listenTo = buildTcpString(opts.Hostname, opts.Port)
	s.serverName = opts.ServerName
	s.driverFactory = opts.Factory
	s.quiet = opts.Quiet
	s.logger = newFtpLogger("", s.quiet)
	s.passiveOpts = opts.PassiveOpts
	s.cryptoConfig = opts.CryptoConfig
	return s
}

// ListenAndServe asks a new FTPServer to begin accepting client connections. It
// accepts no arguments - all configuration is provided via the NewFTPServer
// function.
//
// If the server fails to start for any reason, an error will be returned. Common
// errors are trying to bind to a privileged port or something else is already
// listening on the same port.
//
func (ftpServer *FTPServer) ListenAndServe() error {
	laddr, err := net.ResolveTCPAddr("tcp", ftpServer.listenTo)
	if err != nil {
		return err
	}
	listener, err := net.ListenTCP("tcp", laddr)
	if err != nil {
		return err
	}
	ftpServer.listener = listener
	for {
		tcpConn, err := listener.AcceptTCP()

		if ftpServer.closed {
			return nil
		}

		if err != nil {
			ftpServer.logger.Print("listening error")
			break
		}
		driver, err := ftpServer.driverFactory.NewDriver()
		if err != nil {
			ftpServer.logger.Print("Error creating driver, aborting client connection")
		} else {
			ftpConn := newftpConn(tcpConn, driver, ftpServer.serverName, ftpServer.passiveOpts, ftpServer.cryptoConfig, ftpServer.quiet)
			go ftpConn.Serve()
		}
	}
	return nil
}

func (ftpServer *FTPServer) Close() error {
	ftpServer.closed = true
	return ftpServer.listener.Close()
}

func buildTcpString(hostname string, port int) (result string) {
	if strings.Contains(hostname, ":") {
		// ipv6
		if port == 0 {
			result = "[" + hostname + "]"
		} else {
			result = "[" + hostname + "]:" + strconv.Itoa(port)
		}
	} else {
		// ipv4
		if port == 0 {
			result = hostname
		} else {
			result = hostname + ":" + strconv.Itoa(port)
		}
	}
	return
}
<|endoftext|>


<commit_before>// replicator.go copies S3 files from one bucket to another.
// This is used to replicate files in one S3 region (Virginia)
// to another region (Oregon).

package workers

import (
	"encoding/json"
	"fmt"
	"github.com/APTrust/bagman/bagman"
	"github.com/bitly/go-nsq"
	"github.com/crowdmob/goamz/aws"
	"github.com/crowdmob/goamz/s3"
	"os"
	"path/filepath"
	"time"
)

type ReplicationObject struct {
	File       *bagman.File
	NsqMessage *nsq.Message
}

type Replicator struct {
	ReplicationChannel  chan *ReplicationObject
	S3ReplicationClient *bagman.S3Client
	ProcUtil            *bagman.ProcessUtil
}

func NewReplicator(procUtil *bagman.ProcessUtil) (*Replicator) {
	replicationClient, _ := bagman.NewS3Client(aws.USWest2)
	replicator := &Replicator{
		ProcUtil: procUtil,
		S3ReplicationClient: replicationClient,
	}
	workerBufferSize := procUtil.Config.StoreWorker.Workers * 10
	replicator.ReplicationChannel = make(chan *ReplicationObject, workerBufferSize)
	for i := 0; i < procUtil.Config.StoreWorker.Workers; i++ {
		go replicator.replicate()
	}
	return replicator
}

// MessageHandler handles messages from the queue, putting each
// item into the replication channel.
func (replicator *Replicator) HandleMessage(message *nsq.Message) error {
	message.DisableAutoResponse()
	var file bagman.File
	err := json.Unmarshal(message.Body, &file)
	if err != nil {
		replicator.ProcUtil.MessageLog.Error("Could not unmarshal JSON data from nsq:",
			string(message.Body))
		message.Finish()
		return fmt.Errorf("Could not unmarshal JSON data from nsq")
	}
	if file.NeedsSave == false &&  replicator.ReplicatedFileExists(file.S3UUID()){
		// This occurs when the depositor has uploaded a new version
		// of an existing bag, but this particiluar file within the
		// bag did not change.
		replicator.ProcUtil.MessageLog.Info("File %s has not changed since it was last saved " +
			"and is confirmed to exist in replication bucket with UUID %s",
			file.Identifier, file.S3UUID())
		message.Finish()
		return nil
	}
	if replicator.ReplicatedFileExists(file.Uuid) {
		replicator.ProcUtil.MessageLog.Info("File %s already exists in replication bucket",
			file.Identifier)
		message.Finish()
		return nil
	}
	replicationObject := &ReplicationObject{
		NsqMessage: message,
		File: &file,
	}

	// Unfortunately, we're probably running on a machine that's
	// also running either ingest or restore. The Volume monitor
	// can know how much disk space those processes are using, but
	// not how much they have reserved. So until we have a volume
	// monitoring service, all we can do is pad the estimate of
	// how much space we might need.
	err = replicator.ProcUtil.Volume.Reserve(uint64(file.Size * 2))
	if err != nil {
		// Not enough room on disk
		replicator.ProcUtil.MessageLog.Warning("Requeueing %s (%d bytes) - not enough disk space",
			file.Identifier, file.Size)
		message.Requeue(10 * time.Minute)
		return nil
	}

	replicator.ReplicationChannel <- replicationObject
	replicator.ProcUtil.MessageLog.Debug("Put %s (%d bytes) into replication queue",
		file.Identifier, file.Size)
	return nil
}

func (replicator *Replicator) ReplicatedFileExists(fileUUID string) (bool) {
	exists, _ := replicator.S3ReplicationClient.Exists(
		replicator.ProcUtil.Config.PreservationBucket,
		fileUUID)
	return exists
}

func (replicator *Replicator) replicate() {
	for replicationObject := range replicator.ReplicationChannel {
		replicator.ProcUtil.MessageLog.Info("Starting %s",
			replicationObject.File.Identifier)
		url, err := replicator.CopyAndSaveEvent(replicationObject)
		if err != nil {
			replicationObject.File.ReplicationError = err.Error()
			// If we failed too many times, send this into the failure
			// queue. Otherwise, just requeue and try again.
			if (replicationObject.NsqMessage.Attempts >=
				uint16(replicator.ProcUtil.Config.ReplicationWorker.MaxAttempts)) {
				replicator.SendToTroubleQueue(replicationObject.File)
				replicationObject.NsqMessage.Finish()
				replicator.ProcUtil.IncrementFailed()
			} else {
				replicator.ProcUtil.MessageLog.Error(
					"Requeuing %s (%s) because copy failed. Error: %v",
					replicationObject.File.Identifier,
					replicationObject.File.StorageURL,
					err)
				replicationObject.NsqMessage.Requeue(5 * time.Minute)
				replicator.ProcUtil.IncrementFailed()
			}
		} else {
			// Success!
			replicator.ProcUtil.MessageLog.Info("Finished %s. Replication " +
				"copy is at %s.",
				replicationObject.File.Identifier,
				url)
			replicationObject.NsqMessage.Finish()
			replicator.ProcUtil.IncrementSucceeded()
		}
		replicator.ProcUtil.MessageLog.Info(
			"**STATS** Succeeded: %d, Failed: %d",
			replicator.ProcUtil.Succeeded(),
			replicator.ProcUtil.Failed())
	}
}

// Copy the file to S3 in Oregon and save the replication
// PremisEvent to Fluctus. Returns the S3 URL of the newly-saved file
// or an error.
func (replicator *Replicator) CopyAndSaveEvent(replicationObject *ReplicationObject) (string, error) {
	url, err := replicator.CopyFile(replicationObject)
	if err != nil {
		return "", err
	}
	event, err := replicator.SaveReplicationEvent(replicationObject.File, url)
	if err != nil {
		return "", err
	}
	replicator.ProcUtil.MessageLog.Info(
		"Saved replication PremisEvent for %s (%s) with event identifier %s",
		replicationObject.File.Identifier,
		replicationObject.File.Uuid,
		event.Identifier)
	return url, nil
}

// Copies a file from one bucket to another, across regions,
// including all of APTrust's custom metadata. Returns the URL
// of the destination file (that should be in the replication
// bucket in Oregon), or an error.
//
// This does NOT use PUT COPY internally because PUT COPY is
// limited to files of 5GB or less, and we'll have many files
// over 5GB. The copy operation downloads data from the S3
// preservation bucket and uploads it to the replication bucket.
//
// As long as we're running in the same region as our S3
// preservation bucket (USEast), the download should be fast
// and free. Running this code outside of USEast will be
// slow and expensive, since we'll have to pay for the bandwidth
// of BOTH download and upload.
func (replicator *Replicator) CopyFile(replicationObject *ReplicationObject) (string, error) {
	replicator.ProcUtil.MessageLog.Info("Starting copy of %s (%s)",
		replicationObject.File.Identifier, replicationObject.File.Uuid)
	// Copy options include the md5 sum of the file we're copying
	// and all of our custom meta data.
	copyOptions, err := replicator.GetCopyOptions(replicationObject.File)
	if err != nil {
		return "", err
	}

	// Touch before dowload, because large files can take a long time!
	replicationObject.NsqMessage.Touch()

	localPath, err := replicator.DownloadFromPreservation(replicationObject.File)
	if err != nil {
		return "", err
	}

	reader, err := os.Open(localPath)
	if err != nil {
		return "", err
	}

	// Touch again before upload, because large files are slow
	replicationObject.NsqMessage.Touch()

	// Replication client is configured to us USWest-2 (Oregon),
	// but the bucket name should be enough.
	url := ""
	if replicationObject.File.Size <= bagman.S3_LARGE_FILE {
		url, err = replicator.S3ReplicationClient.SaveToS3(
			replicator.ProcUtil.Config.ReplicationBucket,
			replicationObject.File.Uuid,
			replicationObject.File.MimeType,
			reader,
			replicationObject.File.Size,
			copyOptions)
	} else {
		url, err = replicator.S3ReplicationClient.SaveLargeFileToS3(
			replicator.ProcUtil.Config.ReplicationBucket,
			replicationObject.File.Uuid,
			replicationObject.File.MimeType,
			reader,
			replicationObject.File.Size,
			copyOptions,
			bagman.S3_CHUNK_SIZE)
	}

	// Touch so NSQ knows we're not dead yet!
	replicationObject.NsqMessage.Touch()

	// Delete the local file.
	delErr := os.Remove(localPath)
	if delErr != nil {
		replicator.ProcUtil.MessageLog.Warning("Could not delete local file %s: %v",
			localPath, delErr)
	} else {
		replicator.ProcUtil.Volume.Release(uint64(replicationObject.File.Size * 2))
	}

	if err == nil {
		replicator.ProcUtil.MessageLog.Info("Finished copy of %s (%s)",
			replicationObject.File.Identifier,
			replicationObject.File.Uuid)
	}

	return url, err
}

// Returns S3 options, including the md5 checksum and APTrust's custom
// metadata. These options must accompany the file copy.
func (replicator *Replicator) GetCopyOptions(file *bagman.File) (s3.Options, error) {
	// Copy all of the meta data
	resp, err := replicator.ProcUtil.S3Client.Head(
		replicator.ProcUtil.Config.PreservationBucket,
		file.Uuid)
	if err != nil {
		detailedErr :=  fmt.Errorf("Head request for %s at %s returned error: %v",
			file.Uuid,
			replicator.ProcUtil.Config.PreservationBucket,
			err)
		return s3.Options{}, detailedErr
	}
	if resp.StatusCode != 200 {
		detailedErr := fmt.Errorf(
			"Can't get S3 metadata for file %s. " +
			"Head request for %s/%s returned HTTP Status Code %d",
			file.Identifier,
			replicator.ProcUtil.Config.PreservationBucket,
			file.Uuid,
			resp.StatusCode)
		return s3.Options{}, detailedErr
	}

	s3Metadata := make(map[string][]string)
	s3Metadata["md5"]         = []string{ file.Md5 }
	s3Metadata["institution"] = []string{ resp.Header["X-Amz-Meta-Institution"][0] }
	s3Metadata["bag"]         = []string{ resp.Header["X-Amz-Meta-Bag"][0] }
	s3Metadata["bagpath"]     = []string{ file.Path }

	base64md5, err := bagman.Base64EncodeMd5(file.Md5)
	if err != nil {
		return s3.Options{}, err
	}

	options := replicator.S3ReplicationClient.MakeOptions(base64md5, s3Metadata)
	return options, nil
}

// Copies a file from the preservation bucket to a local file
// and returns the path to the local file.
func (replicator *Replicator) DownloadFromPreservation(file *bagman.File) (string, error) {
	// Make sure we have a folder to put the file in.
	replicationDir := replicator.ProcUtil.Config.ReplicationDirectory
	if _, err := os.Stat(replicationDir); os.IsNotExist(err) {
		err = os.MkdirAll(replicationDir, 0755)
		if err != nil {
			return "", err
		}
	}

	// Now copy the file from S3 to local path.
	localPath := filepath.Join(replicationDir, file.Uuid)
	err := replicator.ProcUtil.S3Client.FetchToFileWithoutChecksum(
		replicator.ProcUtil.Config.PreservationBucket,
		file.Uuid,
		localPath)
	if err != nil {
		detailedErr := fmt.Errorf("Cannot read file %s from %s: %v",
			file.Uuid,
			replicator.ProcUtil.Config.PreservationBucket,
			err)
		return "", detailedErr
	}

	replicator.ProcUtil.MessageLog.Info("Downloaded %s", localPath)
	return localPath, nil
}

// Saves the replication PremisEvent to Fluctus.
// Param url is the S3 URL we just saved that file to.
// That should be in Oregon.
func (replicator *Replicator) SaveReplicationEvent(file *bagman.File, url string) (*bagman.PremisEvent, error) {
	replicationEvent, err := file.ReplicationEvent(url)
	if err != nil {
		return nil, err
	}
	replicator.ProcUtil.MessageLog.Info("Saving replication PremisEvent for %s (%s)",
		file.Identifier, file.Uuid)
	savedEvent, err := replicator.ProcUtil.FluctusClient.PremisEventSave(
		file.Identifier, "GenericFile", replicationEvent)
	if err != nil {
		return nil, err
	}
	return savedEvent, nil
}

// Puts an item into the trouble queue.
func (replicator *Replicator) SendToTroubleQueue(file *bagman.File) {
	err := bagman.Enqueue(
		replicator.ProcUtil.Config.NsqdHttpAddress,
		replicator.ProcUtil.Config.FailedReplicationWorker.NsqTopic,
		file)
	if err != nil {
		replicator.ProcUtil.MessageLog.Error(
			"Could not send '%s' (%s) to failed replication queue: %v\n",
			file.Identifier,
			file.Uuid,
			err)
	} else {
		message := fmt.Sprintf("Failed to copy '%s' (%s) to replication bucket. " +
			"This item is going into the failed replication queue.",
			file.Identifier,
			file.Uuid)
		replicator.ProcUtil.MessageLog.Warning(message)
	}
}
<commit_msg>Check replication bucket, not preservation bucket. On error, assume files needs replication.<commit_after>// replicator.go copies S3 files from one bucket to another.
// This is used to replicate files in one S3 region (Virginia)
// to another region (Oregon).

package workers

import (
	"encoding/json"
	"fmt"
	"github.com/APTrust/bagman/bagman"
	"github.com/bitly/go-nsq"
	"github.com/crowdmob/goamz/aws"
	"github.com/crowdmob/goamz/s3"
	"os"
	"path/filepath"
	"time"
)

type ReplicationObject struct {
	File       *bagman.File
	NsqMessage *nsq.Message
}

type Replicator struct {
	ReplicationChannel  chan *ReplicationObject
	S3ReplicationClient *bagman.S3Client
	ProcUtil            *bagman.ProcessUtil
}

func NewReplicator(procUtil *bagman.ProcessUtil) (*Replicator) {
	replicationClient, _ := bagman.NewS3Client(aws.USWest2)
	replicator := &Replicator{
		ProcUtil: procUtil,
		S3ReplicationClient: replicationClient,
	}
	workerBufferSize := procUtil.Config.StoreWorker.Workers * 10
	replicator.ReplicationChannel = make(chan *ReplicationObject, workerBufferSize)
	for i := 0; i < procUtil.Config.StoreWorker.Workers; i++ {
		go replicator.replicate()
	}
	return replicator
}

// MessageHandler handles messages from the queue, putting each
// item into the replication channel.
func (replicator *Replicator) HandleMessage(message *nsq.Message) error {
	message.DisableAutoResponse()
	var file bagman.File
	err := json.Unmarshal(message.Body, &file)
	if err != nil {
		replicator.ProcUtil.MessageLog.Error("Could not unmarshal JSON data from nsq:",
			string(message.Body))
		message.Finish()
		return fmt.Errorf("Could not unmarshal JSON data from nsq")
	}
	if file.NeedsSave == false && replicator.ReplicatedFileExists(file.S3UUID()){
		// This occurs when the depositor has uploaded a new version
		// of an existing bag, but this particiluar file within the
		// bag did not change.
		replicator.ProcUtil.MessageLog.Info("File %s has not changed since it was last saved " +
			"and is confirmed to exist in replication bucket with UUID %s",
			file.Identifier, file.S3UUID())
		message.Finish()
		return nil
	}
	if replicator.ReplicatedFileExists(file.Uuid) {
		replicator.ProcUtil.MessageLog.Info("File %s already exists in replication bucket",
			file.Identifier)
		message.Finish()
		return nil
	}
	replicationObject := &ReplicationObject{
		NsqMessage: message,
		File: &file,
	}

	// Unfortunately, we're probably running on a machine that's
	// also running either ingest or restore. The Volume monitor
	// can know how much disk space those processes are using, but
	// not how much they have reserved. So until we have a volume
	// monitoring service, all we can do is pad the estimate of
	// how much space we might need.
	err = replicator.ProcUtil.Volume.Reserve(uint64(file.Size * 2))
	if err != nil {
		// Not enough room on disk
		replicator.ProcUtil.MessageLog.Warning("Requeueing %s (%d bytes) - not enough disk space",
			file.Identifier, file.Size)
		message.Requeue(10 * time.Minute)
		return nil
	}

	replicator.ReplicationChannel <- replicationObject
	replicator.ProcUtil.MessageLog.Debug("Put %s (%d bytes) into replication queue",
		file.Identifier, file.Size)
	return nil
}

func (replicator *Replicator) ReplicatedFileExists(fileUUID string) (bool) {
	exists, err := replicator.S3ReplicationClient.Exists(
		replicator.ProcUtil.Config.ReplicationBucket,
		fileUUID)
	if err != nil {
		replicator.ProcUtil.MessageLog.Warning("Error checking S3 file exists " +
			"in replication bucket: %v", err)
		replicator.ProcUtil.MessageLog.Info("File %s will be sent to replication " +
			"because we're not sure it's already there", fileUUID)
		return false
    }
	return exists
}

func (replicator *Replicator) replicate() {
	for replicationObject := range replicator.ReplicationChannel {
		replicator.ProcUtil.MessageLog.Info("Starting %s",
			replicationObject.File.Identifier)
		url, err := replicator.CopyAndSaveEvent(replicationObject)
		if err != nil {
			replicationObject.File.ReplicationError = err.Error()
			// If we failed too many times, send this into the failure
			// queue. Otherwise, just requeue and try again.
			if (replicationObject.NsqMessage.Attempts >=
				uint16(replicator.ProcUtil.Config.ReplicationWorker.MaxAttempts)) {
				replicator.SendToTroubleQueue(replicationObject.File)
				replicationObject.NsqMessage.Finish()
				replicator.ProcUtil.IncrementFailed()
			} else {
				replicator.ProcUtil.MessageLog.Error(
					"Requeuing %s (%s) because copy failed. Error: %v",
					replicationObject.File.Identifier,
					replicationObject.File.StorageURL,
					err)
				replicationObject.NsqMessage.Requeue(5 * time.Minute)
				replicator.ProcUtil.IncrementFailed()
			}
		} else {
			// Success!
			replicator.ProcUtil.MessageLog.Info("Finished %s. Replication " +
				"copy is at %s.",
				replicationObject.File.Identifier,
				url)
			replicationObject.NsqMessage.Finish()
			replicator.ProcUtil.IncrementSucceeded()
		}
		replicator.ProcUtil.MessageLog.Info(
			"**STATS** Succeeded: %d, Failed: %d",
			replicator.ProcUtil.Succeeded(),
			replicator.ProcUtil.Failed())
	}
}

// Copy the file to S3 in Oregon and save the replication
// PremisEvent to Fluctus. Returns the S3 URL of the newly-saved file
// or an error.
func (replicator *Replicator) CopyAndSaveEvent(replicationObject *ReplicationObject) (string, error) {
	url, err := replicator.CopyFile(replicationObject)
	if err != nil {
		return "", err
	}
	event, err := replicator.SaveReplicationEvent(replicationObject.File, url)
	if err != nil {
		return "", err
	}
	replicator.ProcUtil.MessageLog.Info(
		"Saved replication PremisEvent for %s (%s) with event identifier %s",
		replicationObject.File.Identifier,
		replicationObject.File.Uuid,
		event.Identifier)
	return url, nil
}

// Copies a file from one bucket to another, across regions,
// including all of APTrust's custom metadata. Returns the URL
// of the destination file (that should be in the replication
// bucket in Oregon), or an error.
//
// This does NOT use PUT COPY internally because PUT COPY is
// limited to files of 5GB or less, and we'll have many files
// over 5GB. The copy operation downloads data from the S3
// preservation bucket and uploads it to the replication bucket.
//
// As long as we're running in the same region as our S3
// preservation bucket (USEast), the download should be fast
// and free. Running this code outside of USEast will be
// slow and expensive, since we'll have to pay for the bandwidth
// of BOTH download and upload.
func (replicator *Replicator) CopyFile(replicationObject *ReplicationObject) (string, error) {
	replicator.ProcUtil.MessageLog.Info("Starting copy of %s (%s)",
		replicationObject.File.Identifier, replicationObject.File.Uuid)
	// Copy options include the md5 sum of the file we're copying
	// and all of our custom meta data.
	copyOptions, err := replicator.GetCopyOptions(replicationObject.File)
	if err != nil {
		return "", err
	}

	// Touch before dowload, because large files can take a long time!
	replicationObject.NsqMessage.Touch()

	localPath, err := replicator.DownloadFromPreservation(replicationObject.File)
	if err != nil {
		return "", err
	}

	reader, err := os.Open(localPath)
	if err != nil {
		return "", err
	}

	// Touch again before upload, because large files are slow
	replicationObject.NsqMessage.Touch()

	// Replication client is configured to us USWest-2 (Oregon),
	// but the bucket name should be enough.
	url := ""
	if replicationObject.File.Size <= bagman.S3_LARGE_FILE {
		url, err = replicator.S3ReplicationClient.SaveToS3(
			replicator.ProcUtil.Config.ReplicationBucket,
			replicationObject.File.Uuid,
			replicationObject.File.MimeType,
			reader,
			replicationObject.File.Size,
			copyOptions)
	} else {
		url, err = replicator.S3ReplicationClient.SaveLargeFileToS3(
			replicator.ProcUtil.Config.ReplicationBucket,
			replicationObject.File.Uuid,
			replicationObject.File.MimeType,
			reader,
			replicationObject.File.Size,
			copyOptions,
			bagman.S3_CHUNK_SIZE)
	}

	// Touch so NSQ knows we're not dead yet!
	replicationObject.NsqMessage.Touch()

	// Delete the local file.
	delErr := os.Remove(localPath)
	if delErr != nil {
		replicator.ProcUtil.MessageLog.Warning("Could not delete local file %s: %v",
			localPath, delErr)
	} else {
		replicator.ProcUtil.Volume.Release(uint64(replicationObject.File.Size * 2))
	}

	if err == nil {
		replicator.ProcUtil.MessageLog.Info("Finished copy of %s (%s)",
			replicationObject.File.Identifier,
			replicationObject.File.Uuid)
	}

	return url, err
}

// Returns S3 options, including the md5 checksum and APTrust's custom
// metadata. These options must accompany the file copy.
func (replicator *Replicator) GetCopyOptions(file *bagman.File) (s3.Options, error) {
	// Copy all of the meta data
	resp, err := replicator.ProcUtil.S3Client.Head(
		replicator.ProcUtil.Config.PreservationBucket,
		file.Uuid)
	if err != nil {
		detailedErr :=  fmt.Errorf("Head request for %s at %s returned error: %v",
			file.Uuid,
			replicator.ProcUtil.Config.PreservationBucket,
			err)
		return s3.Options{}, detailedErr
	}
	if resp.StatusCode != 200 {
		detailedErr := fmt.Errorf(
			"Can't get S3 metadata for file %s. " +
			"Head request for %s/%s returned HTTP Status Code %d",
			file.Identifier,
			replicator.ProcUtil.Config.PreservationBucket,
			file.Uuid,
			resp.StatusCode)
		return s3.Options{}, detailedErr
	}

	s3Metadata := make(map[string][]string)
	s3Metadata["md5"]         = []string{ file.Md5 }
	s3Metadata["institution"] = []string{ resp.Header["X-Amz-Meta-Institution"][0] }
	s3Metadata["bag"]         = []string{ resp.Header["X-Amz-Meta-Bag"][0] }
	s3Metadata["bagpath"]     = []string{ file.Path }

	base64md5, err := bagman.Base64EncodeMd5(file.Md5)
	if err != nil {
		return s3.Options{}, err
	}

	options := replicator.S3ReplicationClient.MakeOptions(base64md5, s3Metadata)
	return options, nil
}

// Copies a file from the preservation bucket to a local file
// and returns the path to the local file.
func (replicator *Replicator) DownloadFromPreservation(file *bagman.File) (string, error) {
	// Make sure we have a folder to put the file in.
	replicationDir := replicator.ProcUtil.Config.ReplicationDirectory
	if _, err := os.Stat(replicationDir); os.IsNotExist(err) {
		err = os.MkdirAll(replicationDir, 0755)
		if err != nil {
			return "", err
		}
	}

	// Now copy the file from S3 to local path.
	localPath := filepath.Join(replicationDir, file.Uuid)
	err := replicator.ProcUtil.S3Client.FetchToFileWithoutChecksum(
		replicator.ProcUtil.Config.PreservationBucket,
		file.Uuid,
		localPath)
	if err != nil {
		detailedErr := fmt.Errorf("Cannot read file %s from %s: %v",
			file.Uuid,
			replicator.ProcUtil.Config.PreservationBucket,
			err)
		return "", detailedErr
	}

	replicator.ProcUtil.MessageLog.Info("Downloaded %s", localPath)
	return localPath, nil
}

// Saves the replication PremisEvent to Fluctus.
// Param url is the S3 URL we just saved that file to.
// That should be in Oregon.
func (replicator *Replicator) SaveReplicationEvent(file *bagman.File, url string) (*bagman.PremisEvent, error) {
	replicationEvent, err := file.ReplicationEvent(url)
	if err != nil {
		return nil, err
	}
	replicator.ProcUtil.MessageLog.Info("Saving replication PremisEvent for %s (%s)",
		file.Identifier, file.Uuid)
	savedEvent, err := replicator.ProcUtil.FluctusClient.PremisEventSave(
		file.Identifier, "GenericFile", replicationEvent)
	if err != nil {
		return nil, err
	}
	return savedEvent, nil
}

// Puts an item into the trouble queue.
func (replicator *Replicator) SendToTroubleQueue(file *bagman.File) {
	err := bagman.Enqueue(
		replicator.ProcUtil.Config.NsqdHttpAddress,
		replicator.ProcUtil.Config.FailedReplicationWorker.NsqTopic,
		file)
	if err != nil {
		replicator.ProcUtil.MessageLog.Error(
			"Could not send '%s' (%s) to failed replication queue: %v\n",
			file.Identifier,
			file.Uuid,
			err)
	} else {
		message := fmt.Sprintf("Failed to copy '%s' (%s) to replication bucket. " +
			"This item is going into the failed replication queue.",
			file.Identifier,
			file.Uuid)
		replicator.ProcUtil.MessageLog.Warning(message)
	}
}
<|endoftext|>


<commit_before>package funcs

import (
	"sync"

	"github.com/hairyhenderson/gomplate/aws"
)

var (
	af     *Funcs
	afInit sync.Once
)

// AWSNS - the aws namespace
func AWSNS() *Funcs {
	afInit.Do(func() {
		af = &Funcs{
			awsopts: aws.GetClientOptions(),
		}
	})
	return af
}

// AWSFuncs -
func AWSFuncs(f map[string]interface{}) {
	f["aws"] = AWSNS

	// global aliases - for backwards compatibility
	f["ec2meta"] = AWSNS().EC2Meta
	f["ec2dynamic"] = AWSNS().EC2Dynamic
	f["ec2tag"] = AWSNS().EC2Tag
	f["ec2region"] = AWSNS().EC2Region
	f["kmsEncrypt"] = AWSNS().KMSEncrypt
	f["kmsDecrypt"] = AWSNS().KMSDecrypt
}

// Funcs -
type Funcs struct {
	meta     *aws.Ec2Meta
	info     *aws.Ec2Info
	kms      *aws.KMS
	metaInit sync.Once
	infoInit sync.Once
	kmsInit  sync.Once
	awsopts  aws.ClientOptions
}

// EC2Region -
func (a *Funcs) EC2Region(def ...string) (string, error) {
	a.metaInit.Do(a.initMeta)
	return a.meta.Region(def...)
}

// EC2Meta -
func (a *Funcs) EC2Meta(key string, def ...string) (string, error) {
	a.metaInit.Do(a.initMeta)
	return a.meta.Meta(key, def...)
}

// EC2Dynamic -
func (a *Funcs) EC2Dynamic(key string, def ...string) (string, error) {
	a.metaInit.Do(a.initMeta)
	return a.meta.Dynamic(key, def...)
}

// EC2Tag -
func (a *Funcs) EC2Tag(tag string, def ...string) (string, error) {
	a.infoInit.Do(a.initInfo)
	return a.info.Tag(tag, def...)
}

// KMSEncrypt -
func (a *Funcs) KMSEncrypt(keyID, plaintext string) (string, error) {
	a.kmsInit.Do(a.initKMS)
	return a.kms.Encrypt(keyID, plaintext)
}

// KMSDecrypt -
func (a *Funcs) KMSDecrypt(ciphertext string) (string, error) {
	a.kmsInit.Do(a.initKMS)
	return a.kms.Decrypt(ciphertext)
}

func (a *Funcs) initMeta() {
	if a.meta == nil {
		a.meta = aws.NewEc2Meta(a.awsopts)
	}
}

func (a *Funcs) initInfo() {
	if a.info == nil {
		a.info = aws.NewEc2Info(a.awsopts)
	}
}

func (a *Funcs) initKMS() {
	if a.kms == nil {
		a.kms = aws.NewKMS(a.awsopts)
	}
}
<commit_msg>remove aliases<commit_after>package funcs

import (
	"sync"

	"github.com/hairyhenderson/gomplate/aws"
)

var (
	af     *Funcs
	afInit sync.Once
)

// AWSNS - the aws namespace
func AWSNS() *Funcs {
	afInit.Do(func() {
		af = &Funcs{
			awsopts: aws.GetClientOptions(),
		}
	})
	return af
}

// AWSFuncs -
func AWSFuncs(f map[string]interface{}) {
	f["aws"] = AWSNS

	// global aliases - for backwards compatibility
	f["ec2meta"] = AWSNS().EC2Meta
	f["ec2dynamic"] = AWSNS().EC2Dynamic
	f["ec2tag"] = AWSNS().EC2Tag
	f["ec2region"] = AWSNS().EC2Region
}

// Funcs -
type Funcs struct {
	meta     *aws.Ec2Meta
	info     *aws.Ec2Info
	kms      *aws.KMS
	metaInit sync.Once
	infoInit sync.Once
	kmsInit  sync.Once
	awsopts  aws.ClientOptions
}

// EC2Region -
func (a *Funcs) EC2Region(def ...string) (string, error) {
	a.metaInit.Do(a.initMeta)
	return a.meta.Region(def...)
}

// EC2Meta -
func (a *Funcs) EC2Meta(key string, def ...string) (string, error) {
	a.metaInit.Do(a.initMeta)
	return a.meta.Meta(key, def...)
}

// EC2Dynamic -
func (a *Funcs) EC2Dynamic(key string, def ...string) (string, error) {
	a.metaInit.Do(a.initMeta)
	return a.meta.Dynamic(key, def...)
}

// EC2Tag -
func (a *Funcs) EC2Tag(tag string, def ...string) (string, error) {
	a.infoInit.Do(a.initInfo)
	return a.info.Tag(tag, def...)
}

// KMSEncrypt -
func (a *Funcs) KMSEncrypt(keyID, plaintext string) (string, error) {
	a.kmsInit.Do(a.initKMS)
	return a.kms.Encrypt(keyID, plaintext)
}

// KMSDecrypt -
func (a *Funcs) KMSDecrypt(ciphertext string) (string, error) {
	a.kmsInit.Do(a.initKMS)
	return a.kms.Decrypt(ciphertext)
}

func (a *Funcs) initMeta() {
	if a.meta == nil {
		a.meta = aws.NewEc2Meta(a.awsopts)
	}
}

func (a *Funcs) initInfo() {
	if a.info == nil {
		a.info = aws.NewEc2Info(a.awsopts)
	}
}

func (a *Funcs) initKMS() {
	if a.kms == nil {
		a.kms = aws.NewKMS(a.awsopts)
	}
}
<|endoftext|>


<commit_before>package functions

import (
	"html/template"
	"os"
	"strings"
)

var (
	Bold = func(content string) (template.HTML, error) {
		return template.HTML("" + content + ""), nil
	}
	Italic = func(content string) (template.HTML, error) {
		return template.HTML("" + content + ""), nil
	}
	Underline = func(content string) (template.HTML, error) {
		return template.HTML("" + content + ""), nil
	}
)

func Exists(path string) bool {
	_, err := os.Stat(path)
	if err == nil {
		return true
	}
	if os.IsNotExist(err) {
		return false
	}
	return true
}

func GetFilenameNoExtension(s string) string {
	n := strings.LastIndexByte(s, '.')
	if n >= 0 {
		return s[:n]
	}
	return s
}

func RemoveDuplicates(a []string) []string {
	result := []string{}
	seen := map[string]string{}
	for _, val := range a {
		if _, ok := seen[val]; !ok {
			result = append(result, val)
			seen[val] = val
		}
	}
	return result
}
<commit_msg>Added Encrypt/Decrypt functions<commit_after>package functions

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"html/template"
	"io"
	"os"
	"strings"
)

var (
	Bold = func(content string) (template.HTML, error) {
		return template.HTML("" + content + ""), nil
	}
	Italic = func(content string) (template.HTML, error) {
		return template.HTML("" + content + ""), nil
	}
	Underline = func(content string) (template.HTML, error) {
		return template.HTML("" + content + ""), nil
	}
)

func Exists(path string) bool {
	_, err := os.Stat(path)
	if err == nil {
		return true
	}
	if os.IsNotExist(err) {
		return false
	}
	return true
}

func GetFilenameNoExtension(s string) string {
	n := strings.LastIndexByte(s, '.')
	if n >= 0 {
		return s[:n]
	}
	return s
}

func RemoveDuplicates(a []string) []string {
	result := []string{}
	seen := map[string]string{}
	for _, val := range a {
		if _, ok := seen[val]; !ok {
			result = append(result, val)
			seen[val] = val
		}
	}
	return result
}

func encodeBase64(b []byte) string {
	return base64.StdEncoding.EncodeToString(b)
}

func decodeBase64(s string) []byte {
	data, err := base64.StdEncoding.DecodeString(s)
	if err != nil {
		panic(err)
	}
	return data
}

func Encrypt(key, text []byte) string {
	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}
	ciphertext := make([]byte, aes.BlockSize+len(text))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}
	cfb := cipher.NewCFBEncrypter(block, iv)
	cfb.XORKeyStream(ciphertext[aes.BlockSize:], text)
	return encodeBase64(ciphertext)
}

func Decrypt(key []byte, b64 string) string {
	text := decodeBase64(b64)
	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}
	if len(text) < aes.BlockSize {
		panic("ciphertext too short")
	}
	iv := text[:aes.BlockSize]
	text = text[aes.BlockSize:]
	cfb := cipher.NewCFBDecrypter(block, iv)
	cfb.XORKeyStream(text, text)
	return string(text)
}
<|endoftext|>


<commit_before>package fuzz

import (
	"testing"
)

func TestParse(t *testing.T) {
	Parse("")
}

<commit_msg>Expanding testing stub<commit_after>package fuzz

import (
	"testing"
)

type FuzzTest struct {

}

var fuzzTests = []FuzzTest{}

func TestFuzz(t *testing.T) {

}

type ParseTest struct {

}

var parseTests = []ParseTest{}

func TestParse(t *testing.T) {
	Parse("")
}

<|endoftext|>


<commit_before>// Copyright 2012 SocialCode. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package gelf

type GELFWriter struct {

}

// Used to control GELF chunking
const ChunkSize = 1420

func NewGELFWriter(addr string) (*GELFWriter, error) {
	w := new(GELFWriter)

	return w, nil
}

func (g *GELFWriter) Write(p []byte) (n int, err error) {
	return 0, nil
}
<commit_msg>lots of work on writer<commit_after>// Copyright 2012 SocialCode. All rights reserved.
// Use of this source code is governed by the MIT
// license that can be found in the LICENSE file.

package gelf

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"runtime"
	"strings"
	"sync"
	"time"
)

type Writer struct {
	mu       sync.Mutex
	conn     net.Conn
	facility string
	hostname string
}

type Message struct {
	Version      string `json:"version"`
	Host         string `json:"host"`
	ShortMessage string `json:"short_message"`
	FullMessage  string `json:"full_message"`
	TimeUnix     int64  `json:"timestamp"`
	Level        int32  `json:"level"`
	Facility     string `json:"facility"`
	File         string `json:"file"`
	Line         int    `json:"line"`
}

/*
{
  "version": "1.0",
  "host": "www1",
  "short_message": "Short message",
  "full_message": "Backtrace here\n\nmore stuff",
  "timestamp": 1291899928.412,
  "level": 1,
  "facility": "payment-backend",
  "file": "/var/www/somefile.rb",
  "line": 356,
  "_user_id": 42,
  "_something_else": "foo"
}
*/

// Used to control GELF chunking
const ChunkSize = 1420

func New(addr string) (*Writer, error) {
	var err error
	w := new(Writer)

	if w.conn, err = net.Dial("udp", addr); err != nil {
		return nil, err
	}
	if w.hostname, err = os.Hostname(); err != nil {
		return nil, err
	}

	return w, nil
}

func (w *Writer) WriteMessage(m *Message) error {
	mBytes, err := json.Marshal(m)
	if err != nil {
		return err
	}

	n, err := w.conn.Write(mBytes)
	if err != nil {
		return err
	}
	if n != len(mBytes) {
		return fmt.Errorf("bad write (%d/%d)", n, len(mBytes))
	}

	return nil
}

/*
func (w *Writer) Alert(m string) (err error)
func (w *Writer) Close() error
func (w *Writer) Crit(m string) (err error)
func (w *Writer) Debug(m string) (err error)
func (w *Writer) Emerg(m string) (err error)
func (w *Writer) Err(m string) (err error)
func (w *Writer) Info(m string) (err error)
func (w *Writer) Notice(m string) (err error)
func (w *Writer) Warning(m string) (err error)
*/

// getCallerIgnoringLog returns the filename and the line info
func getCallerIgnoringLog() (file string, line int) {
	var ok bool
	calldepth := 2 // 0 would be this function, 1 would be Writer.Write

	for {
		_, file, line, ok = runtime.Caller(calldepth)
		if !ok {
			file = "???"
			line = 0
			break
		}
		if !strings.HasSuffix(file, "/pkg/log/log.go") {
			break
		}
		calldepth++
	}
	return
}

func (w *Writer) Write(p []byte) (n int, err error) {

	file, line := getCallerIgnoringLog()

	fmt.Printf("f: %s, line %d\n", file, line)

	m := Message{
		Version:      "1.0",
		Host:         w.hostname,
		ShortMessage: "",
		FullMessage:  "",
		TimeUnix:     time.Now().Unix(),
		Level:        1,
		Facility:     w.facility,
		File:         file,
		Line:         line,
	}

	if err = w.WriteMessage(&m); err != nil {
		return 0, err
	}

	return len(p), nil
}
<|endoftext|>


<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/gogo/protobuf/proto"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/auth"
	"github.com/mesos/mesos-go/auth/sasl"
	"github.com/mesos/mesos-go/auth/sasl/mech"
	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	sched "github.com/mesos/mesos-go/scheduler"
	"golang.org/x/net/context"
)

const (
	CPUS_PER_EXECUTOR   = 0.01
	CPUS_PER_TASK       = 1
	MEM_PER_EXECUTOR    = 64
	MEM_PER_TASK        = 64
	defaultArtifactPort = 12345
)

var (
	address      = flag.String("address", "127.0.0.1", "Binding address for artifact server")
	artifactPort = flag.Int("artifactPort", defaultArtifactPort, "Binding port for artifact server")
	authProvider = flag.String("mesos_authentication_provider", sasl.ProviderName,
		fmt.Sprintf("Authentication provider to use, default is SASL that supports mechanisms: %+v", mech.ListSupported()))
	master              = flag.String("master", "127.0.0.1:5050", "Master address <ip:port>")
	executorPath        = flag.String("executor", "./executor", "Path to test executor")
	taskCount           = flag.String("task-count", "5", "Total task count to run.")
	mesosAuthPrincipal  = flag.String("mesos_authentication_principal", "", "Mesos authentication principal.")
	mesosAuthSecretFile = flag.String("mesos_authentication_secret_file", "", "Mesos authentication secret file.")
	slowLaunch          = flag.Bool("slow_launch", false, "When true the ResourceOffers func waits for several seconds before attempting to launch tasks; useful for debugging failover")
)

type ExampleScheduler struct {
	executor      *mesos.ExecutorInfo
	tasksLaunched int
	tasksFinished int
	totalTasks    int
}

func newExampleScheduler(exec *mesos.ExecutorInfo) *ExampleScheduler {
	total, err := strconv.Atoi(*taskCount)
	if err != nil {
		total = 5
	}
	return &ExampleScheduler{
		executor:      exec,
		tasksLaunched: 0,
		tasksFinished: 0,
		totalTasks:    total,
	}
}

func (sched *ExampleScheduler) Registered(driver sched.SchedulerDriver, frameworkId *mesos.FrameworkID, masterInfo *mesos.MasterInfo) {
	log.Infoln("Framework Registered with Master ", masterInfo)
}

func (sched *ExampleScheduler) Reregistered(driver sched.SchedulerDriver, masterInfo *mesos.MasterInfo) {
	log.Infoln("Framework Re-Registered with Master ", masterInfo)
	_, err := driver.ReconcileTasks([]*mesos.TaskStatus{})
	if err != nil {
		log.Errorf("failed to request task reconciliation: %v", err)
	}
}

func (sched *ExampleScheduler) Disconnected(sched.SchedulerDriver) {
	log.Warningf("disconnected from master")
}

func (sched *ExampleScheduler) ResourceOffers(driver sched.SchedulerDriver, offers []*mesos.Offer) {

	if *slowLaunch {
		time.Sleep(3 * time.Second)
	}

	if sched.tasksLaunched >= sched.totalTasks {
		log.Info("decline all of the offers since all of our tasks are already launched")
		ids := make([]*mesos.OfferID, len(offers))
		for i, offer := range offers {
			ids[i] = offer.Id
		}
		driver.LaunchTasks(ids, []*mesos.TaskInfo{}, &mesos.Filters{RefuseSeconds: proto.Float64(120)})
		return
	}
	for _, offer := range offers {
		cpuResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
			return res.GetName() == "cpus"
		})
		cpus := 0.0
		for _, res := range cpuResources {
			cpus += res.GetScalar().GetValue()
		}

		memResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
			return res.GetName() == "mem"
		})
		mems := 0.0
		for _, res := range memResources {
			mems += res.GetScalar().GetValue()
		}

		log.Infoln("Received Offer <", offer.Id.GetValue(), "> with cpus=", cpus, " mem=", mems)

		remainingCpus := cpus
		remainingMems := mems

		// account for executor resources if there's not an executor already running on the slave
		if len(offer.ExecutorIds) == 0 {
			remainingCpus -= CPUS_PER_EXECUTOR
			remainingMems -= MEM_PER_EXECUTOR
		}

		var tasks []*mesos.TaskInfo
		for sched.tasksLaunched < sched.totalTasks &&
			CPUS_PER_TASK <= remainingCpus &&
			MEM_PER_TASK <= remainingMems {

			sched.tasksLaunched++

			taskId := &mesos.TaskID{
				Value: proto.String(strconv.Itoa(sched.tasksLaunched)),
			}

			task := &mesos.TaskInfo{
				Name:     proto.String("go-task-" + taskId.GetValue()),
				TaskId:   taskId,
				SlaveId:  offer.SlaveId,
				Executor: sched.executor,
				Resources: []*mesos.Resource{
					util.NewScalarResource("cpus", CPUS_PER_TASK),
					util.NewScalarResource("mem", MEM_PER_TASK),
				},
			}
			log.Infof("Prepared task: %s with offer %s for launch\n", task.GetName(), offer.Id.GetValue())

			tasks = append(tasks, task)
			remainingCpus -= CPUS_PER_TASK
			remainingMems -= MEM_PER_TASK
		}
		log.Infoln("Launching ", len(tasks), "tasks for offer", offer.Id.GetValue())
		driver.LaunchTasks([]*mesos.OfferID{offer.Id}, tasks, &mesos.Filters{RefuseSeconds: proto.Float64(5)})
	}
}

func (sched *ExampleScheduler) StatusUpdate(driver sched.SchedulerDriver, status *mesos.TaskStatus) {
	log.Infoln("Status update: task", status.TaskId.GetValue(), " is in state ", status.State.Enum().String())
	if status.GetState() == mesos.TaskState_TASK_FINISHED {
		sched.tasksFinished++
		driver.ReviveOffers() // TODO(jdef) rate-limit this
	}

	if sched.tasksFinished >= sched.totalTasks {
		log.Infoln("Total tasks completed, stopping framework.")
		driver.Stop(false)
	}

	if status.GetState() == mesos.TaskState_TASK_LOST ||
		status.GetState() == mesos.TaskState_TASK_KILLED ||
		status.GetState() == mesos.TaskState_TASK_FAILED ||
		status.GetState() == mesos.TaskState_TASK_ERROR {
		log.Infoln(
			"Aborting because task", status.TaskId.GetValue(),
			"is in unexpected state", status.State.String(),
			"with message", status.GetMessage(),
		)
		driver.Abort()
	}
}

func (sched *ExampleScheduler) OfferRescinded(_ sched.SchedulerDriver, oid *mesos.OfferID) {
	log.Errorf("offer rescinded: %v", oid)
}
func (sched *ExampleScheduler) FrameworkMessage(_ sched.SchedulerDriver, eid *mesos.ExecutorID, sid *mesos.SlaveID, msg string) {
	log.Errorf("framework message from executor %q slave %q: %q", eid, sid, msg)
}
func (sched *ExampleScheduler) SlaveLost(_ sched.SchedulerDriver, sid *mesos.SlaveID) {
	log.Errorf("slave lost: %v", sid)
}
func (sched *ExampleScheduler) ExecutorLost(_ sched.SchedulerDriver, eid *mesos.ExecutorID, sid *mesos.SlaveID, code int) {
	log.Errorf("executor %q lost on slave %q code %d", eid, sid, code)
}
func (sched *ExampleScheduler) Error(_ sched.SchedulerDriver, err string) {
	log.Errorf("Scheduler received error: %v", err)
}

// ----------------------- func init() ------------------------- //

func init() {
	flag.Parse()
	log.Infoln("Initializing the Example Scheduler...")
}

// returns (downloadURI, basename(path))
func serveExecutorArtifact(path string) (*string, string) {
	serveFile := func(pattern string, filename string) {
		http.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, filename)
		})
	}

	// Create base path (http://foobar:5000/<base>)
	pathSplit := strings.Split(path, "/")
	var base string
	if len(pathSplit) > 0 {
		base = pathSplit[len(pathSplit)-1]
	} else {
		base = path
	}
	serveFile("/"+base, path)

	hostURI := fmt.Sprintf("http://%s:%d/%s", *address, *artifactPort, base)
	log.V(2).Infof("Hosting artifact '%s' at '%s'", path, hostURI)

	return &hostURI, base
}

func prepareExecutorInfo() *mesos.ExecutorInfo {
	executorUris := []*mesos.CommandInfo_URI{}
	uri, executorCmd := serveExecutorArtifact(*executorPath)
	executorUris = append(executorUris, &mesos.CommandInfo_URI{Value: uri, Executable: proto.Bool(true)})

	// forward the value of the scheduler's -v flag to the executor
	v := 0
	if f := flag.Lookup("v"); f != nil && f.Value != nil {
		if vstr := f.Value.String(); vstr != "" {
			if vi, err := strconv.ParseInt(vstr, 10, 32); err == nil {
				v = int(vi)
			}
		}
	}
	executorCommand := fmt.Sprintf("./%s -logtostderr=true -v=%d", executorCmd, v)

	go http.ListenAndServe(fmt.Sprintf("%s:%d", *address, *artifactPort), nil)
	log.V(2).Info("Serving executor artifacts...")

	// Create mesos scheduler driver.
	return &mesos.ExecutorInfo{
		ExecutorId: util.NewExecutorID("default"),
		Name:       proto.String("Test Executor (Go)"),
		Source:     proto.String("go_test"),
		Command: &mesos.CommandInfo{
			Value: proto.String(executorCommand),
			Uris:  executorUris,
		},
		Resources: []*mesos.Resource{
			util.NewScalarResource("cpus", CPUS_PER_EXECUTOR),
			util.NewScalarResource("mem", MEM_PER_EXECUTOR),
		},
	}
}

func parseIP(address string) net.IP {
	addr, err := net.LookupIP(address)
	if err != nil {
		log.Fatal(err)
	}
	if len(addr) < 1 {
		log.Fatalf("failed to parse IP from address '%v'", address)
	}
	return addr[0]
}

// ----------------------- func main() ------------------------- //

func main() {

	// build command executor
	exec := prepareExecutorInfo()

	// the framework
	fwinfo := &mesos.FrameworkInfo{
		User: proto.String(""), // Mesos-go will fill in user.
		Name: proto.String("Test Framework (Go)"),
	}

	cred := (*mesos.Credential)(nil)
	if *mesosAuthPrincipal != "" {
		fwinfo.Principal = proto.String(*mesosAuthPrincipal)
		cred = &mesos.Credential{
			Principal: proto.String(*mesosAuthPrincipal),
		}
		if *mesosAuthSecretFile != "" {
			_, err := os.Stat(*mesosAuthSecretFile)
			if err != nil {
				log.Fatal("missing secret file: ", err.Error())
			}
			secret, err := ioutil.ReadFile(*mesosAuthSecretFile)
			if err != nil {
				log.Fatal("failed to read secret file: ", err.Error())
			}
			cred.Secret = proto.String(string(secret))
		}
	}
	bindingAddress := parseIP(*address)
	config := sched.DriverConfig{
		Scheduler:      newExampleScheduler(exec),
		Framework:      fwinfo,
		Master:         *master,
		Credential:     cred,
		BindingAddress: bindingAddress,
		WithAuthContext: func(ctx context.Context) context.Context {
			ctx = auth.WithLoginProvider(ctx, *authProvider)
			ctx = sasl.WithBindingAddress(ctx, bindingAddress)
			return ctx
		},
	}
	driver, err := sched.NewMesosSchedulerDriver(config)

	if err != nil {
		log.Errorln("Unable to create a SchedulerDriver ", err.Error())
	}

	if stat, err := driver.Run(); err != nil {
		log.Infof("Framework stopped with status %s and error: %s\n", stat.String(), err.Error())
		time.Sleep(2 * time.Second)
		os.Exit(1)
	}
	log.Infof("framework terminating")
}
<commit_msg>do not abort when tasks error<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/gogo/protobuf/proto"
	log "github.com/golang/glog"
	"github.com/mesos/mesos-go/auth"
	"github.com/mesos/mesos-go/auth/sasl"
	"github.com/mesos/mesos-go/auth/sasl/mech"
	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	sched "github.com/mesos/mesos-go/scheduler"
	"golang.org/x/net/context"
)

const (
	CPUS_PER_EXECUTOR   = 0.01
	CPUS_PER_TASK       = 1
	MEM_PER_EXECUTOR    = 64
	MEM_PER_TASK        = 64
	defaultArtifactPort = 12345
)

var (
	address      = flag.String("address", "127.0.0.1", "Binding address for artifact server")
	artifactPort = flag.Int("artifactPort", defaultArtifactPort, "Binding port for artifact server")
	authProvider = flag.String("mesos_authentication_provider", sasl.ProviderName,
		fmt.Sprintf("Authentication provider to use, default is SASL that supports mechanisms: %+v", mech.ListSupported()))
	master              = flag.String("master", "127.0.0.1:5050", "Master address <ip:port>")
	executorPath        = flag.String("executor", "./executor", "Path to test executor")
	taskCount           = flag.String("task-count", "5", "Total task count to run.")
	mesosAuthPrincipal  = flag.String("mesos_authentication_principal", "", "Mesos authentication principal.")
	mesosAuthSecretFile = flag.String("mesos_authentication_secret_file", "", "Mesos authentication secret file.")
	slowLaunch          = flag.Bool("slow_launch", false, "When true the ResourceOffers func waits for several seconds before attempting to launch tasks; useful for debugging failover")
)

type ExampleScheduler struct {
	executor      *mesos.ExecutorInfo
	tasksLaunched int
	tasksFinished int
	tasksErrored  int
	totalTasks    int
}

func newExampleScheduler(exec *mesos.ExecutorInfo) *ExampleScheduler {
	total, err := strconv.Atoi(*taskCount)
	if err != nil {
		total = 5
	}
	return &ExampleScheduler{
		executor:   exec,
		totalTasks: total,
	}
}

func (sched *ExampleScheduler) Registered(driver sched.SchedulerDriver, frameworkId *mesos.FrameworkID, masterInfo *mesos.MasterInfo) {
	log.Infoln("Framework Registered with Master ", masterInfo)
}

func (sched *ExampleScheduler) Reregistered(driver sched.SchedulerDriver, masterInfo *mesos.MasterInfo) {
	log.Infoln("Framework Re-Registered with Master ", masterInfo)
	_, err := driver.ReconcileTasks([]*mesos.TaskStatus{})
	if err != nil {
		log.Errorf("failed to request task reconciliation: %v", err)
	}
}

func (sched *ExampleScheduler) Disconnected(sched.SchedulerDriver) {
	log.Warningf("disconnected from master")
}

func (sched *ExampleScheduler) ResourceOffers(driver sched.SchedulerDriver, offers []*mesos.Offer) {

	if *slowLaunch {
		time.Sleep(3 * time.Second)
	}

	if (sched.tasksLaunched - sched.tasksErrored) >= sched.totalTasks {
		log.Info("decline all of the offers since all of our tasks are already launched")
		ids := make([]*mesos.OfferID, len(offers))
		for i, offer := range offers {
			ids[i] = offer.Id
		}
		driver.LaunchTasks(ids, []*mesos.TaskInfo{}, &mesos.Filters{RefuseSeconds: proto.Float64(120)})
		return
	}
	for _, offer := range offers {
		cpuResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
			return res.GetName() == "cpus"
		})
		cpus := 0.0
		for _, res := range cpuResources {
			cpus += res.GetScalar().GetValue()
		}

		memResources := util.FilterResources(offer.Resources, func(res *mesos.Resource) bool {
			return res.GetName() == "mem"
		})
		mems := 0.0
		for _, res := range memResources {
			mems += res.GetScalar().GetValue()
		}

		log.Infoln("Received Offer <", offer.Id.GetValue(), "> with cpus=", cpus, " mem=", mems)

		remainingCpus := cpus
		remainingMems := mems

		// account for executor resources if there's not an executor already running on the slave
		if len(offer.ExecutorIds) == 0 {
			remainingCpus -= CPUS_PER_EXECUTOR
			remainingMems -= MEM_PER_EXECUTOR
		}

		var tasks []*mesos.TaskInfo
		for (sched.tasksLaunched-sched.tasksErrored) < sched.totalTasks &&
			CPUS_PER_TASK <= remainingCpus &&
			MEM_PER_TASK <= remainingMems {

			sched.tasksLaunched++

			taskId := &mesos.TaskID{
				Value: proto.String(strconv.Itoa(sched.tasksLaunched)),
			}

			task := &mesos.TaskInfo{
				Name:     proto.String("go-task-" + taskId.GetValue()),
				TaskId:   taskId,
				SlaveId:  offer.SlaveId,
				Executor: sched.executor,
				Resources: []*mesos.Resource{
					util.NewScalarResource("cpus", CPUS_PER_TASK),
					util.NewScalarResource("mem", MEM_PER_TASK),
				},
			}
			log.Infof("Prepared task: %s with offer %s for launch\n", task.GetName(), offer.Id.GetValue())

			tasks = append(tasks, task)
			remainingCpus -= CPUS_PER_TASK
			remainingMems -= MEM_PER_TASK
		}
		log.Infoln("Launching ", len(tasks), "tasks for offer", offer.Id.GetValue())
		driver.LaunchTasks([]*mesos.OfferID{offer.Id}, tasks, &mesos.Filters{RefuseSeconds: proto.Float64(5)})
	}
}

func (sched *ExampleScheduler) StatusUpdate(driver sched.SchedulerDriver, status *mesos.TaskStatus) {
	log.Infoln("Status update: task", status.TaskId.GetValue(), " is in state ", status.State.Enum().String())
	if status.GetState() == mesos.TaskState_TASK_FINISHED {
		sched.tasksFinished++
		driver.ReviveOffers() // TODO(jdef) rate-limit this
	}

	if sched.tasksFinished >= sched.totalTasks {
		log.Infoln("Total tasks completed, stopping framework.")
		driver.Stop(false)
	}

	if status.GetState() == mesos.TaskState_TASK_LOST ||
		status.GetState() == mesos.TaskState_TASK_KILLED ||
		status.GetState() == mesos.TaskState_TASK_FAILED ||
		status.GetState() == mesos.TaskState_TASK_ERROR {
		sched.tasksErrored++
	}
}

func (sched *ExampleScheduler) OfferRescinded(_ sched.SchedulerDriver, oid *mesos.OfferID) {
	log.Errorf("offer rescinded: %v", oid)
}
func (sched *ExampleScheduler) FrameworkMessage(_ sched.SchedulerDriver, eid *mesos.ExecutorID, sid *mesos.SlaveID, msg string) {
	log.Errorf("framework message from executor %q slave %q: %q", eid, sid, msg)
}
func (sched *ExampleScheduler) SlaveLost(_ sched.SchedulerDriver, sid *mesos.SlaveID) {
	log.Errorf("slave lost: %v", sid)
}
func (sched *ExampleScheduler) ExecutorLost(_ sched.SchedulerDriver, eid *mesos.ExecutorID, sid *mesos.SlaveID, code int) {
	log.Errorf("executor %q lost on slave %q code %d", eid, sid, code)
}
func (sched *ExampleScheduler) Error(_ sched.SchedulerDriver, err string) {
	log.Errorf("Scheduler received error: %v", err)
}

// ----------------------- func init() ------------------------- //

func init() {
	flag.Parse()
	log.Infoln("Initializing the Example Scheduler...")
}

// returns (downloadURI, basename(path))
func serveExecutorArtifact(path string) (*string, string) {
	serveFile := func(pattern string, filename string) {
		http.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, filename)
		})
	}

	// Create base path (http://foobar:5000/<base>)
	pathSplit := strings.Split(path, "/")
	var base string
	if len(pathSplit) > 0 {
		base = pathSplit[len(pathSplit)-1]
	} else {
		base = path
	}
	serveFile("/"+base, path)

	hostURI := fmt.Sprintf("http://%s:%d/%s", *address, *artifactPort, base)
	log.V(2).Infof("Hosting artifact '%s' at '%s'", path, hostURI)

	return &hostURI, base
}

func prepareExecutorInfo() *mesos.ExecutorInfo {
	executorUris := []*mesos.CommandInfo_URI{}
	uri, executorCmd := serveExecutorArtifact(*executorPath)
	executorUris = append(executorUris, &mesos.CommandInfo_URI{Value: uri, Executable: proto.Bool(true)})

	// forward the value of the scheduler's -v flag to the executor
	v := 0
	if f := flag.Lookup("v"); f != nil && f.Value != nil {
		if vstr := f.Value.String(); vstr != "" {
			if vi, err := strconv.ParseInt(vstr, 10, 32); err == nil {
				v = int(vi)
			}
		}
	}
	executorCommand := fmt.Sprintf("./%s -logtostderr=true -v=%d", executorCmd, v)

	go http.ListenAndServe(fmt.Sprintf("%s:%d", *address, *artifactPort), nil)
	log.V(2).Info("Serving executor artifacts...")

	// Create mesos scheduler driver.
	return &mesos.ExecutorInfo{
		ExecutorId: util.NewExecutorID("default"),
		Name:       proto.String("Test Executor (Go)"),
		Source:     proto.String("go_test"),
		Command: &mesos.CommandInfo{
			Value: proto.String(executorCommand),
			Uris:  executorUris,
		},
		Resources: []*mesos.Resource{
			util.NewScalarResource("cpus", CPUS_PER_EXECUTOR),
			util.NewScalarResource("mem", MEM_PER_EXECUTOR),
		},
	}
}

func parseIP(address string) net.IP {
	addr, err := net.LookupIP(address)
	if err != nil {
		log.Fatal(err)
	}
	if len(addr) < 1 {
		log.Fatalf("failed to parse IP from address '%v'", address)
	}
	return addr[0]
}

// ----------------------- func main() ------------------------- //

func main() {

	// build command executor
	exec := prepareExecutorInfo()

	// the framework
	fwinfo := &mesos.FrameworkInfo{
		User: proto.String(""), // Mesos-go will fill in user.
		Name: proto.String("Test Framework (Go)"),
	}

	cred := (*mesos.Credential)(nil)
	if *mesosAuthPrincipal != "" {
		fwinfo.Principal = proto.String(*mesosAuthPrincipal)
		cred = &mesos.Credential{
			Principal: proto.String(*mesosAuthPrincipal),
		}
		if *mesosAuthSecretFile != "" {
			_, err := os.Stat(*mesosAuthSecretFile)
			if err != nil {
				log.Fatal("missing secret file: ", err.Error())
			}
			secret, err := ioutil.ReadFile(*mesosAuthSecretFile)
			if err != nil {
				log.Fatal("failed to read secret file: ", err.Error())
			}
			cred.Secret = proto.String(string(secret))
		}
	}
	bindingAddress := parseIP(*address)
	config := sched.DriverConfig{
		Scheduler:      newExampleScheduler(exec),
		Framework:      fwinfo,
		Master:         *master,
		Credential:     cred,
		BindingAddress: bindingAddress,
		WithAuthContext: func(ctx context.Context) context.Context {
			ctx = auth.WithLoginProvider(ctx, *authProvider)
			ctx = sasl.WithBindingAddress(ctx, bindingAddress)
			return ctx
		},
	}
	driver, err := sched.NewMesosSchedulerDriver(config)

	if err != nil {
		log.Errorln("Unable to create a SchedulerDriver ", err.Error())
	}

	if stat, err := driver.Run(); err != nil {
		log.Infof("Framework stopped with status %s and error: %s\n", stat.String(), err.Error())
		time.Sleep(2 * time.Second)
		os.Exit(1)
	}
	log.Infof("framework terminating")
}
<|endoftext|>