<commit_before>/*
*	Author: Taha Shaikh
*	=== Ant Colony Optimization for the Traveling Salesman Problem ===
*
*	=== Usage ===
*	Comment out whichever file/data not needed and uncomment file/data at
*	line numbers 60 and 90 to execute ACO on
*
*	=== Implementation ===
*	- Reads the DIMENSION field in the file provided for the total number of cities
*	- Reads the coordinates of each city in the NODE_COORD_SECTION and saves them in
*	  a list of cities
*	- Creates an adjacency matrix for the cities, computing edge weights by the
*	  calculating the euclidean distance between every two cities
*	- Initialize the tau matrix (pheromone levels) for each edge
*	- initialize each ant and provide them a random city to start
*	- Each ant traveses the path and chooses each city using a probability which
*	  is computed using a formula
*	- Update the tau matrix (pheromone levels) after one tour has end
*	- Run 500 tours 10 times to obtain average and best tour length values
 */
package main

import (
	"bufio"
	"fmt"
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"

	"github.com/gonum/plot"
	"github.com/gonum/plot/plotter"
	"github.com/gonum/plot/plotutil"
	"github.com/gonum/plot/vg"
)

//ant implementation
type ant_t struct {
	tabulist    []int
	currentCity int
	nextCity    int
	tour        []int
	tourIndex   int
	tourlength  float64
}

//city implementation
type city struct {
	x, y float64
}

//declaration and initialization
var ants []ant_t
var cities []city
var tauMatrix [][]float64
var adjMatrix [][]float64
var besttour []int
var currentIndex int
var numCities = 0
var rho = 0.6
var qval = 1.0
var alpha = 0.8
var beta = 0.8
var numAnts = 10

//prints error and exits on abnormal conditions
func printError(err error) {
	fmt.Print(err)
	os.Exit(2)
}

func main() {
	//comment other to use the other one
	//initGraph("TSP_D")
	initGraph("TSP_WS") //other TSP file

	avgavg := make([]float64, 10)
	avgbest := make([]float64, 10)
	besttourlens := make([]float64, 10)
	avgtourlens := make([]float64, 10)
	iterations := 0
	i := 0
	//initializes pheromone levels with base pheromone i.e. 1
	initTrail()
	for i < 10 {
		iterations = 0
		besttour = nil
		for iterations < 500 {
			//initializes each ant
			initAnts()
			//tour for each ant
			moveAnts()
			//intensify pheromone levels
			intensifyTrail()
			//compute best and avg tour length for every tour in
			//500 and add them for average
			besttourlens[i] += calculateBest()
			avgtourlens[i] += calculateAvg()
			iterations++
		}
		//evaporate pheromone to obtain better results
		evaporatePheromone()
		fmt.Println("Iteration: ", i)
		fmt.Println("Optimal Path: ", besttour)
		avgavg[i] = avgtourlens[i] / 500.0
		avgbest[i] = besttourlens[i] / 500.0
		i++
	}
	fmt.Println("Average Average:", avgavg)
	fmt.Println("Average Best:", avgbest)

	p, err := plot.New()
	if err != nil {
		printError(err)
	}

	//comment other to use the other one
	//p.Title.Text = "Dijibouti TSP"
	p.Title.Text = "Western Sahara TSP"
	p.X.Label.Text = "X"
	p.Y.Label.Text = "Y"

	avgpts := make(plotter.XYs, 10)
	for i := range avgpts {
		avgpts[i].Y = avgavg[i]
		avgpts[i].X = float64(i)
	}

	bestpts := make(plotter.XYs, 10)
	for i := range bestpts {
		bestpts[i].Y = avgbest[i]
		bestpts[i].X = float64(i)
	}
	err = plotutil.AddLinePoints(p,
		"Average So Far", avgpts,
		"Best So Far", bestpts)
	if err != nil {
		printError(err)
	}

	// Save the plot to a PNG file.
	/*if err := p.Save(4*vg.Inch, 4*vg.Inch, "djibouti.png"); err != nil {
		printError(err)
	}*/
	if err := p.Save(4*vg.Inch, 4*vg.Inch, "westernsahara.png"); err != nil {
		printError(err)
	}
}

//calculate average tour length of all length for one tour
func calculateAvg() float64 {
	avglength := ants[0].tourlength
	total := 0.0
	for i := range ants {
		total += ants[i].tourlength
	}
	avglength = total / 10.0
	return avglength
}

//calculate best tour length of all length for one tour
func calculateBest() float64 {
	var bestlength float64
	bestlength = ants[0].tourlength
	if besttour == nil {
		besttour = make([]int, numCities)
		copy(besttour, ants[0].tour)
	}
	for i := range ants {
		if ants[i].tourlength < bestlength {

			bestlength = ants[i].tourlength
			copy(besttour, ants[i].tour)
		}
	}
	return bestlength
}

//initialize pheromone levels
func initTrail() {
	tauMatrix = make([][]float64, numCities)
	for i := range tauMatrix {
		tauMatrix[i] = make([]float64, numCities)
		for j := range tauMatrix[i] {
			if i != j {
				tauMatrix[i][j] = 1.0 //initialize to base pheromone = 1
			} else {
				tauMatrix[i][j] = 0.0
			}
		}
	}
}

//initialize ants
func initAnts() {
	ants = nil
	ants = make([]ant_t, numAnts)
	for i := range ants {
		ants[i].tabulist = make([]int, numCities)
		ants[i].currentCity = rand.Intn(numCities) //randomly assigns ant a city to s
		ants[i].nextCity = 0
		ants[i].tour = make([]int, numCities)
		ants[i].tourIndex = 0
		ants[i].tourlength = 0.0
	}
}

//move all ants to visit the whole graph
func moveAnts() {
	for i := range ants {
		currentIndex = 0
		for currentIndex < numCities {
			goToNewCity(&ants[i])
			currentIndex++
		}
	}
}

//choosing next city
func goToNewCity(ant *ant_t) {
	var from, to int
	var p float64
	denom := 0.0
	from = ant.currentCity
	for to = 0; to < numCities; to++ {
		if from != to {
			if ant.tabulist[to] == 0 && tauMatrix[from][to] != 0 && adjMatrix[from][to] != 0 {
				denom += math.Pow(tauMatrix[from][to], alpha) * math.Pow((1.0/adjMatrix[from][to]), beta)
			}
		} else {
			continue
		}

	}
	to = 0
	for {
		if from != to {
			if ant.tabulist[to] == 0 {
				p = (math.Pow(tauMatrix[from][to], alpha) * math.Pow((1.0/adjMatrix[from][to]), beta)) / denom

				if rand.Float64() < p {
					break
				}
			}
		} else {
			to = ((to + 1) % numCities)
			continue
		}
		to = ((to + 1) % numCities)
	}
	ant.nextCity = to
	ant.tabulist[ant.nextCity] = 1
	ant.tour[ant.tourIndex] = ant.nextCity
	ant.tourIndex++
	ant.tourlength += adjMatrix[ant.currentCity][ant.nextCity]
	if ant.tourIndex == numCities {
		ant.tourlength += adjMatrix[ant.tour[numCities-1]][ant.tour[0]]
	}
	ant.currentCity = ant.nextCity
}

//reads from file and creates a list of all the cities coordinates
func readFile(name string) []city {
	var dim, i int
	var cities []city
	i, dim = 1, 0
	var startFlag bool
	startFlag = false
	if file, err := os.Open(name); err == nil {
		// make sure it gets closed
		defer file.Close()
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			str := scanner.Text()
			if strings.Contains(str, "DIMENSION") {
				dim = getDim(str)
				break
			}
		}
		cities = make([]city, dim)
		for scanner.Scan() {
			str := scanner.Text()
			if strings.Contains(str, "EOF") {
				break
			} else if startFlag {
				x, y := tokenize(str)
				if i <= dim {
					cities[i-1] = city{x, y}
					i++
				} else {
					startFlag = false
				}
			} else if strings.Contains(str, "NODE_COORD_SECTION") {
				startFlag = true
			}
		}
		// check for errors
		if err = scanner.Err(); err != nil {
			printError(err)
		}
	} else {
		printError(err)
	}
	return cities
}

//evaporating pheromone after each iteration of the algorithm
func evaporatePheromone() {
	var from, to int
	for from = 0; from < numCities; from++ {
		for to = 0; to < numCities; to++ {
			tauMatrix[from][to] = tauMatrix[from][to] * (1.0 - rho)
			if tauMatrix[from][to] < 0.0 {
				tauMatrix[from][to] = 1.0
			}
		}
	}
}

//intensifying pheromone levels
func intensifyTrail() {
	var from, to, i, c int
	for i = 0; i < numAnts; i++ {
		for c = 0; c < numCities; c++ {
			from = ants[i].tour[c]
			to = ants[i].tour[((c + 1) % numCities)]
			deltatau := (qval / ants[i].tourlength)
			tauMatrix[from][to] = tauMatrix[from][to] + deltatau
			tauMatrix[to][from] = tauMatrix[from][to]
		}
	}
}

//making graph
func initGraph(name string) {
	cities = readFile(name)
	numCities = len(cities)
	adjMatrix = make([][]float64, numCities)
	for i := range adjMatrix {
		adjMatrix[i] = make([]float64, numCities)
		for j := range adjMatrix[i] {
			adjMatrix[i][j] = calEdge(cities[i], cities[j])
		}
	}
}

//calculates edge weight (euclidiean distance)
func calEdge(c1, c2 city) float64 {
	return math.Pow((math.Pow((c2.y-c1.y), 2) + math.Pow((c2.y-c1.y), 2)), 0.5)
}

//tokenizes and converts to float
func tokenize(str string) (x, y float64) {
	s := strings.Split(str, " ")
	strX, strY := s[1], s[2]
	x, err := strconv.ParseFloat(strX, 64) //converts string to float64
	if err != nil {
		printError(err)
	}
	y, err = strconv.ParseFloat(strY, 64) //converts string to float64
	if err != nil {
		printError(err)
	}
	return x, y
}

//gets number of cities from the file
func getDim(str string) (dim int) {
	s := strings.Split(str, ":")
	num := strings.TrimLeft(s[1], " ")
	if dim, err := strconv.Atoi(num); err == nil {
		return dim
	} else {
		fmt.Print(err)
		os.Exit(2)
	}
	return 0
}
<commit_msg>Calculating distances correctly<commit_after>/*
*	Author: Taha Shaikh
*	=== Ant Colony Optimization for the Traveling Salesman Problem ===
*
*	=== Usage ===
*	Comment out whichever file/data not needed and uncomment file/data at
*	line numbers 60 and 90 to execute ACO on
*
*	=== Implementation ===
*	- Reads the DIMENSION field in the file provided for the total number of cities
*	- Reads the coordinates of each city in the NODE_COORD_SECTION and saves them in
*	  a list of cities
*	- Creates an adjacency matrix for the cities, computing edge weights by the
*	  calculating the euclidean distance between every two cities
*	- Initialize the tau matrix (pheromone levels) for each edge
*	- initialize each ant and provide them a random city to start
*	- Each ant traveses the path and chooses each city using a probability which
*	  is computed using a formula
*	- Update the tau matrix (pheromone levels) after one tour has end
*	- Run 500 tours 10 times to obtain average and best tour length values
 */
package main

import (
	"bufio"
	"fmt"
	"math"
	"math/rand"
	"os"
	"strconv"
	"strings"

	"github.com/gonum/plot"
	"github.com/gonum/plot/plotter"
	"github.com/gonum/plot/plotutil"
	"github.com/gonum/plot/vg"
)

//ant implementation
type ant_t struct {
	tabulist    []int
	currentCity int
	nextCity    int
	tour        []int
	tourIndex   int
	tourlength  float64
}

//city implementation
type city struct {
	x, y float64
}

//declaration and initialization
var ants []ant_t
var cities []city
var tauMatrix [][]float64
var adjMatrix [][]float64
var besttour []int
var currentIndex int
var numCities = 0
var rho = 0.6
var qval = 1.0
var alpha = 0.8
var beta = 0.8
var numAnts = 10

//prints error and exits on abnormal conditions
func printError(err error) {
	fmt.Print(err)
	os.Exit(2)
}

