<commit_before>package main

import (
	"bufio"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"github.com/monnand/dhkx"
)

// struct for holding room info
type Room struct {
	name        string
	discription string
	password    string
	maxClients  int
	welcomeMsg  string
}

// struct for holding client info
type Client struct {
	connection net.Conn
	nick       string
	clientKey  string
}

var rooms []Room
var clientRoom map[Client]Room
var publicKeyCode string
var nickCode string

// clear the terminal screen
func clear() {
	switch runtime.GOOS {
	case "linux":
		cmd := exec.Command("clear")
		cmd.Stdout = os.Stdout
		cmd.Run()
	case "windows":
		cmd := exec.Command("cmd", "/c", "cls")
		cmd.Stdout = os.Stdout
		cmd.Run()
	default:
		fmt.Println("Attempted to clear terminal, but OS is not supported.")
	}
}

// init
func setup() {
	publicKeyCode = "ssd990=+?¡][ªs)(sdª]ßð=S)]"
	nickCode = "!#28jKas>zzx'**!+?,>lzc012"
	clientRoom = make(map[Client]Room)
	makeRoom("Lobby", "Welcome to Lobby")
	makeRoom("TestRoom", "Welcome to TestRoom")
}

// check if []byte contains specific byte
func contains(s []byte, e byte) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}

// exchange key with client
func exchangeKeys(c Client) {

	var clientPublicKey []byte

	// generate private key
	g, _ := dhkx.GetGroup(0)
	serverPrivateKey, _ := g.GeneratePrivateKey(nil)

	// make sure the key does not contain '\n' or '%'
	for {
		if contains(serverPrivateKey.Bytes(), byte('\n')) || contains(serverPrivateKey.Bytes(), byte('%')) {
			newKey, _ := g.GeneratePrivateKey(nil)
			serverPrivateKey = newKey
		} else {
			break
		}
	}

	// generate public key
	serverPublicKey := serverPrivateKey.Bytes()

	// listening for client public key
	for {
		message, _ := bufio.NewReader(c.connection).ReadString('\n')
		if len(message) > len(publicKeyCode) {
			if message[0:len(publicKeyCode)] == publicKeyCode {
				clientPublicKey = []byte(message[len(publicKeyCode) : len(message)-1])
				break
			}
		}
	}

	// sending server public key
	msg := publicKeyCode + string(serverPublicKey) + "\n"
	fmt.Fprintf(c.connection, msg)

	// finding common key
	pubKey := dhkx.NewPublicKey(clientPublicKey)
	k, _ := g.ComputeKey(pubKey, serverPrivateKey)
	c.clientKey = string(k.Bytes()[0:32])

	var nick string
	// waiting for nickname
	for {
		msg, _ := bufio.NewReader(c.connection).ReadString('\n')
		message := decrypt([]byte(c.clientKey), msg)
		if len(message) > len(nickCode) {
			if message[0:len(nickCode)] == nickCode {
				nick = message[len(nickCode) : len(message)-1]
				fmt.Println("NICKNAME RECIEVED")
				break
			}
		}
	}

	// start the client thread and place in room 0 (Lobby)
	c.nick = nick
	c.startThread()
	fmt.Println(c.connection.RemoteAddr().String() + " connected as " + c.nick)
	switchRoom(c, rooms[0])
}

func main() {
	setup()
	startServer()
}

// start server listening
func startServer() {
	port := ":8081"
	ln, _ := net.Listen("tcp", port)
	clear()
	fmt.Println("Server is listening on " + port)

	// accept incomming connectoins, make new client struct and start thread for exchanging keys
	for {
		conn, _ := ln.Accept()
		newClient := Client{connection: conn}
		go exchangeKeys(newClient)
	}
}

// listen for messages from client
func (c Client) listener() {
	quitting := false

	for {
		message, _ := bufio.NewReader(c.connection).ReadString('\n')
		if len(message) > 0 {
			msg := message[0 : len(message)-1]

			// decrypt message
			msgDecrypted := decrypt([]byte(c.clientKey), msg)

			// if client wants to quit
			if msgDecrypted == "!quit" {
				quitting = true
			}

			// if not quitting, print message, check for commands
			if !quitting {
				fmt.Print(c.connection.RemoteAddr().String() + " (" + c.nick + "): " + msgDecrypted)

				if !checkForCmd(c, msgDecrypted) {
					// for each client in the same room, encrypt the message
					// with their key and send
					for mapKey, value := range clientRoom {
						if mapKey != c && value == clientRoom[c] {
							mapKey.sendEncrypted(makeBold(c.nick) + " > " + msgDecrypted)
						}
					}
				}
			} else {
				// delete client from the map, close connection and print
				delete(clientRoom, c)
				c.connection.Close()
				fmt.Println(c.connection.RemoteAddr().String() + " has disconnected")
				break
			}
		}
	}
}

// send message to client
func (c *Client) send(message []byte) {
	c.connection.Write(message)
}

// send encrypted message to client
func (c *Client) sendEncrypted(message string) {
	cryptMsg := encrypt([]byte(c.clientKey), message)
	c.connection.Write([]byte(cryptMsg + "\n"))
}

// start a listener thread for a client
func (c *Client) startThread() {
	go c.listener()
}

// check messsage for commands
func checkForCmd(client Client, msg string) bool {
	if len(msg) > 1 {
		words := strings.Split(msg[0:len(msg)-1], " ")
		switch words[0] {
		// !room displays a list of available rooms
		case "!room":
			// if more than 1 word = we have arguments
			if len(words) > 1 {
				// if attempting to join the room client is already in
				if words[1] == clientRoom[client].name {
					message := "You are already in this room\nType" + makeBold(" !room ") + "to get a list of other available chatrooms\n"
					client.sendEncrypted(message)
				} else {
					// look for the room the client wants to join, and join if found
					for _, element := range rooms {
						if element.name == words[1] {
							client.sendEncrypted("Switching room: " + makeBold(element.name) + "\n")
							switchRoom(client, element)
						}
					}
				}
			} else {
				roomMsg := ""
				for _, element := range rooms {
					roomMsg += " - " + element.name + "\n"
				}
				client.sendEncrypted(roomMsg)
			}
			return true
		}
	}
	return false
}

// client switch room
func switchRoom(client Client, room Room) {
	// tell clients in the room that we are leaving that
	// and clients in the new room that we are joining
	for k, v := range clientRoom {
		if k != client && v == room {
			k.sendEncrypted(makeBold(client.nick) + " has joined " + makeBold(room.name) + "\n")
		}
		if k != client && v != room {
			k.sendEncrypted(makeBold(client.nick) + " has left " + makeBold(clientRoom[client].name) + "\n")
		}
	}

	// change room
	clientRoom[client] = room
	// send the rooms welcome message
	client.sendEncrypted(makeBold(room.welcomeMsg) + "\n")
}

// make a new room
func makeRoom(name string, welcomeMsg string) {
	newRoom := Room{name: name, welcomeMsg: welcomeMsg}
	rooms = append(rooms, newRoom)
}

// make string bold
func makeBold(text string) string {
	return "\033[1m" + text + "\033[0m"
}

// create key
func createKey() []byte {
	randKey := make([]byte, 32)
	_, err := rand.Read(randKey)
	if err != nil {
		panic(err)
	}
	return randKey
}

// encrypt message
func encrypt(key []byte, text string) string {
	plaintext := []byte(text)

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)

	return base64.URLEncoding.EncodeToString(ciphertext)
}

// decrypt message
func decrypt(key []byte, cryptoText string) string {
	ciphertext, _ := base64.URLEncoding.DecodeString(cryptoText)

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}
	if len(ciphertext) < aes.BlockSize {
		panic("Ciphertext too short")
	}

	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)

	stream.XORKeyStream(ciphertext, ciphertext)

	return fmt.Sprintf("%s", ciphertext)
}
<commit_msg>clients in the same room are notified when a client leaves<commit_after>package main

import (
	"bufio"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strings"

	"github.com/monnand/dhkx"
)

// struct for holding room info
type Room struct {
	name        string
	discription string
	password    string
	maxClients  int
	welcomeMsg  string
}

// struct for holding client info
type Client struct {
	connection net.Conn
	nick       string
	clientKey  string
}

var rooms []Room
var clientRoom map[Client]Room
var publicKeyCode string
var nickCode string

// clear the terminal screen
func clear() {
	switch runtime.GOOS {
	case "linux":
		cmd := exec.Command("clear")
		cmd.Stdout = os.Stdout
		cmd.Run()
	case "windows":
		cmd := exec.Command("cmd", "/c", "cls")
		cmd.Stdout = os.Stdout
		cmd.Run()
	default:
		fmt.Println("Attempted to clear terminal, but OS is not supported.")
	}
}

// init
func setup() {
	publicKeyCode = "ssd990=+?¡][ªs)(sdª]ßð=S)]"
	nickCode = "!#28jKas>zzx'**!+?,>lzc012"
	clientRoom = make(map[Client]Room)
	makeRoom("Lobby", "Welcome to Lobby")
	makeRoom("TestRoom", "Welcome to TestRoom")
}

// check if []byte contains specific byte
func contains(s []byte, e byte) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}

// exchange key with client
func exchangeKeys(c Client) {

	var clientPublicKey []byte

	// generate private key
	g, _ := dhkx.GetGroup(0)
	serverPrivateKey, _ := g.GeneratePrivateKey(nil)

	// make sure the key does not contain '\n' or '%'
	for {
		if contains(serverPrivateKey.Bytes(), byte('\n')) || contains(serverPrivateKey.Bytes(), byte('%')) {
			newKey, _ := g.GeneratePrivateKey(nil)
			serverPrivateKey = newKey
		} else {
			break
		}
	}

	// generate public key
	serverPublicKey := serverPrivateKey.Bytes()

	// listening for client public key
	for {
		message, _ := bufio.NewReader(c.connection).ReadString('\n')
		if len(message) > len(publicKeyCode) {
			if message[0:len(publicKeyCode)] == publicKeyCode {
				clientPublicKey = []byte(message[len(publicKeyCode) : len(message)-1])
				break
			}
		}
	}

	// sending server public key
	msg := publicKeyCode + string(serverPublicKey) + "\n"
	fmt.Fprintf(c.connection, msg)

	// finding common key
	pubKey := dhkx.NewPublicKey(clientPublicKey)
	k, _ := g.ComputeKey(pubKey, serverPrivateKey)
	c.clientKey = string(k.Bytes()[0:32])

	var nick string
	// waiting for nickname
	for {
		msg, _ := bufio.NewReader(c.connection).ReadString('\n')
		message := decrypt([]byte(c.clientKey), msg)
		if len(message) > len(nickCode) {
			if message[0:len(nickCode)] == nickCode {
				nick = message[len(nickCode) : len(message)-1]
				fmt.Println("NICKNAME RECIEVED")
				break
			}
		}
	}

	// start the client thread and place in room 0 (Lobby)
	c.nick = nick
	c.startThread()
	fmt.Println(c.connection.RemoteAddr().String() + " connected as " + c.nick)
	switchRoom(c, rooms[0])
}

func main() {
	setup()
	startServer()
}

// start server listening
func startServer() {
	port := ":8081"
	ln, _ := net.Listen("tcp", port)
	clear()
	fmt.Println("Server is listening on " + port)

	// accept incomming connectoins, make new client struct and start thread for exchanging keys
	for {
		conn, _ := ln.Accept()
		newClient := Client{connection: conn}
		go exchangeKeys(newClient)
	}
}

// listen for messages from client
func (c Client) listener() {
	quitting := false

	for {
		message, _ := bufio.NewReader(c.connection).ReadString('\n')
		if len(message) > 0 {
			msg := message[0 : len(message)-1]

			// decrypt message
			msgDecrypted := decrypt([]byte(c.clientKey), msg)

			// if client wants to quit
			if msgDecrypted == "!quit" {
				quitting = true
			}

			// if not quitting, print message, check for commands
			if !quitting {
				fmt.Print(c.connection.RemoteAddr().String() + " (" + c.nick + "): " + msgDecrypted)

				if !checkForCmd(c, msgDecrypted) {
					// for each client in the same room, encrypt the message
					// with their key and send
					for mapKey, value := range clientRoom {
						if mapKey != c && value == clientRoom[c] {
							mapKey.sendEncrypted(makeBold(c.nick) + " > " + msgDecrypted)
						}
					}
				}
			} else {
				// delete client from the map, close connection and print
				for m, v := range clientRoom {
					if m != c && v == clientRoom[c] {
						m.sendEncrypted(makeBold(c.nick) + " disconnected")
					}
				}
				delete(clientRoom, c)
				c.connection.Close()
				fmt.Println(c.connection.RemoteAddr().String() + " has disconnected")
				break
			}
		}
	}
}

// send message to client
func (c *Client) send(message []byte) {
	c.connection.Write(message)
}

// send encrypted message to client
func (c *Client) sendEncrypted(message string) {
	cryptMsg := encrypt([]byte(c.clientKey), message)
	c.connection.Write([]byte(cryptMsg + "\n"))
}

