<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"golang.org/x/net/proxy"

	"github.com/olekukonko/tablewriter"

	"gopkg.in/alecthomas/kingpin.v2"
	"gopkg.in/yaml.v2"
)

type config struct {
	BaseURL  string             `yaml:"baseURL"`
	Formats  map[string]format  `yaml:"formats,flow"`
	Elements map[string]element `yaml:"elements,flow"`
}

type element struct {
	ID     string   `yaml:"id"`
	File   string   `yaml:"file"`
	Meta   bool     `yaml:"meta"`
	Name   string   `yaml:"name"`
	Files  []string `yaml:"files"`
	Parent string   `yaml:"parent"`
}

type format struct {
	ID  string `yaml:"ext"`
	Loc string `yaml:"loc"`
}

var (
	app         = kingpin.New("download-geofabrik", "A command-line tool for downloading OSM files.")
	fConfig     = app.Flag("config", "Set Config file.").Default("./geofabrik.yml").Short('c').String()
	nodownload  = app.Flag("nodownload", "Do not download file (test only)").Short('n').Bool()
	verbose     = app.Flag("verbose", "Be verbose").Short('v').Bool()
	fProxyHTTP  = app.Flag("proxy-http", "Use http proxy, format: proxy_address:port").Default("").String()
	fProxySock5 = app.Flag("proxy-sock5", "Use Sock5 proxy, format: proxy_address:port").Default("").String()
	fProxyUser  = app.Flag("proxy-user", "Proxy user").Default("").String()
	fProxyPass  = app.Flag("proxy-pass", "Proxy password").Default("").String()

	update = app.Command("update", "Update geofabrik.yml from github")
	fURL   = update.Flag("url", "Url for config source").Default("https://raw.githubusercontent.com/julien-noblet/download-geofabrik/master/geofabrik.yml").String()

	list = app.Command("list", "Show elements available")
	lmd  = list.Flag("markdown", "generate list in Markdown format").Bool()

	download = app.Command("download", "Download element") //TODO : add d as command
	delement = download.Arg("element", "OSM element").Required().String()
	dosmBz2  = download.Flag("osm.bz2", "Download osm.bz2 if available").Short('B').Bool()
	dshpZip  = download.Flag("shp.zip", "Download shp.zip if available").Short('S').Bool()
	dosmPbf  = download.Flag("osm.pbf", "Download osm.pbf (default)").Short('P').Bool()
	doshPbf  = download.Flag("osh.pbf", "Download osh.pbf (default)").Short('H').Bool()
	dstate   = download.Flag("state", "Download state.txt file").Short('s').Bool()
	dpoly    = download.Flag("poly", "Download poly file").Short('p').Bool()
)

func (e *element) hasParent() bool {
	return len(e.Parent) != 0
}

func miniFormats(s []string) string {
	res := make([]string, 6)
	for _, item := range s {
		switch item {
		case "state":
			res[0] = "s"
		case "osm.pbf":
			res[1] = "P"
		case "osm.bz2":
			res[2] = "B"
		case "osh.pbf":
			res[3] = "H"
		case "poly":
			res[4] = "p"
		case "shp.zip":
			res[5] = "S"
		}
	}

	return strings.Join(res, "")
}

func downloadFromURL(myURL string, fileName string) {
	if *verbose == true {
		log.Println(" Downloading", myURL, "to", fileName)
	}

	if *nodownload == false {
		// TODO: check file existence first with io.IsExist
		output, err := os.Create(fileName)
		if err != nil {
			log.Fatalln(" Error while creating ", fileName, "-", err)
			return
		}
		defer output.Close()
		transport := &http.Transport{}
		if *fProxyHTTP != "" {
			u, err := url.Parse(myURL)
			//log.Println(u.Scheme +"://"+ *fProxyHTTP)
			proxyURL, err := url.Parse(u.Scheme + "://" + *fProxyHTTP)
			if err != nil {
				log.Fatalln(" Wrong proxy url, please use format proxy_address:port")
				return
			}
			transport = &http.Transport{Proxy: http.ProxyURL(proxyURL)}
		}
		client := &http.Client{Transport: transport}
		if *fProxySock5 != "" {
			auth := proxy.Auth{*fProxyUser, *fProxyPass}
			dialer, err := proxy.SOCKS5("tcp", *fProxySock5, &auth, proxy.Direct)
			if err != nil {
				log.Fatalln(" can't connect to the proxy:", err)
				return
			}
			transport.Dial = dialer.Dial
		}
		response, err := client.Get(myURL)
		if err != nil {
			log.Fatalln(" Error while downloading ", myURL, "-", err)
			return
		}
		defer response.Body.Close()

		n, err := io.Copy(output, response.Body)
		if err != nil {
			log.Fatalln(" Error while downloading ", myURL, "-", err)
			return
		}

		if *verbose == true {
			log.Println(" ", n, "bytes downloaded.")
		}
	}
}
func elem2preURL(c config, e element) string {
	var res string
	if e.hasParent() {
		res = elem2preURL(c, findElem(c, e.Parent)) + "/"
		if e.File != "" { //TODO use file in config???
			res = res + e.File
		} else {
			res = res + e.ID
		}
	} else {
		res = c.BaseURL + "/" + e.ID
	}
	return res
}

func elem2URL(c config, e element, ext string) string {
	res := elem2preURL(c, e)
	res += c.Formats[ext].Loc
	if !stringInSlice(ext, e.Files) {
		log.Fatalln(" Error!!! " + res + " not exist")
	}

	return res
}

func findElem(c config, e string) element {
	res := c.Elements[e]
	if res.ID == "" {
		log.Fatalln(" " + e + " is not in config! Please use \"list\" command!")
	}
	return res
}
func getFormats() []string {
	var formatFile []string
	if *dosmPbf {
		formatFile = append(formatFile, "osm.pbf")
	}
	if *doshPbf {
		formatFile = append(formatFile, "osh.pbf")
	}
	if *dosmBz2 {
		formatFile = append(formatFile, "osm.bz2")
	}
	if *dshpZip {
		formatFile = append(formatFile, "shp.zip")
	}
	if *dstate {
		formatFile = append(formatFile, "state")
	}
	if *dpoly {
		formatFile = append(formatFile, "poly")
	}
	if len(formatFile) == 0 {
		formatFile = append(formatFile, "osm.pbf")
	}
	return formatFile
}

func listAllRegions(c config, format string) {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.SetHeader([]string{"ShortName", "Is in", "Long Name", "formats"})
	if format == "Markdown" {
		table.SetBorders(tablewriter.Border{Left: true, Top: false, Right: true, Bottom: false})
		table.SetCenterSeparator("|")
	}
	keys := make(sort.StringSlice, len(c.Elements))
	i := 0
	for k := range c.Elements {
		keys[i] = k
		i++
	}
	keys.Sort()
	for _, item := range keys {
		table.Append([]string{item, c.Elements[c.Elements[item].Parent].Name, c.Elements[item].Name, miniFormats(c.Elements[item].Files)})
	}
	table.Render()
	fmt.Printf("Total elements: %#v\n", len(c.Elements))
}

func loadConfig(configFile string) config {
	filename, _ := filepath.Abs(configFile)
	file, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalln(" File error: %v ", err)
		os.Exit(1)
	}
	var myConfig config
	err = yaml.Unmarshal(file, &myConfig)
	if err != nil {
		panic(err)
	}
	return myConfig

}
func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

// UpdateConfig : simple script to download lastest config from repo
func UpdateConfig(myURL string, myconfig string) {
	downloadFromURL(myURL, myconfig)
	fmt.Println("Congratulation, you have the latest geofabrik.yml")
}

func main() {
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {

	case list.FullCommand():
		var format = ""
		if *lmd {
			format = "Markdown"
		}
		listAllRegions(loadConfig(*fConfig), format)
	case update.FullCommand():
		UpdateConfig(*fURL, *fConfig)
	case download.FullCommand():
		formatFile := getFormats()
		for _, format := range formatFile {
			downloadFromURL(elem2URL(loadConfig(*fConfig), findElem(loadConfig(*fConfig), *delement), format), *delement+"."+format)
		}
	}
}
<commit_msg>feat: Add user/pass support for http proxy (#8)<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"golang.org/x/net/proxy"

	"github.com/olekukonko/tablewriter"

	"gopkg.in/alecthomas/kingpin.v2"
	"gopkg.in/yaml.v2"
)

type config struct {
	BaseURL  string             `yaml:"baseURL"`
	Formats  map[string]format  `yaml:"formats,flow"`
	Elements map[string]element `yaml:"elements,flow"`
}

type element struct {
	ID     string   `yaml:"id"`
	File   string   `yaml:"file"`
	Meta   bool     `yaml:"meta"`
	Name   string   `yaml:"name"`
	Files  []string `yaml:"files"`
	Parent string   `yaml:"parent"`
}

type format struct {
	ID  string `yaml:"ext"`
	Loc string `yaml:"loc"`
}

var (
	app         = kingpin.New("download-geofabrik", "A command-line tool for downloading OSM files.")
	fConfig     = app.Flag("config", "Set Config file.").Default("./geofabrik.yml").Short('c').String()
	nodownload  = app.Flag("nodownload", "Do not download file (test only)").Short('n').Bool()
	verbose     = app.Flag("verbose", "Be verbose").Short('v').Bool()
	fProxyHTTP  = app.Flag("proxy-http", "Use http proxy, format: proxy_address:port").Default("").String()
	fProxySock5 = app.Flag("proxy-sock5", "Use Sock5 proxy, format: proxy_address:port").Default("").String()
	fProxyUser  = app.Flag("proxy-user", "Proxy user").Default("").String()
	fProxyPass  = app.Flag("proxy-pass", "Proxy password").Default("").String()

	update = app.Command("update", "Update geofabrik.yml from github")
	fURL   = update.Flag("url", "Url for config source").Default("https://raw.githubusercontent.com/julien-noblet/download-geofabrik/master/geofabrik.yml").String()

	list = app.Command("list", "Show elements available")
	lmd  = list.Flag("markdown", "generate list in Markdown format").Bool()

	download = app.Command("download", "Download element") //TODO : add d as command
	delement = download.Arg("element", "OSM element").Required().String()
	dosmBz2  = download.Flag("osm.bz2", "Download osm.bz2 if available").Short('B').Bool()
	dshpZip  = download.Flag("shp.zip", "Download shp.zip if available").Short('S').Bool()
	dosmPbf  = download.Flag("osm.pbf", "Download osm.pbf (default)").Short('P').Bool()
	doshPbf  = download.Flag("osh.pbf", "Download osh.pbf (default)").Short('H').Bool()
	dstate   = download.Flag("state", "Download state.txt file").Short('s').Bool()
	dpoly    = download.Flag("poly", "Download poly file").Short('p').Bool()
)

func (e *element) hasParent() bool {
	return len(e.Parent) != 0
}

func miniFormats(s []string) string {
	res := make([]string, 6)
	for _, item := range s {
		switch item {
		case "state":
			res[0] = "s"
		case "osm.pbf":
			res[1] = "P"
		case "osm.bz2":
			res[2] = "B"
		case "osh.pbf":
			res[3] = "H"
		case "poly":
			res[4] = "p"
		case "shp.zip":
			res[5] = "S"
		}
	}

	return strings.Join(res, "")
}