func main() {
	//comment other to use the other one
	//initGraph("TSP_D")
	initGraph("TSP_WS") //other TSP file

	avgavg := make([]float64, 10)
	avgbest := make([]float64, 10)
	besttourlens := make([]float64, 10)
	avgtourlens := make([]float64, 10)
	iterations := 0
	i := 0
	//initializes pheromone levels with base pheromone i.e. 1
	initTrail()
	for i < 10 {
		iterations = 0
		besttour = nil
		for iterations < 500 {
			//initializes each ant
			initAnts()
			//tour for each ant
			moveAnts()
			//intensify pheromone levels
			intensifyTrail()
			//compute best and avg tour length for every tour in
			//500 and add them for average
			besttourlens[i] += calculateBest()
			avgtourlens[i] += calculateAvg()
			iterations++
		}
		//evaporate pheromone to obtain better results
		evaporatePheromone()
		fmt.Println("Iteration: ", i)
		fmt.Println("Optimal Path: ", besttour)
		avgavg[i] = avgtourlens[i] / 500.0
		avgbest[i] = besttourlens[i] / 500.0
		i++
	}
	fmt.Println("Average Average:", avgavg)
	fmt.Println("Average Best:", avgbest)

	p, err := plot.New()
	if err != nil {
		printError(err)
	}

	//comment other to use the other one
	//p.Title.Text = "Dijibouti TSP"
	p.Title.Text = "Western Sahara TSP"
	p.X.Label.Text = "X"
	p.Y.Label.Text = "Y"

	avgpts := make(plotter.XYs, 10)
	for i := range avgpts {
		avgpts[i].Y = avgavg[i]
		avgpts[i].X = float64(i)
	}

	bestpts := make(plotter.XYs, 10)
	for i := range bestpts {
		bestpts[i].Y = avgbest[i]
		bestpts[i].X = float64(i)
	}
	err = plotutil.AddLinePoints(p,
		"Average So Far", avgpts,
		"Best So Far", bestpts)
	if err != nil {
		printError(err)
	}

	// Save the plot to a PNG file.
	/*if err := p.Save(4*vg.Inch, 4*vg.Inch, "djibouti.png"); err != nil {
		printError(err)
	}*/
	if err := p.Save(4*vg.Inch, 4*vg.Inch, "westernsahara.png"); err != nil {
		printError(err)
	}
}

//calculate average tour length of all length for one tour
func calculateAvg() float64 {
	avglength := ants[0].tourlength
	total := 0.0
	for i := range ants {
		total += ants[i].tourlength
	}
	avglength = total / 10.0
	return avglength
}

//calculate best tour length of all length for one tour
func calculateBest() float64 {
	var bestlength float64
	bestlength = ants[0].tourlength
	if besttour == nil {
		besttour = make([]int, numCities)
		copy(besttour, ants[0].tour)
	}
	for i := range ants {
		if ants[i].tourlength < bestlength {

			bestlength = ants[i].tourlength
			copy(besttour, ants[i].tour)
		}
	}
	return bestlength
}

//initialize pheromone levels
func initTrail() {
	tauMatrix = make([][]float64, numCities)
	for i := range tauMatrix {
		tauMatrix[i] = make([]float64, numCities)
		for j := range tauMatrix[i] {
			if i != j {
				tauMatrix[i][j] = 1.0 //initialize to base pheromone = 1
			} else {
				tauMatrix[i][j] = 0.0
			}
		}
	}
}

//initialize ants
func initAnts() {
	ants = nil
	ants = make([]ant_t, numAnts)
	for i := range ants {
		ants[i].tabulist = make([]int, numCities)
		ants[i].currentCity = rand.Intn(numCities) //randomly assigns ant a city to s
		ants[i].nextCity = 0
		ants[i].tour = make([]int, numCities)
		ants[i].tourIndex = 0
		ants[i].tourlength = 0.0
	}
}

//move all ants to visit the whole graph
func moveAnts() {
	for i := range ants {
		currentIndex = 0
		for currentIndex < numCities {
			goToNewCity(&ants[i])
			currentIndex++
		}
	}
}

//choosing next city
func goToNewCity(ant *ant_t) {
	var from, to int
	var p float64
	denom := 0.0
	from = ant.currentCity
	for to = 0; to < numCities; to++ {
		if from != to {
			if ant.tabulist[to] == 0 && tauMatrix[from][to] != 0 && adjMatrix[from][to] != 0 {
				denom += math.Pow(tauMatrix[from][to], alpha) * math.Pow((1.0/adjMatrix[from][to]), beta)
			}
		} else {
			continue
		}

	}
	to = 0
	for {
		if from != to {
			if ant.tabulist[to] == 0 {
				p = (math.Pow(tauMatrix[from][to], alpha) * math.Pow((1.0/adjMatrix[from][to]), beta)) / denom

				if rand.Float64() < p {
					break
				}
			}
		} else {
			to = ((to + 1) % numCities)
			continue
		}
		to = ((to + 1) % numCities)
	}
	ant.nextCity = to
	ant.tabulist[ant.nextCity] = 1
	ant.tour[ant.tourIndex] = ant.nextCity
	ant.tourIndex++
	ant.tourlength += adjMatrix[ant.currentCity][ant.nextCity]
	if ant.tourIndex == numCities {
		ant.tourlength += adjMatrix[ant.tour[numCities-1]][ant.tour[0]]
	}
	ant.currentCity = ant.nextCity
}

//reads from file and creates a list of all the cities coordinates
func readFile(name string) []city {
	var dim, i int
	var cities []city
	i, dim = 1, 0
	var startFlag bool
	startFlag = false
	if file, err := os.Open(name); err == nil {
		// make sure it gets closed
		defer file.Close()
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			str := scanner.Text()
			if strings.Contains(str, "DIMENSION") {
				dim = getDim(str)
				break
			}
		}
		cities = make([]city, dim)
		for scanner.Scan() {
			str := scanner.Text()
			if strings.Contains(str, "EOF") {
				break
			} else if startFlag {
				x, y := tokenize(str)
				if i <= dim {
					cities[i-1] = city{x, y}
					i++
				} else {
					startFlag = false
				}
			} else if strings.Contains(str, "NODE_COORD_SECTION") {
				startFlag = true
			}
		}
		// check for errors
		if err = scanner.Err(); err != nil {
			printError(err)
		}
	} else {
		printError(err)
	}
	return cities
}

//evaporating pheromone after each iteration of the algorithm
func evaporatePheromone() {
	var from, to int
	for from = 0; from < numCities; from++ {
		for to = 0; to < numCities; to++ {
			tauMatrix[from][to] = tauMatrix[from][to] * (1.0 - rho)
			if tauMatrix[from][to] < 0.0 {
				tauMatrix[from][to] = 1.0
			}
		}
	}
}

//intensifying pheromone levels
func intensifyTrail() {
	var from, to, i, c int
	for i = 0; i < numAnts; i++ {
		for c = 0; c < numCities; c++ {
			from = ants[i].tour[c]
			to = ants[i].tour[((c + 1) % numCities)]
			deltatau := (qval / ants[i].tourlength)
			tauMatrix[from][to] = tauMatrix[from][to] + deltatau
			tauMatrix[to][from] = tauMatrix[from][to]
		}
	}
}

//making graph
func initGraph(name string) {
	cities = readFile(name)
	numCities = len(cities)
	adjMatrix = make([][]float64, numCities)
	for i := range adjMatrix {
		adjMatrix[i] = make([]float64, numCities)
		for j := range adjMatrix[i] {
			adjMatrix[i][j] = calEdge(cities[i], cities[j])
		}
	}
}

//calculates edge weight (euclidiean distance)
func calEdge(c1, c2 city) float64 {
	return math.Pow((math.Pow((c2.x-c1.x), 2) + math.Pow((c2.y-c1.y), 2)), 0.5)
}

//tokenizes and converts to float
func tokenize(str string) (x, y float64) {
	s := strings.Split(str, " ")
	strX, strY := s[1], s[2]
	x, err := strconv.ParseFloat(strX, 64) //converts string to float64
	if err != nil {
		printError(err)
	}
	y, err = strconv.ParseFloat(strY, 64) //converts string to float64
	if err != nil {
		printError(err)
	}
	return x, y
}

//gets number of cities from the file
func getDim(str string) (dim int) {
	s := strings.Split(str, ":")
	num := strings.TrimLeft(s[1], " ")
	if dim, err := strconv.Atoi(num); err == nil {
		return dim
	} else {
		fmt.Print(err)
		os.Exit(2)
	}
	return 0
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"net"
	"strings"
	"time"
)

type User struct {
	nick       string
	user       string
	ident      string
	dead       bool
	nickset    bool
	waiting    bool
	connection net.Conn
	id         int
	realname   string
	userset    bool
	registered bool
	ip         string
	host       string
	epoch      time.Time
	lastrcv    time.Time
	nextcheck  time.Time
	chanlist   map[string]*Channel
}

func (user *User) PingChecker() {
	for {
		if user.dead {
			break
		}
		if time.Now().After(user.nextcheck) {
			if user.waiting {
				since := time.Since(user.lastrcv).Seconds()
				user.Quit(fmt.Sprintf("Ping Timeout: %.0f seconds", since))
				break
			} else {
				user.SendLine(fmt.Sprintf("PING :%s", sname))
				user.waiting = true
				user.nextcheck.Add(ping_time * time.Second)
			}
		}
		time.Sleep(ping_check_time * time.Second)
	}
}

func (user *User) QuitCommandHandler(args []string) {
	var reason string
	if len(args) > 1 {
		args[1] = StripLeading(args[1], ":")
		var buffer bytes.Buffer
		for i := 1; i < len(args); i++ {
			buffer.WriteString(args[i])
			buffer.WriteString(" ")
		}
		reason = strings.TrimSpace(buffer.String())
	} else {
		reason = "Leaving"
	}
	user.Quit(reason)
}

func (user *User) Quit(reason string) {
	targets := []*User{user}
	for _, k := range user.chanlist {
		targets = append(targets, k.GetUserList()...)
		delete(k.userlist, user.id)
		delete(user.chanlist, k.name)
		k.ShouldIDie()
	}
	SendToMany(fmt.Sprintf(":%s QUIT :%s", user.GetHostMask(), reason), targets)
	user.SendLine(fmt.Sprintf("ERROR :Closing Link: %s (%s)", user.host, reason))
	user.dead = true
	if user.connection != nil {
		user.connection.Close()
	}
	delete(userlist, user.id)
}

func (user *User) FireNumeric(numeric int, args ...interface{}) {
	msg := strcat(fmt.Sprintf(":%s %.3d %s ", sname, numeric, user.nick), fmt.Sprintf(NUM[numeric], args...))
	user.SendLine(msg)
}

func NewUser(conn net.Conn) *User {
	userip := GetIpFromConn(conn)
	log.Println("New connection from", userip)
	counter = counter + 1
	user := &User{id: counter, connection: conn, ip: userip, nick: "*"}
	user.chanlist = make(map[string]*Channel)
	user.host = user.ip
	user.epoch = time.Now()
	user.lastrcv = time.Now()
	user.nextcheck = time.Now().Add(ping_time * time.Second)
	userlist[user.id] = user
	go user.UserHostLookup()
	go user.PingChecker()
	return user
}

func (user *User) SendLine(msg string) {
	msg = fmt.Sprintf("%s\n", msg)
	if user.dead {
		return
	}
	_, err := user.connection.Write([]byte(msg))
	if err != nil {
		user.dead = true
		user.Quit("Error")
		log.Printf("Error sending message to %s, disconnecting\n", user.nick)
	}
	log.Printf("Send to %s: %s", user.nick, msg)
}