// start a listener thread for a client
func (c *Client) startThread() {
	go c.listener()
}

// check messsage for commands
func checkForCmd(client Client, msg string) bool {
	if len(msg) > 1 {
		words := strings.Split(msg[0:len(msg)-1], " ")
		switch words[0] {
		// !room displays a list of available rooms
		case "!room":
			// if more than 1 word = we have arguments
			if len(words) > 1 {
				// if attempting to join the room client is already in
				if words[1] == clientRoom[client].name {
					message := "You are already in this room\nType" + makeBold(" !room ") + "to get a list of other available chatrooms\n"
					client.sendEncrypted(message)
				} else {
					// look for the room the client wants to join, and join if found
					for _, element := range rooms {
						if element.name == words[1] {
							client.sendEncrypted("Switching room: " + makeBold(element.name) + "\n")
							switchRoom(client, element)
						}
					}
				}
			} else {
				roomMsg := ""
				for _, element := range rooms {
					roomMsg += " - " + element.name + "\n"
				}
				client.sendEncrypted(roomMsg)
			}
			return true
		}
	}
	return false
}

// client switch room
func switchRoom(client Client, room Room) {
	// tell clients in the room that we are leaving that
	// and clients in the new room that we are joining
	for k, v := range clientRoom {
		if k != client && v == room {
			k.sendEncrypted(makeBold(client.nick) + " has joined " + makeBold(room.name) + "\n")
		}
		if k != client && v != room {
			k.sendEncrypted(makeBold(client.nick) + " has left " + makeBold(clientRoom[client].name) + "\n")
		}
	}

	// change room
	clientRoom[client] = room
	// send the rooms welcome message
	client.sendEncrypted(makeBold(room.welcomeMsg) + "\n")
}

// make a new room
func makeRoom(name string, welcomeMsg string) {
	newRoom := Room{name: name, welcomeMsg: welcomeMsg}
	rooms = append(rooms, newRoom)
}

// make string bold
func makeBold(text string) string {
	return "\033[1m" + text + "\033[0m"
}

// create key
func createKey() []byte {
	randKey := make([]byte, 32)
	_, err := rand.Read(randKey)
	if err != nil {
		panic(err)
	}
	return randKey
}

// encrypt message
func encrypt(key []byte, text string) string {
	plaintext := []byte(text)

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err)
	}

	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(ciphertext[aes.BlockSize:], plaintext)

	return base64.URLEncoding.EncodeToString(ciphertext)
}

// decrypt message
func decrypt(key []byte, cryptoText string) string {
	ciphertext, _ := base64.URLEncoding.DecodeString(cryptoText)

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}
	if len(ciphertext) < aes.BlockSize {
		panic("Ciphertext too short")
	}

	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	stream := cipher.NewCFBDecrypter(block, iv)

	stream.XORKeyStream(ciphertext, ciphertext)

	return fmt.Sprintf("%s", ciphertext)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/pkg"
	"github.com/coreos/coreos-cloudinit/system"
)

const (
	version               = "0.8.2+git"
	datasourceInterval    = 100 * time.Millisecond
	datasourceMaxInterval = 30 * time.Second
	datasourceTimeout     = 5 * time.Minute
)

var (
	printVersion  bool
	ignoreFailure bool
	sources       struct {
		file            string
		configDrive     string
		metadataService bool
		url             string
		procCmdLine     bool
	}
	convertNetconf string
	workspace      string
	sshKeyName     string
)