func downloadFromURL(myURL string, fileName string) {
	if *verbose == true {
		log.Println(" Downloading", myURL, "to", fileName)
	}

	if *nodownload == false {
		// TODO: check file existence first with io.IsExist
		output, err := os.Create(fileName)
		if err != nil {
			log.Fatalln(" Error while creating ", fileName, "-", err)
			return
		}
		defer output.Close()
		transport := &http.Transport{}
		if *fProxyHTTP != "" {
			u, err := url.Parse(myURL)
			//log.Println(u.Scheme +"://"+ *fProxyHTTP)
			proxyURL, err := url.Parse(u.Scheme + "://" + *fProxyHTTP)
			if *fProxyUser != "" && *fProxyPass != "" {
				proxyURL, err = url.Parse(u.Scheme + "://" + *fProxyUser + ":" + *fProxyPass + *fProxyHTTP)
			}
			if err != nil {
				log.Fatalln(" Wrong proxy url, please use format proxy_address:port")
				return
			}
			transport = &http.Transport{Proxy: http.ProxyURL(proxyURL)}
		}
		client := &http.Client{Transport: transport}
		if *fProxySock5 != "" {
			auth := proxy.Auth{*fProxyUser, *fProxyPass}
			dialer, err := proxy.SOCKS5("tcp", *fProxySock5, &auth, proxy.Direct)
			if err != nil {
				log.Fatalln(" can't connect to the proxy:", err)
				return
			}
			transport.Dial = dialer.Dial
		}
		response, err := client.Get(myURL)
		if err != nil {
			log.Fatalln(" Error while downloading ", myURL, "-", err)
			return
		}
		defer response.Body.Close()

		n, err := io.Copy(output, response.Body)
		if err != nil {
			log.Fatalln(" Error while downloading ", myURL, "-", err)
			return
		}

		if *verbose == true {
			log.Println(" ", n, "bytes downloaded.")
		}
	}
}
func elem2preURL(c config, e element) string {
	var res string
	if e.hasParent() {
		res = elem2preURL(c, findElem(c, e.Parent)) + "/"
		if e.File != "" { //TODO use file in config???
			res = res + e.File
		} else {
			res = res + e.ID
		}
	} else {
		res = c.BaseURL + "/" + e.ID
	}
	return res
}

func elem2URL(c config, e element, ext string) string {
	res := elem2preURL(c, e)
	res += c.Formats[ext].Loc
	if !stringInSlice(ext, e.Files) {
		log.Fatalln(" Error!!! " + res + " not exist")
	}

	return res
}

func findElem(c config, e string) element {
	res := c.Elements[e]
	if res.ID == "" {
		log.Fatalln(" " + e + " is not in config! Please use \"list\" command!")
	}
	return res
}
func getFormats() []string {
	var formatFile []string
	if *dosmPbf {
		formatFile = append(formatFile, "osm.pbf")
	}
	if *doshPbf {
		formatFile = append(formatFile, "osh.pbf")
	}
	if *dosmBz2 {
		formatFile = append(formatFile, "osm.bz2")
	}
	if *dshpZip {
		formatFile = append(formatFile, "shp.zip")
	}
	if *dstate {
		formatFile = append(formatFile, "state")
	}
	if *dpoly {
		formatFile = append(formatFile, "poly")
	}
	if len(formatFile) == 0 {
		formatFile = append(formatFile, "osm.pbf")
	}
	return formatFile
}

func listAllRegions(c config, format string) {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.SetHeader([]string{"ShortName", "Is in", "Long Name", "formats"})
	if format == "Markdown" {
		table.SetBorders(tablewriter.Border{Left: true, Top: false, Right: true, Bottom: false})
		table.SetCenterSeparator("|")
	}
	keys := make(sort.StringSlice, len(c.Elements))
	i := 0
	for k := range c.Elements {
		keys[i] = k
		i++
	}
	keys.Sort()
	for _, item := range keys {
		table.Append([]string{item, c.Elements[c.Elements[item].Parent].Name, c.Elements[item].Name, miniFormats(c.Elements[item].Files)})
	}
	table.Render()
	fmt.Printf("Total elements: %#v\n", len(c.Elements))
}

func loadConfig(configFile string) config {
	filename, _ := filepath.Abs(configFile)
	file, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalln(" File error: %v ", err)
		os.Exit(1)
	}
	var myConfig config
	err = yaml.Unmarshal(file, &myConfig)
	if err != nil {
		panic(err)
	}
	return myConfig

}
func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

// UpdateConfig : simple script to download lastest config from repo
func UpdateConfig(myURL string, myconfig string) {
	downloadFromURL(myURL, myconfig)
	fmt.Println("Congratulation, you have the latest geofabrik.yml")
}

func main() {
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {

	case list.FullCommand():
		var format = ""
		if *lmd {
			format = "Markdown"
		}
		listAllRegions(loadConfig(*fConfig), format)
	case update.FullCommand():
		UpdateConfig(*fURL, *fConfig)
	case download.FullCommand():
		formatFile := getFormats()
		for _, format := range formatFile {
			downloadFromURL(elem2URL(loadConfig(*fConfig), findElem(loadConfig(*fConfig), *delement), format), *delement+"."+format)
		}
	}
}
<|endoftext|>


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

// Package mobileproxy instantiates a Martian Proxy.
// This package is a reference implementation of Martian Proxy intended to
// be cross compiled with gomobile for use on Android and iOS.
package mobileproxy

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"github.com/google/martian"
	"github.com/google/martian/har"
	"github.com/google/martian/httpspec"
	mlog "github.com/google/martian/log"
	"github.com/google/martian/martianhttp"
	"github.com/google/martian/mitm"
	"github.com/google/martian/verify"
	"log"
	"net"
	"net/http"
	"os/signal"
	"syscall"
	"time"

	// side-effect importing modifiers to register them with the proxy
	_ "github.com/google/martian/body"
	_ "github.com/google/martian/cookie"
	_ "github.com/google/martian/fifo"
	_ "github.com/google/martian/header"
	_ "github.com/google/martian/martianurl"
	_ "github.com/google/martian/method"
	_ "github.com/google/martian/pingback"
	_ "github.com/google/martian/priority"
	_ "github.com/google/martian/querystring"
	_ "github.com/google/martian/skip"
	_ "github.com/google/martian/status"
)

// Martian is a wrapper for the initialized Martian proxy
type Martian struct {
	proxy    *martian.Proxy
	listener net.Listener
	mux      *http.ServeMux
}

// Start runs a martian.Proxy on addr
func Start(proxyAddr string) (*Martian, error) {
	return StartWithCertificate(proxyAddr, "", "")
}

// StartWithCertificate runs a proxy on addr and configures a cert for MITM
func StartWithCertificate(proxyAddr string, cert string, key string) (*Martian, error) {
	flag.Set("logtostderr", "true")

	signal.Ignore(syscall.SIGPIPE)

	l, err := net.Listen("tcp", proxyAddr)
	if err != nil {
		return nil, err
	}

	mlog.Debugf("mobileproxy: started listener: %v", l.Addr())

	p := martian.NewProxy()

	mux := http.NewServeMux()
	p.SetMux(mux)

	if cert != "" && key != "" {
		tlsc, err := tls.X509KeyPair([]byte(cert), []byte(key))
		if err != nil {
			log.Fatal(err)
		}

		mlog.Debugf("mobileproxy: loaded cert and key")

		x509c, err := x509.ParseCertificate(tlsc.Certificate[0])
		if err != nil {
			log.Fatal(err)
		}

		mlog.Debugf("mobileproxy: parsed cert")

		mc, err := mitm.NewConfig(x509c, tlsc.PrivateKey)
		if err != nil {
			log.Fatal(err)
		}

		mc.SetValidity(12 * time.Hour)
		mc.SetOrganization("Martian Proxy")

		p.SetMITM(mc)

		mux.Handle("martian.proxy/authority.cer", martianhttp.NewAuthorityHandler(x509c))
		mlog.Debugf("mobileproxy: install cert from http://martian.proxy/authority.cer")
	}

	stack, fg := httpspec.NewStack("martian.mobileproxy")
	p.SetRequestModifier(stack)
	p.SetResponseModifier(stack)

	// add HAR logger
	hl := har.NewLogger()
	stack.AddRequestModifier(hl)
	stack.AddResponseModifier(hl)

	m := martianhttp.NewModifier()
	fg.AddRequestModifier(m)
	fg.AddResponseModifier(m)

	mlog.Debugf("mobileproxy: set martianhttp modifier")

	// Proxy specific handlers.
	// These handlers take precendence over proxy traffic and will not be intercepted.

	// Retrieve HAR logs
	mux.Handle("martian.proxy/logs", har.NewExportHandler(hl))
	mux.Handle("martian.proxy/logs/reset", har.NewResetHandler(hl))

	// Update modifiers.
	mux.Handle("martian.proxy/configure", m)
	mlog.Debugf("mobileproxy: configure with requests to http://martian.proxy/configure")

	// Verify assertions.
	vh := verify.NewHandler()
	vh.SetRequestVerifier(m)
	vh.SetResponseVerifier(m)
	mux.Handle("martian.proxy/verify", vh)
	mlog.Debugf("mobileproxy: check verifications with requests to http://martian.proxy/verify")

	// Reset verifications.
	rh := verify.NewResetHandler()
	rh.SetRequestVerifier(m)
	rh.SetResponseVerifier(m)
	mux.Handle("martian.proxy/verify/reset", rh)
	mlog.Debugf("mobileproxy: reset verifications with requests to http://martian.proxy/verify/reset")

	go p.Serve(l)
	mlog.Infof("mobileproxy: started proxy on listener")

	return &Martian{
		proxy:    p,
		listener: l,
		mux:      mux,
	}, nil
}

// Shutdown tells the Proxy to close. The proxy will stay alive until all connections through it
// have closed or timed out.
func (p *Martian) Shutdown() {
	mlog.Infof("mobileproxy: telling proxy to close")
	p.proxy.Close()
	mlog.Infof("mobileproxy: proxy closed")
}

// Sets the Martian log level (Silent = 0, Error, Info, Debug), controlling which Martian
// log calls are displayed in the console
func SetLogLevel(l int) {
	mlog.SetLevel(l)
}
<commit_msg>Updating mobileproxy.go (#113)<commit_after>// Copyright 2016 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package mobileproxy instantiates a Martian Proxy.
// This package is a reference implementation of Martian Proxy intended to
// be cross compiled with gomobile for use on Android and iOS.
package mobileproxy

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"log"
	"net"
	"net/http"
	"os/signal"
	"syscall"
	"time"

	"github.com/google/martian"
	// side-effect importing to register with JSON API
	_ "github.com/google/martian/body"
	_ "github.com/google/martian/cookie"
	_ "github.com/google/martian/fifo"
	"github.com/google/martian/har"
	// side-effect importing to register with JSON API
	_ "github.com/google/martian/header"
	"github.com/google/martian/httpspec"
	mlog "github.com/google/martian/log"
	"github.com/google/martian/martianhttp"
	// side-effect importing to register with JSON API
	_ "github.com/google/martian/martianurl"
	_ "github.com/google/martian/method"
	"github.com/google/martian/mitm"
	// side-effect importing to register with JSON API
	_ "github.com/google/martian/pingback"
	_ "github.com/google/martian/priority"
	_ "github.com/google/martian/querystring"
	_ "github.com/google/martian/skip"
	_ "github.com/google/martian/status"
	"github.com/google/martian/verify"
)

