<commit_before>/*
 * Copyright (c) Clinton Freeman 2014
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
 * associated documentation files (the "Software"), to deal in the Software without restriction,
 * including without limitation the rights to use, copy, modify, merge, publish, distribute,
 * sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies or
 * substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
 * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"strings"
	"unicode/utf8"
)

type messageFn func(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error)

func announce(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	source, destination, err := ParsePeerAndRoom(message)
	if err != nil {
		return state, err
	}

	peer, exists := state.Peers[source.Id]
	if !exists {
		log.Printf("INFO - Adding Peer: %s\n", source.Id)
		state.Peers[source.Id] = new(Peer)
		state.Peers[source.Id].Id = source.Id
		state.Peers[source.Id].socket = sourceSocket // Inject a reference to the websocket within the new peer.
		peer = state.Peers[source.Id]
	}

	room, exists := state.Rooms[destination.Room]
	if !exists {
		log.Printf("INFO - Adding Room: %s\n", destination.Room)
		state.Rooms[destination.Room] = new(Room)
		state.Rooms[destination.Room].Room = destination.Room
		room = state.Rooms[destination.Room]
	}

	if state.PeerIsIn[peer.Id] == nil {
		state.PeerIsIn[peer.Id] = make(map[string]*Room)
	}
	state.PeerIsIn[peer.Id][room.Room] = room

	if state.RoomContains[room.Room] == nil {
		state.RoomContains[room.Room] = make(map[string]*Peer)
	}
	state.RoomContains[room.Room][peer.Id] = peer

	// Annouce the arrival to all the peers currently in the room.
	for _, p := range state.RoomContains[room.Room] {
		if p.Id != peer.Id && p.socket != nil {
			writeMessage(p.socket, message)
		}
	}

	return state, nil
}

func leave(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	source, destination, err := ParsePeerAndRoom(message)
	if err != nil {
		return state, err
	}

	peer, exists := state.Peers[source.Id]
	if !exists {
		return state, errors.New(fmt.Sprintf("Unable to leave, peer %s doesn't exist", source.Id))
	}

	room, exists := state.Rooms[destination.Room]
	if !exists {
		return state, errors.New(fmt.Sprintf("Unable to leave, room %s doesn't exist", destination.Room))
	}

	return removePeer(peer, room, message, state)
}

func closePeer(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	source := findPeerBySocket(sourceSocket, state)
	if source == nil {
		return state, errors.New("Unable to close - no Peer matching socket.")
	}

	// Announce to everyone that the peer belonging to sourceSocket
	// has closed and bailed out of their rooms.
	for _, r := range state.PeerIsIn[source.Id] {
		for _, p := range state.RoomContains[r.Room] {
			if p.Id != source.Id && p.socket != nil {
				src := fmt.Sprintf("{\"id\":\"%s\"}", source.Id)
				rm := fmt.Sprintf("{\"room\":\"%s\"}", r.Room)

				state, err = removePeer(source, r, []string{"/leave", src, rm}, state)
				if err != nil {
					return state, err
				}
			}
		}
	}

	// Make sure the socket is closed from this end.
	sourceSocket.Close()

	return state, nil
}

func removePeer(source *Peer, destination *Room, message []string, state SignalBox) (newState SignalBox, err error) {
	delete(state.PeerIsIn[source.Id], destination.Room)
	if len(state.PeerIsIn[source.Id]) == 0 {
		log.Printf("INFO - Removing Peer: %s\n", source.Id)
		delete(state.Peers, source.Id)
		delete(state.PeerIsIn, source.Id)
	}

	delete(state.RoomContains[destination.Room], source.Id)
	if len(state.RoomContains[destination.Room]) == 0 {
		log.Printf("INFO - Removing Room: %s\n", destination.Room)
		delete(state.Rooms, destination.Room)
		delete(state.RoomContains, destination.Room)
	} else {
		// Broadcast the departure to everyone else still in the room
		for _, p := range state.RoomContains[destination.Room] {
			if p.socket != nil {
				writeMessage(p.socket, message)
			}
		}
	}

	return state, nil
}

func to(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	if len(message) < 3 {
		return state, errors.New("Not enouth parts for personalised 'to' message")
	}

	d, exists := state.Peers[message[1]]
	if !exists {
		return state, nil
	}

	if d.socket != nil {
		writeMessage(d.socket, message)
	}

	return state, nil
}

func writeMessage(ws *websocket.Conn, message []string) {
	b, err := json.Marshal(strings.Join(message, "|"))
	if err == nil {
		log.Printf("INFO - Writing %s to %p\n", string(b), ws)
		ws.WriteMessage(websocket.TextMessage, b)
	}
}

func custom(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	source := Peer{}
	if len(message) < 2 {
		return state, errors.New("Not enough parts to custom message")
	}

	err = json.Unmarshal([]byte(message[1]), &source)
	if err != nil {
		return state, err
	}

	peer, exists := state.Peers[source.Id]
	if !exists {
		return state, nil
	}

	for _, r := range state.PeerIsIn[peer.Id] {
		for _, p := range state.RoomContains[r.Room] {
			if p.Id != peer.Id && p.socket != nil {
				writeMessage(p.socket, message)
			}
		}
	}

	return state, nil
}

func ignore(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {
	return state, nil
}

func findPeerBySocket(sourceSocket *websocket.Conn, state SignalBox) *Peer {
	for _, p := range state.Peers {
		if p.socket == sourceSocket {
			return p
		}
	}

	return nil
}

func ParsePeerAndRoom(message []string) (source Peer, destination Room, err error) {
	if len(message) < 3 {
		return Peer{}, Room{}, errors.New("Not enough parts in the message body to parse peer and room.")
	}

	err = json.Unmarshal([]byte(message[1]), &source)
	if err != nil {
		log.Print(err)
		return Peer{}, Room{}, err
	}

	err = json.Unmarshal([]byte(message[2]), &destination)
	if err != nil {
		return Peer{}, Room{}, err
	}

	return source, destination, nil
}

func ParseMessage(message string) (action messageFn, messageBody []string, err error) {
	// All messages are text (utf-8 encoded at present)
	if !utf8.Valid([]byte(message)) {
		return nil, nil, errors.New("Message is not utf-8 encoded")
	}

	parts := strings.Split(message, "|")

	// rtc.io commands start with "/" - ignore everything else.
	if len(message) > 0 && message[0:1] == "/" {
		switch parts[0] {
		case "/announce":
			return announce, parts, nil

		case "/leave":
			return leave, parts, nil

		case "/to":
			return to, parts, nil

		case "/close":
			return closePeer, parts, nil

		default:
			return custom, parts, nil
		}
	}

	return ignore, parts, nil
}
<commit_msg>Removed socket close.<commit_after>/*
 * Copyright (c) Clinton Freeman 2014
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
 * associated documentation files (the "Software"), to deal in the Software without restriction,
 * including without limitation the rights to use, copy, modify, merge, publish, distribute,
 * sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies or
 * substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
 * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"strings"
	"unicode/utf8"
)

type messageFn func(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error)

func announce(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	source, destination, err := ParsePeerAndRoom(message)
	if err != nil {
		return state, err
	}

	peer, exists := state.Peers[source.Id]
	if !exists {
		log.Printf("INFO - Adding Peer: %s\n", source.Id)
		state.Peers[source.Id] = new(Peer)
		state.Peers[source.Id].Id = source.Id
		state.Peers[source.Id].socket = sourceSocket // Inject a reference to the websocket within the new peer.
		peer = state.Peers[source.Id]
	}

	room, exists := state.Rooms[destination.Room]
	if !exists {
		log.Printf("INFO - Adding Room: %s\n", destination.Room)
		state.Rooms[destination.Room] = new(Room)
		state.Rooms[destination.Room].Room = destination.Room
		room = state.Rooms[destination.Room]
	}

	if state.PeerIsIn[peer.Id] == nil {
		state.PeerIsIn[peer.Id] = make(map[string]*Room)
	}
	state.PeerIsIn[peer.Id][room.Room] = room

	if state.RoomContains[room.Room] == nil {
		state.RoomContains[room.Room] = make(map[string]*Peer)
	}
	state.RoomContains[room.Room][peer.Id] = peer

	// Annouce the arrival to all the peers currently in the room.
	for _, p := range state.RoomContains[room.Room] {
		if p.Id != peer.Id && p.socket != nil {
			writeMessage(p.socket, message)
		}
	}

	return state, nil
}

func leave(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	source, destination, err := ParsePeerAndRoom(message)
	if err != nil {
		return state, err
	}

	peer, exists := state.Peers[source.Id]
	if !exists {
		return state, errors.New(fmt.Sprintf("Unable to leave, peer %s doesn't exist", source.Id))
	}

	room, exists := state.Rooms[destination.Room]
	if !exists {
		return state, errors.New(fmt.Sprintf("Unable to leave, room %s doesn't exist", destination.Room))
	}

	return removePeer(peer, room, message, state)
}

func closePeer(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	source := findPeerBySocket(sourceSocket, state)
	if source == nil {
		return state, errors.New("Unable to close - no Peer matching socket.")
	}

	// Announce to everyone that the peer belonging to sourceSocket
	// has closed and bailed out of their rooms.
	for _, r := range state.PeerIsIn[source.Id] {
		for _, p := range state.RoomContains[r.Room] {
			if p.Id != source.Id && p.socket != nil {
				src := fmt.Sprintf("{\"id\":\"%s\"}", source.Id)
				rm := fmt.Sprintf("{\"room\":\"%s\"}", r.Room)

				state, err = removePeer(source, r, []string{"/leave", src, rm}, state)
				if err != nil {
					return state, err
				}
			}
		}
	}

	// Make sure the socket is closed from this end.
	//sourceSocket.Close()

	return state, nil
}

func removePeer(source *Peer, destination *Room, message []string, state SignalBox) (newState SignalBox, err error) {
	delete(state.PeerIsIn[source.Id], destination.Room)
	if len(state.PeerIsIn[source.Id]) == 0 {
		log.Printf("INFO - Removing Peer: %s\n", source.Id)
		delete(state.Peers, source.Id)
		delete(state.PeerIsIn, source.Id)
	}

	delete(state.RoomContains[destination.Room], source.Id)
	if len(state.RoomContains[destination.Room]) == 0 {
		log.Printf("INFO - Removing Room: %s\n", destination.Room)
		delete(state.Rooms, destination.Room)
		delete(state.RoomContains, destination.Room)
	} else {
		// Broadcast the departure to everyone else still in the room
		for _, p := range state.RoomContains[destination.Room] {
			if p.socket != nil {
				writeMessage(p.socket, message)
			}
		}
	}

	return state, nil
}

func to(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	if len(message) < 3 {
		return state, errors.New("Not enouth parts for personalised 'to' message")
	}

	d, exists := state.Peers[message[1]]
	if !exists {
		return state, nil
	}

	if d.socket != nil {
		writeMessage(d.socket, message)
	}

	return state, nil
}

func writeMessage(ws *websocket.Conn, message []string) {
	b, err := json.Marshal(strings.Join(message, "|"))
	if err == nil {
		log.Printf("INFO - Writing %s to %p\n", string(b), ws)
		ws.WriteMessage(websocket.TextMessage, b)
	}
}

func custom(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {

	source := Peer{}
	if len(message) < 2 {
		return state, errors.New("Not enough parts to custom message")
	}

	err = json.Unmarshal([]byte(message[1]), &source)
	if err != nil {
		return state, err
	}

	peer, exists := state.Peers[source.Id]
	if !exists {
		return state, nil
	}

	for _, r := range state.PeerIsIn[peer.Id] {
		for _, p := range state.RoomContains[r.Room] {
			if p.Id != peer.Id && p.socket != nil {
				writeMessage(p.socket, message)
			}
		}
	}

	return state, nil
}

func ignore(message []string,
	sourceSocket *websocket.Conn,
	state SignalBox) (newState SignalBox, err error) {
	return state, nil
}

func findPeerBySocket(sourceSocket *websocket.Conn, state SignalBox) *Peer {
	for _, p := range state.Peers {
		if p.socket == sourceSocket {
			return p
		}
	}

	return nil
}

func ParsePeerAndRoom(message []string) (source Peer, destination Room, err error) {
	if len(message) < 3 {
		return Peer{}, Room{}, errors.New("Not enough parts in the message body to parse peer and room.")
	}

	err = json.Unmarshal([]byte(message[1]), &source)
	if err != nil {
		log.Print(err)
		return Peer{}, Room{}, err
	}

	err = json.Unmarshal([]byte(message[2]), &destination)
	if err != nil {
		return Peer{}, Room{}, err
	}

	return source, destination, nil
}

func ParseMessage(message string) (action messageFn, messageBody []string, err error) {
	// All messages are text (utf-8 encoded at present)
	if !utf8.Valid([]byte(message)) {
		return nil, nil, errors.New("Message is not utf-8 encoded")
	}

	parts := strings.Split(message, "|")

	// rtc.io commands start with "/" - ignore everything else.
	if len(message) > 0 && message[0:1] == "/" {
		switch parts[0] {
		case "/announce":
			return announce, parts, nil

		case "/leave":
			return leave, parts, nil

		case "/to":
			return to, parts, nil

		case "/close":
			return closePeer, parts, nil

		default:
			return custom, parts, nil
		}
	}

	return ignore, parts, nil
}
<|endoftext|>


<commit_before>package commands

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/net/context"

	shlex "github.com/flynn/go-shlex"
	"github.com/urfave/cli"

	"github.com/moul/advanced-ssh-config/pkg/config"
	. "github.com/moul/advanced-ssh-config/pkg/logger"
)

func cmdProxy(c *cli.Context) error {
	Logger.Debugf("assh args: %s", c.Args())

	if len(c.Args()) < 1 {
		Logger.Fatalf("assh: \"connect\" requires 1 argument. See 'assh connect --help'.")
	}

	// dry-run option
	// Setting the 'ASSH_DRYRUN=1' environment variable,
	// so 'assh' can use gateways using sub-SSH commands.
	if c.Bool("dry-run") {
		os.Setenv("ASSH_DRYRUN", "1")
	}
	dryRun := os.Getenv("ASSH_DRYRUN") == "1"

	conf, err := config.Open(c.GlobalString("config"))
	if err != nil {
		Logger.Fatalf("Cannot open configuration file: %v", err)
	}

	if err = conf.LoadKnownHosts(); err != nil {
		Logger.Debugf("Failed to load assh known_hosts: %v", err)
	}

	target := c.Args()[0]

	isOutdated, err := conf.IsConfigOutdated(target)
	if err != nil {
		Logger.Warnf("Cannot check if ~/.ssh/config is outdated.")
	}
	if isOutdated {
		Logger.Debugf("The configuration file is outdated, rebuilding it before calling ssh")
		Logger.Warnf("'~/.ssh/config' has been rewritten.  SSH needs to be restarted.  See https://github.com/moul/advanced-ssh-config/issues/122 for more information.")
		Logger.Debugf("Saving SSH config")
		err = conf.SaveSSHConfig()
		if err != nil {
			Logger.Fatalf("Cannot save SSH config file: %v", err)
		}
	}

	// FIXME: handle complete host with json

	host, err := computeHost(target, c.Int("port"), conf)
	if err != nil {
		Logger.Fatalf("Cannot get host '%s': %v", target, err)
	}
	w := Logger.Writer()
	host.WriteSSHConfigTo(w)
	w.Close()

	hostJson, err := json.Marshal(host)
	if err != nil {
		Logger.Warnf("Failed to marshal host: %v", err)
	}
	Logger.Debugf("Host: %s", hostJson)

	Logger.Debugf("Proxying")
	err = proxy(host, conf, dryRun)
	if err != nil {
		Logger.Fatalf("Proxy error: %v", err)
	}

	return nil
}

func computeHost(dest string, portOverride int, conf *config.Config) (*config.Host, error) {
	host := conf.GetHostSafe(dest)

	if portOverride > 0 {
		host.Port = strconv.Itoa(portOverride)
	}

	return host, nil
}

func prepareHostControlPath(host, gateway *config.Host) error {
	controlPathDir := path.Dir(os.ExpandEnv(strings.Replace(host.ControlPath, "~", "$HOME", -1)))
	gatewayControlPath := path.Join(controlPathDir, gateway.Name())
	if config.BoolVal(host.NoControlMasterMkdir) {
		return nil
	}
	return os.MkdirAll(gatewayControlPath, 0700)
}

func proxy(host *config.Host, conf *config.Config, dryRun bool) error {
	if len(host.Gateways) > 0 {
		Logger.Debugf("Trying gateways: %s", host.Gateways)
		for _, gateway := range host.Gateways {
			if gateway == "direct" {
				err := proxyDirect(host, dryRun)
				if err != nil {
					Logger.Errorf("Failed to use 'direct' connection: %v", err)
				}
			} else {
				hostCopy := host.Clone()
				gatewayHost := conf.GetGatewaySafe(gateway)

				err := prepareHostControlPath(hostCopy, gatewayHost)
				if err != nil {
					return err
				}

				// FIXME: dynamically add "-v" flags

				var command string

				// FIXME: detect ssh client version and use netcat if too old
				// for now, the workaround is to configure the ProxyCommand of the host to "nc %h %p"

				if err = hostPrepare(hostCopy); err != nil {
					return err
				}

				if hostCopy.ProxyCommand != "" {
					command = "ssh %name -- " + hostCopy.ExpandString(hostCopy.ProxyCommand)
				} else {
					command = hostCopy.ExpandString("ssh -W %h:%p ") + "%name"
				}

				Logger.Debugf("Using gateway '%s': %s", gateway, command)
				err = proxyCommand(gatewayHost, command, dryRun)
				if err == nil {
					return nil
				}
				Logger.Errorf("Cannot use gateway '%s': %v", gateway, err)
			}
		}
		return fmt.Errorf("No such available gateway")
	}

	Logger.Debugf("Connecting without gateway")
	return proxyDirect(host, dryRun)
}

func proxyDirect(host *config.Host, dryRun bool) error {
	if host.ProxyCommand != "" {
		return proxyCommand(host, host.ProxyCommand, dryRun)
	}
	return proxyGo(host, dryRun)
}

func proxyCommand(host *config.Host, command string, dryRun bool) error {
	command = host.ExpandString(command)
	Logger.Debugf("ProxyCommand: %s", command)
	args, err := shlex.Split(command)
	if err != nil {
		return err
	}

	if dryRun {
		return fmt.Errorf("dry-run: Execute %s", args)
	}

	spawn := exec.Command(args[0], args[1:]...)
	spawn.Stdout = os.Stdout
	spawn.Stdin = os.Stdin
	spawn.Stderr = os.Stderr
	return spawn.Run()
}

func hostPrepare(host *config.Host) error {
	if host.HostName == "" {
		host.HostName = host.Name()
	}

	if len(host.ResolveNameservers) > 0 {
		Logger.Debugf("Resolving host: '%s' using nameservers %s", host.HostName, host.ResolveNameservers)
		// FIXME: resolve using custom dns server
		results, err := net.LookupAddr(host.HostName)
		if err != nil {
			return err
		}
		if len(results) > 0 {
			host.HostName = results[0]
		}
		Logger.Debugf("Resolved host is: %s", host.HostName)
	}

	if host.ResolveCommand != "" {
		command := host.ExpandString(host.ResolveCommand)
		Logger.Debugf("Resolving host: %q using command: %q", host.HostName, command)

		args, err := shlex.Split(command)
		if err != nil {
			return err
		}

		cmd := exec.Command(args[0], args[1:]...)
		var stdout bytes.Buffer
		var stderr bytes.Buffer
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		if err := cmd.Run(); err != nil {
			Logger.Errorf("ResolveCommand failed: %s", stderr.String())
			return err
		}

		host.HostName = strings.TrimSpace(fmt.Sprintf("%s", stdout.String()))
		Logger.Debugf("Resolved host is: %s", host.HostName)
	}
	return nil
}

type exportReadWrite struct {
	written uint64
	err     error
}

// ConnectionStats contains network and timing informations about a connection
type ConnectionStats struct {
	WrittenBytes       uint64
	CreatedAt          time.Time
	ConnectedAt        time.Time
	DisconnectedAt     time.Time
	ConnectionDuration time.Duration
	AverageSpeed       float64
}

// ConnectHookArgs is the struture sent to the hooks and used in Go templates by the hook drivers
type ConnectHookArgs struct {
	Host  *config.Host
	Stats *ConnectionStats
}

func proxyGo(host *config.Host, dryRun bool) error {
	stats := ConnectionStats{
		CreatedAt: time.Now(),
	}
	connectHookArgs := ConnectHookArgs{
		Host:  host,
		Stats: &stats,
	}

	Logger.Debugf("Preparing host object")
	if err := hostPrepare(host); err != nil {
		return err
	}

	if dryRun {
		return fmt.Errorf("dry-run: Golang native TCP connection to '%s:%s'", host.HostName, host.Port)
	}

	Logger.Debugf("Connecting to %s:%s", host.HostName, host.Port)
	conn, err := net.Dial("tcp", fmt.Sprintf("%s:%s", host.HostName, host.Port))
	if err != nil {
		return err
	}
	Logger.Debugf("Connected to %s:%s", host.HostName, host.Port)
	stats.ConnectedAt = time.Now()

	// OnConnect hook
	Logger.Debugf("Calling OnConnect hooks")
	if err := host.Hooks.OnConnect.InvokeAll(connectHookArgs); err != nil {
		Logger.Errorf("OnConnect hook failed: %v", err)
	}

	// Ignore SIGHUP
	signal.Ignore(syscall.SIGHUP)

	waitGroup := sync.WaitGroup{}
	result := exportReadWrite{}

	ctx, cancel := context.WithCancel(context.Background())
	ctx = context.WithValue(ctx, "sync", &waitGroup)

	waitGroup.Add(2)
	c1 := readAndWrite(ctx, conn, os.Stdout)
	c2 := readAndWrite(ctx, os.Stdin, conn)
	select {
	case result = <-c1:
		stats.WrittenBytes = result.written
	case result = <-c2:
	}
	if result.err != nil && result.err == io.EOF {
		result.err = nil
	}

	stats.DisconnectedAt = time.Now()
	stats.ConnectionDuration = stats.DisconnectedAt.Sub(stats.ConnectedAt)
	averageSpeed := float64(stats.WrittenBytes) / stats.ConnectionDuration.Seconds()
	stats.AverageSpeed = math.Ceil(averageSpeed*1000) / 1000
	conn.Close()
	cancel()
	waitGroup.Wait()
	select {
	case res := <-c1:
		stats.WrittenBytes = res.written
	default:
	}
	// OnDisconnect hook
	Logger.Debugf("Calling OnDisconnect hooks")
	if err := host.Hooks.OnDisconnect.InvokeAll(connectHookArgs); err != nil {
		Logger.Errorf("OnDisconnect hook failed: %v", err)
	}
	Logger.Debugf("Byte written %v", stats.WrittenBytes)
	return result.err
}

func readAndWrite(ctx context.Context, r io.Reader, w io.Writer) <-chan exportReadWrite {
	buff := make([]byte, 1024)
	c := make(chan exportReadWrite, 1)

	go func() {
		defer ctx.Value("sync").(*sync.WaitGroup).Done()

		export := exportReadWrite{}
		for {
			select {
			case <-ctx.Done():
				c <- export
				return
			default:
				nr, err := r.Read(buff)
				if err != nil {
					export.err = err
					c <- export
					return
				}
				if nr > 0 {
					wr, err := w.Write(buff[:nr])
					if err != nil {
						export.err = err
						c <- export
						return
					}
					if wr > 0 {
						export.written += uint64(wr)
					}
				}
			}
		}
	}()
	return c
}
<commit_msg>Move some statistics computing<commit_after>package commands

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/net/context"

	shlex "github.com/flynn/go-shlex"
	"github.com/urfave/cli"

	"github.com/moul/advanced-ssh-config/pkg/config"
	. "github.com/moul/advanced-ssh-config/pkg/logger"
)

func cmdProxy(c *cli.Context) error {
	Logger.Debugf("assh args: %s", c.Args())

	if len(c.Args()) < 1 {
		Logger.Fatalf("assh: \"connect\" requires 1 argument. See 'assh connect --help'.")
	}

	// dry-run option
	// Setting the 'ASSH_DRYRUN=1' environment variable,
	// so 'assh' can use gateways using sub-SSH commands.
	if c.Bool("dry-run") {
		os.Setenv("ASSH_DRYRUN", "1")
	}
	dryRun := os.Getenv("ASSH_DRYRUN") == "1"

	conf, err := config.Open(c.GlobalString("config"))
	if err != nil {
		Logger.Fatalf("Cannot open configuration file: %v", err)
	}

	if err = conf.LoadKnownHosts(); err != nil {
		Logger.Debugf("Failed to load assh known_hosts: %v", err)
	}

	target := c.Args()[0]

	isOutdated, err := conf.IsConfigOutdated(target)
	if err != nil {
		Logger.Warnf("Cannot check if ~/.ssh/config is outdated.")
	}
	if isOutdated {
		Logger.Debugf("The configuration file is outdated, rebuilding it before calling ssh")
		Logger.Warnf("'~/.ssh/config' has been rewritten.  SSH needs to be restarted.  See https://github.com/moul/advanced-ssh-config/issues/122 for more information.")
		Logger.Debugf("Saving SSH config")
		err = conf.SaveSSHConfig()
		if err != nil {
			Logger.Fatalf("Cannot save SSH config file: %v", err)
		}
	}

	// FIXME: handle complete host with json

	host, err := computeHost(target, c.Int("port"), conf)
	if err != nil {
		Logger.Fatalf("Cannot get host '%s': %v", target, err)
	}
	w := Logger.Writer()
	host.WriteSSHConfigTo(w)
	w.Close()

	hostJson, err := json.Marshal(host)
	if err != nil {
		Logger.Warnf("Failed to marshal host: %v", err)
	}
	Logger.Debugf("Host: %s", hostJson)

	Logger.Debugf("Proxying")
	err = proxy(host, conf, dryRun)
	if err != nil {
		Logger.Fatalf("Proxy error: %v", err)
	}

	return nil
}

func computeHost(dest string, portOverride int, conf *config.Config) (*config.Host, error) {
	host := conf.GetHostSafe(dest)

	if portOverride > 0 {
		host.Port = strconv.Itoa(portOverride)
	}

	return host, nil
}

func prepareHostControlPath(host, gateway *config.Host) error {
	controlPathDir := path.Dir(os.ExpandEnv(strings.Replace(host.ControlPath, "~", "$HOME", -1)))
	gatewayControlPath := path.Join(controlPathDir, gateway.Name())
	if config.BoolVal(host.NoControlMasterMkdir) {
		return nil
	}
	return os.MkdirAll(gatewayControlPath, 0700)
}

func proxy(host *config.Host, conf *config.Config, dryRun bool) error {
	if len(host.Gateways) > 0 {
		Logger.Debugf("Trying gateways: %s", host.Gateways)
		for _, gateway := range host.Gateways {
			if gateway == "direct" {
				err := proxyDirect(host, dryRun)
				if err != nil {
					Logger.Errorf("Failed to use 'direct' connection: %v", err)
				}
			} else {
				hostCopy := host.Clone()
				gatewayHost := conf.GetGatewaySafe(gateway)

				err := prepareHostControlPath(hostCopy, gatewayHost)
				if err != nil {
					return err
				}

				// FIXME: dynamically add "-v" flags

				var command string

				// FIXME: detect ssh client version and use netcat if too old
				// for now, the workaround is to configure the ProxyCommand of the host to "nc %h %p"

				if err = hostPrepare(hostCopy); err != nil {
					return err
				}

				if hostCopy.ProxyCommand != "" {
					command = "ssh %name -- " + hostCopy.ExpandString(hostCopy.ProxyCommand)
				} else {
					command = hostCopy.ExpandString("ssh -W %h:%p ") + "%name"
				}

				Logger.Debugf("Using gateway '%s': %s", gateway, command)
				err = proxyCommand(gatewayHost, command, dryRun)
				if err == nil {
					return nil
				}
				Logger.Errorf("Cannot use gateway '%s': %v", gateway, err)
			}
		}
		return fmt.Errorf("No such available gateway")
	}

	Logger.Debugf("Connecting without gateway")
	return proxyDirect(host, dryRun)
}

func proxyDirect(host *config.Host, dryRun bool) error {
	if host.ProxyCommand != "" {
		return proxyCommand(host, host.ProxyCommand, dryRun)
	}
	return proxyGo(host, dryRun)
}

func proxyCommand(host *config.Host, command string, dryRun bool) error {
	command = host.ExpandString(command)
	Logger.Debugf("ProxyCommand: %s", command)
	args, err := shlex.Split(command)
	if err != nil {
		return err
	}

	if dryRun {
		return fmt.Errorf("dry-run: Execute %s", args)
	}

	spawn := exec.Command(args[0], args[1:]...)
	spawn.Stdout = os.Stdout
	spawn.Stdin = os.Stdin
	spawn.Stderr = os.Stderr
	return spawn.Run()
}

func hostPrepare(host *config.Host) error {
	if host.HostName == "" {
		host.HostName = host.Name()
	}

	if len(host.ResolveNameservers) > 0 {
		Logger.Debugf("Resolving host: '%s' using nameservers %s", host.HostName, host.ResolveNameservers)
		// FIXME: resolve using custom dns server
		results, err := net.LookupAddr(host.HostName)
		if err != nil {
			return err
		}
		if len(results) > 0 {
			host.HostName = results[0]
		}
		Logger.Debugf("Resolved host is: %s", host.HostName)
	}

	if host.ResolveCommand != "" {
		command := host.ExpandString(host.ResolveCommand)
		Logger.Debugf("Resolving host: %q using command: %q", host.HostName, command)

		args, err := shlex.Split(command)
		if err != nil {
			return err
		}

		cmd := exec.Command(args[0], args[1:]...)
		var stdout bytes.Buffer
		var stderr bytes.Buffer
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		if err := cmd.Run(); err != nil {
			Logger.Errorf("ResolveCommand failed: %s", stderr.String())
			return err
		}

		host.HostName = strings.TrimSpace(fmt.Sprintf("%s", stdout.String()))
		Logger.Debugf("Resolved host is: %s", host.HostName)
	}
	return nil
}

type exportReadWrite struct {
	written uint64
	err     error
}

// ConnectionStats contains network and timing informations about a connection
type ConnectionStats struct {
	WrittenBytes       uint64
	CreatedAt          time.Time
	ConnectedAt        time.Time
	DisconnectedAt     time.Time
	ConnectionDuration time.Duration
	AverageSpeed       float64
}

// ConnectHookArgs is the struture sent to the hooks and used in Go templates by the hook drivers
type ConnectHookArgs struct {
	Host  *config.Host
	Stats *ConnectionStats
}

func proxyGo(host *config.Host, dryRun bool) error {
	stats := ConnectionStats{
		CreatedAt: time.Now(),
	}
	connectHookArgs := ConnectHookArgs{
		Host:  host,
		Stats: &stats,
	}

	Logger.Debugf("Preparing host object")
	if err := hostPrepare(host); err != nil {
		return err
	}

	if dryRun {
		return fmt.Errorf("dry-run: Golang native TCP connection to '%s:%s'", host.HostName, host.Port)
	}

	Logger.Debugf("Connecting to %s:%s", host.HostName, host.Port)
	conn, err := net.Dial("tcp", fmt.Sprintf("%s:%s", host.HostName, host.Port))
	if err != nil {
		return err
	}
	Logger.Debugf("Connected to %s:%s", host.HostName, host.Port)
	stats.ConnectedAt = time.Now()

	// OnConnect hook
	Logger.Debugf("Calling OnConnect hooks")
	if err := host.Hooks.OnConnect.InvokeAll(connectHookArgs); err != nil {
		Logger.Errorf("OnConnect hook failed: %v", err)
	}

	// Ignore SIGHUP
	signal.Ignore(syscall.SIGHUP)

	waitGroup := sync.WaitGroup{}
	result := exportReadWrite{}

	ctx, cancel := context.WithCancel(context.Background())
	ctx = context.WithValue(ctx, "sync", &waitGroup)

	waitGroup.Add(2)
	c1 := readAndWrite(ctx, conn, os.Stdout)
	c2 := readAndWrite(ctx, os.Stdin, conn)
	select {
	case result = <-c1:
		stats.WrittenBytes = result.written
	case result = <-c2:
	}
	if result.err != nil && result.err == io.EOF {
		result.err = nil
	}

	conn.Close()
	cancel()
	waitGroup.Wait()
	select {
	case res := <-c1:
		stats.WrittenBytes = res.written
	default:
	}

	stats.DisconnectedAt = time.Now()
	stats.ConnectionDuration = stats.DisconnectedAt.Sub(stats.ConnectedAt)
	averageSpeed := float64(stats.WrittenBytes) / stats.ConnectionDuration.Seconds()
	stats.AverageSpeed = math.Ceil(averageSpeed*1000) / 1000

	// OnDisconnect hook
	Logger.Debugf("Calling OnDisconnect hooks")
	if err := host.Hooks.OnDisconnect.InvokeAll(connectHookArgs); err != nil {
		Logger.Errorf("OnDisconnect hook failed: %v", err)
	}
	Logger.Debugf("Byte written %v", stats.WrittenBytes)
	return result.err
}

func readAndWrite(ctx context.Context, r io.Reader, w io.Writer) <-chan exportReadWrite {
	buff := make([]byte, 1024)
	c := make(chan exportReadWrite, 1)

	go func() {
		defer ctx.Value("sync").(*sync.WaitGroup).Done()

		export := exportReadWrite{}
		for {
			select {
			case <-ctx.Done():
				c <- export
				return
			default:
				nr, err := r.Read(buff)
				if err != nil {
					export.err = err
					c <- export
					return
				}
				if nr > 0 {
					wr, err := w.Write(buff[:nr])
					if err != nil {
						export.err = err
						c <- export
						return
					}
					if wr > 0 {
						export.written += uint64(wr)
					}
				}
			}
		}
	}()
	return c
}
<|endoftext|>


<commit_before>package converter

import (
	"math"
	"testing"
)

var (
	dataUint8   = []uint8{0, 64, 128, 192, 255}
	dataInt16   = []int16{-32768, -16384, 0, 16383, 32767}
	dataInt24   = []int32{-2147483648, -1073741824, 0, 1073741823, 2147483647}
	dataInt32   = []int32{-2147483648, -1073741824, 0, 1073741823, 2147483647}
	dataFloat32 = []float32{-1, -0.5, 0, 0.5, 1}
	dataFloat64 = []float64{-1, -0.5, 0, 0.5, 1}
	dataLen     = len(dataFloat32)
)

func testResultUint8(a []uint8, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataUint8 {
		if math.Abs(float64(s)-float64(a[i])) > 1 {
			t.Fail()
		}
	}
}

func testResultInt16(a []int16, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataInt16 {
		if math.Abs(float64(s)-float64(a[i])) > 1 {
			t.Fail()
		}
	}
}

func testResultInt24(a []int32, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataInt32 {
		if math.Abs(float64(s)-float64(a[i])) > 0x100 {
			t.Fail()
		}
	}
}

func testResultInt32(a []int32, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataInt32 {
		if math.Abs(float64(s)-float64(a[i])) > 1 {
			t.Fail()
		}
	}
}

func testResultFloat32(a []float32, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataFloat32 {
		if math.Abs(float64(s)-float64(a[i])) > 0.01 {
			t.Fail()
		}
	}
}

func testResultFloat64(a []float64, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataFloat64 {
		if math.Abs(s-a[i]) > 0.01 {
			t.Fail()
		}
	}
}
<commit_msg>fix test<commit_after>package converter

import (
	"math"
	"testing"
)

var (
	dataUint8   = []uint8{0, 64, 128, 192, 255}
	dataInt16   = []int16{-32768, -16384, 0, 16383, 32767}
	dataInt24   = []int32{-2147483648, -1073741824, 0, 1073741823, 2147483647}
	dataInt32   = []int32{-2147483648, -1073741824, 0, 1073741823, 2147483647}
	dataFloat32 = []float32{-1, -0.5, 0, 0.5, 1}
	dataFloat64 = []float64{-1, -0.5, 0, 0.5, 1}
	dataLen     = len(dataFloat32)
)

func testResultUint8(a []uint8, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataUint8 {
		if math.Abs(float64(s)-float64(a[i])) > 1 {
			t.Fail()
		}
	}
}

func testResultInt16(a []int16, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataInt16 {
		if math.Abs(float64(s)-float64(a[i])) > 1 {
			t.Fail()
		}
	}
}

func testResultInt24(a []int32, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataInt24 {
		if math.Abs(float64(s)-float64(a[i])) > 0x100 {
			t.Fail()
		}
	}
}

func testResultInt32(a []int32, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataInt32 {
		if math.Abs(float64(s)-float64(a[i])) > 1 {
			t.Fail()
		}
	}
}

func testResultFloat32(a []float32, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataFloat32 {
		if math.Abs(float64(s)-float64(a[i])) > 0.01 {
			t.Fail()
		}
	}
}

func testResultFloat64(a []float64, t *testing.T) {
	t.Log("min:", a[0], "low:", a[1], "zero:", a[2], "high:", a[3], "max:", a[4])
	for i, s := range dataFloat64 {
		if math.Abs(s-a[i]) > 0.01 {
			t.Fail()
		}
	}
}
<|endoftext|>


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

// Package ioutil implements some I/O utility functions.
package ioutil

import (
	"bytes"
	"io"
	"os"
	"sort"
)

// readAll reads from r until an error or EOF and returns the data it read
// from the internal buffer allocated with a specified capacity.
func readAll(r io.Reader, capacity int64) (b []byte, err error) {
	buf := bytes.NewBuffer(make([]byte, 0, capacity))
	// If the buffer overflows, we will get bytes.ErrTooLarge.
	// Return that as an error. Any other panic remains.
	defer func() {
		e := recover()
		if e == nil {
			return
		}
		if panicErr, ok := e.(error); ok && panicErr == bytes.ErrTooLarge {
			err = panicErr
		} else {
			panic(e)
		}
	}()
	_, err = buf.ReadFrom(r)
	return buf.Bytes(), err
}

// ReadAll reads from r until an error or EOF and returns the data it read.
// A successful call returns err == nil, not err == EOF. Because ReadAll is
// defined to read from src until EOF, it does not treat an EOF from Read
// as an error to be reported.
func ReadAll(r io.Reader) ([]byte, error) {
	return readAll(r, bytes.MinRead)
}

// ReadFile reads the file named by filename and returns the contents.
// A successful call returns err == nil, not err == EOF. Because ReadFile
// reads the whole file, it does not treat an EOF from Read as an error
// to be reported.
func ReadFile(filename string) ([]byte, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	// It's a good but not certain bet that FileInfo will tell us exactly how much to
	// read, so let's try it but be prepared for the answer to be wrong.
	var n int64

	if fi, err := f.Stat(); err == nil {
		// Don't preallocate a huge buffer, just in case.
		if size := fi.Size(); size < 1e9 {
			n = size
		}
	}
	// As initial capacity for readAll, use n + a little extra in case Size is zero,
	// and to avoid another allocation after Read has filled the buffer.  The readAll
	// call will read into its allocated internal buffer cheaply.  If the size was
	// wrong, we'll either waste some space off the end or reallocate as needed, but
	// in the overwhelmingly common case we'll get it just right.
	return readAll(f, n+bytes.MinRead)
}

// WriteFile writes data to a file named by filename.
// If the file does not exist, WriteFile creates it with permissions perm;
// otherwise WriteFile truncates it before writing.
func WriteFile(filename string, data []byte, perm os.FileMode) error {
	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
	if err != nil {
		return err
	}
	n, err := f.Write(data)
	if err == nil && n < len(data) {
		err = io.ErrShortWrite
	}
	if err1 := f.Close(); err == nil {
		err = err1
	}
	return err
}

// byName implements sort.Interface.
type byName []os.FileInfo

func (f byName) Len() int           { return len(f) }
func (f byName) Less(i, j int) bool { return f[i].Name() < f[j].Name() }
func (f byName) Swap(i, j int)      { f[i], f[j] = f[j], f[i] }

// ReadDir reads the directory named by dirname and returns
// a list of sorted directory entries.
func ReadDir(dirname string) ([]os.FileInfo, error) {
	f, err := os.Open(dirname)
	if err != nil {
		return nil, err
	}
	list, err := f.Readdir(-1)
	f.Close()
	if err != nil {
		return nil, err
	}
	sort.Sort(byName(list))
	return list, nil
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

// NopCloser returns a ReadCloser with a no-op Close method wrapping
// the provided Reader r.
func NopCloser(r io.Reader) io.ReadCloser {
	return nopCloser{r}
}

type devNull int

// devNull implements ReaderFrom as an optimization so io.Copy to
// ioutil.Discard can avoid doing unnecessary work.
var _ io.ReaderFrom = devNull(0)

func (devNull) Write(p []byte) (int, error) {
	return len(p), nil
}

func (devNull) ReadFrom(r io.Reader) (n int64, err error) {
	buf := blackHole()
	defer blackHolePut(buf)
	readSize := 0
	for {
		readSize, err = r.Read(buf)
		n += int64(readSize)
		if err != nil {
			if err == io.EOF {
				return n, nil
			}
			return
		}
	}
}

// Discard is an io.Writer on which all Write calls succeed
// without doing anything.
var Discard io.Writer = devNull(0)
<commit_msg>io/ioutil: add WriteString to Discard<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package ioutil implements some I/O utility functions.
package ioutil

import (
	"bytes"
	"io"
	"os"
	"sort"
)

// readAll reads from r until an error or EOF and returns the data it read
// from the internal buffer allocated with a specified capacity.
func readAll(r io.Reader, capacity int64) (b []byte, err error) {
	buf := bytes.NewBuffer(make([]byte, 0, capacity))
	// If the buffer overflows, we will get bytes.ErrTooLarge.
	// Return that as an error. Any other panic remains.
	defer func() {
		e := recover()
		if e == nil {
			return
		}
		if panicErr, ok := e.(error); ok && panicErr == bytes.ErrTooLarge {
			err = panicErr
		} else {
			panic(e)
		}
	}()
	_, err = buf.ReadFrom(r)
	return buf.Bytes(), err
}

// ReadAll reads from r until an error or EOF and returns the data it read.
// A successful call returns err == nil, not err == EOF. Because ReadAll is
// defined to read from src until EOF, it does not treat an EOF from Read
// as an error to be reported.
func ReadAll(r io.Reader) ([]byte, error) {
	return readAll(r, bytes.MinRead)
}

// ReadFile reads the file named by filename and returns the contents.
// A successful call returns err == nil, not err == EOF. Because ReadFile
// reads the whole file, it does not treat an EOF from Read as an error
// to be reported.
func ReadFile(filename string) ([]byte, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	// It's a good but not certain bet that FileInfo will tell us exactly how much to
	// read, so let's try it but be prepared for the answer to be wrong.
	var n int64

	if fi, err := f.Stat(); err == nil {
		// Don't preallocate a huge buffer, just in case.
		if size := fi.Size(); size < 1e9 {
			n = size
		}
	}
	// As initial capacity for readAll, use n + a little extra in case Size is zero,
	// and to avoid another allocation after Read has filled the buffer.  The readAll
	// call will read into its allocated internal buffer cheaply.  If the size was
	// wrong, we'll either waste some space off the end or reallocate as needed, but
	// in the overwhelmingly common case we'll get it just right.
	return readAll(f, n+bytes.MinRead)
}

// WriteFile writes data to a file named by filename.
// If the file does not exist, WriteFile creates it with permissions perm;
// otherwise WriteFile truncates it before writing.
func WriteFile(filename string, data []byte, perm os.FileMode) error {
	f, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
	if err != nil {
		return err
	}
	n, err := f.Write(data)
	if err == nil && n < len(data) {
		err = io.ErrShortWrite
	}
	if err1 := f.Close(); err == nil {
		err = err1
	}
	return err
}

// byName implements sort.Interface.
type byName []os.FileInfo

func (f byName) Len() int           { return len(f) }
func (f byName) Less(i, j int) bool { return f[i].Name() < f[j].Name() }
func (f byName) Swap(i, j int)      { f[i], f[j] = f[j], f[i] }

// ReadDir reads the directory named by dirname and returns
// a list of sorted directory entries.
func ReadDir(dirname string) ([]os.FileInfo, error) {
	f, err := os.Open(dirname)
	if err != nil {
		return nil, err
	}
	list, err := f.Readdir(-1)
	f.Close()
	if err != nil {
		return nil, err
	}
	sort.Sort(byName(list))
	return list, nil
}

type nopCloser struct {
	io.Reader
}

func (nopCloser) Close() error { return nil }

// NopCloser returns a ReadCloser with a no-op Close method wrapping
// the provided Reader r.
func NopCloser(r io.Reader) io.ReadCloser {
	return nopCloser{r}
}

type devNull int

// devNull implements ReaderFrom as an optimization so io.Copy to
// ioutil.Discard can avoid doing unnecessary work.
var _ io.ReaderFrom = devNull(0)

func (devNull) Write(p []byte) (int, error) {
	return len(p), nil
}

func (devNull) WriteString(s string) (int, error) {
	return len(s), nil
}

func (devNull) ReadFrom(r io.Reader) (n int64, err error) {
	buf := blackHole()
	defer blackHolePut(buf)
	readSize := 0
	for {
		readSize, err = r.Read(buf)
		n += int64(readSize)
		if err != nil {
			if err == io.EOF {
				return n, nil
			}
			return
		}
	}
}

// Discard is an io.Writer on which all Write calls succeed
// without doing anything.
var Discard io.Writer = devNull(0)
<|endoftext|>


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

// Package websocket implements a client and server for the WebSocket protocol
// as specified in RFC 6455.
package websocket // import "golang.org/x/net/websocket"

import (
	"bufio"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"
)

const (
	ProtocolVersionHybi13    = 13
	ProtocolVersionHybi      = ProtocolVersionHybi13
	SupportedProtocolVersion = "13"

	ContinuationFrame = 0
	TextFrame         = 1
	BinaryFrame       = 2
	CloseFrame        = 8
	PingFrame         = 9
	PongFrame         = 10
	UnknownFrame      = 255

	DefaultMaxPayloadBytes = 32 << 20 // 32MB
)

// ProtocolError represents WebSocket protocol errors.
type ProtocolError struct {
	ErrorString string
}

func (err *ProtocolError) Error() string { return err.ErrorString }

var (
	ErrBadProtocolVersion   = &ProtocolError{"bad protocol version"}
	ErrBadScheme            = &ProtocolError{"bad scheme"}
	ErrBadStatus            = &ProtocolError{"bad status"}
	ErrBadUpgrade           = &ProtocolError{"missing or bad upgrade"}
	ErrBadWebSocketOrigin   = &ProtocolError{"missing or bad WebSocket-Origin"}
	ErrBadWebSocketLocation = &ProtocolError{"missing or bad WebSocket-Location"}
	ErrBadWebSocketProtocol = &ProtocolError{"missing or bad WebSocket-Protocol"}
	ErrBadWebSocketVersion  = &ProtocolError{"missing or bad WebSocket Version"}
	ErrChallengeResponse    = &ProtocolError{"mismatch challenge/response"}
	ErrBadFrame             = &ProtocolError{"bad frame"}
	ErrBadFrameBoundary     = &ProtocolError{"not on frame boundary"}
	ErrNotWebSocket         = &ProtocolError{"not websocket protocol"}
	ErrBadRequestMethod     = &ProtocolError{"bad method"}
	ErrNotSupported         = &ProtocolError{"not supported"}
)

// ErrFrameTooLarge is returned by Codec's Receive method if payload size
// exceeds limit set by Conn.MaxPayloadBytes
var ErrFrameTooLarge = errors.New("websocket: frame payload size exceeds limit")

// Addr is an implementation of net.Addr for WebSocket.
type Addr struct {
	*url.URL
}

// Network returns the network type for a WebSocket, "websocket".
func (addr *Addr) Network() string { return "websocket" }

// Config is a WebSocket configuration
type Config struct {
	// A WebSocket server address.
	Location *url.URL

	// A Websocket client origin.
	Origin *url.URL

	// WebSocket subprotocols.
	Protocol []string

	// WebSocket protocol version.
	Version int

	// TLS config for secure WebSocket (wss).
	TlsConfig *tls.Config

	// Additional header fields to be sent in WebSocket opening handshake.
	Header http.Header

	// Dialer used when opening websocket connections.
	Dialer *net.Dialer

	handshakeData map[string]string
}

// serverHandshaker is an interface to handle WebSocket server side handshake.
type serverHandshaker interface {
	// ReadHandshake reads handshake request message from client.
	// Returns http response code and error if any.
	ReadHandshake(buf *bufio.Reader, req *http.Request) (code int, err error)

	// AcceptHandshake accepts the client handshake request and sends
	// handshake response back to client.
	AcceptHandshake(buf *bufio.Writer) (err error)

	// NewServerConn creates a new WebSocket connection.
	NewServerConn(buf *bufio.ReadWriter, rwc io.ReadWriteCloser, request *http.Request) (conn *Conn)
}

// frameReader is an interface to read a WebSocket frame.
type frameReader interface {
	// Reader is to read payload of the frame.
	io.Reader

	// PayloadType returns payload type.
	PayloadType() byte

	// HeaderReader returns a reader to read header of the frame.
	HeaderReader() io.Reader

	// TrailerReader returns a reader to read trailer of the frame.
	// If it returns nil, there is no trailer in the frame.
	TrailerReader() io.Reader

	// Len returns total length of the frame, including header and trailer.
	Len() int
}

// frameReaderFactory is an interface to creates new frame reader.
type frameReaderFactory interface {
	NewFrameReader() (r frameReader, err error)
}

// frameWriter is an interface to write a WebSocket frame.
type frameWriter interface {
	// Writer is to write payload of the frame.
	io.WriteCloser
}

// frameWriterFactory is an interface to create new frame writer.
type frameWriterFactory interface {
	NewFrameWriter(payloadType byte) (w frameWriter, err error)
}

type frameHandler interface {
	HandleFrame(frame frameReader) (r frameReader, err error)
	WriteClose(status int) (err error)
}

// Conn represents a WebSocket connection.
//
// Multiple goroutines may invoke methods on a Conn simultaneously.
type Conn struct {
	config  *Config
	request *http.Request

	buf *bufio.ReadWriter
	rwc io.ReadWriteCloser

	rio sync.Mutex
	frameReaderFactory
	frameReader

	wio sync.Mutex
	frameWriterFactory

	frameHandler
	PayloadType        byte
	defaultCloseStatus int

	// MaxPayloadBytes limits the size of frame payload received over Conn
	// by Codec's Receive method. If zero, DefaultMaxPayloadBytes is used.
	MaxPayloadBytes int
}

// Read implements the io.Reader interface:
// it reads data of a frame from the WebSocket connection.
// if msg is not large enough for the frame data, it fills the msg and next Read
// will read the rest of the frame data.
// it reads Text frame or Binary frame.
func (ws *Conn) Read(msg []byte) (n int, err error) {
	ws.rio.Lock()
	defer ws.rio.Unlock()
again:
	if ws.frameReader == nil {
		frame, err := ws.frameReaderFactory.NewFrameReader()
		if err != nil {
			return 0, err
		}
		ws.frameReader, err = ws.frameHandler.HandleFrame(frame)
		if err != nil {
			return 0, err
		}
		if ws.frameReader == nil {
			goto again
		}
	}
	n, err = ws.frameReader.Read(msg)
	if err == io.EOF {
		if trailer := ws.frameReader.TrailerReader(); trailer != nil {
			io.Copy(ioutil.Discard, trailer)
		}
		ws.frameReader = nil
		goto again
	}
	return n, err
}

// Write implements the io.Writer interface:
// it writes data as a frame to the WebSocket connection.
func (ws *Conn) Write(msg []byte) (n int, err error) {
	ws.wio.Lock()
	defer ws.wio.Unlock()
	w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
	if err != nil {
		return 0, err
	}
	n, err = w.Write(msg)
	w.Close()
	return n, err
}

// Close implements the io.Closer interface.
func (ws *Conn) Close() error {
	err := ws.frameHandler.WriteClose(ws.defaultCloseStatus)
	err1 := ws.rwc.Close()
	if err != nil {
		return err
	}
	return err1
}

func (ws *Conn) IsClientConn() bool { return ws.request == nil }
func (ws *Conn) IsServerConn() bool { return ws.request != nil }

// LocalAddr returns the WebSocket Origin for the connection for client, or
// the WebSocket location for server.
func (ws *Conn) LocalAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Origin}
	}
	return &Addr{ws.config.Location}
}

// RemoteAddr returns the WebSocket location for the connection for client, or
// the Websocket Origin for server.
func (ws *Conn) RemoteAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Location}
	}
	return &Addr{ws.config.Origin}
}

var errSetDeadline = errors.New("websocket: cannot set deadline: not using a net.Conn")

// SetDeadline sets the connection's network read & write deadlines.
func (ws *Conn) SetDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetDeadline(t)
	}
	return errSetDeadline
}

// SetReadDeadline sets the connection's network read deadline.
func (ws *Conn) SetReadDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetReadDeadline(t)
	}
	return errSetDeadline
}

// SetWriteDeadline sets the connection's network write deadline.
func (ws *Conn) SetWriteDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetWriteDeadline(t)
	}
	return errSetDeadline
}

// Config returns the WebSocket config.
func (ws *Conn) Config() *Config { return ws.config }

// Request returns the http request upgraded to the WebSocket.
// It is nil for client side.
func (ws *Conn) Request() *http.Request { return ws.request }

// Codec represents a symmetric pair of functions that implement a codec.
type Codec struct {
	Marshal   func(v interface{}) (data []byte, payloadType byte, err error)
	Unmarshal func(data []byte, payloadType byte, v interface{}) (err error)
}

// Send sends v marshaled by cd.Marshal as single frame to ws.
func (cd Codec) Send(ws *Conn, v interface{}) (err error) {
	data, payloadType, err := cd.Marshal(v)
	if err != nil {
		return err
	}
	ws.wio.Lock()
	defer ws.wio.Unlock()
	w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
	if err != nil {
		return err
	}
	_, err = w.Write(data)
	w.Close()
	return err
}

// Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores
// in v. The whole frame payload is read to an in-memory buffer; max size of
// payload is defined by ws.MaxPayloadBytes. If frame payload size exceeds
// limit, ErrFrameTooLarge is returned; in this case frame is not read off wire
// completely. The next call to Receive would read and discard leftover data of
// previous oversized frame before processing next frame.
func (cd Codec) Receive(ws *Conn, v interface{}) (err error) {
	ws.rio.Lock()
	defer ws.rio.Unlock()
	if ws.frameReader != nil {
		_, err = io.Copy(ioutil.Discard, ws.frameReader)
		if err != nil {
			return err
		}
		ws.frameReader = nil
	}
again:
	frame, err := ws.frameReaderFactory.NewFrameReader()
	if err != nil {
		return err
	}
	frame, err = ws.frameHandler.HandleFrame(frame)
	if err != nil {
		return err
	}
	if frame == nil {
		goto again
	}
	maxPayloadBytes := ws.MaxPayloadBytes
	if maxPayloadBytes == 0 {
		maxPayloadBytes = DefaultMaxPayloadBytes
	}
	if hf, ok := frame.(*hybiFrameReader); ok && hf.header.Length > int64(maxPayloadBytes) {
		// payload size exceeds limit, no need to call Unmarshal
		//
		// set frameReader to current oversized frame so that
		// the next call to this function can drain leftover
		// data before processing the next frame
		ws.frameReader = frame
		return ErrFrameTooLarge
	}
	payloadType := frame.PayloadType()
	data, err := ioutil.ReadAll(frame)
	if err != nil {
		return err
	}
	return cd.Unmarshal(data, payloadType, v)
}

func marshal(v interface{}) (msg []byte, payloadType byte, err error) {
	switch data := v.(type) {
	case string:
		return []byte(data), TextFrame, nil
	case []byte:
		return data, BinaryFrame, nil
	}
	return nil, UnknownFrame, ErrNotSupported
}

func unmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
	switch data := v.(type) {
	case *string:
		*data = string(msg)
		return nil
	case *[]byte:
		*data = msg
		return nil
	}
	return ErrNotSupported
}

/*
Message is a codec to send/receive text/binary data in a frame on WebSocket connection.
To send/receive text frame, use string type.
To send/receive binary frame, use []byte type.

Trivial usage:

	import "websocket"

	// receive text frame
	var message string
	websocket.Message.Receive(ws, &message)

	// send text frame
	message = "hello"
	websocket.Message.Send(ws, message)

	// receive binary frame
	var data []byte
	websocket.Message.Receive(ws, &data)

	// send binary frame
	data = []byte{0, 1, 2}
	websocket.Message.Send(ws, data)

*/
var Message = Codec{marshal, unmarshal}