func (user *User) HandleRequests() {
	b := bufio.NewReader(user.connection)
	for {
		if user.dead {
			break
		}
		line, err := b.ReadString('\n')
		if err != nil {
			log.Println("Error reading:", err.Error())
			user.dead = true
			user.Quit("Error")
		}
		if line == "" {
			user.dead = true
			user.Quit("Error")
			break
		}
		line = strings.TrimSpace(line)
		log.Println("Receive from", fmt.Sprintf("%s:", user.nick), line)
		go ProcessLine(user, line)
	}
}
func (user *User) NickHandler(args []string) {
	if len(args) < 2 {
		user.FireNumeric(ERR_NONICKNAMEGIVEN)
		return
	}
	if NickHasBadChars(args[1]) {
		user.FireNumeric(ERR_ERRONEOUSNICKNAME, args[1])
		return
	}
	if GetUserByNick(args[1]) != nil {
		user.FireNumeric(ERR_NICKNAMEINUSE, args[1])
		return
	}
	if !user.nickset {
		user.nickset = true
	} else if user.registered {
		targets := []*User{}
		targets = append(targets, user)
		for _, k := range user.chanlist {
			targets = append(targets, k.GetUserList()...)
		}
		SendToMany(fmt.Sprintf(":%s NICK %s", user.GetHostMask(), args[1]), targets)
	}
	user.nick = args[1]
	if !user.registered && user.userset {
		user.UserRegistrationFinished()
	}
}

func (user *User) UserHandler(args []string) {
	if len(args) < 5 {
		user.FireNumeric(ERR_NEEDMOREPARAMS, "USER")
		return
	}
	user.ident = args[1]
	args[4] = StripLeading(args[4], ":")
	var buffer bytes.Buffer
	for i := 4; i < len(args); i++ {
		buffer.WriteString(args[i])
		buffer.WriteString(" ")
	}
	user.realname = strings.TrimSpace(buffer.String())
	user.userset = true
	if !user.registered && user.nickset {
		user.UserRegistrationFinished()
	}
}

func (user *User) UserRegistrationFinished() {
	user.registered = true
	log.Printf("User %d finished registration\n", user.id)
	user.FireNumeric(RPL_WELCOME, user.nick, user.ident, user.host)
	user.FireNumeric(RPL_YOURHOST, sname, software, softwarev)
	user.FireNumeric(RPL_CREATED, epoch)
	//TODO fire RPL_MYINFO when we actually have enough stuff to do it
	user.FireLusers()
}

func (user *User) UserHostLookup() {
	user.SendLine(fmt.Sprintf(":%s NOTICE %s :*** Looking up your hostname...", sname, user.nick))
	adds, err := net.LookupAddr(user.ip)
	if err != nil {
		user.SendLine(fmt.Sprintf("%s NOTICE %s :*** Unable to resolve your hostname", sname, user.nick))
		return
	}
	addstring := adds[0]
	adds, err = net.LookupHost(addstring)
	if err != nil {
		user.SendLine(fmt.Sprintf("%s NOTICE %s :*** Unable to resolve your hostname", sname, user.nick))
		return
	}
	for _, k := range adds {
		if user.ip == k {
			user.host = addstring
			user.SendLine(fmt.Sprintf(":%s NOTICE %s :*** Found your hostname", sname, user.nick))
			return
		}
	}
	user.SendLine(fmt.Sprintf(":%s NOTICE %s :*** Your forward and reverse DNS do not match, ignoring hostname", sname, user.nick))
}

func (user *User) CommandNotFound(args []string) {
	user.FireNumeric(ERR_UNKNOWNCOMMAND, args[0])
}

func (user *User) GetHostMask() string {
	return fmt.Sprintf("%s!%s@%s", user.nick, user.ident, user.host)
}

func (user *User) JoinHandler(args []string) {
	if len(args) < 2 {
		user.FireNumeric(ERR_NEEDMOREPARAMS, "JOIN")
		return
	}
	if !ValidChanName(args[1]) {
		user.FireNumeric(ERR_NOSUCHCHANNEL, args[1])
		return
	}
	_, channel := GetChannelByName(args[1])
	channel.JoinUser(user)
	user.chanlist[channel.name] = channel
}

func (user *User) FireLusers() {
	user.FireNumeric(RPL_LUSERCLIENT, len(userlist), 0, 1) //0 services and 1 servers for now
	user.FireNumeric(RPL_LUSEROP, 0)                       //also 0 for now
	user.FireNumeric(RPL_LUSERUNKNOWN, 0)                  //also 0...
	user.FireNumeric(RPL_LUSERCHANNELS, len(chanlist))
	user.FireNumeric(RPL_LUSERME, len(userlist), 1)
}

func (user *User) PrivmsgHandler(args []string) {
	if len(args) < 3 {
		user.FireNumeric(ERR_NEEDMOREPARAMS, "PRIVMSG")
		return
	}
	if ValidChanName(args[1]) { //TODO part of this should be sent to the channel "object"
		//presumably a channel
		k, j := GetChannelByName(args[1])
		if k {
			//channel exists, send the message
			msg := FormatMessageArgs(args)
			list := j.GetUserList()
			for _, l := range list {
				if l != user {
					l.SendLine(fmt.Sprintf(":%s PRIVMSG %s :%s", user.GetHostMask(), j.name, msg))
				}
			}
			return
		} else {
			//channel didnt exist but get channel by name makes one anyways, lets kill it...
			user.FireNumeric(ERR_NOSUCHCHANNEL, args[1])
			j.ShouldIDie()
			return
		}
	} else {
		//maybe its a user
		target := GetUserByNick(args[1])
		if target != nil {
			msg := FormatMessageArgs(args)
			target.SendLine(fmt.Sprint(":%s PRIVMSG %s :%s", user.GetHostMask(), target.nick, msg))
		}
	}
}
<commit_msg>remove that go routine, race detector complains :[<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"net"
	"strings"
	"time"
)

type User struct {
	nick       string
	user       string
	ident      string
	dead       bool
	nickset    bool
	waiting    bool
	connection net.Conn
	id         int
	realname   string
	userset    bool
	registered bool
	ip         string
	host       string
	epoch      time.Time
	lastrcv    time.Time
	nextcheck  time.Time
	chanlist   map[string]*Channel
}

func (user *User) PingChecker() {
	for {
		if user.dead {
			break
		}
		if time.Now().After(user.nextcheck) {
			if user.waiting {
				since := time.Since(user.lastrcv).Seconds()
				user.Quit(fmt.Sprintf("Ping Timeout: %.0f seconds", since))
				break
			} else {
				user.SendLine(fmt.Sprintf("PING :%s", sname))
				user.waiting = true
				user.nextcheck.Add(ping_time * time.Second)
			}
		}
		time.Sleep(ping_check_time * time.Second)
	}
}

func (user *User) QuitCommandHandler(args []string) {
	var reason string
	if len(args) > 1 {
		args[1] = StripLeading(args[1], ":")
		var buffer bytes.Buffer
		for i := 1; i < len(args); i++ {
			buffer.WriteString(args[i])
			buffer.WriteString(" ")
		}
		reason = strings.TrimSpace(buffer.String())
	} else {
		reason = "Leaving"
	}
	user.Quit(reason)
}

func (user *User) Quit(reason string) {
	targets := []*User{user}
	for _, k := range user.chanlist {
		targets = append(targets, k.GetUserList()...)
		delete(k.userlist, user.id)
		delete(user.chanlist, k.name)
		k.ShouldIDie()
	}
	SendToMany(fmt.Sprintf(":%s QUIT :%s", user.GetHostMask(), reason), targets)
	user.SendLine(fmt.Sprintf("ERROR :Closing Link: %s (%s)", user.host, reason))
	user.dead = true
	if user.connection != nil {
		user.connection.Close()
	}
	delete(userlist, user.id)
}

func (user *User) FireNumeric(numeric int, args ...interface{}) {
	msg := strcat(fmt.Sprintf(":%s %.3d %s ", sname, numeric, user.nick), fmt.Sprintf(NUM[numeric], args...))
	user.SendLine(msg)
}

func NewUser(conn net.Conn) *User {
	userip := GetIpFromConn(conn)
	log.Println("New connection from", userip)
	counter = counter + 1
	user := &User{id: counter, connection: conn, ip: userip, nick: "*"}
	user.chanlist = make(map[string]*Channel)
	user.host = user.ip
	user.epoch = time.Now()
	user.lastrcv = time.Now()
	user.nextcheck = time.Now().Add(ping_time * time.Second)
	userlist[user.id] = user
	go user.UserHostLookup()
	go user.PingChecker()
	return user
}

func (user *User) SendLine(msg string) {
	msg = fmt.Sprintf("%s\n", msg)
	if user.dead {
		return
	}
	_, err := user.connection.Write([]byte(msg))
	if err != nil {
		user.dead = true
		user.Quit("Error")
		log.Printf("Error sending message to %s, disconnecting\n", user.nick)
	}
	log.Printf("Send to %s: %s", user.nick, msg)
}

func (user *User) HandleRequests() {
	b := bufio.NewReader(user.connection)
	for {
		if user.dead {
			break
		}
		line, err := b.ReadString('\n')
		if err != nil {
			log.Println("Error reading:", err.Error())
			user.dead = true
			user.Quit("Error")
		}
		if line == "" {
			user.dead = true
			user.Quit("Error")
			break
		}
		line = strings.TrimSpace(line)
		log.Println("Receive from", fmt.Sprintf("%s:", user.nick), line)
		ProcessLine(user, line)
	}
}
func (user *User) NickHandler(args []string) {
	if len(args) < 2 {
		user.FireNumeric(ERR_NONICKNAMEGIVEN)
		return
	}
	if NickHasBadChars(args[1]) {
		user.FireNumeric(ERR_ERRONEOUSNICKNAME, args[1])
		return
	}
	if GetUserByNick(args[1]) != nil {
		user.FireNumeric(ERR_NICKNAMEINUSE, args[1])
		return
	}
	if !user.nickset {
		user.nickset = true
	} else if user.registered {
		targets := []*User{}
		targets = append(targets, user)
		for _, k := range user.chanlist {
			targets = append(targets, k.GetUserList()...)
		}
		SendToMany(fmt.Sprintf(":%s NICK %s", user.GetHostMask(), args[1]), targets)
	}
	user.nick = args[1]
	if !user.registered && user.userset {
		user.UserRegistrationFinished()
	}
}

func (user *User) UserHandler(args []string) {
	if len(args) < 5 {
		user.FireNumeric(ERR_NEEDMOREPARAMS, "USER")
		return
	}
	user.ident = args[1]
	args[4] = StripLeading(args[4], ":")
	var buffer bytes.Buffer
	for i := 4; i < len(args); i++ {
		buffer.WriteString(args[i])
		buffer.WriteString(" ")
	}
	user.realname = strings.TrimSpace(buffer.String())
	user.userset = true
	if !user.registered && user.nickset {
		user.UserRegistrationFinished()
	}
}

func (user *User) UserRegistrationFinished() {
	user.registered = true
	log.Printf("User %d finished registration\n", user.id)
	user.FireNumeric(RPL_WELCOME, user.nick, user.ident, user.host)
	user.FireNumeric(RPL_YOURHOST, sname, software, softwarev)
	user.FireNumeric(RPL_CREATED, epoch)
	//TODO fire RPL_MYINFO when we actually have enough stuff to do it
	user.FireLusers()
}

func (user *User) UserHostLookup() {
	user.SendLine(fmt.Sprintf(":%s NOTICE %s :*** Looking up your hostname...", sname, user.nick))
	adds, err := net.LookupAddr(user.ip)
	if err != nil {
		user.SendLine(fmt.Sprintf("%s NOTICE %s :*** Unable to resolve your hostname", sname, user.nick))
		return
	}
	addstring := adds[0]
	adds, err = net.LookupHost(addstring)
	if err != nil {
		user.SendLine(fmt.Sprintf("%s NOTICE %s :*** Unable to resolve your hostname", sname, user.nick))
		return
	}
	for _, k := range adds {
		if user.ip == k {
			user.host = addstring
			user.SendLine(fmt.Sprintf(":%s NOTICE %s :*** Found your hostname", sname, user.nick))
			return
		}
	}
	user.SendLine(fmt.Sprintf(":%s NOTICE %s :*** Your forward and reverse DNS do not match, ignoring hostname", sname, user.nick))
}

func (user *User) CommandNotFound(args []string) {
	user.FireNumeric(ERR_UNKNOWNCOMMAND, args[0])
}

func (user *User) GetHostMask() string {
	return fmt.Sprintf("%s!%s@%s", user.nick, user.ident, user.host)
}