func init() {
	flag.BoolVar(&printVersion, "version", false, "Print the version and exit")
	flag.BoolVar(&ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")
	flag.StringVar(&sources.file, "from-file", "", "Read user-data from provided file")
	flag.StringVar(&sources.configDrive, "from-configdrive", "", "Read data from provided cloud-drive directory")
	flag.BoolVar(&sources.metadataService, "from-metadata-service", false, "Download data from metadata service")
	flag.StringVar(&sources.url, "from-url", "", "Download user-data from provided url")
	flag.BoolVar(&sources.procCmdLine, "from-proc-cmdline", false, fmt.Sprintf("Parse %s for '%s=<url>', using the cloud-config served by an HTTP GET to <url>", datasource.ProcCmdlineLocation, datasource.ProcCmdlineCloudConfigFlag))
	flag.StringVar(&convertNetconf, "convert-netconf", "", "Read the network config provided in cloud-drive and translate it from the specified format into networkd unit files (requires the -from-configdrive flag)")
	flag.StringVar(&workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")
	flag.StringVar(&sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")
}

func main() {
	flag.Parse()

	die := func() {
		if ignoreFailure {
			os.Exit(0)
		}
		os.Exit(1)
	}

	if printVersion == true {
		fmt.Printf("coreos-cloudinit version %s\n", version)
		os.Exit(0)
	}

	if convertNetconf != "" && sources.configDrive == "" {
		fmt.Println("-convert-netconf flag requires -from-configdrive")
		os.Exit(1)
	}

	switch convertNetconf {
	case "":
	case "debian":
	default:
		fmt.Printf("Invalid option to -convert-netconf: '%s'. Supported options: 'debian'\n", convertNetconf)
		os.Exit(1)
	}

	dss := getDatasources()
	if len(dss) == 0 {
		fmt.Println("Provide at least one of --from-file, --from-configdrive, --from-metadata-service, --from-url or --from-proc-cmdline")
		os.Exit(1)
	}

	ds := selectDatasource(dss)
	if ds == nil {
		fmt.Println("No datasources available in time")
		die()
	}

	fmt.Printf("Fetching user-data from datasource of type %q\n", ds.Type())
	userdataBytes, err := ds.FetchUserdata()
	if err != nil {
		fmt.Printf("Failed fetching user-data from datasource: %v\n", err)
		die()
	}

	fmt.Printf("Fetching meta-data from datasource of type %q\n", ds.Type())
	metadataBytes, err := ds.FetchMetadata()
	if err != nil {
		fmt.Printf("Failed fetching meta-data from datasource: %v\n", err)
		die()
	}

	// Extract IPv4 addresses from metadata if possible
	var subs map[string]string
	if len(metadataBytes) > 0 {
		subs, err = initialize.ExtractIPsFromMetadata(metadataBytes)
		if err != nil {
			fmt.Printf("Failed extracting IPs from meta-data: %v\n", err)
			die()
		}
	}

	env := initialize.NewEnvironment("/", ds.ConfigRoot(), workspace, convertNetconf, sshKeyName, subs)

	var ccm, ccu *initialize.CloudConfig
	var script *system.Script
	if ccm, err = initialize.ParseMetaData(string(metadataBytes)); err != nil {
		fmt.Printf("Failed to parse meta-data: %v\n", err)
		die()
	}
	if ud, err := initialize.ParseUserData(string(userdataBytes)); err != nil {
		fmt.Printf("Failed to parse user-data: %v\n", err)
		die()
	} else {
		switch t := ud.(type) {
		case *initialize.CloudConfig:
			ccu = t
		case system.Script:
			script = &t
		}
	}

	var cc *initialize.CloudConfig
	if ccm != nil && ccu != nil {
		fmt.Println("Merging cloud-config from meta-data and user-data")
		merged := mergeCloudConfig(*ccu, *ccm)
		cc = &merged
	} else if ccm != nil && ccu == nil {
		fmt.Println("Processing cloud-config from meta-data")
		cc = ccm
	} else if ccm == nil && ccu != nil {
		fmt.Println("Processing cloud-config from user-data")
		cc = ccu
	} else {
		fmt.Println("No cloud-config data to handle.")
	}

	if cc != nil {
		if err = initialize.Apply(*cc, env); err != nil {
			fmt.Printf("Failed to apply cloud-config: %v\n", err)
			die()
		}
	}

	if script != nil {
		if err = runScript(*script, env); err != nil {
			fmt.Printf("Failed to run script: %v\n", err)
			die()
		}
	}
}

// mergeCloudConfig merges certain options from mdcc (a CloudConfig derived from
// meta-data) onto udcc (a CloudConfig derived from user-data), if they are
// not already set on udcc (i.e. user-data always takes precedence)
// NB: This needs to be kept in sync with ParseMetadata so that it tracks all
// elements of a CloudConfig which that function can populate.
func mergeCloudConfig(mdcc, udcc initialize.CloudConfig) (cc initialize.CloudConfig) {
	if mdcc.Hostname != "" {
		if udcc.Hostname != "" {
			fmt.Printf("Warning: user-data hostname (%s) overrides metadata hostname (%s)", udcc.Hostname, mdcc.Hostname)
		} else {
			udcc.Hostname = mdcc.Hostname
		}

	}
	for _, key := range mdcc.SSHAuthorizedKeys {
		udcc.SSHAuthorizedKeys = append(udcc.SSHAuthorizedKeys, key)
	}
	if mdcc.NetworkConfigPath != "" {
		if udcc.NetworkConfigPath != "" {
			fmt.Printf("Warning: user-data NetworkConfigPath %s overrides metadata NetworkConfigPath %s", udcc.NetworkConfigPath, mdcc.NetworkConfigPath)
		} else {
			udcc.NetworkConfigPath = mdcc.NetworkConfigPath
		}
	}
	return udcc
}

// getDatasources creates a slice of possible Datasources for cloudinit based
// on the different source command-line flags.
func getDatasources() []datasource.Datasource {
	dss := make([]datasource.Datasource, 0, 5)
	if sources.file != "" {
		dss = append(dss, datasource.NewLocalFile(sources.file))
	}
	if sources.url != "" {
		dss = append(dss, datasource.NewRemoteFile(sources.url))
	}
	if sources.configDrive != "" {
		dss = append(dss, datasource.NewConfigDrive(sources.configDrive))
	}
	if sources.metadataService {
		dss = append(dss, datasource.NewMetadataService())
	}
	if sources.procCmdLine {
		dss = append(dss, datasource.NewProcCmdline())
	}
	return dss
}

// selectDatasource attempts to choose a valid Datasource to use based on its
// current availability. The first Datasource to report to be available is
// returned. Datasources will be retried if possible if they are not
// immediately available. If all Datasources are permanently unavailable or
// datasourceTimeout is reached before one becomes available, nil is returned.
func selectDatasource(sources []datasource.Datasource) datasource.Datasource {
	ds := make(chan datasource.Datasource)
	stop := make(chan struct{})
	var wg sync.WaitGroup

	for _, s := range sources {
		wg.Add(1)
		go func(s datasource.Datasource) {
			defer wg.Done()

			duration := datasourceInterval
			for {
				fmt.Printf("Checking availability of %q\n", s.Type())
				if s.IsAvailable() {
					ds <- s
					return
				} else if !s.AvailabilityChanges() {
					return
				}
				select {
				case <-stop:
					return
				case <-time.Tick(duration):
					duration = pkg.ExpBackoff(duration, datasourceMaxInterval)
				}
			}
		}(s)
	}

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	var s datasource.Datasource
	select {
	case s = <-ds:
	case <-done:
	case <-time.Tick(datasourceTimeout):
	}

	close(stop)
	return s
}

// TODO(jonboulle): this should probably be refactored and moved into a different module
func runScript(script system.Script, env *initialize.Environment) error {
	err := initialize.PrepWorkspace(env.Workspace())
	if err != nil {
		fmt.Printf("Failed preparing workspace: %v\n", err)
		return err
	}
	path, err := initialize.PersistScriptInWorkspace(script, env.Workspace())
	if err == nil {
		var name string
		name, err = system.ExecuteScript(path)
		initialize.PersistUnitNameInWorkspace(name, env.Workspace())
	}
	return err
}
<commit_msg>coreos-cloudinit: fix ordering error in merge invocation<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/initialize"
	"github.com/coreos/coreos-cloudinit/pkg"
	"github.com/coreos/coreos-cloudinit/system"
)

const (
	version               = "0.8.2+git"
	datasourceInterval    = 100 * time.Millisecond
	datasourceMaxInterval = 30 * time.Second
	datasourceTimeout     = 5 * time.Minute
)

var (
	printVersion  bool
	ignoreFailure bool
	sources       struct {
		file            string
		configDrive     string
		metadataService bool
		url             string
		procCmdLine     bool
	}
	convertNetconf string
	workspace      string
	sshKeyName     string
)

func init() {
	flag.BoolVar(&printVersion, "version", false, "Print the version and exit")
	flag.BoolVar(&ignoreFailure, "ignore-failure", false, "Exits with 0 status in the event of malformed input from user-data")
	flag.StringVar(&sources.file, "from-file", "", "Read user-data from provided file")
	flag.StringVar(&sources.configDrive, "from-configdrive", "", "Read data from provided cloud-drive directory")
	flag.BoolVar(&sources.metadataService, "from-metadata-service", false, "Download data from metadata service")
	flag.StringVar(&sources.url, "from-url", "", "Download user-data from provided url")
	flag.BoolVar(&sources.procCmdLine, "from-proc-cmdline", false, fmt.Sprintf("Parse %s for '%s=<url>', using the cloud-config served by an HTTP GET to <url>", datasource.ProcCmdlineLocation, datasource.ProcCmdlineCloudConfigFlag))
	flag.StringVar(&convertNetconf, "convert-netconf", "", "Read the network config provided in cloud-drive and translate it from the specified format into networkd unit files (requires the -from-configdrive flag)")
	flag.StringVar(&workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")
	flag.StringVar(&sshKeyName, "ssh-key-name", initialize.DefaultSSHKeyName, "Add SSH keys to the system with the given name")
}

func main() {
	flag.Parse()

	die := func() {
		if ignoreFailure {
			os.Exit(0)
		}
		os.Exit(1)
	}

	if printVersion == true {
		fmt.Printf("coreos-cloudinit version %s\n", version)
		os.Exit(0)
	}

	if convertNetconf != "" && sources.configDrive == "" {
		fmt.Println("-convert-netconf flag requires -from-configdrive")
		os.Exit(1)
	}

	switch convertNetconf {
	case "":
	case "debian":
	default:
		fmt.Printf("Invalid option to -convert-netconf: '%s'. Supported options: 'debian'\n", convertNetconf)
		os.Exit(1)
	}

	dss := getDatasources()
	if len(dss) == 0 {
		fmt.Println("Provide at least one of --from-file, --from-configdrive, --from-metadata-service, --from-url or --from-proc-cmdline")
		os.Exit(1)
	}

	ds := selectDatasource(dss)
	if ds == nil {
		fmt.Println("No datasources available in time")
		die()
	}

	fmt.Printf("Fetching user-data from datasource of type %q\n", ds.Type())
	userdataBytes, err := ds.FetchUserdata()
	if err != nil {
		fmt.Printf("Failed fetching user-data from datasource: %v\n", err)
		die()
	}

	fmt.Printf("Fetching meta-data from datasource of type %q\n", ds.Type())
	metadataBytes, err := ds.FetchMetadata()
	if err != nil {
		fmt.Printf("Failed fetching meta-data from datasource: %v\n", err)
		die()
	}

	// Extract IPv4 addresses from metadata if possible
	var subs map[string]string
	if len(metadataBytes) > 0 {
		subs, err = initialize.ExtractIPsFromMetadata(metadataBytes)
		if err != nil {
			fmt.Printf("Failed extracting IPs from meta-data: %v\n", err)
			die()
		}
	}

	env := initialize.NewEnvironment("/", ds.ConfigRoot(), workspace, convertNetconf, sshKeyName, subs)

	var ccm, ccu *initialize.CloudConfig
	var script *system.Script
	if ccm, err = initialize.ParseMetaData(string(metadataBytes)); err != nil {
		fmt.Printf("Failed to parse meta-data: %v\n", err)
		die()
	}
	if ud, err := initialize.ParseUserData(string(userdataBytes)); err != nil {
		fmt.Printf("Failed to parse user-data: %v\n", err)
		die()
	} else {
		switch t := ud.(type) {
		case *initialize.CloudConfig:
			ccu = t
		case system.Script:
			script = &t
		}
	}

	var cc *initialize.CloudConfig
	if ccm != nil && ccu != nil {
		fmt.Println("Merging cloud-config from meta-data and user-data")
		merged := mergeCloudConfig(*ccm, *ccu)
		cc = &merged
	} else if ccm != nil && ccu == nil {
		fmt.Println("Processing cloud-config from meta-data")
		cc = ccm
	} else if ccm == nil && ccu != nil {
		fmt.Println("Processing cloud-config from user-data")
		cc = ccu
	} else {
		fmt.Println("No cloud-config data to handle.")
	}

	if cc != nil {
		if err = initialize.Apply(*cc, env); err != nil {
			fmt.Printf("Failed to apply cloud-config: %v\n", err)
			die()
		}
	}

	if script != nil {
		if err = runScript(*script, env); err != nil {
			fmt.Printf("Failed to run script: %v\n", err)
			die()
		}
	}
}

// mergeCloudConfig merges certain options from mdcc (a CloudConfig derived from
// meta-data) onto udcc (a CloudConfig derived from user-data), if they are
// not already set on udcc (i.e. user-data always takes precedence)
// NB: This needs to be kept in sync with ParseMetadata so that it tracks all
// elements of a CloudConfig which that function can populate.
func mergeCloudConfig(mdcc, udcc initialize.CloudConfig) (cc initialize.CloudConfig) {
	if mdcc.Hostname != "" {
		if udcc.Hostname != "" {
			fmt.Printf("Warning: user-data hostname (%s) overrides metadata hostname (%s)", udcc.Hostname, mdcc.Hostname)
		} else {
			udcc.Hostname = mdcc.Hostname
		}

	}
	for _, key := range mdcc.SSHAuthorizedKeys {
		udcc.SSHAuthorizedKeys = append(udcc.SSHAuthorizedKeys, key)
	}
	if mdcc.NetworkConfigPath != "" {
		if udcc.NetworkConfigPath != "" {
			fmt.Printf("Warning: user-data NetworkConfigPath %s overrides metadata NetworkConfigPath %s", udcc.NetworkConfigPath, mdcc.NetworkConfigPath)
		} else {
			udcc.NetworkConfigPath = mdcc.NetworkConfigPath
		}
	}
	return udcc
}

// getDatasources creates a slice of possible Datasources for cloudinit based
// on the different source command-line flags.
func getDatasources() []datasource.Datasource {
	dss := make([]datasource.Datasource, 0, 5)
	if sources.file != "" {
		dss = append(dss, datasource.NewLocalFile(sources.file))
	}
	if sources.url != "" {
		dss = append(dss, datasource.NewRemoteFile(sources.url))
	}
	if sources.configDrive != "" {
		dss = append(dss, datasource.NewConfigDrive(sources.configDrive))
	}
	if sources.metadataService {
		dss = append(dss, datasource.NewMetadataService())
	}
	if sources.procCmdLine {
		dss = append(dss, datasource.NewProcCmdline())
	}
	return dss
}

// selectDatasource attempts to choose a valid Datasource to use based on its
// current availability. The first Datasource to report to be available is
// returned. Datasources will be retried if possible if they are not
// immediately available. If all Datasources are permanently unavailable or
// datasourceTimeout is reached before one becomes available, nil is returned.
func selectDatasource(sources []datasource.Datasource) datasource.Datasource {
	ds := make(chan datasource.Datasource)
	stop := make(chan struct{})
	var wg sync.WaitGroup

	for _, s := range sources {
		wg.Add(1)
		go func(s datasource.Datasource) {
			defer wg.Done()

			duration := datasourceInterval
			for {
				fmt.Printf("Checking availability of %q\n", s.Type())
				if s.IsAvailable() {
					ds <- s
					return
				} else if !s.AvailabilityChanges() {
					return
				}
				select {
				case <-stop:
					return
				case <-time.Tick(duration):
					duration = pkg.ExpBackoff(duration, datasourceMaxInterval)
				}
			}
		}(s)
	}

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	var s datasource.Datasource
	select {
	case s = <-ds:
	case <-done:
	case <-time.Tick(datasourceTimeout):
	}

	close(stop)
	return s
}

// TODO(jonboulle): this should probably be refactored and moved into a different module
func runScript(script system.Script, env *initialize.Environment) error {
	err := initialize.PrepWorkspace(env.Workspace())
	if err != nil {
		fmt.Printf("Failed preparing workspace: %v\n", err)
		return err
	}
	path, err := initialize.PersistScriptInWorkspace(script, env.Workspace())
	if err == nil {
		var name string
		name, err = system.ExecuteScript(path)
		initialize.PersistUnitNameInWorkspace(name, env.Workspace())
	}
	return err
}
<|endoftext|>


<commit_before>package runtime

// Version is the current version of the buffalo binary
const Version = "v0.13.4"
<commit_msg>version bump: v0.13.5<commit_after>package runtime

// Version is the current version of the buffalo binary
const Version = "v0.13.5"
<|endoftext|>


<commit_before>package veneur

import (
	"crypto/tls"
	"fmt"
	"net"
	"os"
	"sync"

	"github.com/Sirupsen/logrus"
	flock "github.com/theckman/go-flock"
)

// StartStatsd spawns a goroutine that listens for metrics in statsd
// format on the address a. As this is a setup routine, if any error
// occurs, it panics.
func StartStatsd(s *Server, a net.Addr, packetPool *sync.Pool) {
	switch addr := a.(type) {
	case *net.UDPAddr:
		startStatsdUDP(s, addr, packetPool)
	case *net.TCPAddr:
		startStatsdTCP(s, addr, packetPool)
	default:
		panic(fmt.Sprintf("Can't listen on %v: only TCP and UDP are supported", a))
	}
}

func startStatsdUDP(s *Server, addr *net.UDPAddr, packetPool *sync.Pool) {
	for i := 0; i < s.numReaders; i++ {
		go func() {
			defer func() {
				ConsumePanic(s.Sentry, s.Statsd, s.Hostname, recover())
			}()
			// each goroutine gets its own socket
			// if the sockets support SO_REUSEPORT, then this will cause the
			// kernel to distribute datagrams across them, for better read
			// performance
			sock, err := NewSocket(addr, s.RcvbufBytes, s.numReaders != 1)
			if err != nil {
				// if any goroutine fails to create the socket, we can't really
				// recover, so we just blow up
				// this probably indicates a systemic issue, eg lack of
				// SO_REUSEPORT support
				panic(fmt.Sprintf("couldn't listen on UDP socket %v: %v", addr, err))
			}
			log.WithField("address", addr).Info("Listening for statsd metrics on UDP socket")
			s.ReadMetricSocket(sock, packetPool)
		}()
	}
}

func startStatsdTCP(s *Server, addr *net.TCPAddr, packetPool *sync.Pool) {
	var listener net.Listener
	var err error

	listener, err = net.ListenTCP("tcp", addr)
	if err != nil {
		panic(fmt.Sprintf("couldn't listen on TCP socket %v: %v", addr, err))
	}

	go func() {
		<-s.shutdown
		// TODO: the socket is in use until there are no goroutines blocked in Accept
		// we should wait until the accepting goroutine exits
		err := listener.Close()
		if err != nil {
			log.WithError(err).Warn("Ignoring error closing TCP listener")
		}
	}()

	mode := "unencrypted"
	if s.tlsConfig != nil {
		// wrap the listener with TLS
		listener = tls.NewListener(listener, s.tlsConfig)
		if s.tlsConfig.ClientAuth == tls.RequireAndVerifyClientCert {
			mode = "authenticated"
		} else {
			mode = "encrypted"
		}
	}

	log.WithFields(logrus.Fields{
		"address": addr, "mode": mode,
	}).Info("Listening for statsd metrics on TCP socket")

	go func() {
		defer func() {
			ConsumePanic(s.Sentry, s.Statsd, s.Hostname, recover())
		}()
		s.ReadTCPSocket(listener)
	}()
}

func StartSSF(s *Server, a net.Addr, tracePool *sync.Pool) {
	switch addr := a.(type) {
	case *net.UDPAddr:
		startSSFUDP(s, addr, tracePool)
	case *net.UnixAddr:
		startSSFUnix(s, addr)
	default:
		panic(fmt.Sprintf("Can't listen for SSF on %v: only udp:// & unix:// are supported", a))
	}
	log.WithFields(logrus.Fields{
		"address": a.String(),
		"network": a.Network(),
	}).Info("Listening for SSF traces")
}

func startSSFUDP(s *Server, addr *net.UDPAddr, tracePool *sync.Pool) {
	// if we want to use multiple readers, make reuseport a parameter, like ReadMetricSocket.
	listener, err := NewSocket(addr, s.RcvbufBytes, false)
	if err != nil {
		// if any goroutine fails to create the socket, we can't really
		// recover, so we just blow up
		// this probably indicates a systemic issue, eg lack of
		// SO_REUSEPORT support
		panic(fmt.Sprintf("couldn't listen on UDP socket %v: %v", addr, err))
	}
	go func() {
		defer func() {
			ConsumePanic(s.Sentry, s.Statsd, s.Hostname, recover())
		}()
		s.ReadSSFPacketSocket(listener, tracePool)
	}()
}

// startSSFUnix starts listening for connections that send framed SSF
// spans on a UNIX domain socket address. It does so until the
// server's shutdown socket is closed. startSSFUnix returns a channel
// that is closed once the listener has terminated.
func startSSFUnix(s *Server, addr *net.UnixAddr) <-chan struct{} {
	done := make(chan struct{})
	if addr.Network() != "unix" {
		panic(fmt.Sprintf("Can't listen for SSF on %v: only udp:// and unix:// addresses are supported", addr))
	}
	// ensure we are the only ones locking this socket:
	lockname := fmt.Sprintf("%s.lock", addr.String())
	lock := flock.NewFlock(lockname)
	locked, err := lock.TryLock()
	if err != nil {
		panic(fmt.Sprintf("Could not acquire the lock %q to listen on %v: %v", lockname, addr, err))
	}
	if !locked {
		panic(fmt.Sprintf("Lock file %q for %v is in use by another process already", lockname, addr))
	}
	// We have the exclusive use of the socket, clear away any old sockets and listen:
	_ = os.Remove(addr.String())
	listener, err := net.ListenUnix(addr.Network(), addr)
	if err != nil {
		panic(fmt.Sprintf("Couldn't listen on UNIX socket %v: %v", addr, err))
	}

	// Make the socket connectable by everyone with access to the socket pathname:
	err = os.Chmod(addr.String(), 0666)
	if err != nil {
		panic(fmt.Sprintf("Couldn't set permissions on %v: %v", addr, err))
	}

	go func() {
		conns := make(chan net.Conn)
		go func() {
			defer lock.Unlock()
			defer close(done)
			for {
				conn, err := listener.AcceptUnix()
				if err != nil {
					select {
					case <-s.shutdown:
						// occurs when cleanly shutting down the server e.g. in tests; ignore errors
						log.WithError(err).Info("Ignoring Accept error while shutting down")
						return
					default:
						log.WithError(err).Fatal("Unix accept failed")
					}
				}
				conns <- conn
			}
		}()
		for {
			select {
			case conn := <-conns:
				go s.ReadSSFStreamSocket(conn)
			case <-s.shutdown:
				listener.Close()
				return
			}
		}
	}()
	return done
}
<commit_msg>Use numReaders as a proxy for SSF listener using SO_REUSEPORT<commit_after>package veneur

import (
	"crypto/tls"
	"fmt"
	"net"
	"os"
	"sync"

	"github.com/Sirupsen/logrus"
	flock "github.com/theckman/go-flock"
)

// StartStatsd spawns a goroutine that listens for metrics in statsd
// format on the address a. As this is a setup routine, if any error
// occurs, it panics.
func StartStatsd(s *Server, a net.Addr, packetPool *sync.Pool) {
	switch addr := a.(type) {
	case *net.UDPAddr:
		startStatsdUDP(s, addr, packetPool)
	case *net.TCPAddr:
		startStatsdTCP(s, addr, packetPool)
	default:
		panic(fmt.Sprintf("Can't listen on %v: only TCP and UDP are supported", a))
	}
}

func startStatsdUDP(s *Server, addr *net.UDPAddr, packetPool *sync.Pool) {
	for i := 0; i < s.numReaders; i++ {
		go func() {
			defer func() {
				ConsumePanic(s.Sentry, s.Statsd, s.Hostname, recover())
			}()
			// each goroutine gets its own socket
			// if the sockets support SO_REUSEPORT, then this will cause the
			// kernel to distribute datagrams across them, for better read
			// performance
			sock, err := NewSocket(addr, s.RcvbufBytes, s.numReaders != 1)
			if err != nil {
				// if any goroutine fails to create the socket, we can't really
				// recover, so we just blow up
				// this probably indicates a systemic issue, eg lack of
				// SO_REUSEPORT support
				panic(fmt.Sprintf("couldn't listen on UDP socket %v: %v", addr, err))
			}
			log.WithField("address", addr).Info("Listening for statsd metrics on UDP socket")
			s.ReadMetricSocket(sock, packetPool)
		}()
	}
}

func startStatsdTCP(s *Server, addr *net.TCPAddr, packetPool *sync.Pool) {
	var listener net.Listener
	var err error

	listener, err = net.ListenTCP("tcp", addr)
	if err != nil {
		panic(fmt.Sprintf("couldn't listen on TCP socket %v: %v", addr, err))
	}

	go func() {
		<-s.shutdown
		// TODO: the socket is in use until there are no goroutines blocked in Accept
		// we should wait until the accepting goroutine exits
		err := listener.Close()
		if err != nil {
			log.WithError(err).Warn("Ignoring error closing TCP listener")
		}
	}()

	mode := "unencrypted"
	if s.tlsConfig != nil {
		// wrap the listener with TLS
		listener = tls.NewListener(listener, s.tlsConfig)
		if s.tlsConfig.ClientAuth == tls.RequireAndVerifyClientCert {
			mode = "authenticated"
		} else {
			mode = "encrypted"
		}
	}

	log.WithFields(logrus.Fields{
		"address": addr, "mode": mode,
	}).Info("Listening for statsd metrics on TCP socket")

	go func() {
		defer func() {
			ConsumePanic(s.Sentry, s.Statsd, s.Hostname, recover())
		}()
		s.ReadTCPSocket(listener)
	}()
}

func StartSSF(s *Server, a net.Addr, tracePool *sync.Pool) {
	switch addr := a.(type) {
	case *net.UDPAddr:
		startSSFUDP(s, addr, tracePool)
	case *net.UnixAddr:
		startSSFUnix(s, addr)
	default:
		panic(fmt.Sprintf("Can't listen for SSF on %v: only udp:// & unix:// are supported", a))
	}
	log.WithFields(logrus.Fields{
		"address": a.String(),
		"network": a.Network(),
	}).Info("Listening for SSF traces")
}

func startSSFUDP(s *Server, addr *net.UDPAddr, tracePool *sync.Pool) {
	// TODO: Make this actually use readers / add a predicate
	// function for testing if we should SO_REUSEPORT.
	listener, err := NewSocket(addr, s.RcvbufBytes, s.numReaders > 1)
	if err != nil {
		// if any goroutine fails to create the socket, we can't really
		// recover, so we just blow up
		// this probably indicates a systemic issue, eg lack of
		// SO_REUSEPORT support
		panic(fmt.Sprintf("couldn't listen on UDP socket %v: %v", addr, err))
	}
	go func() {
		defer func() {
			ConsumePanic(s.Sentry, s.Statsd, s.Hostname, recover())
		}()
		s.ReadSSFPacketSocket(listener, tracePool)
	}()
}

// startSSFUnix starts listening for connections that send framed SSF
// spans on a UNIX domain socket address. It does so until the
// server's shutdown socket is closed. startSSFUnix returns a channel
// that is closed once the listener has terminated.
func startSSFUnix(s *Server, addr *net.UnixAddr) <-chan struct{} {
	done := make(chan struct{})
	if addr.Network() != "unix" {
		panic(fmt.Sprintf("Can't listen for SSF on %v: only udp:// and unix:// addresses are supported", addr))
	}
	// ensure we are the only ones locking this socket:
	lockname := fmt.Sprintf("%s.lock", addr.String())
	lock := flock.NewFlock(lockname)
	locked, err := lock.TryLock()
	if err != nil {
		panic(fmt.Sprintf("Could not acquire the lock %q to listen on %v: %v", lockname, addr, err))
	}
	if !locked {
		panic(fmt.Sprintf("Lock file %q for %v is in use by another process already", lockname, addr))
	}
	// We have the exclusive use of the socket, clear away any old sockets and listen:
	_ = os.Remove(addr.String())
	listener, err := net.ListenUnix(addr.Network(), addr)
	if err != nil {
		panic(fmt.Sprintf("Couldn't listen on UNIX socket %v: %v", addr, err))
	}

	// Make the socket connectable by everyone with access to the socket pathname:
	err = os.Chmod(addr.String(), 0666)
	if err != nil {
		panic(fmt.Sprintf("Couldn't set permissions on %v: %v", addr, err))
	}

	go func() {
		conns := make(chan net.Conn)
		go func() {
			defer lock.Unlock()
			defer close(done)
			for {
				conn, err := listener.AcceptUnix()
				if err != nil {
					select {
					case <-s.shutdown:
						// occurs when cleanly shutting down the server e.g. in tests; ignore errors
						log.WithError(err).Info("Ignoring Accept error while shutting down")
						return
					default:
						log.WithError(err).Fatal("Unix accept failed")
					}
				}
				conns <- conn
			}
		}()
		for {
			select {
			case conn := <-conns:
				go s.ReadSSFStreamSocket(conn)
			case <-s.shutdown:
				listener.Close()
				return
			}
		}
	}()
	return done
}
<|endoftext|>


<commit_before>package common

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/user"
	"strconv"
	"strings"
)

// CatFile cats the contents of a file (if it exists)
func CatFile(filename string) {
	slice, err := FileToSlice(filename)
	if err != nil {
		LogDebug(fmt.Sprintf("Error cat'ing file %s: %s", filename, err.Error()))
		return
	}

	for _, line := range slice {
		LogDebug(fmt.Sprintf("line: '%s'", line))
	}
}

// CopyFile copies a file from src to dst. If src and dst files exist, and are
// the same, then return success. Otherise, attempt to create a hard link
// between the two files. If that fail, copy the file contents from src to dst.
// FROM: https://stackoverflow.com/a/21067803/1515875
func CopyFile(src, dst string) (err error) {
	sfi, err := os.Stat(src)
	if err != nil {
		return
	}
	if !sfi.Mode().IsRegular() {
		// cannot copy non-regular files (e.g., directories,
		// symlinks, devices, etc.)
		return fmt.Errorf("CopyFile: non-regular source file %s (%q)", sfi.Name(), sfi.Mode().String())
	}
	dfi, err := os.Stat(dst)
	if err != nil {
		if !os.IsNotExist(err) {
			return
		}
	} else {
		if !(dfi.Mode().IsRegular()) {
			return fmt.Errorf("CopyFile: non-regular destination file %s (%q)", dfi.Name(), dfi.Mode().String())
		}
		if os.SameFile(sfi, dfi) {
			return
		}
	}
	if err = os.Link(src, dst); err == nil {
		return
	}
	err = copyFileContents(src, dst)
	return
}

// copyFileContents copies the contents of the file named src to the file named
// by dst. The file will be created if it does not already exist. If the
// destination file exists, all it's contents will be replaced by the contents
// of the source file.
// FROM: https://stackoverflow.com/a/21067803/1515875
func copyFileContents(src, dst string) (err error) {
	in, err := os.Open(src)
	if err != nil {
		return
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return
	}
	defer func() {
		cerr := out.Close()
		if err == nil {
			err = cerr
		}
	}()
	if _, err = io.Copy(out, in); err != nil {
		return
	}
	err = out.Sync()
	return
}

// DirectoryExists returns if a path exists and is a directory
func DirectoryExists(filename string) bool {
	fi, err := os.Stat(filename)
	if err != nil {
		return false
	}

	return fi.IsDir()
}

// FileToSlice reads in all the lines from a file into a string slice
func FileToSlice(filename string) (lines []string, err error) {
	f, err := os.Open(filename)
	if err != nil {
		return
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		text := strings.TrimSpace(scanner.Text())
		if text == "" {
			continue
		}
		lines = append(lines, text)
	}
	err = scanner.Err()
	return
}

// FileExists returns if a path exists and is a file
func FileExists(filename string) bool {
	fi, err := os.Stat(filename)
	if err != nil {
		return false
	}

	return fi.Mode().IsRegular()
}

// IsAbsPath returns 0 if input path is absolute
func IsAbsPath(path string) bool {
	return strings.HasPrefix(path, "/")
}

// ListFilesWithPrefix lists files within a given path that have a given prefix
func ListFilesWithPrefix(path string, prefix string) []string {
	names, err := ioutil.ReadDir(path)
	if err != nil {
		return []string{}
	}

	files := []string{}
	for _, f := range names {
		if prefix != "" && !strings.HasPrefix(f.Name(), prefix) {
			continue
		}

		if f.Mode().IsRegular() {
			files = append(files, fmt.Sprintf("%s/%s", path, f.Name()))
		}
	}

	return files
}

// ReadFirstLine gets the first line of a file that has contents and returns it
// if there are no contents, an empty string is returned
// will also return an empty string if the file does not exist
func ReadFirstLine(filename string) (text string) {
	if !FileExists(filename) {
		return
	}
	f, err := os.Open(filename)
	if err != nil {
		return
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if text = strings.TrimSpace(scanner.Text()); text == "" {
			continue
		}
		return
	}
	return
}

// SetPermissions sets the proper owner and filemode for a given file
func SetPermissions(path string, fileMode os.FileMode) error {
	if err := os.Chmod(path, fileMode); err != nil {
		return err
	}

	systemGroup := GetenvWithDefault("DOKKU_SYSTEM_GROUP", "dokku")
	systemUser := GetenvWithDefault("DOKKU_SYSTEM_USER", "dokku")
	if strings.HasPrefix("/etc/sudoers.d/", path) {
		systemGroup = "root"
		systemUser = "root"
	}

	group, err := user.LookupGroup(systemGroup)
	if err != nil {
		return err
	}
	user, err := user.Lookup(systemUser)
	if err != nil {
		return err
	}

	uid, err := strconv.Atoi(user.Uid)
	if err != nil {
		return err
	}

	gid, err := strconv.Atoi(group.Gid)
	if err != nil {
		return err
	}
	return os.Chown(path, uid, gid)
}

// WriteSliceToFile writes a slice of strings to a file
func WriteSliceToFile(filename string, lines []string) error {
	mode := os.FileMode(0600)
	if strings.HasPrefix("/etc/sudoers.d/", filename) {
		// sudoers files should be either 0600 (rhel) or 0440 (debian)
		defaultMode := map[string]bool{
			"centos": true,
			"fedora": true,
			"rhel":   true,
		}
		if !defaultMode[os.Getenv("DOKKU_DISTRO")] {
			mode = os.FileMode(0440)
		}
	}

	file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_TRUNC, mode)
	if err != nil {
		return err
	}

	w := bufio.NewWriter(file)
	for _, line := range lines {
		fmt.Fprintln(w, line)
	}
	if err = w.Flush(); err != nil {
		return err
	}

	file.Chmod(mode)
	SetPermissions(filename, mode)

	return nil
}
<commit_msg>fix: correct the prefix check<commit_after>package common

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/user"
	"strconv"
	"strings"
)