func jsonMarshal(v interface{}) (msg []byte, payloadType byte, err error) {
	msg, err = json.Marshal(v)
	return msg, TextFrame, err
}

func jsonUnmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
	return json.Unmarshal(msg, v)
}

/*
JSON is a codec to send/receive JSON data in a frame from a WebSocket connection.

Trivial usage:

	import "websocket"

	type T struct {
		Msg string
		Count int
	}

	// receive JSON type T
	var data T
	websocket.JSON.Receive(ws, &data)

	// send JSON type T
	websocket.JSON.Send(ws, data)
*/
var JSON = Codec{jsonMarshal, jsonUnmarshal}
<commit_msg>websocket: mention the gorilla package<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package websocket implements a client and server for the WebSocket protocol
// as specified in RFC 6455.
//
// This package currently lacks some features found in an alternative
// and more actively maintained WebSocket package:
//
//     https://godoc.org/github.com/gorilla/websocket
//
package websocket // import "golang.org/x/net/websocket"

import (
	"bufio"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"
)

const (
	ProtocolVersionHybi13    = 13
	ProtocolVersionHybi      = ProtocolVersionHybi13
	SupportedProtocolVersion = "13"

	ContinuationFrame = 0
	TextFrame         = 1
	BinaryFrame       = 2
	CloseFrame        = 8
	PingFrame         = 9
	PongFrame         = 10
	UnknownFrame      = 255

	DefaultMaxPayloadBytes = 32 << 20 // 32MB
)