// Martian is a wrapper for the initialized Martian proxy
type Martian struct {
	proxy    *martian.Proxy
	listener net.Listener
	mux      *http.ServeMux
}

// Start runs a martian.Proxy on addr
func Start(proxyAddr string) (*Martian, error) {
	return StartWithCertificate(proxyAddr, "", "")
}

// StartWithCertificate runs a proxy on addr and configures a cert for MITM
func StartWithCertificate(proxyAddr string, cert string, key string) (*Martian, error) {
	flag.Set("logtostderr", "true")

	p := martian.NewProxy()

	mux := http.NewServeMux()
	p.SetMux(mux)

	if cert != "" && key != "" {
		tlsc, err := tls.X509KeyPair([]byte(cert), []byte(key))
		if err != nil {
			log.Fatal(err)
		}

		mlog.Debugf("mobileproxy: loaded cert and key")

		x509c, err := x509.ParseCertificate(tlsc.Certificate[0])
		if err != nil {
			log.Fatal(err)
		}

		mlog.Debugf("mobileproxy: parsed cert")

		mc, err := mitm.NewConfig(x509c, tlsc.PrivateKey)
		if err != nil {
			log.Fatal(err)
		}

		mc.SetValidity(12 * time.Hour)
		mc.SetOrganization("Martian Proxy")

		p.SetMITM(mc)

		mux.Handle("martian.proxy/authority.cer", martianhttp.NewAuthorityHandler(x509c))
		mlog.Debugf("mobileproxy: install cert from http://martian.proxy/authority.cer")
	}

	stack, fg := httpspec.NewStack("martian.mobileproxy")
	p.SetRequestModifier(stack)
	p.SetResponseModifier(stack)

	// add HAR logger
	hl := har.NewLogger()
	stack.AddRequestModifier(hl)
	stack.AddResponseModifier(hl)

	m := martianhttp.NewModifier()
	fg.AddRequestModifier(m)
	fg.AddResponseModifier(m)

	mlog.Debugf("mobileproxy: set martianhttp modifier")

	// Proxy specific handlers.
	// These handlers take precendence over proxy traffic and will not be intercepted.

	// Retrieve HAR logs
	mux.Handle("martian.proxy/logs", har.NewExportHandler(hl))
	mux.Handle("martian.proxy/logs/reset", har.NewResetHandler(hl))

	// Update modifiers.
	mux.Handle("martian.proxy/configure", m)
	mlog.Debugf("mobileproxy: configure with requests to http://martian.proxy/configure")

	// Verify assertions.
	vh := verify.NewHandler()
	vh.SetRequestVerifier(m)
	vh.SetResponseVerifier(m)
	mux.Handle("martian.proxy/verify", vh)
	mlog.Debugf("mobileproxy: check verifications with requests to http://martian.proxy/verify")

	// Reset verifications.
	rh := verify.NewResetHandler()
	rh.SetRequestVerifier(m)
	rh.SetResponseVerifier(m)
	mux.Handle("martian.proxy/verify/reset", rh)
	mlog.Debugf("mobileproxy: reset verifications with requests to http://martian.proxy/verify/reset")

	// Ignore SIGPIPE
	mlog.Debugf("mobileproxy: ignoring SIGPIPE signals")
	signal.Ignore(syscall.SIGPIPE)

	l, err := net.Listen("tcp", proxyAddr)
	if err != nil {
		return nil, err
	}

	mlog.Debugf("mobileproxy: started listener: %v", l.Addr())
	mlog.Infof("mobileproxy: starting proxy")
	go p.Serve(l)
	mlog.Infof("mobileproxy: started proxy on listener")

	return &Martian{
		proxy:    p,
		listener: l,
		mux:      mux,
	}, nil
}

// Shutdown tells the Proxy to close. The proxy will stay alive until all connections through it
// have closed or timed out.
func (p *Martian) Shutdown() {
	mlog.Infof("mobileproxy: shutting down proxy")
	p.proxy.Close()
	mlog.Infof("mobileproxy: proxy shut down")
}

// SetLogLevel sets the Martian log level (Silent = 0, Error, Info, Debug), controlling which Martian
// log calls are displayed in the console
func SetLogLevel(l int) {
	mlog.SetLevel(l)
}
<|endoftext|>


<commit_before>package dos

import (
	"os"
	"strings"
)

// Get %HOME% || %USERPROFILE% || ""
func GetHome() string {
	home := os.Getenv("HOME")
	if home == "" {
		home = os.Getenv("USERPROFILE")
	}
	return home
}

func ReplaceHomeToTilde(wd string) string {
	home := GetHome()
	homeLen := len(home)
	if len(wd) >= homeLen && strings.EqualFold(home, wd[0:homeLen]) {
		wd = "~" + wd[homeLen:]
	}
	return wd
}

func ReplaceHomeToTildeSlash(wd string) string {
	return strings.Replace(ReplaceHomeToTilde(wd), "\\", "/", -1)
}
<commit_msg>dos/home.go: Add comment<commit_after>package dos

import (
	"os"
	"strings"
)

// Get %HOME% || %USERPROFILE% || ""
func GetHome() string {
	home := os.Getenv("HOME")
	if home == "" {
		home = os.Getenv("USERPROFILE")
	}
	return home
}

// C:\users\name\foo\bar -> ~\foo\bar
func ReplaceHomeToTilde(wd string) string {
	home := GetHome()
	homeLen := len(home)
	if len(wd) >= homeLen && strings.EqualFold(home, wd[0:homeLen]) {
		wd = "~" + wd[homeLen:]
	}
	return wd
}

// C:\users\name\foo\bar -> ~/foo/bar
func ReplaceHomeToTildeSlash(wd string) string {
	return strings.Replace(ReplaceHomeToTilde(wd), "\\", "/", -1)
}
<|endoftext|>


<commit_before>//    Copyright 2016 Red Hat, Inc.
//
//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at
//
//        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 download

import (
	"crypto"
	"crypto/md5" // #nosec
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"time"

	pb "gopkg.in/cheggaaa/pb.v1"

	"github.com/pkg/errors"
)

// Options holds the possible configuration options for the Downloader.
type Options struct {
	// HTTPClient is an optional client to perform downloads with. If nil, `http.DefaultClient`
	// will be used.
	HTTPClient *http.Client
	// Checksum is either a checksum string, or a URL or path to a file containing the checksum. The file
	// can either contain the checksum only or contain multiple lines of the format:
	// CHECKSUM FILENAME
	Checksum string
	// Checksum hash is the hash for the checksum. Currently only supports SHA1, SHA256, SHA512 and MD5.
	// If unspecified, defaults to SHA256.
	ChecksumHash crypto.Hash
	// ProgressBars is the configuration of progress bars output. Set to `nil` (default) to disable.
	ProgressBars *ProgressBarOptions
}

// FileOptions holds the possible configuration options to download to a file.
type FileOptions struct {
	// Options is the common set of downloader options.
	Options
	// Mkdirs is the option to create parent directories of target directory if they don't
	// exist. Use `download.MkdirAll` or `download.MkdirNone` (or any `*bool`). Defaults to
	// `download.MkdirAll`.
	Mkdirs Mkdirs
}

// ProgressBarOptions holds the configuration for progress bars if required.
type ProgressBarOptions struct {
	// Writer holds where to output the progress bars to. Defaults to `os.Stdout`.
	Writer io.Writer
	// Width is the maximum width of the progress bar. If output to a narrower terminal then this
	// will be ignored.
	MaxWidth int
}

func newBool(b bool) *bool {
	return &b
}

// Mkdirs is a custom type so we can differentiate between not specified (nil)
// and set.
type Mkdirs *bool

var (
	// MkdirAll is used to create all intermediate directories if required.
	MkdirAll = Mkdirs(newBool(true))
	// MkdirNone is used to create no intermediate directories.
	MkdirNone = Mkdirs(newBool(false))
)

// ToFile downloads the specified `src` URL to `dest` file using
// the specified `FileOptions`.
func ToFile(src, dest string, options FileOptions) error {
	u, err := url.Parse(src)
	if err != nil {
		return errors.Wrap(err, "invalid src URL")
	}

	targetDir := filepath.Dir(dest)
	if _, err = os.Stat(targetDir); err != nil {
		if !os.IsNotExist(err) || (options.Mkdirs != nil && !*options.Mkdirs) {
			return errors.Wrap(err, "failed to check destination directory")
		}
		err = os.MkdirAll(targetDir, 0700)
		if err != nil {
			return errors.Wrap(err, "failed to create destination directory")
		}
	}

	targetName := filepath.Base(dest)
	f, err := ioutil.TempFile(targetDir, ".tmp-"+targetName)
	if err != nil {
		return errors.Wrap(err, "failed to create temp file")
	}

	err = downloadFile(u, f, options.Options)
	if err != nil {
		return errors.Wrap(err, "failed to download")
	}

	err = os.Rename(f.Name(), dest)
	if err != nil {
		_ = os.Remove(f.Name()) // #nosec
		return errors.Wrap(err, "failed to rename temp file to destination")
	}

	return nil
}

func downloadFile(u *url.URL, f *os.File, options Options) error {
	err := FromURL(u, f, options)
	if err != nil {
		_ = os.Remove(f.Name()) // #nosec
		return errors.Wrap(err, "failed to download to temp file")
	}
	err = f.Close()
	if err != nil {
		_ = os.Remove(f.Name()) // #nosec
		return errors.Wrap(err, "failed to close temp file")
	}

	return nil
}

// ToWriter downloads the specified `src` URL to `w` writer using
// the specified `Options`.
func ToWriter(src string, w io.Writer, options Options) error {
	u, err := url.Parse(src)
	if err != nil {
		return errors.Wrap(err, "invalid src URL")
	}
	return FromURL(u, w, options)
}

// FromURL downloads the specified `src` URL to `w` writer using
// the specified `Options`.
func FromURL(src *url.URL, w io.Writer, options Options) error {
	httpClient := getHTTPClient(options)
	var err error
	resp, err := httpClient.Get(src.String())
	if err != nil {
		return errors.Wrap(err, "download failed")
	}
	defer func() { _ = resp.Body.Close() }() // #nosec

	if resp.StatusCode != http.StatusOK {
		return errors.Errorf("received invalid status code: %d (expected %d)", resp.StatusCode, http.StatusOK)
	}

	var (
		validator checksumValidator

		reader io.Reader = resp.Body
	)

	if options.ProgressBars != nil && resp.ContentLength > 0 {
		bar := newProgressBar(resp.ContentLength, options.ProgressBars.MaxWidth, options.ProgressBars.Writer)
		bar.Start()
		reader = bar.NewProxyReader(reader)
		defer func() {
			<-time.After(bar.RefreshRate)
			fmt.Println()
		}()
	}

	if len(options.Checksum) != 0 {
		validator, err = createValidator(options.ChecksumHash, httpClient, options.Checksum, path.Base(src.Path))
		if err != nil {
			return errors.Wrap(err, "failed to create validator")
		}
		reader = io.TeeReader(reader, validator)
	}

	if _, err = io.Copy(w, reader); err != nil {
		return errors.Wrap(err, "failed to copy contents")
	}

	if validator != nil && !validator.validate() {
		return errors.New("checksum validation failed")
	}

	return nil
}