// CatFile cats the contents of a file (if it exists)
func CatFile(filename string) {
	slice, err := FileToSlice(filename)
	if err != nil {
		LogDebug(fmt.Sprintf("Error cat'ing file %s: %s", filename, err.Error()))
		return
	}

	for _, line := range slice {
		LogDebug(fmt.Sprintf("line: '%s'", line))
	}
}

// CopyFile copies a file from src to dst. If src and dst files exist, and are
// the same, then return success. Otherise, attempt to create a hard link
// between the two files. If that fail, copy the file contents from src to dst.
// FROM: https://stackoverflow.com/a/21067803/1515875
func CopyFile(src, dst string) (err error) {
	sfi, err := os.Stat(src)
	if err != nil {
		return
	}
	if !sfi.Mode().IsRegular() {
		// cannot copy non-regular files (e.g., directories,
		// symlinks, devices, etc.)
		return fmt.Errorf("CopyFile: non-regular source file %s (%q)", sfi.Name(), sfi.Mode().String())
	}
	dfi, err := os.Stat(dst)
	if err != nil {
		if !os.IsNotExist(err) {
			return
		}
	} else {
		if !(dfi.Mode().IsRegular()) {
			return fmt.Errorf("CopyFile: non-regular destination file %s (%q)", dfi.Name(), dfi.Mode().String())
		}
		if os.SameFile(sfi, dfi) {
			return
		}
	}
	if err = os.Link(src, dst); err == nil {
		return
	}
	err = copyFileContents(src, dst)
	return
}