// ProtocolError represents WebSocket protocol errors.
type ProtocolError struct {
	ErrorString string
}

func (err *ProtocolError) Error() string { return err.ErrorString }

var (
	ErrBadProtocolVersion   = &ProtocolError{"bad protocol version"}
	ErrBadScheme            = &ProtocolError{"bad scheme"}
	ErrBadStatus            = &ProtocolError{"bad status"}
	ErrBadUpgrade           = &ProtocolError{"missing or bad upgrade"}
	ErrBadWebSocketOrigin   = &ProtocolError{"missing or bad WebSocket-Origin"}
	ErrBadWebSocketLocation = &ProtocolError{"missing or bad WebSocket-Location"}
	ErrBadWebSocketProtocol = &ProtocolError{"missing or bad WebSocket-Protocol"}
	ErrBadWebSocketVersion  = &ProtocolError{"missing or bad WebSocket Version"}
	ErrChallengeResponse    = &ProtocolError{"mismatch challenge/response"}
	ErrBadFrame             = &ProtocolError{"bad frame"}
	ErrBadFrameBoundary     = &ProtocolError{"not on frame boundary"}
	ErrNotWebSocket         = &ProtocolError{"not websocket protocol"}
	ErrBadRequestMethod     = &ProtocolError{"bad method"}
	ErrNotSupported         = &ProtocolError{"not supported"}
)