func (user *User) JoinHandler(args []string) {
	if len(args) < 2 {
		user.FireNumeric(ERR_NEEDMOREPARAMS, "JOIN")
		return
	}
	if !ValidChanName(args[1]) {
		user.FireNumeric(ERR_NOSUCHCHANNEL, args[1])
		return
	}
	_, channel := GetChannelByName(args[1])
	channel.JoinUser(user)
	user.chanlist[channel.name] = channel
}

func (user *User) FireLusers() {
	user.FireNumeric(RPL_LUSERCLIENT, len(userlist), 0, 1) //0 services and 1 servers for now
	user.FireNumeric(RPL_LUSEROP, 0)                       //also 0 for now
	user.FireNumeric(RPL_LUSERUNKNOWN, 0)                  //also 0...
	user.FireNumeric(RPL_LUSERCHANNELS, len(chanlist))
	user.FireNumeric(RPL_LUSERME, len(userlist), 1)
}

func (user *User) PrivmsgHandler(args []string) {
	if len(args) < 3 {
		user.FireNumeric(ERR_NEEDMOREPARAMS, "PRIVMSG")
		return
	}
	if ValidChanName(args[1]) { //TODO part of this should be sent to the channel "object"
		//presumably a channel
		k, j := GetChannelByName(args[1])
		if k {
			//channel exists, send the message
			msg := FormatMessageArgs(args)
			list := j.GetUserList()
			for _, l := range list {
				if l != user {
					l.SendLine(fmt.Sprintf(":%s PRIVMSG %s :%s", user.GetHostMask(), j.name, msg))
				}
			}
			return
		} else {
			//channel didnt exist but get channel by name makes one anyways, lets kill it...
			user.FireNumeric(ERR_NOSUCHCHANNEL, args[1])
			j.ShouldIDie()
			return
		}
	} else {
		//maybe its a user
		target := GetUserByNick(args[1])
		if target != nil {
			msg := FormatMessageArgs(args)
			target.SendLine(fmt.Sprint(":%s PRIVMSG %s :%s", user.GetHostMask(), target.nick, msg))
		}
	}
}
<|endoftext|>


<commit_before>package rglog

import (
	"github.com/realglobe-Inc/go-lib-rg/rglog/handler"
	"github.com/realglobe-Inc/go-lib-rg/rglog/level"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestLog(t *testing.T) {
	rootLabel := "github.com/realglobe-Inc/go-lib-rg"

	loop := 100

	rootLog := GetLogger(rootLabel)
	rootLog.SetLevel(level.DEBUG)
	rootLog.SetUseParent(false)

	hndl := handler.NewConsoleHandler()
	hndl.SetLevel(level.INFO)
	rootLog.AddHandler(hndl)

	path := filepath.Join(os.TempDir(), "log_test.go.log")
	if e := os.Remove(path); e != nil {
		if !os.IsNotExist(e) {
			t.Fatal(e)
		}
	}
	hndl, err := handler.NewFileHandler(path)
	if err != nil {
		t.Fatal(err)
	}
	hndl.SetLevel(level.DEBUG)
	rootLog.AddHandler(hndl)

	start := time.Now()

	for i := 0; i < loop; i++ {
		GetLogger(rootLabel + "/" + strconv.Itoa(i)).Info(i)
	}

	// 遅過ぎ検知。
	// 1 回 100 マイクロ秒も掛かってるのは遅い。
	limit := start.Add(time.Duration(int64(loop*100) * int64(time.Microsecond)))
	if time.Now().After(limit) {
		t.Error("Too slow")
	}

	Flush()

	// ファイルに書き込めているかどうか検査。
	buff, err := ioutil.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}

	if len(buff) > 0 && buff[len(buff)-1] == '\n' { // 最後の空行は抜かしとく。
		buff = buff[:len(buff)-1]
	}

	lines := strings.Split(string(buff), "\n")
	if len(lines) != loop {
		t.Error(len(lines), loop)
	}

}

func TestConcurrent(t *testing.T) {
	rootLabel := "github.com/realglobe-Inc/go-lib-rg"

	n := 100
	loop := 1000

	rootLog := GetLogger(rootLabel)
	rootLog.SetLevel(level.DEBUG)
	rootLog.SetUseParent(false)

	hndl := handler.NewConsoleHandler()
	hndl.SetLevel(level.INFO)
	rootLog.AddHandler(hndl)

	path := filepath.Join(os.TempDir(), "log_test.go.log")
	if e := os.Remove(path); e != nil {
		if !os.IsNotExist(e) {
			t.Fatal(e)
		}
	}
	hndl, err := handler.NewFileHandler(path)
	if err != nil {
		t.Fatal(err)
	}
	hndl.SetLevel(level.DEBUG)
	rootLog.AddHandler(hndl)

	var lock sync.Mutex
	end := false

	start := time.Now()

	c := make(chan bool)

	for i := 0; i < n; i++ {
		id := i
		go func() {
			for j := 0; j < loop; j++ {
				GetLogger(rootLabel+"/"+strconv.Itoa(id)).Info(id, j)
			}

			c <- true
		}()
	}

	go func() {
		for i := 0; i < n; i++ {
			<-c
		}

		lock.Lock()
		end = true
		lock.Unlock()
	}()

	// 遅過ぎ検知。
	// 1 回 100 マイクロ秒も掛かってるのは遅い。
	limit := start.Add(time.Duration(int64(n*loop*100) * int64(time.Microsecond)))
	for time.Now().Before(limit) {
		lock.Lock()
		flag := end
		lock.Unlock()

		if flag {
			break
		}

		time.Sleep(time.Millisecond)
	}

	if !end {
		t.Fatal("Too slow")
	}

	Flush()

	// ファイルに書き込めているかどうか検査。
	buff, err := ioutil.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}

	if len(buff) > 0 && buff[len(buff)-1] == '\n' { // 最後の空行は抜かしとく。
		buff = buff[:len(buff)-1]
	}

	lines := strings.Split(string(buff), "\n")
	if len(lines) != n*loop {
		t.Error(len(lines), n*loop)
	}

}
<commit_msg>繰り返し回数の修正等<commit_after>package rglog

import (
	"github.com/realglobe-Inc/go-lib-rg/rglog/handler"
	"github.com/realglobe-Inc/go-lib-rg/rglog/level"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestLog(t *testing.T) {
	rootLabel := "github.com/realglobe-Inc/go-lib-rg"

	loop := 100000
	n := 100

	rootLog := GetLogger(rootLabel)
	rootLog.SetLevel(level.DEBUG)
	rootLog.SetUseParent(false)

	hndl := handler.NewConsoleHandler()
	hndl.SetLevel(level.INFO)
	rootLog.AddHandler(hndl)

	path := filepath.Join(os.TempDir(), "log_test.go.log")
	if e := os.Remove(path); e != nil {
		if !os.IsNotExist(e) {
			t.Fatal(e)
		}
	}
	hndl, err := handler.NewFileHandler(path)
	if err != nil {
		t.Fatal(err)
	}
	hndl.SetLevel(level.DEBUG)
	rootLog.AddHandler(hndl)

	start := time.Now()
	for i := 0; i < loop; i++ {
		GetLogger(rootLabel + "/" + strconv.Itoa(i%n)).Info(i)
	}
	end := time.Now()

	// 遅過ぎ検知。
	// 1 回 100 マイクロ秒も掛かってるのは遅い。
	limit := start.Add(time.Duration(int64(loop*100) * int64(time.Microsecond)))
	if end.After(limit) {
		t.Error("Too slow ", end.Sub(start))
	} else {
		//t.Error(end.Sub(start))
	}

	Flush()

	// ファイルに書き込めているかどうか検査。
	buff, err := ioutil.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}

	if len(buff) > 0 && buff[len(buff)-1] == '\n' { // 最後の空行は抜かしとく。
		buff = buff[:len(buff)-1]
	}

	lines := strings.Split(string(buff), "\n")
	if len(lines) != loop {
		t.Error(len(lines), loop)
	}

}

func TestConcurrent(t *testing.T) {
	rootLabel := "github.com/realglobe-Inc/go-lib-rg"

	n := 100
	loop := 1000

	rootLog := GetLogger(rootLabel)
	rootLog.SetLevel(level.DEBUG)
	rootLog.SetUseParent(false)

	hndl := handler.NewConsoleHandler()
	hndl.SetLevel(level.INFO)
	rootLog.AddHandler(hndl)

	path := filepath.Join(os.TempDir(), "log_test.go.log")
	if e := os.Remove(path); e != nil {
		if !os.IsNotExist(e) {
			t.Fatal(e)
		}
	}
	hndl, err := handler.NewFileHandler(path)
	if err != nil {
		t.Fatal(err)
	}
	hndl.SetLevel(level.DEBUG)
	rootLog.AddHandler(hndl)

	var lock sync.Mutex

	var zero time.Time
	start := time.Now()
	end := zero

	c := make(chan bool)

	for i := 0; i < n; i++ {
		id := i
		go func() {
			for j := 0; j < loop; j++ {
				GetLogger(rootLabel+"/"+strconv.Itoa(id)).Info(id, j)
			}

			c <- true
		}()
	}

	go func() {
		for i := 0; i < n; i++ {
			<-c
		}

		lock.Lock()
		end = time.Now()
		lock.Unlock()
	}()

	// 遅過ぎ検知。
	// 1 回 100 マイクロ秒も掛かってるのは遅い。
	limit := start.Add(time.Duration(int64(n*loop*100) * int64(time.Microsecond)))
	for time.Now().Before(limit) {
		lock.Lock()
		curEnd := end
		lock.Unlock()

		if curEnd != zero {
			break
		}

		time.Sleep(time.Millisecond)
	}

	if end == zero {
		t.Fatal("Too slow ", time.Now().Sub(start))
	} else {
		//t.Error(end.Sub(start))
	}

	Flush()

	// ファイルに書き込めているかどうか検査。
	buff, err := ioutil.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}

	if len(buff) > 0 && buff[len(buff)-1] == '\n' { // 最後の空行は抜かしとく。
		buff = buff[:len(buff)-1]
	}

	lines := strings.Split(string(buff), "\n")
	if len(lines) != n*loop {
		t.Error(len(lines), n*loop)
	}

}
<|endoftext|>


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

package main

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/appc/spec/schema/types"
	"github.com/coreos/rkt/common/apps"
	"github.com/coreos/rkt/rkt/config"
	"github.com/coreos/rkt/rkt/image"
	"github.com/coreos/rkt/store"
	"github.com/hashicorp/errwrap"
	"github.com/spf13/pflag"
)

// stage1ImageLocationKind describes the stage1 image location
type stage1ImageLocationKind int

const (
	// location unset, it is not a valid kind to be used
	stage1ImageLocationUnset stage1ImageLocationKind = iota
	// a URL with a scheme
	stage1ImageLocationURL
	// an absolute or a relative path
	stage1ImageLocationPath
	// an image name
	stage1ImageLocationName
	// an image hash
	stage1ImageLocationHash
	// an image in the default dir
	stage1ImageLocationFromDir
)

// stage1FlagData is used for creating the flags for each valid location kind
type stage1FlagData struct {
	kind stage1ImageLocationKind
	flag string
	name string
	help string
}

// stage1ImageLocation is used to store the user's choice of stage1 image via flags
type stage1ImageLocation struct {
	kind     stage1ImageLocationKind
	location string
}

// stage1ImageLocationFlag is an implementation of a pflag.Value
// interface, which handles all the valid location kinds
type stage1ImageLocationFlag struct {
	loc  *stage1ImageLocation
	kind stage1ImageLocationKind
}

func (f *stage1ImageLocationFlag) Set(location string) error {
	if f.loc.kind != stage1ImageLocationUnset {
		wanted := stage1FlagsData[f.kind]
		current := stage1FlagsData[f.loc.kind]
		if f.loc.kind == f.kind {
			return fmt.Errorf("--%s already used", current.flag)
		}
		return fmt.Errorf("flags --%s and --%s are mutually exclusive",
			wanted.flag, current.flag)
	}
	f.loc.kind = f.kind
	f.loc.location = location
	return nil
}

func (f *stage1ImageLocationFlag) String() string {
	return f.loc.location
}

func (f *stage1ImageLocationFlag) Type() string {
	return stage1FlagsData[f.kind].name
}