// copyFileContents copies the contents of the file named src to the file named
// by dst. The file will be created if it does not already exist. If the
// destination file exists, all it's contents will be replaced by the contents
// of the source file.
// FROM: https://stackoverflow.com/a/21067803/1515875
func copyFileContents(src, dst string) (err error) {
	in, err := os.Open(src)
	if err != nil {
		return
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return
	}
	defer func() {
		cerr := out.Close()
		if err == nil {
			err = cerr
		}
	}()
	if _, err = io.Copy(out, in); err != nil {
		return
	}
	err = out.Sync()
	return
}

// DirectoryExists returns if a path exists and is a directory
func DirectoryExists(filename string) bool {
	fi, err := os.Stat(filename)
	if err != nil {
		return false
	}

	return fi.IsDir()
}

// FileToSlice reads in all the lines from a file into a string slice
func FileToSlice(filename string) (lines []string, err error) {
	f, err := os.Open(filename)
	if err != nil {
		return
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		text := strings.TrimSpace(scanner.Text())
		if text == "" {
			continue
		}
		lines = append(lines, text)
	}
	err = scanner.Err()
	return
}

// FileExists returns if a path exists and is a file
func FileExists(filename string) bool {
	fi, err := os.Stat(filename)
	if err != nil {
		return false
	}

	return fi.Mode().IsRegular()
}

// IsAbsPath returns 0 if input path is absolute
func IsAbsPath(path string) bool {
	return strings.HasPrefix(path, "/")
}

// ListFilesWithPrefix lists files within a given path that have a given prefix
func ListFilesWithPrefix(path string, prefix string) []string {
	names, err := ioutil.ReadDir(path)
	if err != nil {
		return []string{}
	}

	files := []string{}
	for _, f := range names {
		if prefix != "" && !strings.HasPrefix(f.Name(), prefix) {
			continue
		}

		if f.Mode().IsRegular() {
			files = append(files, fmt.Sprintf("%s/%s", path, f.Name()))
		}
	}

	return files
}

// ReadFirstLine gets the first line of a file that has contents and returns it
// if there are no contents, an empty string is returned
// will also return an empty string if the file does not exist
func ReadFirstLine(filename string) (text string) {
	if !FileExists(filename) {
		return
	}
	f, err := os.Open(filename)
	if err != nil {
		return
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		if text = strings.TrimSpace(scanner.Text()); text == "" {
			continue
		}
		return
	}
	return
}

// SetPermissions sets the proper owner and filemode for a given file
func SetPermissions(path string, fileMode os.FileMode) error {
	if err := os.Chmod(path, fileMode); err != nil {
		return err
	}

	systemGroup := GetenvWithDefault("DOKKU_SYSTEM_GROUP", "dokku")
	systemUser := GetenvWithDefault("DOKKU_SYSTEM_USER", "dokku")
	if strings.HasPrefix(path, "/etc/sudoers.d/") {
		systemGroup = "root"
		systemUser = "root"
	}

	group, err := user.LookupGroup(systemGroup)
	if err != nil {
		return err
	}
	user, err := user.Lookup(systemUser)
	if err != nil {
		return err
	}

	uid, err := strconv.Atoi(user.Uid)
	if err != nil {
		return err
	}

	gid, err := strconv.Atoi(group.Gid)
	if err != nil {
		return err
	}
	return os.Chown(path, uid, gid)
}

// WriteSliceToFile writes a slice of strings to a file
func WriteSliceToFile(filename string, lines []string) error {
	mode := os.FileMode(0600)
	if strings.HasPrefix(filename, "/etc/sudoers.d/") {
		// sudoers files should be either 0600 (rhel) or 0440 (debian)
		defaultMode := map[string]bool{
			"centos": true,
			"fedora": true,
			"rhel":   true,
		}
		if !defaultMode[os.Getenv("DOKKU_DISTRO")] {
			mode = os.FileMode(0440)
		}
	}

	file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_TRUNC, mode)
	if err != nil {
		return err
	}

	w := bufio.NewWriter(file)
	for _, line := range lines {
		fmt.Fprintln(w, line)
	}
	if err = w.Flush(); err != nil {
		return err
	}

	file.Chmod(mode)
	SetPermissions(filename, mode)

	return nil
}
<|endoftext|>


<commit_before>package hots

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"regexp"
	"strconv"

	"github.com/StalkR/goircbot/lib/transport"
)

type Score struct {
	Medal string
	Score int
	MMR   int
}

type Stats struct {
	PlayerID   int
	TeamLeague Score
	HeroLeague Score
	QuickMatch Score
}

func (s *Stats) String() string {
	tl := fmt.Sprintf("Team League: %s %d (MMR %d)", s.TeamLeague.Medal, s.TeamLeague.Score, s.TeamLeague.MMR)
	if s.TeamLeague.Medal == "" {
		tl = "Team League: n/a"
	}
	hl := fmt.Sprintf("Hero League: %s %d (MMR %d)", s.HeroLeague.Medal, s.HeroLeague.Score, s.HeroLeague.MMR)
	if s.HeroLeague.Medal == "" {
		hl = "Hero League: n/a"
	}
	qm := fmt.Sprintf("Quick Match: %s %d (MMR %d)", s.QuickMatch.Medal, s.QuickMatch.Score, s.QuickMatch.MMR)
	if s.QuickMatch.Medal == "" {
		hl = "Quick Match: n/a"
	}
	return fmt.Sprintf("%s, %s, %s - %s?PlayerID=%d", tl, hl, qm, statsURL, s.PlayerID)
}

func NewStats(id int) (*Stats, error) {
	b, err := get(id)
	if err != nil {
		return nil, err
	}
	return parse(string(b))
}

const statsURL = "https://www.hotslogs.com/Player/Profile"

func get(id int) ([]byte, error) {
	c, err := transport.Client(statsURL)
	if err != nil {
		return nil, err
	}
	v := url.Values{}
	v.Add("PlayerID", fmt.Sprintf("%d", id))
	resp, err := c.Get(fmt.Sprintf("%s?%s", statsURL, v.Encode()))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return b, nil
}

var (
	playerIDRE   = regexp.MustCompile(`<link href="https://www\.hotslogs\.com/Player/Profile\?PlayerID=(\d+)"`)
	teamLeagueRE = regexp.MustCompile(`<th>Team League</th>.*?<span>(.*?) (\d+) \(Current MMR: (\d+)\)</span>`)
	heroLeagueRE = regexp.MustCompile(`<th>Hero League</th>.*?<span>(.*?) (\d+) \(Current MMR: (\d+)\)</span>`)
	quickMatchRE = regexp.MustCompile(`<th>Quick Match</th>.*?<span>(.*?) (\d+) \(Current MMR: (\d+)\)</span>`)
)

func parse(page string) (*Stats, error) {
	var s Stats
	if m := playerIDRE.FindStringSubmatch(page); m != nil {
		playerID, err := strconv.Atoi(m[1])
		if err != nil {
			return nil, err
		}
		s.PlayerID = playerID
	} else {
		return nil, errors.New("hots: could not find player ID")
	}
	if m := teamLeagueRE.FindStringSubmatch(page); m != nil {
		s.TeamLeague.Medal = m[1]
		score, err := strconv.Atoi(m[2])
		if err != nil {
			return nil, err
		}
		s.TeamLeague.Score = score
		mmr, err := strconv.Atoi(m[3])
		if err != nil {
			return nil, err
		}
		s.TeamLeague.MMR = mmr
	}
	if m := heroLeagueRE.FindStringSubmatch(page); m != nil {
		s.HeroLeague.Medal = m[1]
		score, err := strconv.Atoi(m[2])
		if err != nil {
			return nil, err
		}
		s.HeroLeague.Score = score
		mmr, err := strconv.Atoi(m[3])
		if err != nil {
			return nil, err
		}
		s.HeroLeague.MMR = mmr
	}
	if m := quickMatchRE.FindStringSubmatch(page); m != nil {
		s.QuickMatch.Medal = m[1]
		score, err := strconv.Atoi(m[2])
		if err != nil {
			return nil, err
		}
		s.QuickMatch.Score = score
		mmr, err := strconv.Atoi(m[3])
		if err != nil {
			return nil, err
		}
		s.QuickMatch.MMR = mmr
	}
	return &s, nil
}
<commit_msg>plugins/hots: fix quick match if n/a<commit_after>package hots

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"regexp"
	"strconv"

	"github.com/StalkR/goircbot/lib/transport"
)

type Score struct {
	Medal string
	Score int
	MMR   int
}

type Stats struct {
	PlayerID   int
	TeamLeague Score
	HeroLeague Score
	QuickMatch Score
}

func (s *Stats) String() string {
	tl := fmt.Sprintf("Team League: %s %d (MMR %d)", s.TeamLeague.Medal, s.TeamLeague.Score, s.TeamLeague.MMR)
	if s.TeamLeague.Medal == "" {
		tl = "Team League: n/a"
	}
	hl := fmt.Sprintf("Hero League: %s %d (MMR %d)", s.HeroLeague.Medal, s.HeroLeague.Score, s.HeroLeague.MMR)
	if s.HeroLeague.Medal == "" {
		hl = "Hero League: n/a"
	}
	qm := fmt.Sprintf("Quick Match: %s %d (MMR %d)", s.QuickMatch.Medal, s.QuickMatch.Score, s.QuickMatch.MMR)
	if s.QuickMatch.Medal == "" {
		qm = "Quick Match: n/a"
	}
	return fmt.Sprintf("%s, %s, %s - %s?PlayerID=%d", tl, hl, qm, statsURL, s.PlayerID)
}

func NewStats(id int) (*Stats, error) {
	b, err := get(id)
	if err != nil {
		return nil, err
	}
	return parse(string(b))
}

const statsURL = "https://www.hotslogs.com/Player/Profile"

func get(id int) ([]byte, error) {
	c, err := transport.Client(statsURL)
	if err != nil {
		return nil, err
	}
	v := url.Values{}
	v.Add("PlayerID", fmt.Sprintf("%d", id))
	resp, err := c.Get(fmt.Sprintf("%s?%s", statsURL, v.Encode()))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return b, nil
}

var (
	playerIDRE   = regexp.MustCompile(`<link href="https://www\.hotslogs\.com/Player/Profile\?PlayerID=(\d+)"`)
	teamLeagueRE = regexp.MustCompile(`<th>Team League</th>.*?<span>(.*?) (\d+) \(Current MMR: (\d+)\)</span>`)
	heroLeagueRE = regexp.MustCompile(`<th>Hero League</th>.*?<span>(.*?) (\d+) \(Current MMR: (\d+)\)</span>`)
	quickMatchRE = regexp.MustCompile(`<th>Quick Match</th>.*?<span>(.*?) (\d+) \(Current MMR: (\d+)\)</span>`)
)