// ErrFrameTooLarge is returned by Codec's Receive method if payload size
// exceeds limit set by Conn.MaxPayloadBytes
var ErrFrameTooLarge = errors.New("websocket: frame payload size exceeds limit")

// Addr is an implementation of net.Addr for WebSocket.
type Addr struct {
	*url.URL
}

// Network returns the network type for a WebSocket, "websocket".
func (addr *Addr) Network() string { return "websocket" }

// Config is a WebSocket configuration
type Config struct {
	// A WebSocket server address.
	Location *url.URL

	// A Websocket client origin.
	Origin *url.URL

	// WebSocket subprotocols.
	Protocol []string

	// WebSocket protocol version.
	Version int

	// TLS config for secure WebSocket (wss).
	TlsConfig *tls.Config

	// Additional header fields to be sent in WebSocket opening handshake.
	Header http.Header

	// Dialer used when opening websocket connections.
	Dialer *net.Dialer

	handshakeData map[string]string
}

// serverHandshaker is an interface to handle WebSocket server side handshake.
type serverHandshaker interface {
	// ReadHandshake reads handshake request message from client.
	// Returns http response code and error if any.
	ReadHandshake(buf *bufio.Reader, req *http.Request) (code int, err error)

	// AcceptHandshake accepts the client handshake request and sends
	// handshake response back to client.
	AcceptHandshake(buf *bufio.Writer) (err error)

	// NewServerConn creates a new WebSocket connection.
	NewServerConn(buf *bufio.ReadWriter, rwc io.ReadWriteCloser, request *http.Request) (conn *Conn)
}