func createValidator(hashType crypto.Hash, httpClient *http.Client, checksum, filename string) (checksumValidator, error) {
	var hasher hash.Hash
	switch hashType {
	case crypto.SHA256, 0:
		hasher = sha256.New()
	case crypto.SHA1:
		hasher = sha1.New()
	case crypto.SHA512:
		hasher = sha512.New()
	case crypto.MD5:
		hasher = md5.New() // #nosec
	default:
		return nil, errors.New("invalid hash function")
	}

	validator, err := newValidator(hasher, httpClient, checksum, filename)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create validator")
	}

	return validator, nil
}

func getHTTPClient(options Options) *http.Client {
	httpClient := options.HTTPClient
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	return httpClient
}

func getBarWriter(w io.Writer) io.Writer {
	if w == nil {
		w = os.Stdout
	}
	return w
}

func newProgressBar(length int64, maxWidth int, w io.Writer) *pb.ProgressBar {
	bar := pb.New64(length).SetUnits(pb.U_BYTES)
	if maxWidth > 0 {
		bar.SetMaxWidth(maxWidth)
	}
	barWriter := getBarWriter(w)
	bar.Output = barWriter
	return bar
}
<commit_msg>Refactor temp file closing<commit_after>//    Copyright 2016 Red Hat, Inc.
//
//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at
//
//        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 download

import (
	"crypto"
	"crypto/md5" // #nosec
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"time"

	pb "gopkg.in/cheggaaa/pb.v1"

	"github.com/pkg/errors"
)

// Options holds the possible configuration options for the Downloader.
type Options struct {
	// HTTPClient is an optional client to perform downloads with. If nil, `http.DefaultClient`
	// will be used.
	HTTPClient *http.Client
	// Checksum is either a checksum string, or a URL or path to a file containing the checksum. The file
	// can either contain the checksum only or contain multiple lines of the format:
	// CHECKSUM FILENAME
	Checksum string
	// Checksum hash is the hash for the checksum. Currently only supports SHA1, SHA256, SHA512 and MD5.
	// If unspecified, defaults to SHA256.
	ChecksumHash crypto.Hash
	// ProgressBars is the configuration of progress bars output. Set to `nil` (default) to disable.
	ProgressBars *ProgressBarOptions
}

// FileOptions holds the possible configuration options to download to a file.
type FileOptions struct {
	// Options is the common set of downloader options.
	Options
	// Mkdirs is the option to create parent directories of target directory if they don't
	// exist. Use `download.MkdirAll` or `download.MkdirNone` (or any `*bool`). Defaults to
	// `download.MkdirAll`.
	Mkdirs Mkdirs
}

// ProgressBarOptions holds the configuration for progress bars if required.
type ProgressBarOptions struct {
	// Writer holds where to output the progress bars to. Defaults to `os.Stdout`.
	Writer io.Writer
	// Width is the maximum width of the progress bar. If output to a narrower terminal then this
	// will be ignored.
	MaxWidth int
}

func newBool(b bool) *bool {
	return &b
}

// Mkdirs is a custom type so we can differentiate between not specified (nil)
// and set.
type Mkdirs *bool

var (
	// MkdirAll is used to create all intermediate directories if required.
	MkdirAll = Mkdirs(newBool(true))
	// MkdirNone is used to create no intermediate directories.
	MkdirNone = Mkdirs(newBool(false))
)

// ToFile downloads the specified `src` URL to `dest` file using
// the specified `FileOptions`.
func ToFile(src, dest string, options FileOptions) error {
	u, err := url.Parse(src)
	if err != nil {
		return errors.Wrap(err, "invalid src URL")
	}

	targetDir := filepath.Dir(dest)
	if _, err = os.Stat(targetDir); err != nil {
		if !os.IsNotExist(err) || (options.Mkdirs != nil && !*options.Mkdirs) {
			return errors.Wrap(err, "failed to check destination directory")
		}
		err = os.MkdirAll(targetDir, 0700)
		if err != nil {
			return errors.Wrap(err, "failed to create destination directory")
		}
	}

	targetName := filepath.Base(dest)
	f, err := ioutil.TempFile(targetDir, ".tmp-"+targetName)
	if err != nil {
		return errors.Wrap(err, "failed to create temp file")
	}

	err = downloadFile(u, f, options.Options)
	if err != nil {
		_ = f.Close()           // #nosec
		_ = os.Remove(f.Name()) // #nosec
		return errors.Wrap(err, "failed to download")
	}
	err = f.Close()
	if err != nil {
		_ = os.Remove(f.Name()) // #nosec
		return errors.Wrap(err, "failed to close temp file")
	}

	err = os.Rename(f.Name(), dest)
	if err != nil {
		_ = os.Remove(f.Name()) // #nosec
		return errors.Wrap(err, "failed to rename temp file to destination")
	}

	return nil
}

func downloadFile(u *url.URL, f *os.File, options Options) error {
	err := FromURL(u, f, options)
	if err != nil {
		return errors.Wrap(err, "failed to download to temp file")
	}

	return nil
}

// ToWriter downloads the specified `src` URL to `w` writer using
// the specified `Options`.
func ToWriter(src string, w io.Writer, options Options) error {
	u, err := url.Parse(src)
	if err != nil {
		return errors.Wrap(err, "invalid src URL")
	}
	return FromURL(u, w, options)
}

// FromURL downloads the specified `src` URL to `w` writer using
// the specified `Options`.
func FromURL(src *url.URL, w io.Writer, options Options) error {
	httpClient := getHTTPClient(options)
	var err error
	resp, err := httpClient.Get(src.String())
	if err != nil {
		return errors.Wrap(err, "download failed")
	}
	defer func() { _ = resp.Body.Close() }() // #nosec

	if resp.StatusCode != http.StatusOK {
		return errors.Errorf("received invalid status code: %d (expected %d)", resp.StatusCode, http.StatusOK)
	}

	var (
		validator checksumValidator

		reader io.Reader = resp.Body
	)

	if options.ProgressBars != nil && resp.ContentLength > 0 {
		bar := newProgressBar(resp.ContentLength, options.ProgressBars.MaxWidth, options.ProgressBars.Writer)
		bar.Start()
		reader = bar.NewProxyReader(reader)
		defer func() {
			<-time.After(bar.RefreshRate)
			fmt.Println()
		}()
	}

	if len(options.Checksum) != 0 {
		validator, err = createValidator(options.ChecksumHash, httpClient, options.Checksum, path.Base(src.Path))
		if err != nil {
			return errors.Wrap(err, "failed to create validator")
		}
		reader = io.TeeReader(reader, validator)
	}

	if _, err = io.Copy(w, reader); err != nil {
		return errors.Wrap(err, "failed to copy contents")
	}

	if validator != nil && !validator.validate() {
		return errors.New("checksum validation failed")
	}

	return nil
}

func createValidator(hashType crypto.Hash, httpClient *http.Client, checksum, filename string) (checksumValidator, error) {
	var hasher hash.Hash
	switch hashType {
	case crypto.SHA256, 0:
		hasher = sha256.New()
	case crypto.SHA1:
		hasher = sha1.New()
	case crypto.SHA512:
		hasher = sha512.New()
	case crypto.MD5:
		hasher = md5.New() // #nosec
	default:
		return nil, errors.New("invalid hash function")
	}

	validator, err := newValidator(hasher, httpClient, checksum, filename)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create validator")
	}

	return validator, nil
}

func getHTTPClient(options Options) *http.Client {
	httpClient := options.HTTPClient
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	return httpClient
}

func getBarWriter(w io.Writer) io.Writer {
	if w == nil {
		w = os.Stdout
	}
	return w
}

func newProgressBar(length int64, maxWidth int, w io.Writer) *pb.ProgressBar {
	bar := pb.New64(length).SetUnits(pb.U_BYTES)
	if maxWidth > 0 {
		bar.SetMaxWidth(maxWidth)
	}
	barWriter := getBarWriter(w)
	bar.Output = barWriter
	return bar
}
<|endoftext|>


<commit_before>package rain

import (
	"errors"
	"math/rand"
	"net"
	"runtime"
	"sort"
	"time"

	"github.com/cenkalti/rain/internal/connection"
	"github.com/cenkalti/rain/internal/logger"
)

// peerManager receives from t.peersC and keeps most recent peer addresses in t.peerC.
func (t *transfer) peerManager() {
	t.log.Debug("Started peerManager")
	for {
		select {
		case <-t.stopC:
			return
		case peers := <-t.peersC:
			for _, p := range peers {
				t.log.Debugln("Peer:", p)
				go func(addr *net.TCPAddr) {
					select {
					case t.peerC <- addr:
					case <-t.stopC:
					}
				}(p)
			}
		}
	}
}

// connecter connects to peers coming from t. peerC.
func (t *transfer) connecter() {
	limit := make(chan struct{}, maxPeerPerTorrent)
	for {
		select {
		case p := <-t.peerC:
			if p.Port == 0 {
				break
			}
			if p.IP.IsLoopback() && p.Port == int(t.rain.Port()) {
				break
			}

			limit <- struct{}{}
			go func(addr *net.TCPAddr) {
				defer func() {
					if err := recover(); err != nil {
						buf := make([]byte, 10000)
						t.log.Critical(err, "\n", string(buf[:runtime.Stack(buf, false)]))
					}
					<-limit
				}()
				t.connect(addr)
			}(p)
		case <-t.stopC:
			return
		}
	}
}

func (t *transfer) connect(addr *net.TCPAddr) {
	log := logger.New("peer -> " + addr.String())

	conn, cipher, extensions, peerID, err := connection.Dial(addr, !t.rain.config.Encryption.DisableOutgoing, t.rain.config.Encryption.ForceOutgoing, [8]byte{}, t.torrent.Info.Hash, t.rain.peerID)
	if err != nil {
		if err == connection.ErrOwnConnection {
			log.Debug(err)
		} else {
			log.Error(err)
		}
		return
	}
	log.Infof("Connected to peer. (cipher=%s, extensions=%x)", cipher, extensions)
	defer conn.Close()

	p := t.newPeer(conn, peerID, log)

	t.m.Lock()
	t.peers[peerID] = p
	t.m.Unlock()
	defer func() {
		t.m.Lock()
		delete(t.peers, peerID)
		t.m.Unlock()
	}()

	if err = p.SendBitfield(); err != nil {
		log.Error(err)
		return
	}

	p.Run()
}

func (p *Peer) downloader() {
	t := p.transfer
	for {
		t.m.Lock()
		if t.bitfield.All() {
			t.onceFinished.Do(func() { close(t.finished) })
			t.m.Unlock()
			return
		}
		candidates := p.candidates()
		if len(candidates) == 0 {
			t.m.Unlock()
			if err := p.BeNotInterested(); err != nil {
				p.log.Error(err)
				return
			}
			if err := p.waitForHaveMessage(); err != nil {
				p.log.Error(err)
				return
			}
			continue
		}
		piece := selectPiece(candidates)
		// Save selected piece so other downloaders do not try to download the same piece.
		// TODO remove from requests when downloader exited with error.
		t.requests[piece.Index] = &pieceRequest{p, time.Now()}
		t.m.Unlock()

		if err := p.BeInterested(); err != nil {
			p.log.Error(err)
			return
		}

		// TODO queue max 10 requests
		go func() {
			if err := p.requestBlocks(piece); err != nil {
				p.log.Error(err)
			}
		}()

		// TODO handle choke while receiving pieces. Re-reqeust, etc..

		// Read blocks from peer.
		pieceData := make([]byte, piece.Length)
		for i := 0; i < len(piece.Blocks); i++ {
			peerBlock := <-p.pieceC
			data := <-peerBlock.Data
			if data == nil {
				p.log.Error("peer did not send block completely")
				return
			}
			p.log.Debugln("Will receive block of length", len(data))
			copy(pieceData[peerBlock.Begin:], data)
		}

		if _, err := piece.Write(pieceData); err != nil {
			t.log.Error(err)
			p.Close()
			return
		}

		t.m.Lock()
		t.bitfield.Set(piece.Index)
		delete(t.requests, piece.Index)
		t.m.Unlock()
	}
}