func parse(page string) (*Stats, error) {
	var s Stats
	if m := playerIDRE.FindStringSubmatch(page); m != nil {
		playerID, err := strconv.Atoi(m[1])
		if err != nil {
			return nil, err
		}
		s.PlayerID = playerID
	} else {
		return nil, errors.New("hots: could not find player ID")
	}
	if m := teamLeagueRE.FindStringSubmatch(page); m != nil {
		s.TeamLeague.Medal = m[1]
		score, err := strconv.Atoi(m[2])
		if err != nil {
			return nil, err
		}
		s.TeamLeague.Score = score
		mmr, err := strconv.Atoi(m[3])
		if err != nil {
			return nil, err
		}
		s.TeamLeague.MMR = mmr
	}
	if m := heroLeagueRE.FindStringSubmatch(page); m != nil {
		s.HeroLeague.Medal = m[1]
		score, err := strconv.Atoi(m[2])
		if err != nil {
			return nil, err
		}
		s.HeroLeague.Score = score
		mmr, err := strconv.Atoi(m[3])
		if err != nil {
			return nil, err
		}
		s.HeroLeague.MMR = mmr
	}
	if m := quickMatchRE.FindStringSubmatch(page); m != nil {
		s.QuickMatch.Medal = m[1]
		score, err := strconv.Atoi(m[2])
		if err != nil {
			return nil, err
		}
		s.QuickMatch.Score = score
		mmr, err := strconv.Atoi(m[3])
		if err != nil {
			return nil, err
		}
		s.QuickMatch.MMR = mmr
	}
	return &s, nil
}
<|endoftext|>


<commit_before>package pool

import "net"
import "errors"
import "sync"
import "time"

//blockingPool implements the Pool interface.
//Connestions from blockingPool offer a kind of blocking mechanism that is derived from buffered channel.
type blockingPool struct {
	//mutex is to make closing the pool and recycling the connection an atomic operation
	mutex sync.Mutex

	//timeout to Get, default to 3
	timeout time.Duration

	//storage for net.Conn connections
	conns chan *WrappedConn

	//net.Conn generator
	factory Factory
}

//Factory is a function to create new connections
//which is provided by the user
type Factory func() (net.Conn, error)

//Create a new blocking pool. As no new connections would be made when the pool is busy,
//the number of connections of the pool is kept no more than initCap and maxCap does not
//make sense but the api is reserved. The timeout to block Get() is set to 3 by default
//concerning that it is better to be related with Get() method.
func NewBlockingPool(initCap, maxCap int, factory Factory) (Pool, error) {
	if initCap < 0 || maxCap < 1 || initCap > maxCap {
		return nil, errors.New("invalid capacity settings")
	}

	newPool := &blockingPool{
		timeout: 3,
		conns: make(chan *WrappedConn, maxCap),
		factory: factory,
	}

	for i := 0; i < initCap; i++ {
		conn, _ := factory()
		//It doesn't matter if conn it nil, nil is checked whenever Get() is called,
		//so error counted from factory() is ignored.
		newPool.conns <- newPool.wrap(conn)
	}
	return newPool, nil
}

//Get blocks for an available connection.
func (p *blockingPool) Get() (net.Conn, error) {
	//in case that pool is closed or pool.conns is set to nil
	conns := p.conns
	if conns == nil {
		return nil, ErrClosed
	}

	select {
	case conn := <-conns:
		if conn.Conn == nil {
			var err error
			conn.Conn, err = p.factory()
			if err != nil {
				p.put(conn)
				return nil, err
			}
		}
		return conn, nil
	case <-time.After(time.Second*p.timeout):
		return nil, ErrTimeout
	}
}

//put puts the connection back to the pool. If the pool is closed, put simply close
//any connections received and return immediately. A nil net.Conn is illegal and will be rejected.
func (p *blockingPool) put(conn *WrappedConn) error {
	//in case that pool is closed and pool.conns is set to nil
	conns := p.conns
	if conns == nil {
		//conn.Conn is possibly nil coz factory() may fail, in which case conn is immediately
		//put back to the pool
		if conn.Conn != nil {
			return conn.Conn.Close()
		} else {
			return nil
		}
	}

	//if conn is marked unusable, underlying net.Conn is set to nil
	if conn.unusable {
		conn.Conn = nil
	}

	//It is impossible to block as number of connections is never more than length of channel
	conns <-conn
	return nil
}

//TODO
//Close set connection channel to nil and close all the relative connections.
//Yet not implemented.
func (p *blockingPool) Close() {}

//TODO
//Len return the number of current active(in use or available) connections.
func (p *blockingPool) Len() int {
	return 0;
}

<commit_msg>fix bug<commit_after>package pool

import "net"
import "errors"
import "sync"
import "time"

//blockingPool implements the Pool interface.
//Connestions from blockingPool offer a kind of blocking mechanism that is derived from buffered channel.
type blockingPool struct {
	//mutex is to make closing the pool and recycling the connection an atomic operation
	mutex sync.Mutex

	//timeout to Get, default to 3
	timeout time.Duration

	//storage for net.Conn connections
	conns chan *WrappedConn

	//net.Conn generator
	factory Factory
}

//Factory is a function to create new connections
//which is provided by the user
type Factory func() (net.Conn, error)

//Create a new blocking pool. As no new connections would be made when the pool is busy,
//the number of connections of the pool is kept no more than initCap and maxCap does not
//make sense but the api is reserved. The timeout to block Get() is set to 3 by default
//concerning that it is better to be related with Get() method.
func NewBlockingPool(initCap, maxCap int, factory Factory) (Pool, error) {
	if initCap < 0 || maxCap < 1 || initCap > maxCap {
		return nil, errors.New("invalid capacity settings")
	}

	newPool := &blockingPool{
		timeout: 3,
		conns: make(chan *WrappedConn, maxCap),
		factory: factory,
	}

	for i := 0; i < initCap; i++ {
		conn, _ := factory()
		//It doesn't matter if conn it nil, nil is checked whenever Get() is called,
		//so error counted from factory() is ignored.
		newPool.conns <- newPool.wrap(conn)
	}
	return newPool, nil
}

//Get blocks for an available connection.
func (p *blockingPool) Get() (net.Conn, error) {
	//in case that pool is closed or pool.conns is set to nil
	conns := p.conns
	if conns == nil {
		return nil, ErrClosed
	}

	select {
	case conn := <-conns:
		if conn.Conn == nil {
			var err error
			conn.Conn, err = p.factory()
			if err != nil {
				p.put(conn)
				return nil, err
			}
		}
		return conn, nil
	case <-time.After(time.Second*p.timeout):
		return nil, ErrTimeout
	}
}

//put puts the connection back to the pool. If the pool is closed, put simply close
//any connections received and return immediately. A nil net.Conn is illegal and will be rejected.
func (p *blockingPool) put(conn *WrappedConn) error {
	//in case that pool is closed and pool.conns is set to nil
	conns := p.conns
	if conns == nil {
		//conn.Conn is possibly nil coz factory() may fail, in which case conn is immediately
		//put back to the pool
		if conn.Conn != nil {
			conn.Conn.Close()
			conn.Conn = nil
		}
		return ErrClosed
	}

	//if conn is marked unusable, underlying net.Conn is set to nil
	if conn.unusable {
		if conn.Conn != nil {
			conn.Conn.Close()
			conn.Conn = nil
		}
	}

	//It is impossible to block as number of connections is never more than length of channel
	conns <-conn
	return nil
}

//TODO
//Close set connection channel to nil and close all the relative connections.
//Yet not implemented.
func (p *blockingPool) Close() {}

//TODO
//Len return the number of current active(in use or available) connections.
func (p *blockingPool) Len() int {
	return 0;
}

<|endoftext|>


<commit_before>package scanner

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"os"
	"path/filepath"
	"strings"

	"github.com/src-d/protogo/report"
)

// Package holds information about a single Go package and
// a reference of all defined structs and type aliases.
// A Package is only safe to use once it is resolved.
type Package struct {
	Resolved bool
	Path     string
	Name     string
	Structs  []*Struct
	Enums    []*Enum
	Aliases  map[string]Type
	values   map[string][]string
}

// Type is the common interface for all possible types supported in protogo.
// Type is neither a representation of a Go type nor a representation of a
// protobuf type. Is an intermediate representation to ease future steps in
// the conversion from Go to protobuf.
// All types can be nullable (or not) or repeated (or not).
type Type interface {
	SetRepeated(bool)
	SetNullable(bool)
	IsRepeated() bool
	IsNullable() bool
}

// BaseType contains the common fields for all the types.
type BaseType struct {
	Nullable bool
	Repeated bool
}

func newBaseType() *BaseType {
	return &BaseType{
		Nullable: true,
		Repeated: false,
	}
}

func (t *BaseType) IsRepeated() bool   { return t.Repeated }
func (t *BaseType) IsNullable() bool   { return t.Nullable }
func (t *BaseType) SetRepeated(v bool) { t.Repeated = v }
func (t *BaseType) SetNullable(v bool) { t.Nullable = v }

// Basic is a basic type, which only is identified by its name.
type Basic struct {
	*BaseType
	Name string
}

func NewBasic(name string) Type {
	return &Basic{
		newBaseType(),
		name,
	}
}

// Named is non-basic type identified by a name on some package.
type Named struct {
	*BaseType
	Path string
	Name string
}

func NewNamed(path, name string) Type {
	return &Named{
		newBaseType(),
		path,
		name,
	}
}

// Map is a map type with a key and a value type.
type Map struct {
	*BaseType
	Key   Type
	Value Type
}

func NewMap(key, val Type) Type {
	return &Map{
		newBaseType(),
		key,
		val,
	}
}

// Enum consists of a list of possible values.
type Enum struct {
	Name   string
	Values []string
}

// Struct represents a Go struct with its name and fields.
type Struct struct {
	Name   string
	Fields []*Field
}

func (s *Struct) HasField(name string) bool {
	for _, f := range s.Fields {
		if f.Name == name {
			return true
		}
	}
	return false
}

// Field contains name and type of a struct field.
type Field struct {
	Name string
	Type Type
}

// Scanner scans paths looking for Go source files to parse
// and extract types and structs from.
type Scanner struct {
	paths []string
}

// New creates a new Scanner that will look for types and structs
// only in the given paths.
func New(paths ...string) (*Scanner, error) {
	for _, p := range paths {
		fi, err := os.Stat(p)
		switch {
		case err != nil:
			return nil, err
		case !fi.IsDir():
			return nil, fmt.Errorf("path is not directory: %s", p)
		}
	}

	return &Scanner{paths: paths}, nil
}

// Scan retrieves the scanned packages containing the extracted
// go types and structs.
func (s *Scanner) Scan() ([]*Package, error) {
	var pkgs []*Package
	for _, p := range s.paths {
		pkg, err := s.scanPackage(p)
		if err != nil {
			return nil, fmt.Errorf("error scanning package %q: %s", p, err)
		}

		pkgs = append(pkgs, pkg)
	}
	return pkgs, nil
}

func (s *Scanner) scanPackage(path string) (*Package, error) {
	files, err := getSourceFiles(path)
	if err != nil {
		return nil, err
	}

	gopkg, err := parseSourceFiles(path, files)
	if err != nil {
		return nil, err
	}

	return buildPackage(gopkg)
}

func (p *Package) processObject(o types.Object) {
	n, ok := o.Type().(*types.Named)
	if !ok || !o.Exported() {
		return
	}

	switch o.(type) {
	case *types.Var, *types.Const:
		if _, ok := n.Underlying().(*types.Basic); ok {
			p.processEnumValue(o.Name(), n)
		}
		return
	}

	if s, ok := n.Underlying().(*types.Struct); ok {

		st := processStruct(&Struct{Name: o.Name()}, s)
		p.Structs = append(p.Structs, st)
		return
	}

	name := fmt.Sprintf("%s.%s", n.Obj().Pkg().Path(), n.Obj().Name())
	p.Aliases[name] = processType(n.Underlying())
}

func processType(typ types.Type) (t Type) {
	switch u := typ.(type) {
	case *types.Named:
		t = NewNamed(
			u.Obj().Pkg().Path(),
			u.Obj().Name(),
		)
	case *types.Basic:
		t = NewBasic(u.Name())
	case *types.Slice:
		t = processType(u.Elem())
		t.SetRepeated(true)
	case *types.Array:
		t = processType(u.Elem())
		t.SetRepeated(true)
	case *types.Pointer:
		t = processType(u.Elem())
	case *types.Map:
		key := processType(u.Key())
		val := processType(u.Elem())
		t = NewMap(key, val)
	default:
		report.Warn("ignoring type %s", typ.String())
		return nil
	}

	return
}

func (p *Package) processEnumValue(name string, named *types.Named) {
	typ := fmt.Sprintf("%s.%s", named.Obj().Pkg().Path(), named.Obj().Name())
	p.values[typ] = append(p.values[typ], name)
}