// frameReader is an interface to read a WebSocket frame.
type frameReader interface {
	// Reader is to read payload of the frame.
	io.Reader

	// PayloadType returns payload type.
	PayloadType() byte

	// HeaderReader returns a reader to read header of the frame.
	HeaderReader() io.Reader

	// TrailerReader returns a reader to read trailer of the frame.
	// If it returns nil, there is no trailer in the frame.
	TrailerReader() io.Reader

	// Len returns total length of the frame, including header and trailer.
	Len() int
}

// frameReaderFactory is an interface to creates new frame reader.
type frameReaderFactory interface {
	NewFrameReader() (r frameReader, err error)
}

// frameWriter is an interface to write a WebSocket frame.
type frameWriter interface {
	// Writer is to write payload of the frame.
	io.WriteCloser
}

// frameWriterFactory is an interface to create new frame writer.
type frameWriterFactory interface {
	NewFrameWriter(payloadType byte) (w frameWriter, err error)
}

type frameHandler interface {
	HandleFrame(frame frameReader) (r frameReader, err error)
	WriteClose(status int) (err error)
}

// Conn represents a WebSocket connection.
//
// Multiple goroutines may invoke methods on a Conn simultaneously.
type Conn struct {
	config  *Config
	request *http.Request

	buf *bufio.ReadWriter
	rwc io.ReadWriteCloser

	rio sync.Mutex
	frameReaderFactory
	frameReader

	wio sync.Mutex
	frameWriterFactory

	frameHandler
	PayloadType        byte
	defaultCloseStatus int

	// MaxPayloadBytes limits the size of frame payload received over Conn
	// by Codec's Receive method. If zero, DefaultMaxPayloadBytes is used.
	MaxPayloadBytes int
}