func (p *Peer) requestBlocks(piece *Piece) error {
	for _, b := range piece.Blocks {
		// Send requests only when unchoked.
		if err := p.waitForUnchoke(); err != nil {
			return err
		}
		if err := p.Request(piece.Index, b.Begin, b.Length); err != nil {
			return err
		}
	}
	return nil
}

func (p *Peer) waitForHaveMessage() error {
	p.m.Lock()
	defer p.m.Unlock()
	count := p.bitfield.Count()
	for count == p.bitfield.Count() && !p.disconnected {
		p.cond.Wait()
	}
	if p.disconnected {
		return errors.New("disconnected while waiting for new have message")
	}
	return nil
}

func (p *Peer) waitForUnchoke() error {
	p.m.Lock()
	defer p.m.Unlock()
	for p.peerChoking && !p.disconnected {
		p.cond.Wait()
	}
	if p.disconnected {
		return errors.New("disconnected while waiting for unchoke message")
	}
	return nil
}

// candidates returns list of piece indexes which is available on the peer but not available on the client.
func (p *Peer) candidates() (candidates []*Piece) {
	p.m.Lock()
	for i := uint32(0); i < p.transfer.bitfield.Len(); i++ {
		if !p.transfer.bitfield.Test(i) && p.bitfield.Test(i) {
			piece := p.transfer.pieces[i]
			if _, ok := p.transfer.requests[piece.Index]; !ok {
				candidates = append(candidates, piece)
			}
		}
	}
	p.m.Unlock()
	return
}

// selectPiece returns the index of the selected piece from candidates.
func selectPiece(candidates []*Piece) *Piece {
	sort.Sort(rarestFirst(candidates))
	minAvailability := candidates[0].availability()
	var i int
	for _, piece := range candidates {
		if piece.availability() > minAvailability {
			break
		}
		i++
	}
	candidates = candidates[:i]
	return candidates[rand.Intn(len(candidates))]
}

// rarestFirst implements sort.Interface based on availability of piece.
type rarestFirst []*Piece

func (r rarestFirst) Len() int           { return len(r) }
func (r rarestFirst) Swap(i, j int)      { r[i], r[j] = r[j], r[i] }
func (r rarestFirst) Less(i, j int) bool { return r[i].availability() < r[j].availability() }
<commit_msg>remove request<commit_after>package rain

import (
	"errors"
	"math/rand"
	"net"
	"runtime"
	"sort"
	"time"

	"github.com/cenkalti/rain/internal/connection"
	"github.com/cenkalti/rain/internal/logger"
)

// peerManager receives from t.peersC and keeps most recent peer addresses in t.peerC.
func (t *transfer) peerManager() {
	t.log.Debug("Started peerManager")
	for {
		select {
		case <-t.stopC:
			return
		case peers := <-t.peersC:
			for _, p := range peers {
				t.log.Debugln("Peer:", p)
				go func(addr *net.TCPAddr) {
					select {
					case t.peerC <- addr:
					case <-t.stopC:
					}
				}(p)
			}
		}
	}
}

// connecter connects to peers coming from t. peerC.
func (t *transfer) connecter() {
	limit := make(chan struct{}, maxPeerPerTorrent)
	for {
		select {
		case p := <-t.peerC:
			if p.Port == 0 {
				break
			}
			if p.IP.IsLoopback() && p.Port == int(t.rain.Port()) {
				break
			}

			limit <- struct{}{}
			go func(addr *net.TCPAddr) {
				defer func() {
					if err := recover(); err != nil {
						buf := make([]byte, 10000)
						t.log.Critical(err, "\n", string(buf[:runtime.Stack(buf, false)]))
					}
					<-limit
				}()
				t.connect(addr)
			}(p)
		case <-t.stopC:
			return
		}
	}
}

func (t *transfer) connect(addr *net.TCPAddr) {
	log := logger.New("peer -> " + addr.String())

	conn, cipher, extensions, peerID, err := connection.Dial(addr, !t.rain.config.Encryption.DisableOutgoing, t.rain.config.Encryption.ForceOutgoing, [8]byte{}, t.torrent.Info.Hash, t.rain.peerID)
	if err != nil {
		if err == connection.ErrOwnConnection {
			log.Debug(err)
		} else {
			log.Error(err)
		}
		return
	}
	log.Infof("Connected to peer. (cipher=%s, extensions=%x)", cipher, extensions)
	defer conn.Close()

	p := t.newPeer(conn, peerID, log)

	t.m.Lock()
	t.peers[peerID] = p
	t.m.Unlock()
	defer func() {
		t.m.Lock()
		delete(t.peers, peerID)
		t.m.Unlock()
	}()

	if err = p.SendBitfield(); err != nil {
		log.Error(err)
		return
	}

	p.Run()
}

func (p *Peer) downloader() {
	t := p.transfer
	for {
		t.m.Lock()
		if t.bitfield.All() {
			t.onceFinished.Do(func() { close(t.finished) })
			t.m.Unlock()
			return
		}
		candidates := p.candidates()
		if len(candidates) == 0 {
			t.m.Unlock()
			if err := p.BeNotInterested(); err != nil {
				p.log.Error(err)
				return
			}
			if err := p.waitForHaveMessage(); err != nil {
				p.log.Error(err)
				return
			}
			continue
		}
		piece := selectPiece(candidates)
		// Save selected piece so other downloaders do not try to download the same piece.
		t.requests[piece.Index] = &pieceRequest{p, time.Now()}
		t.m.Unlock()

		if err := p.BeInterested(); err != nil {
			p.log.Error(err)
			t.m.Lock()
			delete(t.requests, piece.Index)
			t.m.Unlock()
			return
		}

		// TODO queue max 10 requests
		go func() {
			if err := p.requestBlocks(piece); err != nil {
				p.log.Error(err)
			}
		}()

		// TODO handle choke while receiving pieces. Re-reqeust, etc..

		// Read blocks from peer.
		pieceData := make([]byte, piece.Length)
		for i := 0; i < len(piece.Blocks); i++ {
			peerBlock := <-p.pieceC
			data := <-peerBlock.Data
			if data == nil {
				p.log.Error("peer did not send block completely")
				t.m.Lock()
				delete(t.requests, piece.Index)
				t.m.Unlock()
				return
			}
			p.log.Debugln("Will receive block of length", len(data))
			copy(pieceData[peerBlock.Begin:], data)
		}

		if _, err := piece.Write(pieceData); err != nil {
			t.log.Error(err)
			p.Close()
			t.m.Lock()
			delete(t.requests, piece.Index)
			t.m.Unlock()
			return
		}

		t.m.Lock()
		t.bitfield.Set(piece.Index)
		delete(t.requests, piece.Index)
		t.m.Unlock()
	}
}

func (p *Peer) requestBlocks(piece *Piece) error {
	for _, b := range piece.Blocks {
		// Send requests only when unchoked.
		if err := p.waitForUnchoke(); err != nil {
			return err
		}
		if err := p.Request(piece.Index, b.Begin, b.Length); err != nil {
			return err
		}
	}
	return nil
}

func (p *Peer) waitForHaveMessage() error {
	p.m.Lock()
	defer p.m.Unlock()
	count := p.bitfield.Count()
	for count == p.bitfield.Count() && !p.disconnected {
		p.cond.Wait()
	}
	if p.disconnected {
		return errors.New("disconnected while waiting for new have message")
	}
	return nil
}

func (p *Peer) waitForUnchoke() error {
	p.m.Lock()
	defer p.m.Unlock()
	for p.peerChoking && !p.disconnected {
		p.cond.Wait()
	}
	if p.disconnected {
		return errors.New("disconnected while waiting for unchoke message")
	}
	return nil
}

// candidates returns list of piece indexes which is available on the peer but not available on the client.
func (p *Peer) candidates() (candidates []*Piece) {
	p.m.Lock()
	for i := uint32(0); i < p.transfer.bitfield.Len(); i++ {
		if !p.transfer.bitfield.Test(i) && p.bitfield.Test(i) {
			piece := p.transfer.pieces[i]
			if _, ok := p.transfer.requests[piece.Index]; !ok {
				candidates = append(candidates, piece)
			}
		}
	}
	p.m.Unlock()
	return
}

// selectPiece returns the index of the selected piece from candidates.
func selectPiece(candidates []*Piece) *Piece {
	sort.Sort(rarestFirst(candidates))
	minAvailability := candidates[0].availability()
	var i int
	for _, piece := range candidates {
		if piece.availability() > minAvailability {
			break
		}
		i++
	}
	candidates = candidates[:i]
	return candidates[rand.Intn(len(candidates))]
}

// rarestFirst implements sort.Interface based on availability of piece.
type rarestFirst []*Piece

func (r rarestFirst) Len() int           { return len(r) }
func (r rarestFirst) Swap(i, j int)      { r[i], r[j] = r[j], r[i] }
func (r rarestFirst) Less(i, j int) bool { return r[i].availability() < r[j].availability() }
<|endoftext|>


<commit_before>package lbot

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
)

func (b *Bot) SetConfig(c Config) {
	b.config = &c
}

func (b *Bot) SendTextMessage(m mid, s string) error {
	if b.config == nil {
		return errors.New("Config have not been set")
	}
	if b.config.Debug {
		log.Println("Start to Set Message")
	}

	var payload Request
	payload.SetDefaults()
	payload.SetText(s)
	payload.AddTargetUser(mid(m))

	out, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	if b.config.Debug {
		log.Println("Output json: " + string(out))
	}

	req, err := http.NewRequest("POST", b.config.ServerHost+"/v1/events", strings.NewReader(string(out)))
	if err != nil {
		return err
	}

	b.addAuthHeader(req)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if b.config.Debug {
		log.Println("Result: ", string(result))
	}

	return nil
}

func (b *Bot) GetUserProfile(m mid) ([]UserProfileResponse, error) {
	if b.config == nil {
		return nil, errors.New("Config have not been set")
	}

	if b.config.Debug {
		log.Println("Start to Set Message")
	}

	req, err := http.NewRequest("GET", b.config.ServerHost+"/v1/profiles?mids="+string(m), nil)
	if err != nil {
		return nil, err
	}

	b.addAuthHeader(req)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if b.config.Debug {
		log.Println("Result: ", string(result))
	}

	return nil, nil

}

func (b *Bot) addAuthHeader(r *http.Request) {

	r.Header.Set("Content-type", "application/json; charset=UTF-8")
	r.Header.Set("X-Line-ChannelID", b.config.ChannelID)
	r.Header.Set("X-Line-ChannelSecret", b.config.ChannelSecret)
	r.Header.Set("X-Line-Trusted-User-With-ACL", b.config.MID)

}
<commit_msg>modify return type<commit_after>package lbot

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
)