var (
	// defaults defined by configure, set by linker
	// default stage1 image name
	// (e.g. coreos.com/rkt/stage1-coreos)
	buildDefaultStage1Name string
	// default stage1 image version (e.g. 0.15.0)
	buildDefaultStage1Version string
	// an absolute path or a URL to the default stage1 image file
	buildDefaultStage1ImageLoc string
	// filename of the default stage1 image file in the default
	// stage1 images directory
	buildDefaultStage1ImageInRktDir string
	// an absolute path to the stage1 images directory
	buildDefaultStage1ImagesDir string

	// this holds necessary data to generate the --stage1-* flags
	// for each location kind
	stage1FlagsData = map[stage1ImageLocationKind]*stage1FlagData{
		stage1ImageLocationURL: &stage1FlagData{
			kind: stage1ImageLocationURL,
			flag: "stage1-url",
			name: "stage1URL",
			help: "a URL to an image to use as stage1",
		},

		stage1ImageLocationPath: &stage1FlagData{
			kind: stage1ImageLocationPath,
			flag: "stage1-path",
			name: "stage1Path",
			help: "an absolute or a relative path to an image to use as stage1",
		},

		stage1ImageLocationName: &stage1FlagData{
			kind: stage1ImageLocationName,
			flag: "stage1-name",
			name: "stage1Name",
			help: "a name of an image to use as stage1",
		},

		stage1ImageLocationHash: &stage1FlagData{
			kind: stage1ImageLocationHash,
			flag: "stage1-hash",
			name: "stage1Hash",
			help: "a hash of an image to use as stage1",
		},

		stage1ImageLocationFromDir: &stage1FlagData{
			kind: stage1ImageLocationFromDir,
			flag: "stage1-from-dir",
			name: "stage1FromDir",
			help: "a filename of an image in stage1 images directory to use as stage1",
		},
	}
	// location to stage1 image overridden by one of --stage1-*
	// flags
	overriddenStage1Location = stage1ImageLocation{
		kind:     stage1ImageLocationUnset,
		location: "",
	}
)

// addStage1ImageFlags adds flags for specifying custom stage1 image
func addStage1ImageFlags(flags *pflag.FlagSet) {
	for _, data := range stage1FlagsData {
		wrapper := &stage1ImageLocationFlag{
			loc:  &overriddenStage1Location,
			kind: data.kind,
		}
		flags.Var(wrapper, data.flag, data.help)
	}
}

// getStage1Hash will try to get the hash of stage1 to use.
//
// Before getting inside this rats nest, let's try to write up the
// expected behaviour.
//
// If the user passed --stage1-url, --stage1-path, --stage1-name,
// --stage1-hash, or --stage1-from-dir then we take what was passed
// and try to load it. If it failed, we bail out. No second chances
// and whatnot. The details about how each location type should be
// handled are below.
//
// If the user passed none of the above flags, we try to get the name,
// the version and the location from the configuration. The name and
// the version must be defined in pair, that is - either both of them
// are defined in configuration or none. Values from the configuration
// override the values taken from the configure script. We search for
// an image with the default name and version in the store. If it is
// there, then woo, we are done. Otherwise we get the location and try
// to load it. Depending on location type, we bail out immediately or
// get a second chance.
//
// Details about the handling of different location types follow.
//
// If location is a URL, we do no discovery, just try to fetch it
// directly into the store instead.
//
// If location is a path, we do no discovery, just try to fetch it
// directly into the store instead. If the file is not found and we
// have a second chance, we try to fetch the file in the same
// directory as the rkt binary itself into the store.
//
// If location is a name, we do the discovery, fetch the discovered
// image into the store.
//
// If location is an image hash then we make sure that it exists in
// the store.
func getStage1Hash(s *store.Store, c *config.Config) (*types.Hash, error) {
	imgDir := getStage1ImagesDirectory(c)
	if overriddenStage1Location.kind != stage1ImageLocationUnset {
		// we passed a --stage-{url,path,name,hash,from-dir} flag
		return getStage1HashFromFlag(s, overriddenStage1Location, imgDir)
	}

	imgRef, imgLoc, imgFileName := getStage1DataFromConfig(c)
	return getConfiguredStage1Hash(s, imgRef, imgLoc, imgFileName)
}

func getStage1ImagesDirectory(c *config.Config) string {
	if c.Paths.Stage1ImagesDir != "" {
		return c.Paths.Stage1ImagesDir
	}
	return buildDefaultStage1ImagesDir
}

func getStage1HashFromFlag(s *store.Store, loc stage1ImageLocation, dir string) (*types.Hash, error) {
	imgType := apps.AppImageGuess
	switch loc.kind {
	case stage1ImageLocationURL:
		imgType = apps.AppImageURL
	case stage1ImageLocationPath:
		imgType = apps.AppImagePath
	case stage1ImageLocationName:
		imgType = apps.AppImageName
	case stage1ImageLocationHash:
		imgType = apps.AppImageHash
	case stage1ImageLocationFromDir:
		loc.location = filepath.Join(dir, loc.location)
		imgType = apps.AppImagePath
	}

	fn := getStage1Finder(s)
	return fn.FindImage(loc.location, "", imgType)
}

func getStage1DataFromConfig(c *config.Config) (string, string, string) {
	imgName := c.Stage1.Name
	imgVersion := c.Stage1.Version
	// if the name in the configuration is empty, then the version
	// is empty too, but let's better be safe now then sorry later
	// - if either one is empty we take build defaults for both
	if imgName == "" || imgVersion == "" {
		imgName = buildDefaultStage1Name
		imgVersion = buildDefaultStage1Version
	}
	imgRef := fmt.Sprintf("%s:%s", imgName, imgVersion)

	imgLoc := c.Stage1.Location
	imgFileName := getFileNameFromLocation(imgLoc)
	if imgLoc == "" {
		imgLoc = buildDefaultStage1ImageLoc
		imgFileName = buildDefaultStage1ImageInRktDir
	}

	return imgRef, imgLoc, imgFileName
}

func getFileNameFromLocation(imgLoc string) string {
	if !filepath.IsAbs(imgLoc) {
		return ""
	}
	return filepath.Base(imgLoc)
}

func getConfiguredStage1Hash(s *store.Store, imgRef, imgLoc, imgFileName string) (*types.Hash, error) {
	fn := getStage1Finder(s)
	if !strings.HasSuffix(imgRef, "-dirty") {
		fn.StoreOnly = true
		if hash, err := fn.FindImage(imgRef, "", apps.AppImageName); err == nil {
			return hash, nil
		}
		fn.StoreOnly = false
	}
	if imgLoc == "" && imgFileName == "" {
		return nil, fmt.Errorf("neither the location of the default stage1 image nor its filename are set, use --stage1-{path,url,name,hash,from-dir} flag")
	}
	// If imgLoc is not an absolute path, then it is a URL
	imgLocIsURL := imgLoc != "" && !filepath.IsAbs(imgLoc)
	if imgLocIsURL {
		return fn.FindImage(imgLoc, "", apps.AppImageURL)
	}
	return getStage1HashFromPath(fn, imgLoc, imgFileName)
}

func getStage1Finder(s *store.Store) *image.Finder {
	return &image.Finder{
		S:                  s,
		InsecureFlags:      globalFlags.InsecureFlags,
		TrustKeysFromHTTPS: globalFlags.TrustKeysFromHTTPS,

		StoreOnly: false,
		NoStore:   false,
		WithDeps:  false,
	}
}

func getStage1HashFromPath(fn *image.Finder, imgLoc, imgFileName string) (*types.Hash, error) {
	var fetchErr error
	var fallbackErr error
	if imgLoc != "" {
		hash, err := fn.FindImage(imgLoc, "", apps.AppImagePath)
		if err == nil {
			return hash, nil
		}
		fetchErr = err
	}
	if imgFileName != "" {
		rktDir := getRktBinaryDir()
		imgPath := filepath.Join(rktDir, imgFileName)
		hash, err := fn.FindImage(imgPath, "", apps.AppImagePath)
		if err == nil {
			return hash, nil
		}
		fallbackErr = err
	}
	return nil, mergeStage1Errors(fetchErr, fallbackErr)
}

func mergeStage1Errors(fetchErr, fallbackErr error) error {
	if fetchErr != nil && fallbackErr != nil {
		innerErr := errwrap.Wrap(fallbackErr, fetchErr)
		return errwrap.Wrap(errors.New("failed to fetch stage1 image and failed to fall back to stage1 image in the rkt directory"), innerErr)
	} else if fetchErr != nil {
		return errwrap.Wrap(errors.New("failed to fetch stage1 image"), fetchErr)
	}
	return errwrap.Wrap(errors.New("failed to fall back to stage1 image in rkt directory (default stage1 image location is not specified)"), fallbackErr)
}

func getRktBinaryDir() string {
	return filepath.Dir(os.Args[0])
}
<commit_msg>rkt: fix rkt binary directory detection<commit_after>// Copyright 2015 The rkt Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/appc/spec/schema/types"
	"github.com/coreos/rkt/common/apps"
	"github.com/coreos/rkt/rkt/config"
	"github.com/coreos/rkt/rkt/image"
	"github.com/coreos/rkt/store"
	"github.com/hashicorp/errwrap"
	"github.com/spf13/pflag"
)

// stage1ImageLocationKind describes the stage1 image location
type stage1ImageLocationKind int

const (
	// location unset, it is not a valid kind to be used
	stage1ImageLocationUnset stage1ImageLocationKind = iota
	// a URL with a scheme
	stage1ImageLocationURL
	// an absolute or a relative path
	stage1ImageLocationPath
	// an image name
	stage1ImageLocationName
	// an image hash
	stage1ImageLocationHash
	// an image in the default dir
	stage1ImageLocationFromDir
)

// stage1FlagData is used for creating the flags for each valid location kind
type stage1FlagData struct {
	kind stage1ImageLocationKind
	flag string
	name string
	help string
}

// stage1ImageLocation is used to store the user's choice of stage1 image via flags
type stage1ImageLocation struct {
	kind     stage1ImageLocationKind
	location string
}

// stage1ImageLocationFlag is an implementation of a pflag.Value
// interface, which handles all the valid location kinds
type stage1ImageLocationFlag struct {
	loc  *stage1ImageLocation
	kind stage1ImageLocationKind
}

func (f *stage1ImageLocationFlag) Set(location string) error {
	if f.loc.kind != stage1ImageLocationUnset {
		wanted := stage1FlagsData[f.kind]
		current := stage1FlagsData[f.loc.kind]
		if f.loc.kind == f.kind {
			return fmt.Errorf("--%s already used", current.flag)
		}
		return fmt.Errorf("flags --%s and --%s are mutually exclusive",
			wanted.flag, current.flag)
	}
	f.loc.kind = f.kind
	f.loc.location = location
	return nil
}

func (f *stage1ImageLocationFlag) String() string {
	return f.loc.location
}

func (f *stage1ImageLocationFlag) Type() string {
	return stage1FlagsData[f.kind].name
}

var (
	// defaults defined by configure, set by linker
	// default stage1 image name
	// (e.g. coreos.com/rkt/stage1-coreos)
	buildDefaultStage1Name string
	// default stage1 image version (e.g. 0.15.0)
	buildDefaultStage1Version string
	// an absolute path or a URL to the default stage1 image file
	buildDefaultStage1ImageLoc string
	// filename of the default stage1 image file in the default
	// stage1 images directory
	buildDefaultStage1ImageInRktDir string
	// an absolute path to the stage1 images directory
	buildDefaultStage1ImagesDir string

	// this holds necessary data to generate the --stage1-* flags
	// for each location kind
	stage1FlagsData = map[stage1ImageLocationKind]*stage1FlagData{
		stage1ImageLocationURL: &stage1FlagData{
			kind: stage1ImageLocationURL,
			flag: "stage1-url",
			name: "stage1URL",
			help: "a URL to an image to use as stage1",
		},

		stage1ImageLocationPath: &stage1FlagData{
			kind: stage1ImageLocationPath,
			flag: "stage1-path",
			name: "stage1Path",
			help: "an absolute or a relative path to an image to use as stage1",
		},

		stage1ImageLocationName: &stage1FlagData{
			kind: stage1ImageLocationName,
			flag: "stage1-name",
			name: "stage1Name",
			help: "a name of an image to use as stage1",
		},

		stage1ImageLocationHash: &stage1FlagData{
			kind: stage1ImageLocationHash,
			flag: "stage1-hash",
			name: "stage1Hash",
			help: "a hash of an image to use as stage1",
		},

		stage1ImageLocationFromDir: &stage1FlagData{
			kind: stage1ImageLocationFromDir,
			flag: "stage1-from-dir",
			name: "stage1FromDir",
			help: "a filename of an image in stage1 images directory to use as stage1",
		},
	}
	// location to stage1 image overridden by one of --stage1-*
	// flags
	overriddenStage1Location = stage1ImageLocation{
		kind:     stage1ImageLocationUnset,
		location: "",
	}
)