// Read implements the io.Reader interface:
// it reads data of a frame from the WebSocket connection.
// if msg is not large enough for the frame data, it fills the msg and next Read
// will read the rest of the frame data.
// it reads Text frame or Binary frame.
func (ws *Conn) Read(msg []byte) (n int, err error) {
	ws.rio.Lock()
	defer ws.rio.Unlock()
again:
	if ws.frameReader == nil {
		frame, err := ws.frameReaderFactory.NewFrameReader()
		if err != nil {
			return 0, err
		}
		ws.frameReader, err = ws.frameHandler.HandleFrame(frame)
		if err != nil {
			return 0, err
		}
		if ws.frameReader == nil {
			goto again
		}
	}
	n, err = ws.frameReader.Read(msg)
	if err == io.EOF {
		if trailer := ws.frameReader.TrailerReader(); trailer != nil {
			io.Copy(ioutil.Discard, trailer)
		}
		ws.frameReader = nil
		goto again
	}
	return n, err
}

// Write implements the io.Writer interface:
// it writes data as a frame to the WebSocket connection.
func (ws *Conn) Write(msg []byte) (n int, err error) {
	ws.wio.Lock()
	defer ws.wio.Unlock()
	w, err := ws.frameWriterFactory.NewFrameWriter(ws.PayloadType)
	if err != nil {
		return 0, err
	}
	n, err = w.Write(msg)
	w.Close()
	return n, err
}

// Close implements the io.Closer interface.
func (ws *Conn) Close() error {
	err := ws.frameHandler.WriteClose(ws.defaultCloseStatus)
	err1 := ws.rwc.Close()
	if err != nil {
		return err
	}
	return err1
}

func (ws *Conn) IsClientConn() bool { return ws.request == nil }
func (ws *Conn) IsServerConn() bool { return ws.request != nil }

// LocalAddr returns the WebSocket Origin for the connection for client, or
// the WebSocket location for server.
func (ws *Conn) LocalAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Origin}
	}
	return &Addr{ws.config.Location}
}

// RemoteAddr returns the WebSocket location for the connection for client, or
// the Websocket Origin for server.
func (ws *Conn) RemoteAddr() net.Addr {
	if ws.IsClientConn() {
		return &Addr{ws.config.Location}
	}
	return &Addr{ws.config.Origin}
}

var errSetDeadline = errors.New("websocket: cannot set deadline: not using a net.Conn")

// SetDeadline sets the connection's network read & write deadlines.
func (ws *Conn) SetDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetDeadline(t)
	}
	return errSetDeadline
}

// SetReadDeadline sets the connection's network read deadline.
func (ws *Conn) SetReadDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetReadDeadline(t)
	}
	return errSetDeadline
}

// SetWriteDeadline sets the connection's network write deadline.
func (ws *Conn) SetWriteDeadline(t time.Time) error {
	if conn, ok := ws.rwc.(net.Conn); ok {
		return conn.SetWriteDeadline(t)
	}
	return errSetDeadline
}

// Config returns the WebSocket config.
func (ws *Conn) Config() *Config { return ws.config }

// Request returns the http request upgraded to the WebSocket.
// It is nil for client side.
func (ws *Conn) Request() *http.Request { return ws.request }

// Codec represents a symmetric pair of functions that implement a codec.
type Codec struct {
	Marshal   func(v interface{}) (data []byte, payloadType byte, err error)
	Unmarshal func(data []byte, payloadType byte, v interface{}) (err error)
}

// Send sends v marshaled by cd.Marshal as single frame to ws.
func (cd Codec) Send(ws *Conn, v interface{}) (err error) {
	data, payloadType, err := cd.Marshal(v)
	if err != nil {
		return err
	}
	ws.wio.Lock()
	defer ws.wio.Unlock()
	w, err := ws.frameWriterFactory.NewFrameWriter(payloadType)
	if err != nil {
		return err
	}
	_, err = w.Write(data)
	w.Close()
	return err
}

// Receive receives single frame from ws, unmarshaled by cd.Unmarshal and stores
// in v. The whole frame payload is read to an in-memory buffer; max size of
// payload is defined by ws.MaxPayloadBytes. If frame payload size exceeds
// limit, ErrFrameTooLarge is returned; in this case frame is not read off wire
// completely. The next call to Receive would read and discard leftover data of
// previous oversized frame before processing next frame.
func (cd Codec) Receive(ws *Conn, v interface{}) (err error) {
	ws.rio.Lock()
	defer ws.rio.Unlock()
	if ws.frameReader != nil {
		_, err = io.Copy(ioutil.Discard, ws.frameReader)
		if err != nil {
			return err
		}
		ws.frameReader = nil
	}
again:
	frame, err := ws.frameReaderFactory.NewFrameReader()
	if err != nil {
		return err
	}
	frame, err = ws.frameHandler.HandleFrame(frame)
	if err != nil {
		return err
	}
	if frame == nil {
		goto again
	}
	maxPayloadBytes := ws.MaxPayloadBytes
	if maxPayloadBytes == 0 {
		maxPayloadBytes = DefaultMaxPayloadBytes
	}
	if hf, ok := frame.(*hybiFrameReader); ok && hf.header.Length > int64(maxPayloadBytes) {
		// payload size exceeds limit, no need to call Unmarshal
		//
		// set frameReader to current oversized frame so that
		// the next call to this function can drain leftover
		// data before processing the next frame
		ws.frameReader = frame
		return ErrFrameTooLarge
	}
	payloadType := frame.PayloadType()
	data, err := ioutil.ReadAll(frame)
	if err != nil {
		return err
	}
	return cd.Unmarshal(data, payloadType, v)
}

func marshal(v interface{}) (msg []byte, payloadType byte, err error) {
	switch data := v.(type) {
	case string:
		return []byte(data), TextFrame, nil
	case []byte:
		return data, BinaryFrame, nil
	}
	return nil, UnknownFrame, ErrNotSupported
}

func unmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
	switch data := v.(type) {
	case *string:
		*data = string(msg)
		return nil
	case *[]byte:
		*data = msg
		return nil
	}
	return ErrNotSupported
}

/*
Message is a codec to send/receive text/binary data in a frame on WebSocket connection.
To send/receive text frame, use string type.
To send/receive binary frame, use []byte type.

Trivial usage:

	import "websocket"

	// receive text frame
	var message string
	websocket.Message.Receive(ws, &message)

	// send text frame
	message = "hello"
	websocket.Message.Send(ws, message)

	// receive binary frame
	var data []byte
	websocket.Message.Receive(ws, &data)

	// send binary frame
	data = []byte{0, 1, 2}
	websocket.Message.Send(ws, data)

*/
var Message = Codec{marshal, unmarshal}

func jsonMarshal(v interface{}) (msg []byte, payloadType byte, err error) {
	msg, err = json.Marshal(v)
	return msg, TextFrame, err
}

func jsonUnmarshal(msg []byte, payloadType byte, v interface{}) (err error) {
	return json.Unmarshal(msg, v)
}

/*
JSON is a codec to send/receive JSON data in a frame from a WebSocket connection.

Trivial usage:

	import "websocket"

	type T struct {
		Msg string
		Count int
	}

	// receive JSON type T
	var data T
	websocket.JSON.Receive(ws, &data)

	// send JSON type T
	websocket.JSON.Send(ws, data)
*/
var JSON = Codec{jsonMarshal, jsonUnmarshal}
<|endoftext|>


<commit_before>// Copyright (c) 2017, CodeBoy. All rights reserved.
//
// This Source Code Form is subject to the terms of the
// license that can be found in the LICENSE file.

package fileapp

import (
	log "github.com/Sirupsen/logrus"
	"os/exec"
	"strings"

	"github.com/filemaps/filemaps/pkg/config"
)

func init() {
	register(NewCustom())
}

type Custom struct {
}

func NewCustom() *Custom {
	return &Custom{}
}

func (a *Custom) getInfo() FileAppInfo {
	return FileAppInfo{
		ID:   "custom1",
		Name: "Custom",
	}
}

func (a *Custom) open(path string) int {
	log.WithFields(log.Fields{
		"path": path,
	}).Info("Custom: open")

	cfg := config.GetConfiguration()
	cmd := strings.Split(cfg.TextEditorCustom1Cmd, " ")
	cmd = append(cmd, path)

	out, err := exec.Command(cmd[0], cmd[1:]...).Output()
	if err != nil {
		log.WithFields(log.Fields{
			"err": err,
		}).Error("Custom open error")
		return -1
	}
	log.WithFields(log.Fields{
		"out": out,
	}).Info("Custom")
	return 0
}
<commit_msg>Use comma as argument separator in custom command Space is usual in command for Windows<commit_after>// Copyright (c) 2017, CodeBoy. All rights reserved.
//
// This Source Code Form is subject to the terms of the
// license that can be found in the LICENSE file.

package fileapp

import (
	log "github.com/Sirupsen/logrus"
	"os/exec"
	"strings"

	"github.com/filemaps/filemaps/pkg/config"
)

const (
	// argSeparator separates arguments from custom command string
	argSeparator = ","
)

func init() {
	register(NewCustom())
}

type Custom struct {
}

func NewCustom() *Custom {
	return &Custom{}
}

func (a *Custom) getInfo() FileAppInfo {
	return FileAppInfo{
		ID:   "custom1",
		Name: "Custom",
	}
}

func (a *Custom) open(path string) int {
	log.WithFields(log.Fields{
		"path": path,
	}).Info("Custom: open")

	cfg := config.GetConfiguration()
	cmd := strings.Split(cfg.TextEditorCustom1Cmd, argSeparator)
	cmd = append(cmd, path)

	out, err := exec.Command(cmd[0], cmd[1:]...).Output()
	if err != nil {
		log.WithFields(log.Fields{
			"err": err,
		}).Error("Custom open error")
		return -1
	}
	log.WithFields(log.Fields{
		"out": out,
	}).Info("Custom")
	return 0
}
<|endoftext|>


<commit_before>package storage

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/chrislusf/seaweedfs/weed/images"
	. "github.com/chrislusf/seaweedfs/weed/storage/types"
	"io/ioutil"
)

const (
	NeedleChecksumSize = 4
	PairNamePrefix     = "Seaweed-"
)

/*
* A Needle means a uploaded and stored file.
* Needle file size is limited to 4GB for now.
 */
type Needle struct {
	Cookie Cookie   `comment:"random number to mitigate brute force lookups"`
	Id     NeedleId `comment:"needle id"`
	Size   uint32   `comment:"sum of DataSize,Data,NameSize,Name,MimeSize,Mime"`

	DataSize     uint32 `comment:"Data size"` //version2
	Data         []byte `comment:"The actual file data"`
	Flags        byte   `comment:"boolean flags"`          //version2
	NameSize     uint8                                     //version2
	Name         []byte `comment:"maximum 256 characters"` //version2
	MimeSize     uint8                                     //version2
	Mime         []byte `comment:"maximum 256 characters"` //version2
	PairsSize    uint16                                    //version2
	Pairs        []byte `comment:"additional name value pairs, json format, maximum 64kB"`
	LastModified uint64 //only store LastModifiedBytesLength bytes, which is 5 bytes to disk
	Ttl          *TTL

	Checksum   CRC    `comment:"CRC32 to check integrity"`
	AppendAtNs uint64 `comment:"append timestamp in nano seconds"` //version3
	Padding    []byte `comment:"Aligned to 8 bytes"`
}

func (n *Needle) String() (str string) {
	str = fmt.Sprintf("%s Size:%d, DataSize:%d, Name:%s, Mime:%s", formatNeedleIdCookie(n.Id, n.Cookie), n.Size, n.DataSize, n.Name, n.Mime)
	return
}