func (b *Bot) SetConfig(c Config) {
	b.config = &c
}

func (b *Bot) SendTextMessage(m mid, s string) error {
	if b.config == nil {
		return errors.New("Config have not been set")
	}
	if b.config.Debug {
		log.Println("Start to Set Message")
	}

	var payload Request
	payload.SetDefaults()
	payload.SetText(s)
	payload.AddTargetUser(mid(m))

	out, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	if b.config.Debug {
		log.Println("Output json: " + string(out))
	}

	req, err := http.NewRequest("POST", b.config.ServerHost+"/v1/events", strings.NewReader(string(out)))
	if err != nil {
		return err
	}

	b.addAuthHeader(req)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if b.config.Debug {
		log.Println("Result: ", string(result))
	}

	return nil
}

func (b *Bot) GetUserProfile(m mid) ([]ProfileInfo, error) {
	if b.config == nil {
		return nil, errors.New("Config have not been set")
	}

	if b.config.Debug {
		log.Println("Start to Set Message")
	}

	req, err := http.NewRequest("GET", b.config.ServerHost+"/v1/profiles?mids="+string(m), nil)
	if err != nil {
		return nil, err
	}

	b.addAuthHeader(req)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	result, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if b.config.Debug {
		log.Println("Result: ", string(result))
	}

	return nil, nil

}

func (b *Bot) addAuthHeader(r *http.Request) {

	r.Header.Set("Content-type", "application/json; charset=UTF-8")
	r.Header.Set("X-Line-ChannelID", b.config.ChannelID)
	r.Header.Set("X-Line-ChannelSecret", b.config.ChannelSecret)
	r.Header.Set("X-Line-Trusted-User-With-ACL", b.config.MID)

}
<|endoftext|>


<commit_before>package learn

import (
	"davidb.org/x/stenome/timelearn"
)

// The UI is responsible for asking the user a question, and
// determining if they got it right.  The prob argument will always
// point to the problem to learn.  'next' may point to an upcoming
// problem.
type UI interface {
	Single(prob, next *timelearn.Problem) (int, error)
	Close() error
}

// Mapping holding all of the UIs that can be learned.  The mapping
// holds a generator to generate the UI.
var allUI map[string]func() (UI, error)

func init() {
	allUI = make(map[string]func() (UI, error))

	Register("simple", newSimpleUI)
}

// Register adds a new user interface.  Any database with the given
// kind will use this function to make a new UI for it.
func Register(kind string, gen func() (UI, error)) {
	allUI[kind] = gen
}
<commit_msg>Use map literal instead of init func<commit_after>package learn

import (
	"davidb.org/x/stenome/timelearn"
)

// The UI is responsible for asking the user a question, and
// determining if they got it right.  The prob argument will always
// point to the problem to learn.  'next' may point to an upcoming
// problem.
type UI interface {
	Single(prob, next *timelearn.Problem) (int, error)
	Close() error
}

// Mapping holding all of the UIs that can be learned.  The mapping
// holds a generator to generate the UI.
var allUI = map[string]func() (UI, error){
	"simple": newSimpleUI,
}

// Register adds a new user interface.  Any database with the given
// kind will use this function to make a new UI for it.
func Register(kind string, gen func() (UI, error)) {
	allUI[kind] = gen
}
<|endoftext|>


<commit_before>package huffman

import (
	"encoding/binary"
	"io"

	"github.com/pkg/errors"
)

type BitReader struct {
	reader io.Reader
	buf    uint8
	rest   int
}

func NewBitReader(rdr io.Reader) *BitReader {
	return &BitReader{reader: rdr}
}

func (r *BitReader) ReadBit() (bool, error) {
	if r.rest == 0 {
		err := binary.Read(r.reader, binary.LittleEndian, &r.buf)
		if err != nil {
			return false, errors.WithStack(err)
		}
		r.rest = 8
	}
	/*
	   buf 01234567  rest=8
	   buf 1234567-  rest=7
	   buf 234567--  rest=6
	   ..
	   buf 7-------  rest=1
	   buf --------  rest=0
	*/
	result := r.buf&0x80 != 0
	r.buf <<= 1
	r.rest--
	return result, nil
}

func (r *BitReader) ReadUint8() (uint8, error) {
	var b uint8
	err := binary.Read(r.reader, binary.LittleEndian, &b)
	if err != nil {
		return 0, errors.WithStack(err)
	}
	result := r.buf
	if r.rest != 0 {
		result |= b >> uint(r.rest)
		r.buf = b << uint(8-r.rest)
	}
	return result, nil
}
<commit_msg>BitReaderの問題修正<commit_after>package huffman

import (
	"encoding/binary"
	"io"

	"github.com/pkg/errors"
)

type BitReader struct {
	reader io.Reader
	buf    uint8
	rest   int
}

func NewBitReader(rdr io.Reader) *BitReader {
	return &BitReader{reader: rdr}
}

func (r *BitReader) ReadBit() (bool, error) {
	if r.rest == 0 {
		err := binary.Read(r.reader, binary.LittleEndian, &r.buf)
		if err != nil {
			return false, errors.WithStack(err)
		}
		r.rest = 8
	}
	/*
	   buf 01234567  rest=8
	   buf 1234567-  rest=7
	   buf 234567--  rest=6
	   ..
	   buf 7-------  rest=1
	   buf --------  rest=0
	*/
	result := r.buf&0x80 != 0
	r.buf <<= 1
	r.rest--
	return result, nil
}

func (r *BitReader) ReadUint8() (uint8, error) {
	var buf2 uint8
	err := binary.Read(r.reader, binary.LittleEndian, &buf2)
	if err != nil {
		return 0, errors.WithStack(err)
	}
	result := r.buf | buf2 >> uint(r.rest)
	r.buf = buf2 << uint(8-r.rest)
	return result, nil
}
<|endoftext|>


<commit_before>package vm

import (
	"encoding/json"

	bosherr "github.com/cloudfoundry/bosh-utils/errors"
)

type AgentEnv struct {
	AgentID string `json:"agent_id"`

	VM VMSpec `json:"vm"`

	Mbus string   `json:"mbus"`
	NTP  []string `json:"ntp"`

	Blobstore BlobstoreSpec `json:"blobstore"`

	Networks Networks `json:"networks"`

	Disks DisksSpec `json:"disks"`

	Env EnvSpec `json:"env"`
}

type VMSpec struct {
	Name string `json:"name"`
	ID   string `json:"id"`
}

type DisksSpec struct {
	Ephemeral  string         `json:"ephemeral"`
	Persistent PersistentSpec `json:"persistent"`
}

type PersistentSpec map[string]string

type EnvSpec map[string]interface{}

const (
	BlobstoreTypeDav   = "dav"
	BlobstoreTypeLocal = "local"
)

type BlobstoreSpec struct {
	Provider string                 `json:"provider"`
	Options  map[string]interface{} `json:"options"`
}

type DavConfig map[string]interface{}

func NewAgentEnvFromJSON(bytes []byte) (AgentEnv, error) {
	var agentEnv AgentEnv

	err := json.Unmarshal(bytes, &agentEnv)
	if err != nil {
		return agentEnv, bosherr.WrapError(err, "Unmarshalling agent env")
	}

	return agentEnv, nil
}

func NewAgentEnvForVM(agentID, vmCID string, networks Networks, disksSpec DisksSpec, env Environment, agentOptions AgentOptions) AgentEnv {
	networksSpec := Networks{}

	for netName, network := range networks {
		networksSpec[netName] = Network{
			Type: network.Type,

			IP:      network.IP,
			Netmask: network.Netmask,
			Gateway: network.Gateway,

			DNS:           network.DNS,
			Default:       network.Default,
			Preconfigured: network.Preconfigured,

			MAC: "",

			Preconfigured: true,

			CloudProperties: network.CloudProperties,
		}
	}

	agentEnv := AgentEnv{
		AgentID: agentID,

		VM: VMSpec{
			Name: vmCID, // id for name and id
			ID:   vmCID,
		},

		Mbus: agentOptions.Mbus,
		NTP:  agentOptions.NTP,

		Blobstore: BlobstoreSpec{
			Provider: agentOptions.Blobstore.Type,
			Options:  agentOptions.Blobstore.Options,
		},

		Disks: disksSpec,

		Networks: networksSpec,

		// todo deep copy env?
		Env: EnvSpec(env),
	}

	return agentEnv
}

func (ae AgentEnv) AttachPersistentDisk(diskID, path string) AgentEnv {
	spec := PersistentSpec{}

	if ae.Disks.Persistent != nil {
		for k, v := range ae.Disks.Persistent {
			spec[k] = v
		}
	}

	spec[diskID] = path

	ae.Disks.Persistent = spec

	return ae
}

func (ae AgentEnv) DetachPersistentDisk(diskID string) AgentEnv {
	spec := PersistentSpec{}

	if ae.Disks.Persistent != nil {
		for k, v := range ae.Disks.Persistent {
			spec[k] = v
		}
	}

	delete(spec, diskID)

	ae.Disks.Persistent = spec

	return ae
}
<commit_msg>Fixed issue<commit_after>package vm

import (
	"encoding/json"

	bosherr "github.com/cloudfoundry/bosh-utils/errors"
)

type AgentEnv struct {
	AgentID string `json:"agent_id"`

	VM VMSpec `json:"vm"`

	Mbus string   `json:"mbus"`
	NTP  []string `json:"ntp"`

	Blobstore BlobstoreSpec `json:"blobstore"`

	Networks Networks `json:"networks"`

	Disks DisksSpec `json:"disks"`

	Env EnvSpec `json:"env"`
}

type VMSpec struct {
	Name string `json:"name"`
	ID   string `json:"id"`
}

type DisksSpec struct {
	Ephemeral  string         `json:"ephemeral"`
	Persistent PersistentSpec `json:"persistent"`
}

type PersistentSpec map[string]string

type EnvSpec map[string]interface{}

const (
	BlobstoreTypeDav   = "dav"
	BlobstoreTypeLocal = "local"
)

type BlobstoreSpec struct {
	Provider string                 `json:"provider"`
	Options  map[string]interface{} `json:"options"`
}

type DavConfig map[string]interface{}

func NewAgentEnvFromJSON(bytes []byte) (AgentEnv, error) {
	var agentEnv AgentEnv

	err := json.Unmarshal(bytes, &agentEnv)
	if err != nil {
		return agentEnv, bosherr.WrapError(err, "Unmarshalling agent env")
	}

	return agentEnv, nil
}

func NewAgentEnvForVM(agentID, vmCID string, networks Networks, disksSpec DisksSpec, env Environment, agentOptions AgentOptions) AgentEnv {
	networksSpec := Networks{}

	for netName, network := range networks {
		networksSpec[netName] = Network{
			Type: network.Type,

			IP:      network.IP,
			Netmask: network.Netmask,
			Gateway: network.Gateway,

			DNS:           network.DNS,
			Default:       network.Default,
			Preconfigured: true,

			MAC: "",

			CloudProperties: network.CloudProperties,
		}
	}

	agentEnv := AgentEnv{
		AgentID: agentID,

		VM: VMSpec{
			Name: vmCID, // id for name and id
			ID:   vmCID,
		},

		Mbus: agentOptions.Mbus,
		NTP:  agentOptions.NTP,

		Blobstore: BlobstoreSpec{
			Provider: agentOptions.Blobstore.Type,
			Options:  agentOptions.Blobstore.Options,
		},

		Disks: disksSpec,

		Networks: networksSpec,

		// todo deep copy env?
		Env: EnvSpec(env),
	}

	return agentEnv
}