func processStruct(s *Struct, elem *types.Struct) *Struct {
	for i := 0; i < elem.NumFields(); i++ {
		v := elem.Field(i)
		tags := findProtoTags(elem.Tag(i))

		if isIgnoredField(v, tags) {
			continue
		}

		if s.HasField(v.Name()) {
			report.Warn("struct %q already has a field %q", s.Name, v.Name())
			continue
		}

		if v.Anonymous() {
			embedded := findStruct(v.Type())
			if embedded == nil {
				report.Warn("field %q with type %q is not a valid embedded type", v.Name(), v.Type())
			} else {
				s = processStruct(s, embedded)
			}
			continue
		}

		f := &Field{
			Name: v.Name(),
			Type: processType(v.Type()),
		}
		if f.Type == nil {
			continue
		}

		s.Fields = append(s.Fields, f)
	}

	return s
}

func findStruct(t types.Type) *types.Struct {
	switch elem := t.(type) {
	case *types.Pointer:
		return findStruct(elem.Elem())
	case *types.Named:
		return findStruct(elem.Underlying())
	case *types.Struct:
		return elem
	default:
		return nil
	}
}

func (p *Package) collectEnums() {
	for k := range p.Aliases {
		if vals, ok := p.values[k]; ok {
			idx := strings.LastIndex(k, ".")
			name := k[idx+1:]

			p.Enums = append(p.Enums, &Enum{
				Name:   name,
				Values: vals,
			})

			delete(p.Aliases, k)
		}
	}
}

func isIgnoredField(f *types.Var, tags []string) bool {
	return !f.Exported() || (len(tags) > 0 && tags[0] == "-")
}

func buildPackage(gopkg *types.Package) (*Package, error) {
	objs := objectsInScope(gopkg.Scope())

	pkg := &Package{
		Path:    gopkg.Path(),
		Name:    gopkg.Name(),
		values:  make(map[string][]string),
		Aliases: make(map[string]Type),
	}

	for _, o := range objs {
		pkg.processObject(o)
	}

	pkg.collectEnums()
	return pkg, nil
}

func objectsInScope(scope *types.Scope) (objs []types.Object) {
	for _, n := range scope.Names() {
		objs = append(objs, scope.Lookup(n))
	}
	return
}

func getSourceFiles(path string) ([]string, error) {
	pkg, err := build.ImportDir(path, 0)
	if err != nil {
		return nil, err
	}

	var filenames []string
	filenames = append(filenames, pkg.GoFiles...)
	filenames = append(filenames, pkg.CgoFiles...)

	if len(filenames) == 0 {
		return nil, fmt.Errorf("no go source files in path: %s", path)
	}

	var paths []string
	for _, f := range filenames {
		paths = append(paths, filepath.Join(path, f))
	}

	return paths, nil
}

func parseSourceFiles(root string, paths []string) (*types.Package, error) {
	var files []*ast.File
	fs := token.NewFileSet()
	for _, p := range paths {
		f, err := parser.ParseFile(fs, p, nil, 0)
		if err != nil {
			return nil, err
		}

		files = append(files, f)
	}

	config := types.Config{Importer: importer.For("gc", nil)}

	return config.Check(root, fs, files, new(types.Info))
}
<commit_msg>Make scan of several package happen in parallel<commit_after>package scanner

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/src-d/protogo/report"
)

// Package holds information about a single Go package and
// a reference of all defined structs and type aliases.
// A Package is only safe to use once it is resolved.
type Package struct {
	Resolved bool
	Path     string
	Name     string
	Structs  []*Struct
	Enums    []*Enum
	Aliases  map[string]Type
	values   map[string][]string
}

// Type is the common interface for all possible types supported in protogo.
// Type is neither a representation of a Go type nor a representation of a
// protobuf type. Is an intermediate representation to ease future steps in
// the conversion from Go to protobuf.
// All types can be nullable (or not) or repeated (or not).
type Type interface {
	SetRepeated(bool)
	SetNullable(bool)
	IsRepeated() bool
	IsNullable() bool
}

// BaseType contains the common fields for all the types.
type BaseType struct {
	Nullable bool
	Repeated bool
}

func newBaseType() *BaseType {
	return &BaseType{
		Nullable: true,
		Repeated: false,
	}
}

func (t *BaseType) IsRepeated() bool   { return t.Repeated }
func (t *BaseType) IsNullable() bool   { return t.Nullable }
func (t *BaseType) SetRepeated(v bool) { t.Repeated = v }
func (t *BaseType) SetNullable(v bool) { t.Nullable = v }

// Basic is a basic type, which only is identified by its name.
type Basic struct {
	*BaseType
	Name string
}

func NewBasic(name string) Type {
	return &Basic{
		newBaseType(),
		name,
	}
}

// Named is non-basic type identified by a name on some package.
type Named struct {
	*BaseType
	Path string
	Name string
}

func NewNamed(path, name string) Type {
	return &Named{
		newBaseType(),
		path,
		name,
	}
}

// Map is a map type with a key and a value type.
type Map struct {
	*BaseType
	Key   Type
	Value Type
}

func NewMap(key, val Type) Type {
	return &Map{
		newBaseType(),
		key,
		val,
	}
}

// Enum consists of a list of possible values.
type Enum struct {
	Name   string
	Values []string
}

// Struct represents a Go struct with its name and fields.
type Struct struct {
	Name   string
	Fields []*Field
}

func (s *Struct) HasField(name string) bool {
	for _, f := range s.Fields {
		if f.Name == name {
			return true
		}
	}
	return false
}

// Field contains name and type of a struct field.
type Field struct {
	Name string
	Type Type
}

// Scanner scans paths looking for Go source files to parse
// and extract types and structs from.
type Scanner struct {
	paths []string
}

// New creates a new Scanner that will look for types and structs
// only in the given paths.
func New(paths ...string) (*Scanner, error) {
	for _, p := range paths {
		fi, err := os.Stat(p)
		switch {
		case err != nil:
			return nil, err
		case !fi.IsDir():
			return nil, fmt.Errorf("path is not directory: %s", p)
		}
	}

	return &Scanner{paths: paths}, nil
}

// Scan retrieves the scanned packages containing the extracted
// go types and structs.
func (s *Scanner) Scan() ([]*Package, error) {
	var (
		pkgs   = make([]*Package, len(s.paths))
		errors []error
		mut    sync.Mutex
		wg     = new(sync.WaitGroup)
	)

	wg.Add(len(s.paths))
	for i, p := range s.paths {
		go func(p string, i int) {
			defer wg.Done()

			pkg, err := s.scanPackage(p)
			mut.Lock()
			defer mut.Unlock()
			if err != nil {
				errors = append(errors, fmt.Errorf("error scanning package %q: %s", p, err))
			} else {
				pkgs[i] = pkg
			}
		}(p, i)
	}

	wg.Wait()
	if len(errors) > 0 {
		var lines []string
		for _, err := range errors {
			lines = append(lines, err.Error())
		}
		return nil, fmt.Errorf(strings.Join(lines, "\n"))
	}

	return pkgs, nil
}

func (s *Scanner) scanPackage(path string) (*Package, error) {
	files, err := getSourceFiles(path)
	if err != nil {
		return nil, err
	}

	gopkg, err := parseSourceFiles(path, files)
	if err != nil {
		return nil, err
	}

	return buildPackage(gopkg)
}

func (p *Package) processObject(o types.Object) {
	n, ok := o.Type().(*types.Named)
	if !ok || !o.Exported() {
		return
	}

	switch o.(type) {
	case *types.Var, *types.Const:
		if _, ok := n.Underlying().(*types.Basic); ok {
			p.processEnumValue(o.Name(), n)
		}
		return
	}

	if s, ok := n.Underlying().(*types.Struct); ok {

		st := processStruct(&Struct{Name: o.Name()}, s)
		p.Structs = append(p.Structs, st)
		return
	}

	p.Aliases[objName(n.Obj())] = processType(n.Underlying())
}

func processType(typ types.Type) (t Type) {
	switch u := typ.(type) {
	case *types.Named:
		t = NewNamed(
			u.Obj().Pkg().Path(),
			u.Obj().Name(),
		)
	case *types.Basic:
		t = NewBasic(u.Name())
	case *types.Slice:
		t = processType(u.Elem())
		t.SetRepeated(true)
	case *types.Array:
		t = processType(u.Elem())
		t.SetRepeated(true)
	case *types.Pointer:
		t = processType(u.Elem())
	case *types.Map:
		key := processType(u.Key())
		val := processType(u.Elem())
		t = NewMap(key, val)
	default:
		report.Warn("ignoring type %s", typ.String())
		return nil
	}

	return
}

func (p *Package) processEnumValue(name string, named *types.Named) {
	typ := objName(named.Obj())
	p.values[typ] = append(p.values[typ], name)
}

func processStruct(s *Struct, elem *types.Struct) *Struct {
	for i := 0; i < elem.NumFields(); i++ {
		v := elem.Field(i)
		tags := findProtoTags(elem.Tag(i))

		if isIgnoredField(v, tags) {
			continue
		}

		// TODO: It has not been decided yet what exact behaviour
		// is the intended when a struct overrides a field from
		// a previously embedded type. For now, the field is just
		// completely ignored and a warning is printed to give
		// feedback to the user.
		if s.HasField(v.Name()) {
			report.Warn("struct %q already has a field %q", s.Name, v.Name())
			continue
		}

		if v.Anonymous() {
			embedded := findStruct(v.Type())
			if embedded == nil {
				report.Warn("field %q with type %q is not a valid embedded type", v.Name(), v.Type())
			} else {
				s = processStruct(s, embedded)
			}
			continue
		}

		f := &Field{
			Name: v.Name(),
			Type: processType(v.Type()),
		}
		if f.Type == nil {
			continue
		}

		s.Fields = append(s.Fields, f)
	}

	return s
}

func findStruct(t types.Type) *types.Struct {
	switch elem := t.(type) {
	case *types.Pointer:
		return findStruct(elem.Elem())
	case *types.Named:
		return findStruct(elem.Underlying())
	case *types.Struct:
		return elem
	default:
		return nil
	}
}

func (p *Package) collectEnums() {
	for k := range p.Aliases {
		if vals, ok := p.values[k]; ok {
			idx := strings.LastIndex(k, ".")
			name := k[idx+1:]

			p.Enums = append(p.Enums, &Enum{
				Name:   name,
				Values: vals,
			})

			delete(p.Aliases, k)
		}
	}
}

func isIgnoredField(f *types.Var, tags []string) bool {
	return !f.Exported() || (len(tags) > 0 && tags[0] == "-")
}

func buildPackage(gopkg *types.Package) (*Package, error) {
	objs := objectsInScope(gopkg.Scope())

	pkg := &Package{
		Path:    gopkg.Path(),
		Name:    gopkg.Name(),
		values:  make(map[string][]string),
		Aliases: make(map[string]Type),
	}

	for _, o := range objs {
		pkg.processObject(o)
	}

	pkg.collectEnums()
	return pkg, nil
}

func objectsInScope(scope *types.Scope) (objs []types.Object) {
	for _, n := range scope.Names() {
		objs = append(objs, scope.Lookup(n))
	}
	return
}

func getSourceFiles(path string) ([]string, error) {
	pkg, err := build.ImportDir(path, 0)
	if err != nil {
		return nil, err
	}

	var filenames []string
	filenames = append(filenames, pkg.GoFiles...)
	filenames = append(filenames, pkg.CgoFiles...)

	if len(filenames) == 0 {
		return nil, fmt.Errorf("no go source files in path: %s", path)
	}

	var paths []string
	for _, f := range filenames {
		paths = append(paths, filepath.Join(path, f))
	}

	return paths, nil
}

func parseSourceFiles(root string, paths []string) (*types.Package, error) {
	var files []*ast.File
	fs := token.NewFileSet()
	for _, p := range paths {
		f, err := parser.ParseFile(fs, p, nil, 0)
		if err != nil {
			return nil, err
		}

		files = append(files, f)
	}

	config := types.Config{Importer: importer.For("gc", nil)}

	return config.Check(root, fs, files, new(types.Info))
}

func objName(obj types.Object) string {
	return fmt.Sprintf("%s.%s", obj.Pkg().Path(), obj.Name())
}
<|endoftext|>


<commit_before>package scanner

import (
	"bufio"
	"bytes"
	"github.com/github/git-media/pointer"
	// "github.com/rubyist/tracerx"
	"io"
	"os/exec"
	"strconv"
)

var (
	blobSizeCutoff = 125
	stdoutBufSize  = 16384
)

type wrappedPointer struct {
	Sha1 string
	*pointer.Pointer
}

// Scan takes a ref and returns a slice of pointer.Pointer objects
// for all git media pointers it finds for that ref.
func Scan(ref string) ([]*wrappedPointer, error) {
	revs, _ := revListShas(ref, ref == "")
	smallShas, _ := catFileBatchCheck(revs)
	pointerc, _ := catFileBatch(smallShas)

	pointers := make([]*wrappedPointer, 0)
	for p := range pointerc {
		pointers = append(pointers, p)
	}

	return pointers, nil
}