// addStage1ImageFlags adds flags for specifying custom stage1 image
func addStage1ImageFlags(flags *pflag.FlagSet) {
	for _, data := range stage1FlagsData {
		wrapper := &stage1ImageLocationFlag{
			loc:  &overriddenStage1Location,
			kind: data.kind,
		}
		flags.Var(wrapper, data.flag, data.help)
	}
}

// getStage1Hash will try to get the hash of stage1 to use.
//
// Before getting inside this rats nest, let's try to write up the
// expected behaviour.
//
// If the user passed --stage1-url, --stage1-path, --stage1-name,
// --stage1-hash, or --stage1-from-dir then we take what was passed
// and try to load it. If it failed, we bail out. No second chances
// and whatnot. The details about how each location type should be
// handled are below.
//
// If the user passed none of the above flags, we try to get the name,
// the version and the location from the configuration. The name and
// the version must be defined in pair, that is - either both of them
// are defined in configuration or none. Values from the configuration
// override the values taken from the configure script. We search for
// an image with the default name and version in the store. If it is
// there, then woo, we are done. Otherwise we get the location and try
// to load it. Depending on location type, we bail out immediately or
// get a second chance.
//
// Details about the handling of different location types follow.
//
// If location is a URL, we do no discovery, just try to fetch it
// directly into the store instead.
//
// If location is a path, we do no discovery, just try to fetch it
// directly into the store instead. If the file is not found and we
// have a second chance, we try to fetch the file in the same
// directory as the rkt binary itself into the store.
//
// If location is a name, we do the discovery, fetch the discovered
// image into the store.
//
// If location is an image hash then we make sure that it exists in
// the store.
func getStage1Hash(s *store.Store, c *config.Config) (*types.Hash, error) {
	imgDir := getStage1ImagesDirectory(c)
	if overriddenStage1Location.kind != stage1ImageLocationUnset {
		// we passed a --stage-{url,path,name,hash,from-dir} flag
		return getStage1HashFromFlag(s, overriddenStage1Location, imgDir)
	}

	imgRef, imgLoc, imgFileName := getStage1DataFromConfig(c)
	return getConfiguredStage1Hash(s, imgRef, imgLoc, imgFileName)
}

func getStage1ImagesDirectory(c *config.Config) string {
	if c.Paths.Stage1ImagesDir != "" {
		return c.Paths.Stage1ImagesDir
	}
	return buildDefaultStage1ImagesDir
}

func getStage1HashFromFlag(s *store.Store, loc stage1ImageLocation, dir string) (*types.Hash, error) {
	imgType := apps.AppImageGuess
	switch loc.kind {
	case stage1ImageLocationURL:
		imgType = apps.AppImageURL
	case stage1ImageLocationPath:
		imgType = apps.AppImagePath
	case stage1ImageLocationName:
		imgType = apps.AppImageName
	case stage1ImageLocationHash:
		imgType = apps.AppImageHash
	case stage1ImageLocationFromDir:
		loc.location = filepath.Join(dir, loc.location)
		imgType = apps.AppImagePath
	}

	fn := getStage1Finder(s)
	return fn.FindImage(loc.location, "", imgType)
}

func getStage1DataFromConfig(c *config.Config) (string, string, string) {
	imgName := c.Stage1.Name
	imgVersion := c.Stage1.Version
	// if the name in the configuration is empty, then the version
	// is empty too, but let's better be safe now then sorry later
	// - if either one is empty we take build defaults for both
	if imgName == "" || imgVersion == "" {
		imgName = buildDefaultStage1Name
		imgVersion = buildDefaultStage1Version
	}
	imgRef := fmt.Sprintf("%s:%s", imgName, imgVersion)

	imgLoc := c.Stage1.Location
	imgFileName := getFileNameFromLocation(imgLoc)
	if imgLoc == "" {
		imgLoc = buildDefaultStage1ImageLoc
		imgFileName = buildDefaultStage1ImageInRktDir
	}

	return imgRef, imgLoc, imgFileName
}

func getFileNameFromLocation(imgLoc string) string {
	if !filepath.IsAbs(imgLoc) {
		return ""
	}
	return filepath.Base(imgLoc)
}

func getConfiguredStage1Hash(s *store.Store, imgRef, imgLoc, imgFileName string) (*types.Hash, error) {
	fn := getStage1Finder(s)
	if !strings.HasSuffix(imgRef, "-dirty") {
		fn.StoreOnly = true
		if hash, err := fn.FindImage(imgRef, "", apps.AppImageName); err == nil {
			return hash, nil
		}
		fn.StoreOnly = false
	}
	if imgLoc == "" && imgFileName == "" {
		return nil, fmt.Errorf("neither the location of the default stage1 image nor its filename are set, use --stage1-{path,url,name,hash,from-dir} flag")
	}
	// If imgLoc is not an absolute path, then it is a URL
	imgLocIsURL := imgLoc != "" && !filepath.IsAbs(imgLoc)
	if imgLocIsURL {
		return fn.FindImage(imgLoc, "", apps.AppImageURL)
	}
	return getStage1HashFromPath(fn, imgLoc, imgFileName)
}

func getStage1Finder(s *store.Store) *image.Finder {
	return &image.Finder{
		S:                  s,
		InsecureFlags:      globalFlags.InsecureFlags,
		TrustKeysFromHTTPS: globalFlags.TrustKeysFromHTTPS,

		StoreOnly: false,
		NoStore:   false,
		WithDeps:  false,
	}
}

func getStage1HashFromPath(fn *image.Finder, imgLoc, imgFileName string) (*types.Hash, error) {
	var fetchErr error
	var fallbackErr error
	if imgLoc != "" {
		hash, err := fn.FindImage(imgLoc, "", apps.AppImagePath)
		if err == nil {
			return hash, nil
		}
		fetchErr = err
	}
	if imgFileName != "" {
		exePath, err := os.Readlink("/proc/self/exe")
		if err != nil {
			fallbackErr = err
		} else {
			rktDir := filepath.Dir(exePath)
			imgPath := filepath.Join(rktDir, imgFileName)
			hash, err := fn.FindImage(imgPath, "", apps.AppImagePath)
			if err == nil {
				return hash, nil
			}
			fallbackErr = err
		}
	}
	return nil, mergeStage1Errors(fetchErr, fallbackErr)
}

func mergeStage1Errors(fetchErr, fallbackErr error) error {
	if fetchErr != nil && fallbackErr != nil {
		innerErr := errwrap.Wrap(fallbackErr, fetchErr)
		return errwrap.Wrap(errors.New("failed to fetch stage1 image and failed to fall back to stage1 image in the rkt directory"), innerErr)
	} else if fetchErr != nil {
		return errwrap.Wrap(errors.New("failed to fetch stage1 image"), fetchErr)
	}
	return errwrap.Wrap(errors.New("failed to fall back to stage1 image in rkt directory (default stage1 image location is not specified)"), fallbackErr)
}
<|endoftext|>


<commit_before>package main

import (
	log "github.com/Sirupsen/logrus"
	"github.com/elazarl/goproxy"

	"bufio"
	"flag"
	"fmt"
	"net"
	"net/http"
	"regexp"
)

// VirtualizeMode - default mode when Hoverfly looks for captured requests to respond
const VirtualizeMode = "virtualize"

// SynthesizeMode - all requests are sent to middleware to create response
const SynthesizeMode = "synthesize"

// ModifyMode - middleware is applied to outgoing and incoming traffic
const ModifyMode = "modify"

// CaptureMode - requests are captured and stored in cache
const CaptureMode = "capture"

// orPanic - wrapper for logging errors
func orPanic(err error) {
	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Panic("Got error.")
	}
}

func main() {
	// Output to stderr instead of stdout, could also be a file.
	//	log.SetOutput(os.Stderr)
	//	log.SetFormatter(&log.TextFormatter{})
	log.SetFormatter(&log.JSONFormatter{})

	// getting proxy configuration
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	// modes
	capture := flag.Bool("capture", false, "should proxy capture requests")
	synthesize := flag.Bool("synthesize", false, "should proxy capture requests")
	modify := flag.Bool("modify", false, "should proxy only modify requests")

	destination := flag.String("destination", ".", "destination URI to catch")
	middleware := flag.String("middleware", "", "should proxy use middleware")

	// proxy port
	proxyPort := flag.String("pp", "", "proxy port - run proxy on another port (i.e. '-pp 9999' to run proxy on port 9999)")
	// admin port
	adminPort := flag.String("ap", "", "admin port - run admin interface on another port (i.e. '-ap 1234' to run admin UI on port 1234)")

	flag.Parse()

	// getting settings
	cfg := InitSettings()

	if *verbose {
		// Only log the warning severity or above.
		log.SetLevel(log.DebugLevel)
	}
	cfg.verbose = *verbose

	// overriding environment variables (proxy and admin ports)
	if *proxyPort != "" {
		cfg.proxyPort = *proxyPort
	}
	if *adminPort != "" {
		cfg.adminPort = *adminPort
	}

	// overriding default middleware setting
	cfg.middleware = *middleware

	// setting default mode
	mode := VirtualizeMode

	if *capture {
		mode = CaptureMode
		// checking whether user supplied other modes
		if *synthesize == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *synthesize {
		mode = SynthesizeMode

		if cfg.middleware == "" {
			log.Fatal("Synthesize mode chosen although middleware not supplied")
		}

		if *capture == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *modify {
		mode = ModifyMode

		if cfg.middleware == "" {
			log.Fatal("Modify mode chosen although middleware not supplied")
		}

		if *capture == true || *synthesize == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	}

	// overriding default settings
	cfg.mode = mode

	// overriding destination
	cfg.destination = *destination

	proxy, dbClient := getNewHoverfly(cfg)
	defer dbClient.cache.db.Close()

	log.Warn(http.ListenAndServe(fmt.Sprintf(":%s", cfg.proxyPort), proxy))
}

// getNewHoverfly returns a configured ProxyHttpServer and DBClient, also starts admin interface on configured port
func getNewHoverfly(cfg *Configuration) (*goproxy.ProxyHttpServer, DBClient) {

	// getting boltDB
	db := getDB(cfg.databaseName)

	cache := Cache{
		db:             db,
		requestsBucket: []byte(requestsBucketName),
	}

	// getting connections
	d := DBClient{
		cache: cache,
		http:  &http.Client{},
		cfg:   cfg,
	}

	// creating proxy
	proxy := goproxy.NewProxyHttpServer()

	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(d.cfg.destination))).
		HandleConnect(goproxy.AlwaysMitm)

	// enable curl -p for all hosts on port 80
	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(d.cfg.destination))).
		HijackConnect(func(req *http.Request, client net.Conn, ctx *goproxy.ProxyCtx) {
		defer func() {
			if e := recover(); e != nil {
				ctx.Logf("error connecting to remote: %v", e)
				client.Write([]byte("HTTP/1.1 500 Cannot reach destination\r\n\r\n"))
			}
			client.Close()
		}()
		clientBuf := bufio.NewReadWriter(bufio.NewReader(client), bufio.NewWriter(client))
		remote, err := net.Dial("tcp", req.URL.Host)
		orPanic(err)
		remoteBuf := bufio.NewReadWriter(bufio.NewReader(remote), bufio.NewWriter(remote))
		for {
			req, err := http.ReadRequest(clientBuf.Reader)
			orPanic(err)
			orPanic(req.Write(remoteBuf))
			orPanic(remoteBuf.Flush())
			resp, err := http.ReadResponse(remoteBuf.Reader, req)

			orPanic(err)
			orPanic(resp.Write(clientBuf.Writer))
			orPanic(clientBuf.Flush())
		}
	})

	// processing connections
	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(cfg.destination))).DoFunc(
		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
			return d.processRequest(r)
		})

	go d.startAdminInterface()

	proxy.Verbose = d.cfg.verbose
	// proxy starting message
	log.WithFields(log.Fields{
		"Destination": d.cfg.destination,
		"ProxyPort":   d.cfg.proxyPort,
		"Mode":        d.cfg.GetMode(),
	}).Info("Proxy prepared...")

	return proxy, d
}