func (ae AgentEnv) AttachPersistentDisk(diskID, path string) AgentEnv {
	spec := PersistentSpec{}

	if ae.Disks.Persistent != nil {
		for k, v := range ae.Disks.Persistent {
			spec[k] = v
		}
	}

	spec[diskID] = path

	ae.Disks.Persistent = spec

	return ae
}

func (ae AgentEnv) DetachPersistentDisk(diskID string) AgentEnv {
	spec := PersistentSpec{}

	if ae.Disks.Persistent != nil {
		for k, v := range ae.Disks.Persistent {
			spec[k] = v
		}
	}

	delete(spec, diskID)

	ae.Disks.Persistent = spec

	return ae
}
<|endoftext|>


<commit_before>package editor

import (
	"strings"
	"testing"
)

func TestInsertMode(t *testing.T) {
	i := insert{}
	i.init()
	input := "abc"
	i.in = NewReader(strings.NewReader(input))
	for n := range input {
		c, _, err := i.Run()
		if err != nil {
			t.Errorf("Run: %v", err)
		}
		if c != cont {
			t.Errorf("Run: want %v, but got %v", cont, c)
		}
		if got, want := string(i.Runes()), input[:n+1]; got != want {
			t.Errorf("Run: want %v, but got %v", want, got)
		}
	}
}
<commit_msg>Add test for input of matching characters<commit_after>package editor

import (
	"strings"
	"testing"
)

func TestInsertMode(t *testing.T) {
	i := insert{}
	i.init()
	input := "abc"
	i.in = NewReader(strings.NewReader(input))
	for n := range input {
		c, _, err := i.Run()
		if err != nil {
			t.Errorf("Run: %v", err)
		}
		if c != cont {
			t.Errorf("Run: want %v, but got %v", cont, c)
		}
		if got, want := string(i.Runes()), input[:n+1]; got != want {
			t.Errorf("Run: want %v, but got %v", want, got)
		}
	}
}

func TestInputMatches(t *testing.T) {
	i := insert{}
	i.init()
	input := "a 'b"
	i.in = NewReader(strings.NewReader(input))
	tt := []string{"a", "a ", "a ''", "a 'b'"}
	for n := range input {
		c, _, err := i.Run()
		if err != nil {
			t.Errorf("Run: %v", err)
		}
		if c != cont {
			t.Errorf("Run: want %v, but got %v", cont, c)
		}
		if got, want := string(i.Runes()), tt[n]; got != want {
			t.Errorf("Run/%d: want %q, but got %q", n, want, got)
		}
	}
}
<|endoftext|>


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

// packages used for load balancer setting
package loadbalancer

import (
	"math"
	"sort"

	apiv2 "github.com/envoyproxy/go-control-plane/envoy/api/v2"
	core "github.com/envoyproxy/go-control-plane/envoy/api/v2/core"
	"github.com/golang/protobuf/ptypes/wrappers"

	"istio.io/api/networking/v1alpha3"
	"istio.io/istio/pilot/pkg/networking/util"
)

func GetLocalityLbSetting(
	mesh *v1alpha3.LocalityLoadBalancerSetting,
	destrule *v1alpha3.LocalityLoadBalancerSetting,
) *v1alpha3.LocalityLoadBalancerSetting {
	var enabled bool
	// Locality lb is enabled if its not explicitly disabled in mesh global config
	if mesh != nil && (mesh.Enabled == nil || mesh.Enabled.Value) {
		enabled = true
	}
	// Unless we explicitly override this in destination rule
	if destrule != nil {
		if destrule.Enabled != nil && !destrule.Enabled.Value {
			enabled = false
		} else {
			enabled = true
		}
	}
	if !enabled {
		return nil
	}

	// Destination Rule overrides mesh config. If its defined, use that
	if destrule != nil {
		return destrule
	}
	// Otherwise fall back to mesh default
	return mesh
}

func ApplyLocalityLBSetting(
	locality *core.Locality,
	loadAssignment *apiv2.ClusterLoadAssignment,
	localityLB *v1alpha3.LocalityLoadBalancerSetting,
	enableFailover bool,
) {
	if locality == nil || loadAssignment == nil {
		return
	}

	// one of Distribute or Failover settings can be applied.
	if localityLB.GetDistribute() != nil {
		applyLocalityWeight(locality, loadAssignment, localityLB.GetDistribute())
	} else if enableFailover {
		// Failover needs outlier detection, otherwise Envoy will never drop down to a lower priority.
		applyLocalityFailover(locality, loadAssignment, localityLB.GetFailover())
	}
}

// set locality loadbalancing weight
func applyLocalityWeight(
	locality *core.Locality,
	loadAssignment *apiv2.ClusterLoadAssignment,
	distribute []*v1alpha3.LocalityLoadBalancerSetting_Distribute) {
	if distribute == nil {
		return
	}

	// Support Locality weighted load balancing
	// (https://www.envoyproxy.io/docs/envoy/latest/intro/arch_overview/load_balancing/locality_weight.html)
	// by providing weights in LocalityLbEndpoints via load_balancing_weight.
	// By setting weights across different localities, it can allow
	// Envoy to weight assignments across different zones and geographical locations.
	for _, localityWeightSetting := range distribute {
		if localityWeightSetting != nil &&
			util.LocalityMatch(locality, localityWeightSetting.From) {
			misMatched := map[int]struct{}{}
			for i := range loadAssignment.Endpoints {
				misMatched[i] = struct{}{}
			}
			for locality, weight := range localityWeightSetting.To {
				// index -> original weight
				destLocMap := map[int]uint32{}
				totalWeight := uint32(0)
				for i, ep := range loadAssignment.Endpoints {
					if _, exist := misMatched[i]; exist {
						if util.LocalityMatch(ep.Locality, locality) {
							delete(misMatched, i)
							if ep.LoadBalancingWeight != nil {
								destLocMap[i] = ep.LoadBalancingWeight.Value
							} else {
								destLocMap[i] = 1
							}
							totalWeight += destLocMap[i]
						}
					}
				}
				// in case wildcard dest matching multi groups of endpoints
				// the load balancing weight for a locality is divided by the sum of the weights of all localities
				for index, originalWeight := range destLocMap {
					destWeight := float64(originalWeight*weight) / float64(totalWeight)
					if destWeight > 0 {
						loadAssignment.Endpoints[index].LoadBalancingWeight = &wrappers.UInt32Value{
							Value: uint32(math.Ceil(destWeight)),
						}
					}
				}
			}

			// remove groups of endpoints in a locality that miss matched
			for i := range misMatched {
				loadAssignment.Endpoints[i].LbEndpoints = nil
			}
			break
		}
	}
}

// set locality loadbalancing priority
func applyLocalityFailover(
	locality *core.Locality,
	loadAssignment *apiv2.ClusterLoadAssignment,
	failover []*v1alpha3.LocalityLoadBalancerSetting_Failover) {
	// key is priority, value is the index of the LocalityLbEndpoints in ClusterLoadAssignment
	priorityMap := map[int][]int{}

	// 1. calculate the LocalityLbEndpoints.Priority compared with proxy locality
	for i, localityEndpoint := range loadAssignment.Endpoints {
		// if region/zone/subZone all match, the priority is 0.
		// if region/zone match, the priority is 1.
		// if region matches, the priority is 2.
		// if locality not match, the priority is 3.
		priority := util.LbPriority(locality, localityEndpoint.Locality)
		// region not match, apply failover settings when specified
		// update localityLbEndpoints' priority to 4 if failover not match
		if priority == 3 {
			for _, failoverSetting := range failover {
				if failoverSetting.From == locality.Region {
					if localityEndpoint.Locality == nil || localityEndpoint.Locality.Region != failoverSetting.To {
						priority = 4
					}
					break
				}
			}
		}
		loadAssignment.Endpoints[i].Priority = uint32(priority)
		priorityMap[priority] = append(priorityMap[priority], i)
	}

	// since Priorities should range from 0 (highest) to N (lowest) without skipping.
	// 2. adjust the priorities in order
	// 2.1 sort all priorities in increasing order.
	priorities := []int{}
	for priority := range priorityMap {
		priorities = append(priorities, priority)
	}
	sort.Ints(priorities)
	// 2.2 adjust LocalityLbEndpoints priority
	// if the index and value of priorities array is not equal.
	for i, priority := range priorities {
		if i != priority {
			// the LocalityLbEndpoints index in ClusterLoadAssignment.Endpoints
			for _, index := range priorityMap[priority] {
				loadAssignment.Endpoints[index].Priority = uint32(i)
			}
		}
	}

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

// packages used for load balancer setting
package loadbalancer

import (
	"math"
	"sort"

	apiv2 "github.com/envoyproxy/go-control-plane/envoy/api/v2"
	core "github.com/envoyproxy/go-control-plane/envoy/api/v2/core"
	"github.com/golang/protobuf/ptypes/wrappers"

	"istio.io/api/networking/v1alpha3"
	"istio.io/istio/pilot/pkg/networking/util"
)

func GetLocalityLbSetting(
	mesh *v1alpha3.LocalityLoadBalancerSetting,
	destrule *v1alpha3.LocalityLoadBalancerSetting,
) *v1alpha3.LocalityLoadBalancerSetting {
	var enabled bool
	// Locality lb is enabled if its not explicitly disabled in mesh global config
	if mesh != nil && (mesh.Enabled == nil || mesh.Enabled.Value) {
		enabled = true
	}
	// Unless we explicitly override this in destination rule
	if destrule != nil {
		if destrule.Enabled != nil && !destrule.Enabled.Value {
			enabled = false
		} else {
			enabled = true
		}
	}
	if !enabled {
		return nil
	}

	// Destination Rule overrides mesh config. If its defined, use that
	if destrule != nil {
		return destrule
	}
	// Otherwise fall back to mesh default
	return mesh
}

func ApplyLocalityLBSetting(
	locality *core.Locality,
	loadAssignment *apiv2.ClusterLoadAssignment,
	localityLB *v1alpha3.LocalityLoadBalancerSetting,
	enableFailover bool,
) {
	if locality == nil || loadAssignment == nil {
		return
	}

	// one of Distribute or Failover settings can be applied.
	if localityLB.GetDistribute() != nil {
		applyLocalityWeight(locality, loadAssignment, localityLB.GetDistribute())
	} else if enableFailover {
		// Failover needs outlier detection, otherwise Envoy will never drop down to a lower priority.
		applyLocalityFailover(locality, loadAssignment, localityLB.GetFailover())
	}
}

// set locality loadbalancing weight
func applyLocalityWeight(
	locality *core.Locality,
	loadAssignment *apiv2.ClusterLoadAssignment,
	distribute []*v1alpha3.LocalityLoadBalancerSetting_Distribute) {
	if distribute == nil {
		return
	}

	// Support Locality weighted load balancing
	// (https://www.envoyproxy.io/docs/envoy/latest/intro/arch_overview/upstream/load_balancing/locality_weight#locality-weighted-load-balancing)
	// by providing weights in LocalityLbEndpoints via load_balancing_weight.
	// By setting weights across different localities, it can allow
	// Envoy to weight assignments across different zones and geographical locations.
	for _, localityWeightSetting := range distribute {
		if localityWeightSetting != nil &&
			util.LocalityMatch(locality, localityWeightSetting.From) {
			misMatched := map[int]struct{}{}
			for i := range loadAssignment.Endpoints {
				misMatched[i] = struct{}{}
			}
			for locality, weight := range localityWeightSetting.To {
				// index -> original weight
				destLocMap := map[int]uint32{}
				totalWeight := uint32(0)
				for i, ep := range loadAssignment.Endpoints {
					if _, exist := misMatched[i]; exist {
						if util.LocalityMatch(ep.Locality, locality) {
							delete(misMatched, i)
							if ep.LoadBalancingWeight != nil {
								destLocMap[i] = ep.LoadBalancingWeight.Value
							} else {
								destLocMap[i] = 1
							}
							totalWeight += destLocMap[i]
						}
					}
				}
				// in case wildcard dest matching multi groups of endpoints
				// the load balancing weight for a locality is divided by the sum of the weights of all localities
				for index, originalWeight := range destLocMap {
					destWeight := float64(originalWeight*weight) / float64(totalWeight)
					if destWeight > 0 {
						loadAssignment.Endpoints[index].LoadBalancingWeight = &wrappers.UInt32Value{
							Value: uint32(math.Ceil(destWeight)),
						}
					}
				}
			}

			// remove groups of endpoints in a locality that miss matched
			for i := range misMatched {
				loadAssignment.Endpoints[i].LbEndpoints = nil
			}
			break
		}
	}
}

// set locality loadbalancing priority
func applyLocalityFailover(
	locality *core.Locality,
	loadAssignment *apiv2.ClusterLoadAssignment,
	failover []*v1alpha3.LocalityLoadBalancerSetting_Failover) {
	// key is priority, value is the index of the LocalityLbEndpoints in ClusterLoadAssignment
	priorityMap := map[int][]int{}

	// 1. calculate the LocalityLbEndpoints.Priority compared with proxy locality
	for i, localityEndpoint := range loadAssignment.Endpoints {
		// if region/zone/subZone all match, the priority is 0.
		// if region/zone match, the priority is 1.
		// if region matches, the priority is 2.
		// if locality not match, the priority is 3.
		priority := util.LbPriority(locality, localityEndpoint.Locality)
		// region not match, apply failover settings when specified
		// update localityLbEndpoints' priority to 4 if failover not match
		if priority == 3 {
			for _, failoverSetting := range failover {
				if failoverSetting.From == locality.Region {
					if localityEndpoint.Locality == nil || localityEndpoint.Locality.Region != failoverSetting.To {
						priority = 4
					}
					break
				}
			}
		}
		loadAssignment.Endpoints[i].Priority = uint32(priority)
		priorityMap[priority] = append(priorityMap[priority], i)
	}

	// since Priorities should range from 0 (highest) to N (lowest) without skipping.
	// 2. adjust the priorities in order
	// 2.1 sort all priorities in increasing order.
	priorities := []int{}
	for priority := range priorityMap {
		priorities = append(priorities, priority)
	}
	sort.Ints(priorities)
	// 2.2 adjust LocalityLbEndpoints priority
	// if the index and value of priorities array is not equal.
	for i, priority := range priorities {
		if i != priority {
			// the LocalityLbEndpoints index in ClusterLoadAssignment.Endpoints
			for _, index := range priorityMap[priority] {
				loadAssignment.Endpoints[index].Priority = uint32(i)
			}
		}
	}

}
<|endoftext|>


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

package mutating_webhook

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	k8smetav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/tools/cache"

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/testutils"
)