func ParseUpload(r *http.Request) (
	fileName string, data []byte, mimeType string, pairMap map[string]string, isGzipped bool, originalDataSize int,
	modifiedTime uint64, ttl *TTL, isChunkedFile bool, e error) {
	pairMap = make(map[string]string)
	for k, v := range r.Header {
		if len(v) > 0 && strings.HasPrefix(k, PairNamePrefix) {
			pairMap[k] = v[0]
		}
	}

	if r.Method == "POST" {
		fileName, data, mimeType, isGzipped, originalDataSize, isChunkedFile, e = parseMultipart(r)
	} else {
		isGzipped = false
		mimeType = r.Header.Get("Content-Type")
		fileName = ""
		data, e = ioutil.ReadAll(r.Body)
	}
	if e != nil {
		return
	}

	modifiedTime, _ = strconv.ParseUint(r.FormValue("ts"), 10, 64)
	ttl, _ = ReadTTL(r.FormValue("ttl"))

	return
}
func CreateNeedleFromRequest(r *http.Request, fixJpgOrientation bool) (n *Needle, originalSize int, e error) {
	var pairMap map[string]string
	fname, mimeType, isGzipped, isChunkedFile := "", "", false, false
	n = new(Needle)
	fname, n.Data, mimeType, pairMap, isGzipped, originalSize, n.LastModified, n.Ttl, isChunkedFile, e = ParseUpload(r)
	if e != nil {
		return
	}
	if len(fname) < 256 {
		n.Name = []byte(fname)
		n.SetHasName()
	}
	if len(mimeType) < 256 {
		n.Mime = []byte(mimeType)
		n.SetHasMime()
	}
	if len(pairMap) != 0 {
		trimmedPairMap := make(map[string]string)
		for k, v := range pairMap {
			trimmedPairMap[k[len(PairNamePrefix):]] = v
		}

		pairs, _ := json.Marshal(trimmedPairMap)
		if len(pairs) < 65536 {
			n.Pairs = pairs
			n.PairsSize = uint16(len(pairs))
			n.SetHasPairs()
		}
	}
	if isGzipped {
		n.SetGzipped()
	}
	if n.LastModified == 0 {
		n.LastModified = uint64(time.Now().Unix())
	}
	n.SetHasLastModifiedDate()
	if n.Ttl != EMPTY_TTL {
		n.SetHasTtl()
	}

	if isChunkedFile {
		n.SetIsChunkManifest()
	}

	if fixJpgOrientation {
		loweredName := strings.ToLower(fname)
		if mimeType == "image/jpeg" || strings.HasSuffix(loweredName, ".jpg") || strings.HasSuffix(loweredName, ".jpeg") {
			n.Data = images.FixJpgOrientation(n.Data)
		}
	}

	n.Checksum = NewCRC(n.Data)

	commaSep := strings.LastIndex(r.URL.Path, ",")
	dotSep := strings.LastIndex(r.URL.Path, ".")
	fid := r.URL.Path[commaSep+1:]
	if dotSep > 0 {
		fid = r.URL.Path[commaSep+1 : dotSep]
	}

	e = n.ParsePath(fid)

	return
}
func (n *Needle) ParsePath(fid string) (err error) {
	length := len(fid)
	if length <= CookieSize*2 {
		return fmt.Errorf("Invalid fid: %s", fid)
	}
	delta := ""
	deltaIndex := strings.LastIndex(fid, "_")
	if deltaIndex > 0 {
		fid, delta = fid[0:deltaIndex], fid[deltaIndex+1:]
	}
	n.Id, n.Cookie, err = ParseNeedleIdCookie(fid)
	if err != nil {
		return err
	}
	if delta != "" {
		if d, e := strconv.ParseUint(delta, 10, 64); e == nil {
			n.Id += NeedleId(d)
		} else {
			return e
		}
	}
	return err
}

func ParseNeedleIdCookie(key_hash_string string) (NeedleId, Cookie, error) {
	if len(key_hash_string) <= CookieSize*2 {
		return NeedleIdEmpty, 0, fmt.Errorf("KeyHash is too short.")
	}
	if len(key_hash_string) > (NeedleIdSize+CookieSize)*2 {
		return NeedleIdEmpty, 0, fmt.Errorf("KeyHash is too long.")
	}
	split := len(key_hash_string) - CookieSize*2
	needleId, err := ParseNeedleId(key_hash_string[:split])
	if err != nil {
		return NeedleIdEmpty, 0, fmt.Errorf("Parse needleId error: %v", err)
	}
	cookie, err := ParseCookie(key_hash_string[split:])
	if err != nil {
		return NeedleIdEmpty, 0, fmt.Errorf("Parse cookie error: %v", err)
	}
	return needleId, cookie, nil
}

func (n *Needle) LastModifiedString() string {
	return time.Unix(int64(n.LastModified), 0).Format("2006-01-02T15:04:05")
}
<commit_msg>fix s3cmd put<commit_after>package storage

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/chrislusf/seaweedfs/weed/images"
	. "github.com/chrislusf/seaweedfs/weed/storage/types"
	"io/ioutil"
)

const (
	NeedleChecksumSize = 4
	PairNamePrefix     = "Seaweed-"
)

/*
* A Needle means a uploaded and stored file.
* Needle file size is limited to 4GB for now.
 */
type Needle struct {
	Cookie Cookie   `comment:"random number to mitigate brute force lookups"`
	Id     NeedleId `comment:"needle id"`
	Size   uint32   `comment:"sum of DataSize,Data,NameSize,Name,MimeSize,Mime"`

	DataSize     uint32 `comment:"Data size"` //version2
	Data         []byte `comment:"The actual file data"`
	Flags        byte   `comment:"boolean flags"`          //version2
	NameSize     uint8                                     //version2
	Name         []byte `comment:"maximum 256 characters"` //version2
	MimeSize     uint8                                     //version2
	Mime         []byte `comment:"maximum 256 characters"` //version2
	PairsSize    uint16                                    //version2
	Pairs        []byte `comment:"additional name value pairs, json format, maximum 64kB"`
	LastModified uint64 //only store LastModifiedBytesLength bytes, which is 5 bytes to disk
	Ttl          *TTL

	Checksum   CRC    `comment:"CRC32 to check integrity"`
	AppendAtNs uint64 `comment:"append timestamp in nano seconds"` //version3
	Padding    []byte `comment:"Aligned to 8 bytes"`
}

func (n *Needle) String() (str string) {
	str = fmt.Sprintf("%s Size:%d, DataSize:%d, Name:%s, Mime:%s", formatNeedleIdCookie(n.Id, n.Cookie), n.Size, n.DataSize, n.Name, n.Mime)
	return
}

func ParseUpload(r *http.Request) (
	fileName string, data []byte, mimeType string, pairMap map[string]string, isGzipped bool, originalDataSize int,
	modifiedTime uint64, ttl *TTL, isChunkedFile bool, e error) {
	pairMap = make(map[string]string)
	for k, v := range r.Header {
		if len(v) > 0 && strings.HasPrefix(k, PairNamePrefix) {
			pairMap[k] = v[0]
		}
	}

	if r.Method == "POST" {
		fileName, data, mimeType, isGzipped, originalDataSize, isChunkedFile, e = parseMultipart(r)
	} else {
		isGzipped = false
		mimeType = r.Header.Get("Content-Type")
		fileName = ""
		data, e = ioutil.ReadAll(r.Body)
		originalDataSize = len(data)
	}
	if e != nil {
		return
	}

	modifiedTime, _ = strconv.ParseUint(r.FormValue("ts"), 10, 64)
	ttl, _ = ReadTTL(r.FormValue("ttl"))

	return
}
func CreateNeedleFromRequest(r *http.Request, fixJpgOrientation bool) (n *Needle, originalSize int, e error) {
	var pairMap map[string]string
	fname, mimeType, isGzipped, isChunkedFile := "", "", false, false
	n = new(Needle)
	fname, n.Data, mimeType, pairMap, isGzipped, originalSize, n.LastModified, n.Ttl, isChunkedFile, e = ParseUpload(r)
	if e != nil {
		return
	}
	if len(fname) < 256 {
		n.Name = []byte(fname)
		n.SetHasName()
	}
	if len(mimeType) < 256 {
		n.Mime = []byte(mimeType)
		n.SetHasMime()
	}
	if len(pairMap) != 0 {
		trimmedPairMap := make(map[string]string)
		for k, v := range pairMap {
			trimmedPairMap[k[len(PairNamePrefix):]] = v
		}

		pairs, _ := json.Marshal(trimmedPairMap)
		if len(pairs) < 65536 {
			n.Pairs = pairs
			n.PairsSize = uint16(len(pairs))
			n.SetHasPairs()
		}
	}
	if isGzipped {
		n.SetGzipped()
	}
	if n.LastModified == 0 {
		n.LastModified = uint64(time.Now().Unix())
	}
	n.SetHasLastModifiedDate()
	if n.Ttl != EMPTY_TTL {
		n.SetHasTtl()
	}

	if isChunkedFile {
		n.SetIsChunkManifest()
	}

	if fixJpgOrientation {
		loweredName := strings.ToLower(fname)
		if mimeType == "image/jpeg" || strings.HasSuffix(loweredName, ".jpg") || strings.HasSuffix(loweredName, ".jpeg") {
			n.Data = images.FixJpgOrientation(n.Data)
		}
	}

	n.Checksum = NewCRC(n.Data)

	commaSep := strings.LastIndex(r.URL.Path, ",")
	dotSep := strings.LastIndex(r.URL.Path, ".")
	fid := r.URL.Path[commaSep+1:]
	if dotSep > 0 {
		fid = r.URL.Path[commaSep+1 : dotSep]
	}

	e = n.ParsePath(fid)

	return
}
func (n *Needle) ParsePath(fid string) (err error) {
	length := len(fid)
	if length <= CookieSize*2 {
		return fmt.Errorf("Invalid fid: %s", fid)
	}
	delta := ""
	deltaIndex := strings.LastIndex(fid, "_")
	if deltaIndex > 0 {
		fid, delta = fid[0:deltaIndex], fid[deltaIndex+1:]
	}
	n.Id, n.Cookie, err = ParseNeedleIdCookie(fid)
	if err != nil {
		return err
	}
	if delta != "" {
		if d, e := strconv.ParseUint(delta, 10, 64); e == nil {
			n.Id += NeedleId(d)
		} else {
			return e
		}
	}
	return err
}

func ParseNeedleIdCookie(key_hash_string string) (NeedleId, Cookie, error) {
	if len(key_hash_string) <= CookieSize*2 {
		return NeedleIdEmpty, 0, fmt.Errorf("KeyHash is too short.")
	}
	if len(key_hash_string) > (NeedleIdSize+CookieSize)*2 {
		return NeedleIdEmpty, 0, fmt.Errorf("KeyHash is too long.")
	}
	split := len(key_hash_string) - CookieSize*2
	needleId, err := ParseNeedleId(key_hash_string[:split])
	if err != nil {
		return NeedleIdEmpty, 0, fmt.Errorf("Parse needleId error: %v", err)
	}
	cookie, err := ParseCookie(key_hash_string[split:])
	if err != nil {
		return NeedleIdEmpty, 0, fmt.Errorf("Parse cookie error: %v", err)
	}
	return needleId, cookie, nil
}

func (n *Needle) LastModifiedString() string {
	return time.Unix(int64(n.LastModified), 0).Format("2006-01-02T15:04:05")
}
<|endoftext|>


<commit_before>package vcs

import (
	"fmt"
	"os"
	"path/filepath"
	"text/template"

	"github.com/hashicorp/errwrap"
)

var PostCheckoutTmpl = template.Must(template.New("name").Parse(`#!/bin/sh
# when checkout is a branch, start timer
if [ $3 -eq 1 ]; then
   sourceclock start;
fi
`))

var PrepCommitTmpl = template.Must(template.New("name").Parse(`#!/bin/sh

# only add time to template and message sources
# @see http://git-scm.com/docs/githooks#_prepare_commit_msg
case "$2" in
message,|template) echo normal
	printf "$(cat $1) [$(sourceclock split)]" > "$1" ;;
esac

`))

var PostCommitTmpl = template.Must(template.New("name").Parse(`#!/bin/sh
# @todo handle merge/rebase kind of commits

sourceclock lap
`))

type Git struct {
	dir string
}

func NewGit(dir string) *Git {
	return &Git{
		dir: filepath.Join(dir, ".git"),
	}
}

func (g *Git) Name() string { return "git" }
func (g *Git) Supported() bool {
	fi, err := os.Stat(g.dir)
	if err != nil || !fi.IsDir() {
		return false
	}

	return true
}