// revListShas uses git rev-list to return the list of object sha1s
// for the given ref. If all is true, ref is ignored. It returns a
// channel from which sha1 strings can be read.
func revListShas(ref string, all bool) (chan string, error) {
	refArgs := []string{"rev-list", "--objects"}
	if all {
		refArgs = append(refArgs, "--all")
	} else {
		refArgs = append(refArgs, ref)
	}

	cmd, err := startCommand("git", refArgs...)
	if err != nil {
		return nil, err
	}

	cmd.Stdin.Close()

	revs := make(chan string)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			revs <- scanner.Text()[0:40]
		}
		close(revs)
	}()

	return revs, nil
}

// catFileBatchCheck uses git cat-file --batch-check to get the type
// and size of a git object. Any object that isn't of type blob and
// under the blobSizeCutoff will be ignored. revs is a channel over
// which strings containing git sha1s will be sent. It returns a channel
// from which sha1 strings can be read.
func catFileBatchCheck(revs chan string) (chan string, error) {
	cmd, err := startCommand("git", "cat-file", "--batch-check")
	if err != nil {
		return nil, err
	}

	smallRevs := make(chan string)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := scanner.Text()
			if line[41:45] == "blob" {
				size, err := strconv.Atoi(line[46:len(line)])
				if err != nil {
					continue
				}
				if size < blobSizeCutoff {
					smallRevs <- line[0:40]
				}
			}
		}
		close(smallRevs)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return smallRevs, nil
}

// catFileBatch uses git cat-file --batch to get the object contents
// of a git object, given its sha1. The contents will be decoded into
// a git media pointer. revs is a channel over which strings containing
// git sha1s will be sent. It returns a channel from which point.Pointers
// can be read.
func catFileBatch(revs chan string) (chan *wrappedPointer, error) {
	cmd, err := startCommand("git", "cat-file", "--batch")
	if err != nil {
		return nil, err
	}

	pointers := make(chan *wrappedPointer)

	go func() {
		for {
			l, err := cmd.Stdout.ReadBytes('\n')
			if err != nil { // Probably check for EOF
				break
			}

			// Line is formatted:
			// <sha1> <type> <size>
			fields := bytes.Fields(l)
			s, _ := strconv.Atoi(string(fields[2]))

			nbuf := make([]byte, s)
			_, err = io.ReadFull(cmd.Stdout, nbuf)
			if err != nil {
				break // Legit errors
			}

			p, err := pointer.Decode(bytes.NewBuffer(nbuf))
			if err == nil {
				pointers <- &wrappedPointer{string(fields[0]), p}
			}

			_, err = cmd.Stdout.ReadBytes('\n') // Extra \n inserted by cat-file
			if err != nil {                     // Probably check for EOF
				break
			}
		}
		close(pointers)
	}()

	// writes shas to cat-file stdin
	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return pointers, nil
}

type wrappedCmd struct {
	Stdin  io.WriteCloser
	Stdout *bufio.Reader
	*exec.Cmd
}

// startCommand starts up a command and creates a stdin pipe and a buffered
// stdout pipe, wrapped in a wrappedCmd. The stdout buffer wille be of stdoutBufSize
// bytes.
func startCommand(command string, args ...string) (*wrappedCmd, error) {
	cmd := exec.Command(command, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	if err := cmd.Start(); err != nil {
		return nil, err
	}

	return &wrappedCmd{stdin, bufio.NewReaderSize(stdout, stdoutBufSize), cmd}, nil
}
<commit_msg>When scanner is given a ref, it should not traverse the ancesorts Doing so will e.g. show deleted files in `git ls-files`<commit_after>package scanner

import (
	"bufio"
	"bytes"
	"github.com/github/git-media/pointer"
	// "github.com/rubyist/tracerx"
	"io"
	"os/exec"
	"strconv"
)

var (
	blobSizeCutoff = 125
	stdoutBufSize  = 16384
)

type wrappedPointer struct {
	Sha1 string
	*pointer.Pointer
}

// Scan takes a ref and returns a slice of pointer.Pointer objects
// for all git media pointers it finds for that ref.
func Scan(ref string) ([]*wrappedPointer, error) {
	revs, _ := revListShas(ref, ref == "")
	smallShas, _ := catFileBatchCheck(revs)
	pointerc, _ := catFileBatch(smallShas)

	pointers := make([]*wrappedPointer, 0)
	for p := range pointerc {
		pointers = append(pointers, p)
	}

	return pointers, nil
}

// revListShas uses git rev-list to return the list of object sha1s
// for the given ref. If all is true, ref is ignored. It returns a
// channel from which sha1 strings can be read.
func revListShas(ref string, all bool) (chan string, error) {
	refArgs := []string{"rev-list", "--objects"}
	if all {
		refArgs = append(refArgs, "--all")
	} else {
		refArgs = append(refArgs, "--no-walk")
		refArgs = append(refArgs, ref)
	}

	cmd, err := startCommand("git", refArgs...)
	if err != nil {
		return nil, err
	}

	cmd.Stdin.Close()

	revs := make(chan string)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			revs <- scanner.Text()[0:40]
		}
		close(revs)
	}()

	return revs, nil
}

// catFileBatchCheck uses git cat-file --batch-check to get the type
// and size of a git object. Any object that isn't of type blob and
// under the blobSizeCutoff will be ignored. revs is a channel over
// which strings containing git sha1s will be sent. It returns a channel
// from which sha1 strings can be read.
func catFileBatchCheck(revs chan string) (chan string, error) {
	cmd, err := startCommand("git", "cat-file", "--batch-check")
	if err != nil {
		return nil, err
	}

	smallRevs := make(chan string)

	go func() {
		scanner := bufio.NewScanner(cmd.Stdout)
		for scanner.Scan() {
			line := scanner.Text()
			if line[41:45] == "blob" {
				size, err := strconv.Atoi(line[46:len(line)])
				if err != nil {
					continue
				}
				if size < blobSizeCutoff {
					smallRevs <- line[0:40]
				}
			}
		}
		close(smallRevs)
	}()

	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return smallRevs, nil
}

// catFileBatch uses git cat-file --batch to get the object contents
// of a git object, given its sha1. The contents will be decoded into
// a git media pointer. revs is a channel over which strings containing
// git sha1s will be sent. It returns a channel from which point.Pointers
// can be read.
func catFileBatch(revs chan string) (chan *wrappedPointer, error) {
	cmd, err := startCommand("git", "cat-file", "--batch")
	if err != nil {
		return nil, err
	}

	pointers := make(chan *wrappedPointer)

	go func() {
		for {
			l, err := cmd.Stdout.ReadBytes('\n')
			if err != nil { // Probably check for EOF
				break
			}

			// Line is formatted:
			// <sha1> <type> <size>
			fields := bytes.Fields(l)
			s, _ := strconv.Atoi(string(fields[2]))

			nbuf := make([]byte, s)
			_, err = io.ReadFull(cmd.Stdout, nbuf)
			if err != nil {
				break // Legit errors
			}

			p, err := pointer.Decode(bytes.NewBuffer(nbuf))
			if err == nil {
				pointers <- &wrappedPointer{string(fields[0]), p}
			}

			_, err = cmd.Stdout.ReadBytes('\n') // Extra \n inserted by cat-file
			if err != nil {                     // Probably check for EOF
				break
			}
		}
		close(pointers)
	}()

	// writes shas to cat-file stdin
	go func() {
		for r := range revs {
			cmd.Stdin.Write([]byte(r + "\n"))
		}
		cmd.Stdin.Close()
	}()

	return pointers, nil
}

type wrappedCmd struct {
	Stdin  io.WriteCloser
	Stdout *bufio.Reader
	*exec.Cmd
}

// startCommand starts up a command and creates a stdin pipe and a buffered
// stdout pipe, wrapped in a wrappedCmd. The stdout buffer wille be of stdoutBufSize
// bytes.
func startCommand(command string, args ...string) (*wrappedCmd, error) {
	cmd := exec.Command(command, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	if err := cmd.Start(); err != nil {
		return nil, err
	}

	return &wrappedCmd{stdin, bufio.NewReaderSize(stdout, stdoutBufSize), cmd}, nil
}
<|endoftext|>


<commit_before>package crane

import (
	"fmt"
	"os"
	"text/tabwriter"
)

type Containers []Container

// Lift containers (provision + run).
// When recreate is set, this will re-provision all images
// and recreate all containers.
func (containers Containers) lift(recreate bool, nocache bool) {
	containers.provisionOrSkip(recreate, nocache)
	containers.runOrStart(recreate)
}

// Provision containers.
func (containers Containers) provision(nocache bool) {
	for _, container := range containers {
		container.provision(nocache)
	}
}

// Run containers.
// When recreate is true, removes existing containers first.
func (containers Containers) run(recreate bool) {
	if recreate {
		containers.rm(true)
	}
	for _, container := range containers {
		container.run()
	}
}

// Run or start containers.
// When recreate is true, removes existing containers first.
func (containers Containers) runOrStart(recreate bool) {
	if recreate {
		containers.rm(true)
	}
	for _, container := range containers {
		container.runOrStart()
	}
}

// Provision or skip images.
// When update is true, provisions all images.
func (containers Containers) provisionOrSkip(update bool, nocache bool) {
	for _, container := range containers {
		container.provisionOrSkip(update, nocache)
	}
}

// Start containers.
func (containers Containers) start() {
	for _, container := range containers {
		container.start()
	}
}

// Kill containers.
func (containers Containers) kill() {
	for _, container := range containers {
		container.kill()
	}
}

// Stop containers.
func (containers Containers) stop() {
	for _, container := range containers {
		container.stop()
	}
}

// Pause containers.
func (containers Containers) pause() {
	for _, container := range containers {
		container.pause()
	}
}

// Unpause containers.
func (containers Containers) unpause() {
	for _, container := range containers {
		container.unpause()
	}
}

// Remove containers.
// When kill is true, kills existing containers first.
func (containers Containers) rm(kill bool) {
	if kill {
		containers.kill()
	}
	for _, container := range containers {
		container.rm()
	}
}

// Push containers.
func (containers Containers) push() {
	for _, container := range containers {
		container.push()
	}
}

// Status of containers.
func (containers Containers) status(notrunc bool) {
	w := new(tabwriter.Writer)
	w.Init(os.Stdout, 0, 8, 1, '\t', 0)
	fmt.Fprintln(w, "Name\tRunning\tID\tIP\tPorts")
	for _, container := range containers {
		container.status(w, notrunc)
	}
	w.Flush()
}
<commit_msg>uppercase header to mimic docker CLI behavior<commit_after>package crane

import (
	"fmt"
	"os"
	"text/tabwriter"
)

type Containers []Container

// Lift containers (provision + run).
// When recreate is set, this will re-provision all images
// and recreate all containers.
func (containers Containers) lift(recreate bool, nocache bool) {
	containers.provisionOrSkip(recreate, nocache)
	containers.runOrStart(recreate)
}

// Provision containers.
func (containers Containers) provision(nocache bool) {
	for _, container := range containers {
		container.provision(nocache)
	}
}

// Run containers.
// When recreate is true, removes existing containers first.
func (containers Containers) run(recreate bool) {
	if recreate {
		containers.rm(true)
	}
	for _, container := range containers {
		container.run()
	}
}

// Run or start containers.
// When recreate is true, removes existing containers first.
func (containers Containers) runOrStart(recreate bool) {
	if recreate {
		containers.rm(true)
	}
	for _, container := range containers {
		container.runOrStart()
	}
}

// Provision or skip images.
// When update is true, provisions all images.
func (containers Containers) provisionOrSkip(update bool, nocache bool) {
	for _, container := range containers {
		container.provisionOrSkip(update, nocache)
	}
}

// Start containers.
func (containers Containers) start() {
	for _, container := range containers {
		container.start()
	}
}

// Kill containers.
func (containers Containers) kill() {
	for _, container := range containers {
		container.kill()
	}
}

// Stop containers.
func (containers Containers) stop() {
	for _, container := range containers {
		container.stop()
	}
}

// Pause containers.
func (containers Containers) pause() {
	for _, container := range containers {
		container.pause()
	}
}

// Unpause containers.
func (containers Containers) unpause() {
	for _, container := range containers {
		container.unpause()
	}
}

// Remove containers.
// When kill is true, kills existing containers first.
func (containers Containers) rm(kill bool) {
	if kill {
		containers.kill()
	}
	for _, container := range containers {
		container.rm()
	}
}

// Push containers.
func (containers Containers) push() {
	for _, container := range containers {
		container.push()
	}
}

// Status of containers.
func (containers Containers) status(notrunc bool) {
	w := new(tabwriter.Writer)
	w.Init(os.Stdout, 0, 8, 1, '\t', 0)
	fmt.Fprintln(w, "NAME\tRUNNING\tID\tIP\tPORTS")
	for _, container := range containers {
		container.status(w, notrunc)
	}
	w.Flush()
}
<|endoftext|>