func hoverflyError(req *http.Request, err error, msg string, statusCode int) *http.Response {
	return goproxy.NewResponse(req,
		goproxy.ContentTypeText, statusCode,
		fmt.Sprintf("Hoverfly Error! %s. Got error: %s \n", msg, err.Error()))
}

// processRequest - processes incoming requests and based on proxy state (record/playback)
// returns HTTP response.
func (d *DBClient) processRequest(req *http.Request) (*http.Request, *http.Response) {

	mode := d.cfg.GetMode()

	if mode == CaptureMode {
		newResponse, err := d.captureRequest(req)

		if err != nil {
			return req, hoverflyError(req, err, "Could not capture request", http.StatusServiceUnavailable)
		}
		log.WithFields(log.Fields{
			"mode":        mode,
			"middleware":  d.cfg.middleware,
			"path":        req.URL.Path,
			"rawQuery":    req.URL.RawQuery,
			"method":      req.Method,
			"destination": req.Host,
		}).Info("request and response captured")

		return req, newResponse

	} else if mode == SynthesizeMode {
		response, err := synthesizeResponse(req, d.cfg.middleware)

		if err != nil {
			return req, hoverflyError(req, err, "Could not create synthetic response!", http.StatusServiceUnavailable)
		}

		log.WithFields(log.Fields{
			"mode":        mode,
			"middleware":  d.cfg.middleware,
			"path":        req.URL.Path,
			"rawQuery":    req.URL.RawQuery,
			"method":      req.Method,
			"destination": req.Host,
		}).Info("synthetic response created successfuly")

		return req, response

	} else if mode == ModifyMode {
		response, err := d.modifyRequestResponse(req, d.cfg.middleware)

		if err != nil {
			log.WithFields(log.Fields{
				"error":      err.Error(),
				"middleware": d.cfg.middleware,
			}).Error("Got error when performing request modification")
			return req, nil
		}

		// returning modified response
		return req, response

	}

	newResponse := d.getResponse(req)
	return req, newResponse

}
<commit_msg>failing if middleware failed<commit_after>package main

import (
	log "github.com/Sirupsen/logrus"
	"github.com/elazarl/goproxy"

	"bufio"
	"flag"
	"fmt"
	"net"
	"net/http"
	"regexp"
)

// VirtualizeMode - default mode when Hoverfly looks for captured requests to respond
const VirtualizeMode = "virtualize"

// SynthesizeMode - all requests are sent to middleware to create response
const SynthesizeMode = "synthesize"

// ModifyMode - middleware is applied to outgoing and incoming traffic
const ModifyMode = "modify"

// CaptureMode - requests are captured and stored in cache
const CaptureMode = "capture"

// orPanic - wrapper for logging errors
func orPanic(err error) {
	if err != nil {
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Panic("Got error.")
	}
}

func main() {
	// Output to stderr instead of stdout, could also be a file.
	//	log.SetOutput(os.Stderr)
	//	log.SetFormatter(&log.TextFormatter{})
	log.SetFormatter(&log.JSONFormatter{})

	// getting proxy configuration
	verbose := flag.Bool("v", false, "should every proxy request be logged to stdout")
	// modes
	capture := flag.Bool("capture", false, "should proxy capture requests")
	synthesize := flag.Bool("synthesize", false, "should proxy capture requests")
	modify := flag.Bool("modify", false, "should proxy only modify requests")

	destination := flag.String("destination", ".", "destination URI to catch")
	middleware := flag.String("middleware", "", "should proxy use middleware")

	// proxy port
	proxyPort := flag.String("pp", "", "proxy port - run proxy on another port (i.e. '-pp 9999' to run proxy on port 9999)")
	// admin port
	adminPort := flag.String("ap", "", "admin port - run admin interface on another port (i.e. '-ap 1234' to run admin UI on port 1234)")

	flag.Parse()

	// getting settings
	cfg := InitSettings()

	if *verbose {
		// Only log the warning severity or above.
		log.SetLevel(log.DebugLevel)
	}
	cfg.verbose = *verbose

	// overriding environment variables (proxy and admin ports)
	if *proxyPort != "" {
		cfg.proxyPort = *proxyPort
	}
	if *adminPort != "" {
		cfg.adminPort = *adminPort
	}

	// overriding default middleware setting
	cfg.middleware = *middleware

	// setting default mode
	mode := VirtualizeMode

	if *capture {
		mode = CaptureMode
		// checking whether user supplied other modes
		if *synthesize == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *synthesize {
		mode = SynthesizeMode

		if cfg.middleware == "" {
			log.Fatal("Synthesize mode chosen although middleware not supplied")
		}

		if *capture == true || *modify == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	} else if *modify {
		mode = ModifyMode

		if cfg.middleware == "" {
			log.Fatal("Modify mode chosen although middleware not supplied")
		}

		if *capture == true || *synthesize == true {
			log.Fatal("Two or more modes supplied, check your flags")
		}
	}

	// overriding default settings
	cfg.mode = mode

	// overriding destination
	cfg.destination = *destination

	proxy, dbClient := getNewHoverfly(cfg)
	defer dbClient.cache.db.Close()

	log.Warn(http.ListenAndServe(fmt.Sprintf(":%s", cfg.proxyPort), proxy))
}

// getNewHoverfly returns a configured ProxyHttpServer and DBClient, also starts admin interface on configured port
func getNewHoverfly(cfg *Configuration) (*goproxy.ProxyHttpServer, DBClient) {

	// getting boltDB
	db := getDB(cfg.databaseName)

	cache := Cache{
		db:             db,
		requestsBucket: []byte(requestsBucketName),
	}

	// getting connections
	d := DBClient{
		cache: cache,
		http:  &http.Client{},
		cfg:   cfg,
	}

	// creating proxy
	proxy := goproxy.NewProxyHttpServer()

	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(d.cfg.destination))).
		HandleConnect(goproxy.AlwaysMitm)

	// enable curl -p for all hosts on port 80
	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(d.cfg.destination))).
		HijackConnect(func(req *http.Request, client net.Conn, ctx *goproxy.ProxyCtx) {
		defer func() {
			if e := recover(); e != nil {
				ctx.Logf("error connecting to remote: %v", e)
				client.Write([]byte("HTTP/1.1 500 Cannot reach destination\r\n\r\n"))
			}
			client.Close()
		}()
		clientBuf := bufio.NewReadWriter(bufio.NewReader(client), bufio.NewWriter(client))
		remote, err := net.Dial("tcp", req.URL.Host)
		orPanic(err)
		remoteBuf := bufio.NewReadWriter(bufio.NewReader(remote), bufio.NewWriter(remote))
		for {
			req, err := http.ReadRequest(clientBuf.Reader)
			orPanic(err)
			orPanic(req.Write(remoteBuf))
			orPanic(remoteBuf.Flush())
			resp, err := http.ReadResponse(remoteBuf.Reader, req)

			orPanic(err)
			orPanic(resp.Write(clientBuf.Writer))
			orPanic(clientBuf.Flush())
		}
	})

	// processing connections
	proxy.OnRequest(goproxy.ReqHostMatches(regexp.MustCompile(cfg.destination))).DoFunc(
		func(r *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
			return d.processRequest(r)
		})

	go d.startAdminInterface()

	proxy.Verbose = d.cfg.verbose
	// proxy starting message
	log.WithFields(log.Fields{
		"Destination": d.cfg.destination,
		"ProxyPort":   d.cfg.proxyPort,
		"Mode":        d.cfg.GetMode(),
	}).Info("Proxy prepared...")

	return proxy, d
}

func hoverflyError(req *http.Request, err error, msg string, statusCode int) *http.Response {
	return goproxy.NewResponse(req,
		goproxy.ContentTypeText, statusCode,
		fmt.Sprintf("Hoverfly Error! %s. Got error: %s \n", msg, err.Error()))
}

// processRequest - processes incoming requests and based on proxy state (record/playback)
// returns HTTP response.
func (d *DBClient) processRequest(req *http.Request) (*http.Request, *http.Response) {

	mode := d.cfg.GetMode()

	if mode == CaptureMode {
		newResponse, err := d.captureRequest(req)

		if err != nil {
			return req, hoverflyError(req, err, "Could not capture request", http.StatusServiceUnavailable)
		}
		log.WithFields(log.Fields{
			"mode":        mode,
			"middleware":  d.cfg.middleware,
			"path":        req.URL.Path,
			"rawQuery":    req.URL.RawQuery,
			"method":      req.Method,
			"destination": req.Host,
		}).Info("request and response captured")

		return req, newResponse

	} else if mode == SynthesizeMode {
		response, err := synthesizeResponse(req, d.cfg.middleware)

		if err != nil {
			return req, hoverflyError(req, err, "Could not create synthetic response!", http.StatusServiceUnavailable)
		}

		log.WithFields(log.Fields{
			"mode":        mode,
			"middleware":  d.cfg.middleware,
			"path":        req.URL.Path,
			"rawQuery":    req.URL.RawQuery,
			"method":      req.Method,
			"destination": req.Host,
		}).Info("synthetic response created successfuly")

		return req, response

	} else if mode == ModifyMode {
		response, err := d.modifyRequestResponse(req, d.cfg.middleware)

		if err != nil {
			log.WithFields(log.Fields{
				"error":      err.Error(),
				"middleware": d.cfg.middleware,
			}).Error("Got error when performing request modification")
			return req, hoverflyError(
				req,
				err,
				fmt.Sprintf("Middleware (%s) failed or something else happened!", d.cfg.middleware),
				http.StatusServiceUnavailable)
		}
		// returning modified response
		return req, response
	}

	newResponse := d.getResponse(req)
	return req, newResponse

}
<|endoftext|>


<commit_before>// +build acceptance

package v1

import (
	"testing"

	"github.com/rackspace/gophercloud"
	os "github.com/rackspace/gophercloud/openstack/cdn/v1/services"
	"github.com/rackspace/gophercloud/rackspace/cdn/v1/services"
	th "github.com/rackspace/gophercloud/testhelper"
)

func TestService(t *testing.T) {
	client := newClient(t)

	t.Log("Creating Service")
	loc := testServiceCreate(t, client)
	t.Logf("Created service at location: %s", loc)

	defer testServiceDelete(t, client, loc)

	t.Log("Updating Service")
	testServiceUpdate(t, client, loc)

	t.Log("Retrieving Service")
	testServiceGet(t, client, loc)

	/*
		t.Log("Listing Services")
		testServiceList(t, client)
	*/
}

func testServiceCreate(t *testing.T, client *gophercloud.ServiceClient) string {
	createOpts := os.CreateOpts{
		Name: "gophercloud-test-service",
		Domains: []os.Domain{
			os.Domain{
				Domain: "www.gophercloud-test-service.com",
			},
		},
		Origins: []os.Origin{
			os.Origin{
				Origin: "gophercloud-test-service.com",
				Port:   80,
				SSL:    false,
			},
		},
		FlavorID: "cdn",
	}
	l, err := services.Create(client, createOpts).Extract()
	th.AssertNoErr(t, err)
	return l
}

func testServiceGet(t *testing.T, client *gophercloud.ServiceClient, id string) {
	s, err := services.Get(client, id).Extract()
	th.AssertNoErr(t, err)
	t.Logf("Retrieved service: %+v", *s)
}

func testServiceUpdate(t *testing.T, client *gophercloud.ServiceClient, id string) {
	updateOpts := os.UpdateOpts{
		os.UpdateOpt{
			Op:   os.Add,
			Path: "/domains/-",
			Value: map[string]interface{}{
				"domain":   "newDomain.com",
				"protocol": "http",
			},
		},
	}
	loc, err := services.Update(client, id, updateOpts).Extract()
	th.AssertNoErr(t, err)
	t.Logf("Successfully updated service at location: %s", loc)
}