func (g *Git) Hook() error {
	hpath := filepath.Join(g.dir, "hooks")

	//post checkout: start()
	postchf, err := os.Create(filepath.Join(hpath, "post-checkout"))
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to create post-checkout '%s': {{err}}", postchf.Name()), err)
	}

	err = postchf.Chmod(0766)
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to make post-checkout file '%s' executable: {{err}}", hpath), err)
	}

	err = PostCheckoutTmpl.Execute(postchf, struct{}{})
	if err != nil {
		return errwrap.Wrapf("Failed to run post-checkout template: {{err}}", err)
	}

	//prepare commit msg: split()
	prepcof, err := os.Create(filepath.Join(hpath, "prepare-commit-msg"))
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to create prepare-commit-msg  '%s': {{err}}", postchf.Name()), err)
	}

	err = prepcof.Chmod(0766)
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to make prepare-commit-msg file '%s' executable: {{err}}", hpath), err)
	}

	err = PrepCommitTmpl.Execute(prepcof, struct{}{})
	if err != nil {
		return errwrap.Wrapf("Failed to run post-commit template: {{err}}", err)
	}

	//post commit: lap()
	postcof, err := os.Create(filepath.Join(hpath, "post-commit"))
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to create post-commit  '%s': {{err}}", postchf.Name()), err)
	}

	err = postcof.Chmod(0766)
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to make post-commit file '%s' executable: {{err}}", hpath), err)
	}

	err = PostCommitTmpl.Execute(postcof, struct{}{})
	if err != nil {
		return errwrap.Wrapf("Failed to run post-commit template: {{err}}", err)
	}

	return nil
}
<commit_msg>add to commit only when its source is a message or template [1m10s]<commit_after>package vcs

import (
	"fmt"
	"os"
	"path/filepath"
	"text/template"

	"github.com/hashicorp/errwrap"
)

var PostCheckoutTmpl = template.Must(template.New("name").Parse(`#!/bin/sh
# when checkout is a branch, start timer
if [ $3 -eq 1 ]; then
   sourceclock start;
fi
`))

var PrepCommitTmpl = template.Must(template.New("name").Parse(`#!/bin/sh

# only add time to template and message sources
# @see http://git-scm.com/docs/githooks#_prepare_commit_msg
case "$2" in
message|template) echo normal
	printf "$(cat $1) [$(sourceclock split)]" > "$1" ;;
esac

`))

var PostCommitTmpl = template.Must(template.New("name").Parse(`#!/bin/sh
# @todo handle merge/rebase kind of commits

sourceclock lap
`))

type Git struct {
	dir string
}

func NewGit(dir string) *Git {
	return &Git{
		dir: filepath.Join(dir, ".git"),
	}
}

func (g *Git) Name() string { return "git" }
func (g *Git) Supported() bool {
	fi, err := os.Stat(g.dir)
	if err != nil || !fi.IsDir() {
		return false
	}

	return true
}

func (g *Git) Hook() error {
	hpath := filepath.Join(g.dir, "hooks")

	//post checkout: start()
	postchf, err := os.Create(filepath.Join(hpath, "post-checkout"))
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to create post-checkout '%s': {{err}}", postchf.Name()), err)
	}

	err = postchf.Chmod(0766)
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to make post-checkout file '%s' executable: {{err}}", hpath), err)
	}

	err = PostCheckoutTmpl.Execute(postchf, struct{}{})
	if err != nil {
		return errwrap.Wrapf("Failed to run post-checkout template: {{err}}", err)
	}

	//prepare commit msg: split()
	prepcof, err := os.Create(filepath.Join(hpath, "prepare-commit-msg"))
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to create prepare-commit-msg  '%s': {{err}}", postchf.Name()), err)
	}

	err = prepcof.Chmod(0766)
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to make prepare-commit-msg file '%s' executable: {{err}}", hpath), err)
	}

	err = PrepCommitTmpl.Execute(prepcof, struct{}{})
	if err != nil {
		return errwrap.Wrapf("Failed to run post-commit template: {{err}}", err)
	}

	//post commit: lap()
	postcof, err := os.Create(filepath.Join(hpath, "post-commit"))
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to create post-commit  '%s': {{err}}", postchf.Name()), err)
	}

	err = postcof.Chmod(0766)
	if err != nil {
		return errwrap.Wrapf(fmt.Sprintf("Failed to make post-commit file '%s' executable: {{err}}", hpath), err)
	}

	err = PostCommitTmpl.Execute(postcof, struct{}{})
	if err != nil {
		return errwrap.Wrapf("Failed to run post-commit template: {{err}}", err)
	}

	return nil
}
<|endoftext|>


<commit_before>package wikidump

import (
	"bytes"
	"io/ioutil"
	"os"
	"sync"
	"testing"
)

func assertIntEq(t *testing.T, a, b int) {
	if a != b {
		t.Errorf("%d != %d", a, b)
	}
}

func TestGetPages(t *testing.T) {
	input, err := os.Open("nlwiki-20140927-sample.xml")
	if err != nil {
		panic(err)
	}
	pages, redirs := make(chan *Page), make(chan *Redirect)
	go GetPages(input, pages, redirs)

	var nredirs, npages int
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		for _ = range pages {
			npages++
		}
		wg.Done()
	}()
	go func() {
		for _ = range redirs {
			nredirs++
		}
		wg.Done()
	}()
	wg.Wait()

	assertIntEq(t, npages, 22)
	assertIntEq(t, nredirs, 1)
}

func BenchmarkGetPages(b *testing.B) {
	b.StopTimer()
	f, err := os.Open("nlwiki-20140927-sample.xml")
	if err != nil {
		panic(err)
	}
	content, err := ioutil.ReadAll(f)
	if err != nil {
		panic(err)
	}
	f.Close()

	for i := 0; i < b.N; i++ {
		r := bytes.NewBuffer(content)
		pages, redirs := make(chan *Page), make(chan *Redirect)

		b.StartTimer()
		go GetPages(r, pages, redirs)
		go func() {
			for _ = range pages {
			}
		}()
		for _ = range redirs {
		}
		b.StopTimer()
	}
}
<commit_msg>Improve GetPages test<commit_after>package wikidump

import (
	"bytes"
	"io/ioutil"
	"os"
	"sync"
	"testing"
)

func TestGetPages(t *testing.T) {
	input, err := os.Open("nlwiki-20140927-sample.xml")
	if err != nil {
		panic(err)
	}
	pages, redirs := make(chan *Page), make(chan *Redirect)
	go GetPages(input, pages, redirs)

	var titles []string
	var nredirs int
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		for p := range pages {
			titles = append(titles, p.Title)
		}
		wg.Done()
	}()
	go func() {
		for _ = range redirs {
			nredirs++
		}
		wg.Done()
	}()
	wg.Wait()

	if len(titles) != 22 {
		t.Errorf("expected 22 titles, got %d: %v", len(titles), titles)
	}
	if nredirs != 1 {
		t.Errorf("expected one redirect, got %d", nredirs)
	}
}

func BenchmarkGetPages(b *testing.B) {
	b.StopTimer()
	f, err := os.Open("nlwiki-20140927-sample.xml")
	if err != nil {
		panic(err)
	}
	content, err := ioutil.ReadAll(f)
	if err != nil {
		panic(err)
	}
	f.Close()

	for i := 0; i < b.N; i++ {
		r := bytes.NewBuffer(content)
		pages, redirs := make(chan *Page), make(chan *Redirect)

		b.StartTimer()
		go GetPages(r, pages, redirs)
		go func() {
			for _ = range pages {
			}
		}()
		for _ = range redirs {
		}
		b.StopTimer()
	}
}
<|endoftext|>


<commit_before>package netfilter

import (
	"github.com/42wim/registrator-work/bridge"
	"net/url"
	"strconv"
	"strings"
)

func init() {
	bridge.Register(new(Factory), "netfilter")
}

type Factory struct{}

func (f *Factory) New(uri *url.URL) bridge.RegistryAdapter {
	chain := uri.Host
	set := strings.Replace(uri.Path, "/", "", -1)
	FirewalldInit()
	if firewalldRunning {
		OnReloaded(func() { iptablesInit(chain, set) })
	}
	ipsetInit(set)
	iptablesInit(chain, set)
	return &NetfilterAdapter{Chain: chain, Set: set}
}

type NetfilterAdapter struct {
	Chain string
	Set   string
}

func (r *NetfilterAdapter) Ping() error {
	return nil
}

func (r *NetfilterAdapter) Register(service *bridge.Service) error {
	if strings.Contains(service.IP, ":") {
		return ipsetHost("add", r.Set, service.IP, service.Origin.PortType, strconv.Itoa(service.Port))
	} else {
		return nil
	}
}

func (r *NetfilterAdapter) Deregister(service *bridge.Service) error {
	if strings.Contains(service.IP, ":") {
		return ipsetHost("del", r.Set, service.IP, service.Origin.PortType, strconv.Itoa(service.Port))
	} else {
		return nil
	}
}

func (r *NetfilterAdapter) Refresh(service *bridge.Service) error {
	return nil
}
<commit_msg>Sync with https://github.com/gliderlabs/registrator/pull/234<commit_after>package netfilter

import (
	"github.com/42wim/registrator-work/bridge"
	"net/url"
	"strconv"
	"strings"
)

func init() {
	bridge.Register(new(Factory), "netfilter")
}

type Factory struct{}

func (f *Factory) New(uri *url.URL) bridge.RegistryAdapter {
	var chain, set string
	if uri.Host != "" {
		chain = uri.Host
		set = strings.Replace(uri.Path, "/", "", -1)
	} else {
		chain = "FORWARD_direct"
		set = "containerports"
	}
	FirewalldInit()
	if firewalldRunning {
		OnReloaded(func() { iptablesInit(chain, set) })
	}
	ipsetInit(set)
	iptablesInit(chain, set)
	return &NetfilterAdapter{Chain: chain, Set: set}
}

type NetfilterAdapter struct {
	Chain string
	Set   string
}

func (r *NetfilterAdapter) Ping() error {
	return nil
}

func (r *NetfilterAdapter) Register(service *bridge.Service) error {
	if strings.Contains(service.IP, ":") {
		return ipsetHost("add", r.Set, service.IP, service.Origin.PortType, strconv.Itoa(service.Port))
	} else {
		return nil
	}
}

func (r *NetfilterAdapter) Deregister(service *bridge.Service) error {
	if strings.Contains(service.IP, ":") {
		return ipsetHost("del", r.Set, service.IP, service.Origin.PortType, strconv.Itoa(service.Port))
	} else {
		return nil
	}
}

func (r *NetfilterAdapter) Refresh(service *bridge.Service) error {
	return nil
}
<|endoftext|>


<commit_before>package antibody

import (
	"fmt"
	"sync"
)

// Antibody wraps a list of bundles to be processed.
type Antibody struct {
	bundles []Bundle
}
type bundleFn func(bundle Bundle)

// NewAntibody creates an instance of antibody with the given bundles.
func NewAntibody(bundles []Bundle) Antibody {
	return Antibody{
		bundles: bundles,
	}
}

func (a Antibody) forEach(fn bundleFn) {
	var wg sync.WaitGroup
	for _, bundle := range a.bundles {
		wg.Add(1)
		go func(bundle Bundle, fn bundleFn, wg *sync.WaitGroup) {
			fn(bundle)
			for _, sourceable := range bundle.Sourceables() {
				fmt.Println(sourceable)
			}
			wg.Done()
		}(bundle, fn, &wg)
	}
	wg.Wait()
}

// Download the needed bundles.
func (a Antibody) Download() {
	a.forEach(func(b Bundle) {
		b.Download()
	})
}

// Update all bundles.
func (a Antibody) Update() {
	a.forEach(func(b Bundle) {
		b.Update()
	})
}
<commit_msg>wg do not need to be parameterized<commit_after>package antibody

import (
	"fmt"
	"sync"
)

// Antibody wraps a list of bundles to be processed.
type Antibody struct {
	bundles []Bundle
}
type bundleFn func(bundle Bundle)

// NewAntibody creates an instance of antibody with the given bundles.
func NewAntibody(bundles []Bundle) Antibody {
	return Antibody{
		bundles: bundles,
	}
}

func (a Antibody) forEach(fn bundleFn) {
	var wg sync.WaitGroup
	for _, bundle := range a.bundles {
		wg.Add(1)
		go func(bundle Bundle, fn bundleFn) {
			fn(bundle)
			for _, sourceable := range bundle.Sourceables() {
				fmt.Println(sourceable)
			}
			wg.Done()
		}(bundle, fn)
	}
	wg.Wait()
}

// Download the needed bundles.
func (a Antibody) Download() {
	a.forEach(func(b Bundle) {
		b.Download()
	})
}

// Update all bundles.
func (a Antibody) Update() {
	a.forEach(func(b Bundle) {
		b.Update()
	})
}
<|endoftext|>