var _ = Describe("Mutating Webhook Namespace Limits", func() {
	var vmi v1.VirtualMachineInstance
	var namespaceLimit *k8sv1.LimitRange
	var namespaceLimitInformer cache.SharedIndexInformer

	memory, _ := resource.ParseQuantity("64M")
	limitMemory, _ := resource.ParseQuantity("128M")
	zeroMemory, _ := resource.ParseQuantity("0M")

	BeforeEach(func() {
		vmi = v1.VirtualMachineInstance{
			Spec: v1.VirtualMachineInstanceSpec{
				Domain: v1.DomainSpec{
					Resources: v1.ResourceRequirements{
						Requests: k8sv1.ResourceList{
							"memory": memory,
						},
					},
				},
			},
		}
		namespaceLimit = &k8sv1.LimitRange{
			Spec: k8sv1.LimitRangeSpec{
				Limits: []k8sv1.LimitRangeItem{
					{
						Type: k8sv1.LimitTypeContainer,
						Default: k8sv1.ResourceList{
							k8sv1.ResourceMemory: limitMemory,
						},
					},
				},
			},
		}
		namespaceLimitInformer, _ = testutils.NewFakeInformerFor(&k8sv1.LimitRange{})
		namespaceLimitInformer.GetIndexer().Add(namespaceLimit)
	})

	When("VMI has limits under spec", func() {
		It("should not apply namespace limits", func() {
			vmiCopy := vmi.DeepCopy()
			vmiCopy.Spec.Domain.Resources.Limits = k8sv1.ResourceList{
				k8sv1.ResourceMemory: memory,
			}
			By("Applying namespace range values on the VMI")
			applyNamespaceLimitRangeValues(vmiCopy, namespaceLimitInformer)

			Expect(vmiCopy.Spec.Domain.Resources.Limits.Memory().String()).To(Equal("64M"))
		})
	})

	When("VMI does not have limits under spec", func() {
		It("should apply namespace limits", func() {
			vmiCopy := vmi.DeepCopy()
			By("Applying namespace range values on the VMI")
			applyNamespaceLimitRangeValues(vmiCopy, namespaceLimitInformer)

			Expect(vmiCopy.Spec.Domain.Resources.Limits.Memory().String()).To(Equal("128M"))
		})

		When("namespace limit equals 0", func() {
			It("should not apply namespace limits", func() {
				vmiCopy := vmi.DeepCopy()

				namespaceLimitCopy := namespaceLimit.DeepCopy()
				namespaceLimitCopy.Spec.Limits[0].Default[k8sv1.ResourceMemory] = zeroMemory
				namespaceLimitInformer.GetIndexer().Update(namespaceLimitCopy)

				By("Applying namespace range values on the VMI")
				applyNamespaceLimitRangeValues(vmiCopy, namespaceLimitInformer)
				_, ok := vmiCopy.Spec.Domain.Resources.Limits[k8sv1.ResourceMemory]
				Expect(ok).To(Equal(false))
			})

			AfterEach(func() {
				namespaceLimitInformer.GetIndexer().Update(namespaceLimit)
			})
		})

		When("namespace has more than one limit range", func() {
			var additionalNamespaceLimit *k8sv1.LimitRange

			BeforeEach(func() {
				additionalLimitMemory, _ := resource.ParseQuantity("76M")
				additionalNamespaceLimit = &k8sv1.LimitRange{
					ObjectMeta: k8smetav1.ObjectMeta{
						Name: "additional-limit-range",
					},
					Spec: k8sv1.LimitRangeSpec{
						Limits: []k8sv1.LimitRangeItem{
							{
								Type: k8sv1.LimitTypeContainer,
								Default: k8sv1.ResourceList{
									k8sv1.ResourceMemory: additionalLimitMemory,
								},
							},
						},
					},
				}
				namespaceLimitInformer.GetIndexer().Add(additionalNamespaceLimit)
			})

			It("should apply range with minimal limit", func() {
				vmiCopy := vmi.DeepCopy()
				By("Applying namespace range values on the VMI")
				applyNamespaceLimitRangeValues(vmiCopy, namespaceLimitInformer)

				Expect(vmiCopy.Spec.Domain.Resources.Limits.Memory().String()).To(Equal("76M"))
			})
		})
	})
})
<commit_msg>remove incorrect test in namespace-limits<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Copyright 2018 Red Hat, Inc.
 *
 */

package mutating_webhook

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	k8smetav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/tools/cache"

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	"kubevirt.io/kubevirt/pkg/testutils"
)

var _ = Describe("Mutating Webhook Namespace Limits", func() {
	var vmi v1.VirtualMachineInstance
	var namespaceLimit *k8sv1.LimitRange
	var namespaceLimitInformer cache.SharedIndexInformer

	memory, _ := resource.ParseQuantity("64M")
	limitMemory, _ := resource.ParseQuantity("128M")

	BeforeEach(func() {
		vmi = v1.VirtualMachineInstance{
			Spec: v1.VirtualMachineInstanceSpec{
				Domain: v1.DomainSpec{
					Resources: v1.ResourceRequirements{
						Requests: k8sv1.ResourceList{
							"memory": memory,
						},
					},
				},
			},
		}
		namespaceLimit = &k8sv1.LimitRange{
			Spec: k8sv1.LimitRangeSpec{
				Limits: []k8sv1.LimitRangeItem{
					{
						Type: k8sv1.LimitTypeContainer,
						Default: k8sv1.ResourceList{
							k8sv1.ResourceMemory: limitMemory,
						},
					},
				},
			},
		}
		namespaceLimitInformer, _ = testutils.NewFakeInformerFor(&k8sv1.LimitRange{})
		namespaceLimitInformer.GetIndexer().Add(namespaceLimit)
	})

	When("VMI has limits under spec", func() {
		It("should not apply namespace limits", func() {
			vmiCopy := vmi.DeepCopy()
			vmiCopy.Spec.Domain.Resources.Limits = k8sv1.ResourceList{
				k8sv1.ResourceMemory: memory,
			}
			By("Applying namespace range values on the VMI")
			applyNamespaceLimitRangeValues(vmiCopy, namespaceLimitInformer)

			Expect(vmiCopy.Spec.Domain.Resources.Limits.Memory().String()).To(Equal("64M"))
		})
	})

	When("VMI does not have limits under spec", func() {
		It("should apply namespace limits", func() {
			vmiCopy := vmi.DeepCopy()
			By("Applying namespace range values on the VMI")
			applyNamespaceLimitRangeValues(vmiCopy, namespaceLimitInformer)

			Expect(vmiCopy.Spec.Domain.Resources.Limits.Memory().String()).To(Equal("128M"))
		})

		When("namespace has more than one limit range", func() {
			var additionalNamespaceLimit *k8sv1.LimitRange

			BeforeEach(func() {
				additionalLimitMemory, _ := resource.ParseQuantity("76M")
				additionalNamespaceLimit = &k8sv1.LimitRange{
					ObjectMeta: k8smetav1.ObjectMeta{
						Name: "additional-limit-range",
					},
					Spec: k8sv1.LimitRangeSpec{
						Limits: []k8sv1.LimitRangeItem{
							{
								Type: k8sv1.LimitTypeContainer,
								Default: k8sv1.ResourceList{
									k8sv1.ResourceMemory: additionalLimitMemory,
								},
							},
						},
					},
				}
				namespaceLimitInformer.GetIndexer().Add(additionalNamespaceLimit)
			})

			It("should apply range with minimal limit", func() {
				vmiCopy := vmi.DeepCopy()
				By("Applying namespace range values on the VMI")
				applyNamespaceLimitRangeValues(vmiCopy, namespaceLimitInformer)

				Expect(vmiCopy.Spec.Domain.Resources.Limits.Memory().String()).To(Equal("76M"))
			})
		})
	})
})
<|endoftext|>