/*
func testServiceList(t *testing.T, client *gophercloud.ServiceClient) {
	err := service.List(client).ExtractErr()
	th.AssertNoErr(t, err)
	t.Logf("Successfully pinged root URL")
}
*/

func testServiceDelete(t *testing.T, client *gophercloud.ServiceClient, id string) {
	err := services.Delete(client, id).ExtractErr()
	th.AssertNoErr(t, err)
	t.Logf("Successfully deleted service (%s)", id)
}
<commit_msg>added acceptance test for listing services<commit_after>// +build acceptance

package v1

import (
	"testing"

	"github.com/rackspace/gophercloud"
	os "github.com/rackspace/gophercloud/openstack/cdn/v1/services"
	"github.com/rackspace/gophercloud/pagination"
	"github.com/rackspace/gophercloud/rackspace/cdn/v1/services"
	th "github.com/rackspace/gophercloud/testhelper"
)

func TestService(t *testing.T) {
	client := newClient(t)

	t.Log("Creating Service")
	loc := testServiceCreate(t, client)
	t.Logf("Created service at location: %s", loc)

	defer testServiceDelete(t, client, loc)

	t.Log("Updating Service")
	testServiceUpdate(t, client, loc)

	t.Log("Retrieving Service")
	testServiceGet(t, client, loc)

	t.Log("Listing Services")
	testServiceList(t, client)
}

func testServiceCreate(t *testing.T, client *gophercloud.ServiceClient) string {
	createOpts := os.CreateOpts{
		Name: "gophercloud-test-service",
		Domains: []os.Domain{
			os.Domain{
				Domain: "www.gophercloud-test-service.com",
			},
		},
		Origins: []os.Origin{
			os.Origin{
				Origin: "gophercloud-test-service.com",
				Port:   80,
				SSL:    false,
			},
		},
		FlavorID: "cdn",
	}
	l, err := services.Create(client, createOpts).Extract()
	th.AssertNoErr(t, err)
	return l
}

func testServiceGet(t *testing.T, client *gophercloud.ServiceClient, id string) {
	s, err := services.Get(client, id).Extract()
	th.AssertNoErr(t, err)
	t.Logf("Retrieved service: %+v", *s)
}

func testServiceUpdate(t *testing.T, client *gophercloud.ServiceClient, id string) {
	updateOpts := os.UpdateOpts{
		os.UpdateOpt{
			Op:   os.Add,
			Path: "/domains/-",
			Value: map[string]interface{}{
				"domain":   "newDomain.com",
				"protocol": "http",
			},
		},
	}
	loc, err := services.Update(client, id, updateOpts).Extract()
	th.AssertNoErr(t, err)
	t.Logf("Successfully updated service at location: %s", loc)
}

func testServiceList(t *testing.T, client *gophercloud.ServiceClient) {
	err := services.List(client, nil).EachPage(func(page pagination.Page) (bool, error) {
		serviceList, err := os.ExtractServices(page)
		th.AssertNoErr(t, err)

		for _, service := range serviceList {
			t.Logf("Listing service: %+v", service)
		}

		return true, nil
	})

	th.AssertNoErr(t, err)
}

func testServiceDelete(t *testing.T, client *gophercloud.ServiceClient, id string) {
	err := services.Delete(client, id).ExtractErr()
	th.AssertNoErr(t, err)
	t.Logf("Successfully deleted service (%s)", id)
}
<|endoftext|>


<commit_before>package http

import (
	"fmt"
	"net/http"

	"github.com/VagabondDataNinjas/gizlinebot/domain"

	"github.com/line/line-bot-sdk-go/linebot"

	"github.com/VagabondDataNinjas/gizlinebot/storage"
	"github.com/VagabondDataNinjas/gizlinebot/survey"

	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
)

type ApiConf struct {
	Port int
	// hostname where the API is hosted
	GlobalVars *domain.GlobalTplVars
}

type Api struct {
	Storage storage.Storage
	LineBot *linebot.Client
	Surv    *survey.Survey
	Conf    *ApiConf
}

func NewApi(s storage.Storage, lb *linebot.Client, surv *survey.Survey, conf *ApiConf) *Api {
	return &Api{
		Storage: s,
		LineBot: lb,
		Surv:    surv,
		Conf:    conf,
	}
}

// cache-busting middleware
func NoCacheMW(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		c.Response().Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
		return next(c)
	}
}

func (a *Api) Serve() error {
	e := echo.New()
	// Middleware
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())
	// CORS default
	// Allows requests from any origin wth GET, HEAD, PUT, POST or DELETE method.
	e.Use(middleware.CORS())

	qs, err := a.Storage.GetQuestions()
	if err != nil {
		return err
	}
	e.POST("/linewebhook", LineWebhookHandlerBuilder(a.Surv, a.Storage, a.LineBot, a.Conf.GlobalVars))

	e.GET("/api/webform/questions", func(c echo.Context) error {
		return c.JSON(http.StatusOK, qs)
	})

	e.Group("/", NoCacheMW).Static("", "../gizsurvey/build")

	e.GET("/api/user/wipe/:userid", WipeUserHandlerBuilder(a.Storage, a.LineBot), NoCacheMW)

	e.POST("/api/webform/answer", AnswerHandlerBuilder(a.Storage))
	e.POST("/api/webform/answer-gps", AnswerGpsHandlerBuilder(a.Storage))

	// @TODO add authentication
	e.POST("/api/admin/send-msg", SendLineMsgHandlerBuilder(a.Storage, a.LineBot))

	e.Logger.Fatal(e.Start(fmt.Sprintf(":%d", a.Conf.Port)))
	return nil
}
<commit_msg>Fixed static file assets<commit_after>package http

import (
	"fmt"
	"net/http"

	"github.com/VagabondDataNinjas/gizlinebot/domain"

	"github.com/line/line-bot-sdk-go/linebot"

	"github.com/VagabondDataNinjas/gizlinebot/storage"
	"github.com/VagabondDataNinjas/gizlinebot/survey"

	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
)

type ApiConf struct {
	Port int
	// hostname where the API is hosted
	GlobalVars *domain.GlobalTplVars
}

type Api struct {
	Storage storage.Storage
	LineBot *linebot.Client
	Surv    *survey.Survey
	Conf    *ApiConf
}

func NewApi(s storage.Storage, lb *linebot.Client, surv *survey.Survey, conf *ApiConf) *Api {
	return &Api{
		Storage: s,
		LineBot: lb,
		Surv:    surv,
		Conf:    conf,
	}
}

// cache-busting middleware
func NoCacheMW(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		c.Response().Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
		return next(c)
	}
}

func (a *Api) Serve() error {
	e := echo.New()
	// Middleware
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())
	// CORS default
	// Allows requests from any origin wth GET, HEAD, PUT, POST or DELETE method.
	e.Use(middleware.CORS())

	qs, err := a.Storage.GetQuestions()
	if err != nil {
		return err
	}
	e.POST("/linewebhook", LineWebhookHandlerBuilder(a.Surv, a.Storage, a.LineBot, a.Conf.GlobalVars))

	e.GET("/api/webform/questions", func(c echo.Context) error {
		return c.JSON(http.StatusOK, qs)
	})

	e.Group("/", NoCacheMW).Static("", "../gizsurvey/build")
	e.Group("/static", NoCacheMW).Static("", "../gizsurvey/build/static")

	e.GET("/api/user/wipe/:userid", WipeUserHandlerBuilder(a.Storage, a.LineBot), NoCacheMW)

	e.POST("/api/webform/answer", AnswerHandlerBuilder(a.Storage))
	e.POST("/api/webform/answer-gps", AnswerGpsHandlerBuilder(a.Storage))

	// @TODO add authentication
	e.POST("/api/admin/send-msg", SendLineMsgHandlerBuilder(a.Storage, a.LineBot))

	e.Logger.Fatal(e.Start(fmt.Sprintf(":%d", a.Conf.Port)))
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Tigera, 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 commands

import (
	"context"
	"fmt"
	"os"
	"strings"

	"github.com/docopt/docopt-go"
	"github.com/projectcalico/calicoctl/calicoctl/commands/clientmgr"
	"github.com/projectcalico/calicoctl/calicoctl/commands/constants"
	"github.com/projectcalico/libcalico-go/lib/options"
)

var VERSION, GIT_REVISION string
var VERSION_SUMMARY string

func init() {
	VERSION_SUMMARY = "calicoctl version " + VERSION + ", build " + GIT_REVISION
}

func Version(args []string) {
	doc := `Usage:
  calicoctl version [--config=<CONFIG>]

Options:
  -h --help             Show this screen.
  -c --config=<CONFIG>  Path to the file containing connection configuration in
                        YAML or JSON format.
                        [default: ` + constants.DefaultConfigPath + `]

Description:
  Display the version of calicoctl.
`
	parsedArgs, err := docopt.Parse(doc, args, true, "", false, false)
	if err != nil {
		fmt.Printf("Invalid option: 'calicoctl %s'. Use flag '--help' to read about a specific subcommand.\n", strings.Join(args, " "))
		os.Exit(1)
	}
	if len(parsedArgs) == 0 {
		return
	}

	fmt.Println("Client Version:   ", VERSION)
	fmt.Println("Git commit:       ", GIT_REVISION)

	// Load the client config and connect.
	cf := parsedArgs["--config"].(string)
	client, err := clientmgr.NewClient(cf)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	ctx := context.Background()
	ci, err := client.ClusterInformation().Get(ctx, "default", options.GetOptions{})
	if err != nil {
		fmt.Println("Unable to retrieve Cluster Version or Type: ", err)
		os.Exit(1)
	}

	v := ci.Spec.CalicoVersion
	if v == "" {
		v = "unknown"
	}
	t := ci.Spec.ClusterType
	if t == "" {
		t = "unknown"
	}

	fmt.Println("Cluster Version:  ", v)
	fmt.Println("Cluster Type:     ", t)
}
<commit_msg>Improve 'calicoctl --version' output<commit_after>// Copyright (c) 2016-2017 Tigera, 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 commands

import (
	"context"
	"fmt"
	"os"
	"strings"

	"github.com/docopt/docopt-go"
	"github.com/projectcalico/calicoctl/calicoctl/commands/clientmgr"
	"github.com/projectcalico/calicoctl/calicoctl/commands/constants"
	"github.com/projectcalico/libcalico-go/lib/options"
)

var VERSION, GIT_REVISION string
var VERSION_SUMMARY string

func init() {
	VERSION_SUMMARY = "calicoctl " + VERSION + ", build " + GIT_REVISION + `
Run 'calicoctl version' to see cluster version information as well.`
}

func Version(args []string) {
	doc := `Usage:
  calicoctl version [--config=<CONFIG>]

Options:
  -h --help             Show this screen.
  -c --config=<CONFIG>  Path to the file containing connection configuration in
                        YAML or JSON format.
                        [default: ` + constants.DefaultConfigPath + `]

Description:
  Display the version of calicoctl.
`
	parsedArgs, err := docopt.Parse(doc, args, true, "", false, false)
	if err != nil {
		fmt.Printf("Invalid option: 'calicoctl %s'. Use flag '--help' to read about a specific subcommand.\n", strings.Join(args, " "))
		os.Exit(1)
	}
	if len(parsedArgs) == 0 {
		return
	}

	fmt.Println("Client Version:   ", VERSION)
	fmt.Println("Git commit:       ", GIT_REVISION)

	// Load the client config and connect.
	cf := parsedArgs["--config"].(string)
	client, err := clientmgr.NewClient(cf)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	ctx := context.Background()
	ci, err := client.ClusterInformation().Get(ctx, "default", options.GetOptions{})
	if err != nil {
		fmt.Println("Unable to retrieve Cluster Version or Type: ", err)
		os.Exit(1)
	}

	v := ci.Spec.CalicoVersion
	if v == "" {
		v = "unknown"
	}
	t := ci.Spec.ClusterType
	if t == "" {
		t = "unknown"
	}

	fmt.Println("Cluster Version:  ", v)
	fmt.Println("Cluster Type:     ", t)
}
<|endoftext|>