<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"sync/atomic"

	"github.com/artyom/autoflags"
	"github.com/ttacon/chalk"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/net/proxy"
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] ip1 ip2 ...\n", os.Args[0])
		flag.PrintDefaults()
	}
}

type Config struct {
	Concurrency int    `flag:"n,concurrent ssh sessions"`
	Command     string `flag:"cmd,command to run"`
	Login       string `flag:"l,login to use"`
	Port        int    `flag:"p,default port"`
	StdinFile   string `flag:"stdin,REGULAR (no piping!) file to pass to stdin of remote command"`
	DumpFiles   bool   `flag:"logs,save stdout/stderr to separate per-host logs"`
	StdoutFmt   string `flag:"logs.stdout,format of stdout per-host log name"`
	StderrFmt   string `flag:"logs.stderr,format of stderr per-host log name"`

	stdoutPrefix, stderrPrefix string
}

func main() {
	log.SetFlags(0)
	conf := Config{
		Concurrency: 100,
		Login:       "root",
		Port:        22,
		StdoutFmt:   "/tmp/${" + remoteHostVarname + "}.stdout",
		StderrFmt:   "/tmp/${" + remoteHostVarname + "}.stderr",

		stdoutPrefix: ".",
		stderrPrefix: "E",
	}
	autoflags.Define(&conf)
	flag.Parse()
	if conf.Concurrency < 1 {
		conf.Concurrency = 1
	}
	hosts := flag.Args()
	if len(hosts) == 0 || len(conf.Command) == 0 {
		flag.Usage()
		return
	}
	if conf.DumpFiles {
		if conf.StdoutFmt == conf.StderrFmt {
			log.Fatal("file format for stdout and stderr should differ")
		}
		if err := checkFilenameTemplate(conf.StdoutFmt); err != nil {
			log.Fatal("stdout filename format:", err)
		}
		if err := checkFilenameTemplate(conf.StderrFmt); err != nil {
			log.Fatal("stderr filename format:", err)
		}
	}
	if isTerminal(os.Stdout) {
		conf.stdoutPrefix = chalk.Green.Color(conf.stdoutPrefix)
	}
	if isTerminal(os.Stderr) {
		conf.stderrPrefix = chalk.Red.Color(conf.stderrPrefix)
	}
	var sshAgent agent.Agent
	agentConn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		log.Fatal(err)
	}
	sshAgent = agent.NewClient(agentConn)
	defer agentConn.Close()

	signers, err := sshAgent.Signers()
	if err != nil {
		log.Fatal(err)
	}

	authMethods := []ssh.AuthMethod{ssh.PublicKeys(signers...)}

	var wg sync.WaitGroup

	limit := make(chan struct{}, conf.Concurrency)

	uniqueHosts := make(map[string]struct{})
	for _, host := range hosts {
		uniqueHosts[host] = struct{}{}
	}

	var errCnt int32
	for host := range uniqueHosts {
		limit <- struct{}{}
		wg.Add(1)
		go func(host string) {
			defer wg.Done()
			defer func() { <-limit }()
			switch err := RemoteCommand(host, conf, authMethods); {
			case err == nil && conf.DumpFiles:
				fmt.Println(host, "processed")
			case err == nil:
			default:
				atomic.AddInt32(&errCnt, 1)
				log.Println(host, err)
			}
		}(host)
	}
	wg.Wait()
	if errCnt > 0 {
		os.Exit(123)
	}
}

func RemoteCommand(addr string, conf Config, authMethods []ssh.AuthMethod) error {
	login, addr := loginAndAddr(conf.Login, addr, conf.Port)
	sshConfig := &ssh.ClientConfig{
		User: login,
		Auth: authMethods,
	}
	client, err := sshDial("tcp", addr, sshConfig)
	if err != nil {
		return err
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	host := addr
	if h, _, err := net.SplitHostPort(addr); err == nil {
		host = h
	}

	if conf.StdinFile != "" {
		f, err := os.Open(conf.StdinFile)
		if err != nil {
			return err
		}
		defer f.Close()
		st, err := f.Stat()
		if err != nil {
			return err
		}
		if !st.Mode().IsRegular() {
			return fmt.Errorf("file passed to stdin is not a regular file")
		}
		session.Stdin = f
	}

	switch {
	case conf.DumpFiles:
		stdoutLog, err := os.Create(filepath.Clean(expandHostname(conf.StdoutFmt, host)))
		if err != nil {
			return err
		}
		defer closeAndRemoveIfAt0(stdoutLog)
		session.Stdout = stdoutLog
		stderrLog, err := os.Create(filepath.Clean(expandHostname(conf.StderrFmt, host)))
		if err != nil {
			return err
		}
		defer closeAndRemoveIfAt0(stderrLog)
		session.Stderr = stderrLog
	default:
		stdoutPipe, err := session.StdoutPipe()
		if err != nil {
			return err
		}
		stderrPipe, err := session.StderrPipe()
		if err != nil {
			return err
		}

		go byLineCopy(fmt.Sprintf("%s %s\t", conf.stdoutPrefix, host), os.Stdout, stdoutPipe)
		go byLineCopy(fmt.Sprintf("%s %s\t", conf.stderrPrefix, host), os.Stderr, stderrPipe)
	}

	return session.Run(conf.Command)
}

func byLineCopy(prefix string, sink io.Writer, pipe io.Reader) {
	buf := []byte(prefix)
	scanner := bufio.NewScanner(pipe)
	for scanner.Scan() {
		buf := buf[:len(prefix)]
		buf = append(buf, scanner.Bytes()...)
		buf = append(buf, '\n')
		// this is safe to write a single line to stdout/stderr without
		// additional locking from multiple goroutines as os guarantees
		// those writes are atomic (for stdout/stderr only)
		sink.Write(buf)
	}
	if err := scanner.Err(); err != nil {
		log.Print("string scanner error", err)
	}
}

func closeAndRemoveIfAt0(f *os.File) {
	if n, err := f.Seek(0, os.SEEK_CUR); err == nil && n == 0 {
		os.Remove(f.Name())
	}
	f.Close()
}

func isTerminal(f *os.File) bool {
	st, err := f.Stat()
	if err != nil {
		return false
	}
	return st.Mode()&os.ModeDevice != 0
}

func loginAndAddr(defaultLogin, addr string, defaultPort int) (login, hostPort string) {
	login = defaultLogin
	u, err := url.Parse("//" + addr) // add slashes so that it properly parsed as url
	if err != nil {
		return defaultLogin, addr
	}
	if u.User != nil {
		login = u.User.Username()
	}
	if _, _, err := net.SplitHostPort(u.Host); err == nil {
		return login, u.Host
	}
	return login, net.JoinHostPort(u.Host, strconv.Itoa(defaultPort))
}

func sshDial(network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
	conn, err := proxyDialer.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
	if err != nil {
		return nil, err
	}
	return ssh.NewClient(c, chans, reqs), nil
}

var proxyDialer = proxy.FromEnvironment()

const remoteHostVarname = `REX_REMOTE_HOST`

func expandHostname(s, hostname string) string {
	return os.Expand(s, func(x string) string {
		if x == remoteHostVarname {
			return hostname
		}
		return os.Getenv(s)
	})
}

func checkFilenameTemplate(s string) error {
	varSet := os.Expand(s, func(x string) string {
		if x == remoteHostVarname {
			return "value"
		}
		return ""
	})
	varUnset := os.Expand(s, func(string) string { return "" })
	if varSet == varUnset {
		return fmt.Errorf("no ${%s} in pattern", remoteHostVarname)
	}
	return nil
}
<commit_msg>Paint errors red, stderr marker yellow<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"sync/atomic"

	"github.com/artyom/autoflags"
	"github.com/ttacon/chalk"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/net/proxy"
)

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] ip1 ip2 ...\n", os.Args[0])
		flag.PrintDefaults()
	}
}

type Config struct {
	Concurrency int    `flag:"n,concurrent ssh sessions"`
	Command     string `flag:"cmd,command to run"`
	Login       string `flag:"l,login to use"`
	Port        int    `flag:"p,default port"`
	StdinFile   string `flag:"stdin,REGULAR (no piping!) file to pass to stdin of remote command"`
	DumpFiles   bool   `flag:"logs,save stdout/stderr to separate per-host logs"`
	StdoutFmt   string `flag:"logs.stdout,format of stdout per-host log name"`
	StderrFmt   string `flag:"logs.stderr,format of stderr per-host log name"`

	stdoutPrefix, stderrPrefix string
}

func main() {
	log.SetFlags(0)
	conf := Config{
		Concurrency: 100,
		Login:       "root",
		Port:        22,
		StdoutFmt:   "/tmp/${" + remoteHostVarname + "}.stdout",
		StderrFmt:   "/tmp/${" + remoteHostVarname + "}.stderr",

		stdoutPrefix: ".",
		stderrPrefix: "E",
	}
	autoflags.Define(&conf)
	flag.Parse()
	if conf.Concurrency < 1 {
		conf.Concurrency = 1
	}
	hosts := flag.Args()
	if len(hosts) == 0 || len(conf.Command) == 0 {
		flag.Usage()
		return
	}
	if conf.DumpFiles {
		if conf.StdoutFmt == conf.StderrFmt {
			log.Fatal("file format for stdout and stderr should differ")
		}
		if err := checkFilenameTemplate(conf.StdoutFmt); err != nil {
			log.Fatal("stdout filename format:", err)
		}
		if err := checkFilenameTemplate(conf.StderrFmt); err != nil {
			log.Fatal("stderr filename format:", err)
		}
	}
	if isTerminal(os.Stdout) {
		conf.stdoutPrefix = chalk.Green.Color(conf.stdoutPrefix)
	}
	var colorErrors bool
	if isTerminal(os.Stderr) {
		colorErrors = true
		conf.stderrPrefix = chalk.Yellow.Color(conf.stderrPrefix)
	}
	var sshAgent agent.Agent
	agentConn, err := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
	if err != nil {
		log.Fatal(err)
	}
	sshAgent = agent.NewClient(agentConn)
	defer agentConn.Close()

	signers, err := sshAgent.Signers()
	if err != nil {
		log.Fatal(err)
	}

	authMethods := []ssh.AuthMethod{ssh.PublicKeys(signers...)}

	var wg sync.WaitGroup

	limit := make(chan struct{}, conf.Concurrency)

	uniqueHosts := make(map[string]struct{})
	for _, host := range hosts {
		uniqueHosts[host] = struct{}{}
	}

	var errCnt int32
	for host := range uniqueHosts {
		limit <- struct{}{}
		wg.Add(1)
		go func(host string) {
			defer wg.Done()
			defer func() { <-limit }()
			switch err := RemoteCommand(host, conf, authMethods); {
			case err == nil && conf.DumpFiles:
				fmt.Println(host, "processed")
			case err == nil:
			default:
				atomic.AddInt32(&errCnt, 1)
				if colorErrors {
					host = chalk.Red.Color(host)
				}
				log.Println(host, err)
			}
		}(host)
	}
	wg.Wait()
	if errCnt > 0 {
		os.Exit(123)
	}
}

func RemoteCommand(addr string, conf Config, authMethods []ssh.AuthMethod) error {
	login, addr := loginAndAddr(conf.Login, addr, conf.Port)
	sshConfig := &ssh.ClientConfig{
		User: login,
		Auth: authMethods,
	}
	client, err := sshDial("tcp", addr, sshConfig)
	if err != nil {
		return err
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	host := addr
	if h, _, err := net.SplitHostPort(addr); err == nil {
		host = h
	}

	if conf.StdinFile != "" {
		f, err := os.Open(conf.StdinFile)
		if err != nil {
			return err
		}
		defer f.Close()
		st, err := f.Stat()
		if err != nil {
			return err
		}
		if !st.Mode().IsRegular() {
			return fmt.Errorf("file passed to stdin is not a regular file")
		}
		session.Stdin = f
	}

	switch {
	case conf.DumpFiles:
		stdoutLog, err := os.Create(filepath.Clean(expandHostname(conf.StdoutFmt, host)))
		if err != nil {
			return err
		}
		defer closeAndRemoveIfAt0(stdoutLog)
		session.Stdout = stdoutLog
		stderrLog, err := os.Create(filepath.Clean(expandHostname(conf.StderrFmt, host)))
		if err != nil {
			return err
		}
		defer closeAndRemoveIfAt0(stderrLog)
		session.Stderr = stderrLog
	default:
		stdoutPipe, err := session.StdoutPipe()
		if err != nil {
			return err
		}
		stderrPipe, err := session.StderrPipe()
		if err != nil {
			return err
		}

		go byLineCopy(fmt.Sprintf("%s %s\t", conf.stdoutPrefix, host), os.Stdout, stdoutPipe)
		go byLineCopy(fmt.Sprintf("%s %s\t", conf.stderrPrefix, host), os.Stderr, stderrPipe)
	}

	return session.Run(conf.Command)
}

func byLineCopy(prefix string, sink io.Writer, pipe io.Reader) {
	buf := []byte(prefix)
	scanner := bufio.NewScanner(pipe)
	for scanner.Scan() {
		buf := buf[:len(prefix)]
		buf = append(buf, scanner.Bytes()...)
		buf = append(buf, '\n')
		// this is safe to write a single line to stdout/stderr without
		// additional locking from multiple goroutines as os guarantees
		// those writes are atomic (for stdout/stderr only)
		sink.Write(buf)
	}
	if err := scanner.Err(); err != nil {
		log.Print("string scanner error", err)
	}
}

func closeAndRemoveIfAt0(f *os.File) {
	if n, err := f.Seek(0, os.SEEK_CUR); err == nil && n == 0 {
		os.Remove(f.Name())
	}
	f.Close()
}

func isTerminal(f *os.File) bool {
	st, err := f.Stat()
	if err != nil {
		return false
	}
	return st.Mode()&os.ModeDevice != 0
}

func loginAndAddr(defaultLogin, addr string, defaultPort int) (login, hostPort string) {
	login = defaultLogin
	u, err := url.Parse("//" + addr) // add slashes so that it properly parsed as url
	if err != nil {
		return defaultLogin, addr
	}
	if u.User != nil {
		login = u.User.Username()
	}
	if _, _, err := net.SplitHostPort(u.Host); err == nil {
		return login, u.Host
	}
	return login, net.JoinHostPort(u.Host, strconv.Itoa(defaultPort))
}

func sshDial(network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
	conn, err := proxyDialer.Dial(network, addr)
	if err != nil {
		return nil, err
	}
	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
	if err != nil {
		return nil, err
	}
	return ssh.NewClient(c, chans, reqs), nil
}

var proxyDialer = proxy.FromEnvironment()

const remoteHostVarname = `REX_REMOTE_HOST`

func expandHostname(s, hostname string) string {
	return os.Expand(s, func(x string) string {
		if x == remoteHostVarname {
			return hostname
		}
		return os.Getenv(s)
	})
}

func checkFilenameTemplate(s string) error {
	varSet := os.Expand(s, func(x string) string {
		if x == remoteHostVarname {
			return "value"
		}
		return ""
	})
	varUnset := os.Expand(s, func(string) string { return "" })
	if varSet == varUnset {
		return fmt.Errorf("no ${%s} in pattern", remoteHostVarname)
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package http

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	auth "github.com/abbot/go-http-auth"
	"github.com/gorilla/websocket"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
)

const (
	maxMessages    = 1024
	maxMessageSize = 0
	writeWait      = 10 * time.Second
)

// WSMessage is the interface of a message to send over the wire
type WSMessage interface {
	Bytes() []byte
}

// WSRawMessage represents a raw message (array of bytes)
type WSRawMessage []byte

// Bytes returns the string representation of the raw message
func (m WSRawMessage) Bytes() []byte {
	return m
}

// WSSpeaker is the interface for a websocket speaking client. It is used for outgoing
// or incoming connections.
type WSSpeaker interface {
	GetStatus() WSConnStatus
	GetHost() string
	GetAddrPort() (string, int)
	GetServiceType() common.ServiceType
	GetHeaders() http.Header
	GetURL() *url.URL
	IsConnected() bool
	SendMessage(m WSMessage) error
	Connect()
	Disconnect()
	AddEventHandler(WSSpeakerEventHandler)
}

// WSConnState describes the connection state
type WSConnState int32

// WSConnStatus describes the status of a WebSocket connection
type WSConnStatus struct {
	ServiceType common.ServiceType
	Addr        string
	Port        int
	Host        string       `json:"-"`
	State       *WSConnState `json:"IsConnected"`
	Url         *url.URL     `json:"-"`
	headers     http.Header
}

func (s *WSConnState) MarshalJSON() ([]byte, error) {
	switch *s {
	case common.RunningState:
		return []byte("true"), nil
	case common.StoppedState:
		return []byte("false"), nil
	}
	return nil, fmt.Errorf("Invalid state: %d", s)
}

// UnmarshalJSON deserialize a connection state
func (s *WSConnState) UnmarshalJSON(b []byte) error {
	var state bool
	if err := json.Unmarshal(b, &state); err != nil {
		return err
	}

	if state {
		*s = common.RunningState
	} else {
		*s = common.StoppedState
	}

	return nil
}

// WSConn is the connection object of a WSSpeaker
type WSConn struct {
	sync.RWMutex
	WSConnStatus
	send          chan []byte
	read          chan []byte
	quit          chan bool
	wg            sync.WaitGroup
	conn          *websocket.Conn
	running       atomic.Value
	pingTicker    *time.Ticker // only used by incoming connections
	eventHandlers []WSSpeakerEventHandler
	wsSpeaker     WSSpeaker // speaker owning the connection
}

// wsIncomingClient is only used internally to handle incoming client. It embeds a WSConn.
type wsIncomingClient struct {
	*WSConn
}

// WSClient is a outgoint client meaning a client connected to a remote websocket server.
// It embeds a WSConn.
type WSClient struct {
	*WSConn
	Path       string
	AuthClient *AuthenticationClient
}

// WSSpeakerEventHandler is the interface to be implement by the client events listeners.
type WSSpeakerEventHandler interface {
	OnMessage(c WSSpeaker, m WSMessage)
	OnConnected(c WSSpeaker)
	OnDisconnected(c WSSpeaker)
}

// DefaultWSSpeakerEventHandler implements stubs for the wsIncomingClientEventHandler interface
type DefaultWSSpeakerEventHandler struct {
}

// OnMessage is called when a message is received.
func (d *DefaultWSSpeakerEventHandler) OnMessage(c WSSpeaker, m WSMessage) {
}

// OnConnected is called when the connection is established.
func (d *DefaultWSSpeakerEventHandler) OnConnected(c WSSpeaker) {
}

// OnDisconnected is called when the connection is closed or lost.
func (d *DefaultWSSpeakerEventHandler) OnDisconnected(c WSSpeaker) {
}

// GetHost returns the hostname/host-id of the connection.
func (c *WSConn) GetHost() string {
	return c.Host
}

// GetAddrPort returns the address and the port of the remote end.
func (c *WSConn) GetAddrPort() (string, int) {
	return c.Addr, c.Port
}

// GetURL returns the URL of the connection
func (c *WSConn) GetURL() *url.URL {
	return c.Url
}

// IsConnected returns the connection status.
func (c *WSConn) IsConnected() bool {
	return atomic.LoadInt32((*int32)(c.State)) == common.RunningState
}

// GetStatus returns the status of a WebSocket connection
func (c *WSConn) GetStatus() WSConnStatus {
	c.RLock()
	defer c.RUnlock()

	status := c.WSConnStatus
	status.State = new(WSConnState)
	*status.State = WSConnState(atomic.LoadInt32((*int32)(c.State)))
	return c.WSConnStatus
}

// SendMessage adds a message to sending queue.
func (c *WSConn) SendMessage(m WSMessage) error {
	if !c.IsConnected() {
		return errors.New("Not connected")
	}

	c.send <- m.Bytes()

	return nil
}

// SendRaw adds raw bytes to sending queue.
func (c *WSConn) SendRaw(b []byte) error {
	if !c.IsConnected() {
		return errors.New("Not connected")
	}

	c.send <- b

	return nil
}

// GetServiceType returns the client type.
func (c *WSConn) GetServiceType() common.ServiceType {
	return c.ServiceType
}

// GetHeaders returns the client HTTP headers.
func (c *WSConn) GetHeaders() http.Header {
	return c.headers
}

// SendMessage sends a message directly over the wire.
func (c *WSConn) write(msg []byte) error {
	if !c.IsConnected() {
		return errors.New("Not connected")
	}

	c.conn.SetWriteDeadline(time.Now().Add(writeWait))

	w, err := c.conn.NextWriter(websocket.TextMessage)
	if err != nil {
		return err
	}

	if _, err = w.Write(msg); err != nil {
		return err
	}

	return w.Close()
}

func (c *WSConn) start() {
	c.wg.Add(1)
	go c.run()
}

// main loop to read and send messages
func (c *WSConn) run() {
	defer c.wg.Done()

	go func() {
		for c.running.Load() == true {
			_, m, err := c.conn.ReadMessage()
			if err != nil {
				if c.running.Load() != false {
					c.quit <- true
				}
				break
			}

			c.read <- m
		}
	}()

	defer func() {
		atomic.StoreInt32((*int32)(c.State), common.StoppedState)
		c.RLock()
		for _, l := range c.eventHandlers {
			l.OnDisconnected(c.wsSpeaker)
		}
		c.RUnlock()
	}()

	defer c.pingTicker.Stop()

	for {
		select {
		case <-c.quit:
			return
		case m := <-c.send:
			err := c.write(m)
			if err != nil {
				logging.GetLogger().Errorf("Error while writing to the WebSocket: %s", err.Error())
			}
		case m := <-c.read:
			c.RLock()
			for _, l := range c.eventHandlers {
				l.OnMessage(c.wsSpeaker, WSRawMessage(m))
			}
			c.RUnlock()
		case <-c.pingTicker.C:
			c.sendPing()
		}
	}
}

// sendPing is used for remote connections by the server to send PingMessage
// to remote client.
func (c *WSConn) sendPing() {
	c.conn.SetWriteDeadline(time.Now().Add(writeWait))
	if err := c.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
		logging.GetLogger().Warningf("Error while sending ping to the websocket: %s", err.Error())
	}
}

// AddEventHandler registers a new event handler
func (c *WSConn) AddEventHandler(h WSSpeakerEventHandler) {
	c.Lock()
	c.eventHandlers = append(c.eventHandlers, h)
	c.Unlock()
}

// Connect default implementation doing nothing as for incoming connection it is not used.
func (c *WSConn) Connect() {
}

// Disconnect the WSSpeakers without waiting for termination.
func (c *WSConn) Disconnect() {
	c.running.Store(false)
	if atomic.LoadInt32((*int32)(c.State)) == common.RunningState {
		c.quit <- true
		c.conn.Close()
		close(c.send)
		close(c.read)
	}
}

func newWSConn(host string, clientType common.ServiceType, url *url.URL, headers http.Header) *WSConn {
	if headers == nil {
		headers = http.Header{}
	}

	port, _ := strconv.Atoi(url.Port())
	c := &WSConn{
		WSConnStatus: WSConnStatus{
			Host:        host,
			ServiceType: clientType,
			Addr:        url.Hostname(),
			Port:        port,
			State:       new(WSConnState),
			Url:         url,
			headers:     headers,
		},
		send:       make(chan []byte, maxMessages),
		read:       make(chan []byte, maxMessages),
		quit:       make(chan bool, 2),
		pingTicker: &time.Ticker{},
	}
	*c.State = common.StoppedState
	c.running.Store(true)
	return c
}

func (c *WSClient) scheme() string {
	if config.IsTLSenabled() == true {
		return "wss://"
	}
	return "ws://"
}

func (c *WSClient) connect() {
	var err error
	endpoint := c.Url.String()
	headers := http.Header{
		"X-Host-ID":             {c.Host},
		"Origin":                {endpoint},
		"X-Client-Type":         {c.ServiceType.String()},
		"X-Websocket-Namespace": {WilcardNamespace},
	}

	if c.AuthClient != nil {
		if err = c.AuthClient.Authenticate(); err != nil {
			logging.GetLogger().Errorf("Unable to authenticate %s : %s", endpoint, err.Error())
			return
		}
		c.AuthClient.SetHeaders(headers)
	}

	d := websocket.Dialer{
		Proxy:           http.ProxyFromEnvironment,
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	certPEM := config.GetConfig().GetString("agent.X509_cert")
	keyPEM := config.GetConfig().GetString("agent.X509_key")
	if certPEM != "" && keyPEM != "" {
		d.TLSClientConfig = common.SetupTLSClientConfig(certPEM, keyPEM)
		checkTLSConfig(d.TLSClientConfig)
	}
	c.conn, _, err = d.Dial(endpoint, headers)

	if err != nil {
		logging.GetLogger().Errorf("Unable to create a WebSocket connection %s : %s", endpoint, err.Error())
		return
	}
	defer c.conn.Close()
	c.conn.SetPingHandler(nil)

	atomic.StoreInt32((*int32)(c.State), common.RunningState)
	defer atomic.StoreInt32((*int32)(c.State), common.StoppedState)

	logging.GetLogger().Infof("Connected to %s", endpoint)

	// notify connected
	c.RLock()
	for _, l := range c.eventHandlers {
		l.OnConnected(c)
	}
	c.RUnlock()

	c.wg.Add(1)
	c.run()
}

// Connect to the server - and reconnect if necessary
func (c *WSClient) Connect() {
	go func() {
		for c.running.Load() == true {
			c.connect()
			time.Sleep(1 * time.Second)
		}
	}()
}

// NewWSClient returns a WSClient with a new connection.
func NewWSClient(host string, clientType common.ServiceType, url *url.URL, authClient *AuthenticationClient, headers http.Header) *WSClient {
	wsconn := newWSConn(host, clientType, url, headers)
	c := &WSClient{
		WSConn:     wsconn,
		AuthClient: authClient,
	}
	wsconn.wsSpeaker = c
	return c
}

// NewWSClientFromConfig creates a WSClient based on the configuration
func NewWSClientFromConfig(clientType common.ServiceType, url *url.URL, authClient *AuthenticationClient, headers http.Header) *WSClient {
	host := config.GetConfig().GetString("host_id")
	return NewWSClient(host, clientType, url, authClient, headers)
}

// newIncomingWSClient is called by the server for incoming connections
func newIncomingWSClient(conn *websocket.Conn, r *auth.AuthenticatedRequest) *wsIncomingClient {
	host := getRequestParameter(r, "X-Host-ID")
	if host == "" {
		host = r.RemoteAddr
	}

	clientType := common.ServiceType(getRequestParameter(r, "X-Client-Type"))
	if clientType == "" {
		clientType = common.UnknownService
	}

	svc, _ := common.ServiceAddressFromString(conn.RemoteAddr().String())
	url := config.GetURL("http", svc.Addr, svc.Port, r.URL.Path+"?"+r.URL.RawQuery)
	wsconn := newWSConn(host, clientType, url, r.Header)
	wsconn.conn = conn

	pongTimeout := time.Duration(config.GetConfig().GetInt("ws_pong_timeout")) * time.Second

	conn.SetReadLimit(maxMessageSize)
	conn.SetReadDeadline(time.Now().Add(100 * time.Second))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(pongTimeout))
		return nil
	})

	c := &wsIncomingClient{
		WSConn: wsconn,
	}
	wsconn.wsSpeaker = c

	atomic.StoreInt32((*int32)(c.State), common.RunningState)

	// send a first ping to help firefox and some other client which wait for a
	// first ping before doing something
	c.sendPing()

	wsconn.pingTicker = time.NewTicker(pongTimeout * 8 / 10)

	return c
}
<commit_msg>websocket: move read/write to their own goroutine to avoid ping timeout<commit_after>/*
 * Copyright (C) 2015 Red Hat, Inc.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *  http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 *
 */

package http

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	auth "github.com/abbot/go-http-auth"
	"github.com/gorilla/websocket"

	"github.com/skydive-project/skydive/common"
	"github.com/skydive-project/skydive/config"
	"github.com/skydive-project/skydive/logging"
)

const (
	maxMessages    = 1024
	maxMessageSize = 0
	writeWait      = 10 * time.Second
)

// WSMessage is the interface of a message to send over the wire
type WSMessage interface {
	Bytes() []byte
}

// WSRawMessage represents a raw message (array of bytes)
type WSRawMessage []byte

// Bytes returns the string representation of the raw message
func (m WSRawMessage) Bytes() []byte {
	return m
}

// WSSpeaker is the interface for a websocket speaking client. It is used for outgoing
// or incoming connections.
type WSSpeaker interface {
	GetStatus() WSConnStatus
	GetHost() string
	GetAddrPort() (string, int)
	GetServiceType() common.ServiceType
	GetHeaders() http.Header
	GetURL() *url.URL
	IsConnected() bool
	SendMessage(m WSMessage) error
	Connect()
	Disconnect()
	AddEventHandler(WSSpeakerEventHandler)
}

// WSConnState describes the connection state
type WSConnState int32

// WSConnStatus describes the status of a WebSocket connection
type WSConnStatus struct {
	ServiceType common.ServiceType
	Addr        string
	Port        int
	Host        string       `json:"-"`
	State       *WSConnState `json:"IsConnected"`
	Url         *url.URL     `json:"-"`
	headers     http.Header
}

func (s *WSConnState) MarshalJSON() ([]byte, error) {
	switch *s {
	case common.RunningState:
		return []byte("true"), nil
	case common.StoppedState:
		return []byte("false"), nil
	}
	return nil, fmt.Errorf("Invalid state: %d", s)
}

// UnmarshalJSON deserialize a connection state
func (s *WSConnState) UnmarshalJSON(b []byte) error {
	var state bool
	if err := json.Unmarshal(b, &state); err != nil {
		return err
	}

	if state {
		*s = common.RunningState
	} else {
		*s = common.StoppedState
	}

	return nil
}

// WSConn is the connection object of a WSSpeaker
type WSConn struct {
	sync.RWMutex
	WSConnStatus
	send          chan []byte
	read          chan []byte
	quit          chan bool
	wg            sync.WaitGroup
	conn          *websocket.Conn
	running       atomic.Value
	pingTicker    *time.Ticker // only used by incoming connections
	eventHandlers []WSSpeakerEventHandler
	wsSpeaker     WSSpeaker // speaker owning the connection
}

// wsIncomingClient is only used internally to handle incoming client. It embeds a WSConn.
type wsIncomingClient struct {
	*WSConn
}

// WSClient is a outgoint client meaning a client connected to a remote websocket server.
// It embeds a WSConn.
type WSClient struct {
	*WSConn
	Path       string
	AuthClient *AuthenticationClient
}

// WSSpeakerEventHandler is the interface to be implement by the client events listeners.
type WSSpeakerEventHandler interface {
	OnMessage(c WSSpeaker, m WSMessage)
	OnConnected(c WSSpeaker)
	OnDisconnected(c WSSpeaker)
}

// DefaultWSSpeakerEventHandler implements stubs for the wsIncomingClientEventHandler interface
type DefaultWSSpeakerEventHandler struct {
}

// OnMessage is called when a message is received.
func (d *DefaultWSSpeakerEventHandler) OnMessage(c WSSpeaker, m WSMessage) {
}

// OnConnected is called when the connection is established.
func (d *DefaultWSSpeakerEventHandler) OnConnected(c WSSpeaker) {
}

// OnDisconnected is called when the connection is closed or lost.
func (d *DefaultWSSpeakerEventHandler) OnDisconnected(c WSSpeaker) {
}

// GetHost returns the hostname/host-id of the connection.
func (c *WSConn) GetHost() string {
	return c.Host
}

// GetAddrPort returns the address and the port of the remote end.
func (c *WSConn) GetAddrPort() (string, int) {
	return c.Addr, c.Port
}

// GetURL returns the URL of the connection
func (c *WSConn) GetURL() *url.URL {
	return c.Url
}

// IsConnected returns the connection status.
func (c *WSConn) IsConnected() bool {
	return atomic.LoadInt32((*int32)(c.State)) == common.RunningState
}

// GetStatus returns the status of a WebSocket connection
func (c *WSConn) GetStatus() WSConnStatus {
	c.RLock()
	defer c.RUnlock()

	status := c.WSConnStatus
	status.State = new(WSConnState)
	*status.State = WSConnState(atomic.LoadInt32((*int32)(c.State)))
	return c.WSConnStatus
}

// SendMessage adds a message to sending queue.
func (c *WSConn) SendMessage(m WSMessage) error {
	if !c.IsConnected() {
		return errors.New("Not connected")
	}

	c.send <- m.Bytes()

	return nil
}

// SendRaw adds raw bytes to sending queue.
func (c *WSConn) SendRaw(b []byte) error {
	if !c.IsConnected() {
		return errors.New("Not connected")
	}

	c.send <- b

	return nil
}

// GetServiceType returns the client type.
func (c *WSConn) GetServiceType() common.ServiceType {
	return c.ServiceType
}

// GetHeaders returns the client HTTP headers.
func (c *WSConn) GetHeaders() http.Header {
	return c.headers
}

// SendMessage sends a message directly over the wire.
func (c *WSConn) write(msg []byte) error {
	if !c.IsConnected() {
		return errors.New("Not connected")
	}

	c.conn.SetWriteDeadline(time.Now().Add(writeWait))

	w, err := c.conn.NextWriter(websocket.TextMessage)
	if err != nil {
		return err
	}

	if _, err = w.Write(msg); err != nil {
		return err
	}

	return w.Close()
}

func (c *WSConn) start() {
	c.wg.Add(1)
	go c.run()
}

// main loop to read and send messages
func (c *WSConn) run() {
	defer c.wg.Done()

	go func() {
		for c.running.Load() == true {
			_, m, err := c.conn.ReadMessage()
			if err != nil {
				if c.running.Load() != false {
					c.quit <- true
				}
				break
			}

			c.read <- m
		}
	}()

	defer func() {
		atomic.StoreInt32((*int32)(c.State), common.StoppedState)
		c.RLock()
		for _, l := range c.eventHandlers {
			l.OnDisconnected(c.wsSpeaker)
		}
		c.RUnlock()
	}()

	done := make(chan bool)
	go func() {
		for {
			select {
			case m := <-c.send:
				err := c.write(m)
				if err != nil {
					logging.GetLogger().Errorf("Error while writing to the WebSocket: %s", err.Error())
				}
			case <-c.pingTicker.C:
				c.sendPing()
			case <-done:
				c.pingTicker.Stop()
				return
			}
		}
	}()
	defer func() {
		done <- true
	}()

	for {
		select {
		case <-c.quit:
			return
		case m := <-c.read:
			c.RLock()
			for _, l := range c.eventHandlers {
				l.OnMessage(c.wsSpeaker, WSRawMessage(m))
			}
			c.RUnlock()
		}
	}
}

// sendPing is used for remote connections by the server to send PingMessage
// to remote client.
func (c *WSConn) sendPing() {
	c.conn.SetWriteDeadline(time.Now().Add(writeWait))
	if err := c.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
		logging.GetLogger().Warningf("Error while sending ping to the websocket: %s", err.Error())
	}
}

// AddEventHandler registers a new event handler
func (c *WSConn) AddEventHandler(h WSSpeakerEventHandler) {
	c.Lock()
	c.eventHandlers = append(c.eventHandlers, h)
	c.Unlock()
}

// Connect default implementation doing nothing as for incoming connection it is not used.
func (c *WSConn) Connect() {
}

// Disconnect the WSSpeakers without waiting for termination.
func (c *WSConn) Disconnect() {
	c.running.Store(false)
	if atomic.LoadInt32((*int32)(c.State)) == common.RunningState {
		c.quit <- true
		c.conn.Close()
		close(c.send)
		close(c.read)
	}
}

func newWSConn(host string, clientType common.ServiceType, url *url.URL, headers http.Header) *WSConn {
	if headers == nil {
		headers = http.Header{}
	}

	port, _ := strconv.Atoi(url.Port())
	c := &WSConn{
		WSConnStatus: WSConnStatus{
			Host:        host,
			ServiceType: clientType,
			Addr:        url.Hostname(),
			Port:        port,
			State:       new(WSConnState),
			Url:         url,
			headers:     headers,
		},
		send:       make(chan []byte, maxMessages),
		read:       make(chan []byte, maxMessages),
		quit:       make(chan bool, 2),
		pingTicker: &time.Ticker{},
	}
	*c.State = common.StoppedState
	c.running.Store(true)
	return c
}

func (c *WSClient) scheme() string {
	if config.IsTLSenabled() == true {
		return "wss://"
	}
	return "ws://"
}

func (c *WSClient) connect() {
	var err error
	endpoint := c.Url.String()
	headers := http.Header{
		"X-Host-ID":             {c.Host},
		"Origin":                {endpoint},
		"X-Client-Type":         {c.ServiceType.String()},
		"X-Websocket-Namespace": {WilcardNamespace},
	}

	if c.AuthClient != nil {
		if err = c.AuthClient.Authenticate(); err != nil {
			logging.GetLogger().Errorf("Unable to authenticate %s : %s", endpoint, err.Error())
			return
		}
		c.AuthClient.SetHeaders(headers)
	}

	d := websocket.Dialer{
		Proxy:           http.ProxyFromEnvironment,
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
	certPEM := config.GetConfig().GetString("agent.X509_cert")
	keyPEM := config.GetConfig().GetString("agent.X509_key")
	if certPEM != "" && keyPEM != "" {
		d.TLSClientConfig = common.SetupTLSClientConfig(certPEM, keyPEM)
		checkTLSConfig(d.TLSClientConfig)
	}
	c.conn, _, err = d.Dial(endpoint, headers)

	if err != nil {
		logging.GetLogger().Errorf("Unable to create a WebSocket connection %s : %s", endpoint, err.Error())
		return
	}
	defer c.conn.Close()
	c.conn.SetPingHandler(nil)

	atomic.StoreInt32((*int32)(c.State), common.RunningState)
	defer atomic.StoreInt32((*int32)(c.State), common.StoppedState)

	logging.GetLogger().Infof("Connected to %s", endpoint)

	// notify connected
	c.RLock()
	for _, l := range c.eventHandlers {
		l.OnConnected(c)
	}
	c.RUnlock()

	c.wg.Add(1)
	c.run()
}

// Connect to the server - and reconnect if necessary
func (c *WSClient) Connect() {
	go func() {
		for c.running.Load() == true {
			c.connect()
			time.Sleep(1 * time.Second)
		}
	}()
}

// NewWSClient returns a WSClient with a new connection.
func NewWSClient(host string, clientType common.ServiceType, url *url.URL, authClient *AuthenticationClient, headers http.Header) *WSClient {
	wsconn := newWSConn(host, clientType, url, headers)
	c := &WSClient{
		WSConn:     wsconn,
		AuthClient: authClient,
	}
	wsconn.wsSpeaker = c
	return c
}

// NewWSClientFromConfig creates a WSClient based on the configuration
func NewWSClientFromConfig(clientType common.ServiceType, url *url.URL, authClient *AuthenticationClient, headers http.Header) *WSClient {
	host := config.GetConfig().GetString("host_id")
	return NewWSClient(host, clientType, url, authClient, headers)
}

// newIncomingWSClient is called by the server for incoming connections
func newIncomingWSClient(conn *websocket.Conn, r *auth.AuthenticatedRequest) *wsIncomingClient {
	host := getRequestParameter(r, "X-Host-ID")
	if host == "" {
		host = r.RemoteAddr
	}

	clientType := common.ServiceType(getRequestParameter(r, "X-Client-Type"))
	if clientType == "" {
		clientType = common.UnknownService
	}

	svc, _ := common.ServiceAddressFromString(conn.RemoteAddr().String())
	url := config.GetURL("http", svc.Addr, svc.Port, r.URL.Path+"?"+r.URL.RawQuery)
	wsconn := newWSConn(host, clientType, url, r.Header)
	wsconn.conn = conn

	pongTimeout := time.Duration(config.GetConfig().GetInt("ws_pong_timeout")) * time.Second

	conn.SetReadLimit(maxMessageSize)
	conn.SetReadDeadline(time.Now().Add(100 * time.Second))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(pongTimeout))
		return nil
	})

	c := &wsIncomingClient{
		WSConn: wsconn,
	}
	wsconn.wsSpeaker = c

	atomic.StoreInt32((*int32)(c.State), common.RunningState)

	// send a first ping to help firefox and some other client which wait for a
	// first ping before doing something
	c.sendPing()

	wsconn.pingTicker = time.NewTicker(pongTimeout * 8 / 10)

	return c
}
<|endoftext|>


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

/*
Cacheserver implements a directory and storage cache for Upspin. It is a
long-lived process that interposes itself between the client and the remote
services, presenting itself as a local HTTP server that behaves just like the
remote ones.

In its default mode, cacheserver runs in writeback mode, which means the
writes are asynchronous and appear to complete quickly, but may take longer to
propagate to the servers. A flag sets writethrough mode instead, which operates
synchronously and more slowly, but also more safely. Cacheserver uses local disk
to store data it has read or written. The size of the local disk area is
configurable with a flag.

The 'cache:' key should be set in the config file to enable the cacheserver.
It takes a single value that can be:

	- 'yes' (or 'y') to use a default endpoint for the cacheserver
	- 'no' (or 'n') to specify no cacheserver (the default)
	- a local TCP port (e.g. localhost:9999) to specify a particular port

The cacheserver will be started automatically by the upspin command or upspinfs if it is
not already running, and continues to run once the program that started it
has exited.

Usage:
	cacheserver [flags]

The flags are:

	-log=level
 		Set the log level to 'level'.
	-cachedir=directory
		Cache all state in 'directory'/{storecache,dircache}.
	-writethrough
		Make storage cache writethrough.
	-cachesize=bytes
		Set the maximum bytes usable for the on disk cache to 'bytes'.

Example $HOME/upspin/config entry:

	cache: yes
*/
package main // import "upspin.io/cmd/cacheserver"
<commit_msg>Add cacheserver writethrough config example<commit_after>// Copyright 2017 The Upspin Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Cacheserver implements a directory and storage cache for Upspin. It is a
long-lived process that interposes itself between the client and the remote
services, presenting itself as a local HTTP server that behaves just like the
remote ones.

In its default mode, cacheserver runs in writeback mode, which means the
writes are asynchronous and appear to complete quickly, but may take longer to
propagate to the servers. A flag sets writethrough mode instead, which operates
synchronously and more slowly, but also more safely. Cacheserver uses local disk
to store data it has read or written. The size of the local disk area is
configurable with a flag.

The 'cache:' key should be set in the config file to enable the cacheserver.
It takes a single value that can be:

	- 'yes' (or 'y') to use a default endpoint for the cacheserver
	- 'no' (or 'n') to specify no cacheserver (the default)
	- a local TCP port (e.g. localhost:9999) to specify a particular port

The cacheserver will be started automatically by the upspin command or upspinfs if it is
not already running, and continues to run once the program that started it
has exited.

Usage:
	cacheserver [flags]

The flags are:

	-log=level
 		Set the log level to 'level'.
	-cachedir=directory
		Cache all state in 'directory'/{storecache,dircache}.
	-writethrough
		Make storage cache writethrough.
	-cachesize=bytes
		Set the maximum bytes usable for the on disk cache to 'bytes'.

Example $HOME/upspin/config entry:

	cache: yes
	cmdflags:
	 cacheserver:
	  writethrough: true
*/
package main // import "upspin.io/cmd/cacheserver"
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"

	"golang.org/x/crypto/nacl/box"

	"bitbucket.org/jboverfelt/secure"
)

type secureConn struct {
	io.Reader
	io.Writer
	io.Closer
}

// Dial generates a private/public key pair,
// connects to the server, perform the handshake
// and return a reader/writer.
func dial(addr string) (io.ReadWriteCloser, error) {
	pub, priv, err := box.GenerateKey(rand.Reader)

	if err != nil {
		return nil, err
	}

	conn, err := net.Dial("tcp", addr)

	if err != nil {
		return nil, err
	}

	// first thing we do is send our public key
	_, err = conn.Write(pub[:])

	if err != nil {
		return nil, err
	}

	// wait for the server's public key
	peerPubSlice := make([]byte, secure.KeySize)
	n, err := conn.Read(peerPubSlice)

	if err != nil {
		return nil, err
	}

	peerPubSlice = peerPubSlice[:n]
	var peerPub [secure.KeySize]byte
	copy(peerPub[:], peerPubSlice)

	secCon := secureConn{
		secure.NewReader(conn, priv, &peerPub),
		secure.NewWriter(conn, priv, &peerPub),
		conn,
	}

	return secCon, nil
}

// Serve starts a secure echo server on the given listener.
func serve(l net.Listener) error {
	pub, priv, err := box.GenerateKey(rand.Reader)

	if err != nil {
		return err
	}

	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}

		go handleConnection(conn, pub, priv)
	}
}

func handleConnection(c net.Conn, pub, priv *[32]byte) {
	defer c.Close()
	// first wait for the client's public key
	peerPubSlice := make([]byte, secure.KeySize)
	n, err := c.Read(peerPubSlice)

	if err != nil {
		log.Println(err)
		return
	}

	peerPubSlice = peerPubSlice[:n]
	var peerPub [secure.KeySize]byte
	copy(peerPub[:], peerPubSlice)

	// then, send our public key
	_, err = c.Write(pub[:])

	if err != nil {
		log.Println(err)
		return
	}

	// now session is "secure"
	sr := secure.NewReader(c, priv, &peerPub)
	sw := secure.NewWriter(c, priv, &peerPub)

	// echo
	_, err = io.Copy(sw, sr)

	if err != nil {
		log.Println(err)
	}
}

func main() {
	port := flag.Int("l", 0, "Listen mode. Specify port")
	flag.Parse()

	// Server mode
	if *port != 0 {
		l, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
		if err != nil {
			log.Fatal(err)
		}
		defer l.Close()
		log.Fatal(serve(l))
	}

	// Client mode
	if len(os.Args) != 3 {
		log.Fatalf("Usage: %s <port> <message>", os.Args[0])
	}
	conn, err := dial("localhost:" + os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	if _, err := conn.Write([]byte(os.Args[2])); err != nil {
		log.Fatal(err)
	}
	buf := make([]byte, len(os.Args[2]))
	n, err := conn.Read(buf)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s\n", buf[:n])
}
<commit_msg>rewrite import after repo move<commit_after>package main

import (
	"crypto/rand"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"

	"golang.org/x/crypto/nacl/box"

	"github.com/jboverfelt/secure"
)

type secureConn struct {
	io.Reader
	io.Writer
	io.Closer
}

// Dial generates a private/public key pair,
// connects to the server, perform the handshake
// and return a reader/writer.
func dial(addr string) (io.ReadWriteCloser, error) {
	pub, priv, err := box.GenerateKey(rand.Reader)

	if err != nil {
		return nil, err
	}

	conn, err := net.Dial("tcp", addr)

	if err != nil {
		return nil, err
	}

	// first thing we do is send our public key
	_, err = conn.Write(pub[:])

	if err != nil {
		return nil, err
	}

	// wait for the server's public key
	peerPubSlice := make([]byte, secure.KeySize)
	n, err := conn.Read(peerPubSlice)

	if err != nil {
		return nil, err
	}

	peerPubSlice = peerPubSlice[:n]
	var peerPub [secure.KeySize]byte
	copy(peerPub[:], peerPubSlice)

	secCon := secureConn{
		secure.NewReader(conn, priv, &peerPub),
		secure.NewWriter(conn, priv, &peerPub),
		conn,
	}

	return secCon, nil
}

// Serve starts a secure echo server on the given listener.
func serve(l net.Listener) error {
	pub, priv, err := box.GenerateKey(rand.Reader)

	if err != nil {
		return err
	}

	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}

		go handleConnection(conn, pub, priv)
	}
}

func handleConnection(c net.Conn, pub, priv *[32]byte) {
	defer c.Close()
	// first wait for the client's public key
	peerPubSlice := make([]byte, secure.KeySize)
	n, err := c.Read(peerPubSlice)

	if err != nil {
		log.Println(err)
		return
	}

	peerPubSlice = peerPubSlice[:n]
	var peerPub [secure.KeySize]byte
	copy(peerPub[:], peerPubSlice)

	// then, send our public key
	_, err = c.Write(pub[:])

	if err != nil {
		log.Println(err)
		return
	}

	// now session is "secure"
	sr := secure.NewReader(c, priv, &peerPub)
	sw := secure.NewWriter(c, priv, &peerPub)

	// echo
	_, err = io.Copy(sw, sr)

	if err != nil {
		log.Println(err)
	}
}

func main() {
	port := flag.Int("l", 0, "Listen mode. Specify port")
	flag.Parse()

	// Server mode
	if *port != 0 {
		l, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
		if err != nil {
			log.Fatal(err)
		}
		defer l.Close()
		log.Fatal(serve(l))
	}

	// Client mode
	if len(os.Args) != 3 {
		log.Fatalf("Usage: %s <port> <message>", os.Args[0])
	}
	conn, err := dial("localhost:" + os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	if _, err := conn.Write([]byte(os.Args[2])); err != nil {
		log.Fatal(err)
	}
	buf := make([]byte, len(os.Args[2]))
	n, err := conn.Read(buf)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s\n", buf[:n])
}
<|endoftext|>


<commit_before>/*
    Player structures and functions
*/
package main

import (
    "net/textproto"
    "container/list"
    "os"
    "log"
    "fmt"
    "io"
//    "time"
//    "syscall"
//    "bytes"
    "bufio"
    "ataxia/handler"
)

// The Connection structure wraps all the lower networking details for each connected player
type connection struct {
    socket      io.ReadWriteCloser
    buffer      *bufio.ReadWriter
    server      *Server
    handler     handler.Handler
    remoteAddr  string
    state       string
}

// Account
type Account struct {
    Email string
    Password string
    Name string
    Characters *list.List
}

// Player
type Player struct {
    account Account
    conn *connection
    In chan string
    Out chan string
}

// Player factory
func NewPlayer(conn *connection) (player *Player) {
    player = new(Player)
    player.conn = conn
    player.In = make(chan string, 1024)
    player.Out = make(chan string, 1024)
    return player
}

func (player *Player) Run() {
    buf := make([]byte, 1024)

    // Setup the player here.
    player.Write([]byte("Hello, welcome to Ataxia. What is your account name?\n"))
    if _, err := player.Read(buf); err != nil {
        return
    }
    player.Write([]byte(fmt.Sprintf("Hello %s.\n", string(buf))))
    player.account.Name = string(buf)

    // Create an anonymous goroutine for reading
    go func() {
        for {
            if player.conn.socket == nil {
                break
            }

            data := make([]byte, 1024)
            var n int
            var err os.Error

            if n, err = player.Read(data); err != nil {
                log.Println(n)
                log.Println(err)
                player.Close()
                break
            }

            if n > 0 {
                player.conn.server.SendToAll(fmt.Sprintf("<%s> %s", player.account.Name, string(data)))
            }
        }
    }()

    // Create an anonymous goroutine for writing
    go func() {
        for {
            if player.conn.socket == nil {
                break
            }
            buf := <-player.In
            if _, err := player.Write([]byte(buf)); err != nil {
                log.Println(err)
                player.Close()
                break
            }
        }
    }()
}

func (player *Player) Close() {
    if (player.conn.socket != nil) {
        player.conn.socket.Close()
        player.conn.socket = nil
        player.conn.buffer = nil
        player.conn.server.RemovePlayer(player)
        log.Println("Player disconnected:", player.account.Name)
    }
}

func (player *Player) Write(buf []byte) (n int, err os.Error) {
    n, err = player.conn.buffer.Write(buf)
    player.conn.buffer.Flush()
    return
}

func (player *Player) Read(buf []byte) (n int, err os.Error) {
    if player.conn.buffer == nil {
        return
    }

    if player.conn.socket == nil {
        return
    }

    tp := textproto.NewReader(player.conn.buffer.Reader)

    var data []byte
    if data, err = tp.ReadLineBytes(); err != nil {
        if err == os.EOF {
            log.Println("Read EOF, disconnecting player")
            player.Close()
            return 0, nil
        }
        return 0, err
    }

    copy(buf, data)
    return len(buf), err
}
<commit_msg>Fix a crash bug where it tries to write to a closed socket if the player disconnects.<commit_after>/*
    Player structures and functions
*/
package main

import (
    "net/textproto"
    "container/list"
    "os"
    "log"
    "fmt"
    "io"
//    "time"
//    "syscall"
//    "bytes"
    "bufio"
    "ataxia/handler"
)

// The Connection structure wraps all the lower networking details for each connected player
type connection struct {
    socket      io.ReadWriteCloser
    buffer      *bufio.ReadWriter
    server      *Server
    handler     handler.Handler
    remoteAddr  string
    state       string
}

// Account
type Account struct {
    Email string
    Password string
    Name string
    Characters *list.List
}

// Player
type Player struct {
    account Account
    conn *connection
    In chan string
    Out chan string
}

// Player factory
func NewPlayer(conn *connection) (player *Player) {
    player = new(Player)
    player.conn = conn
    player.In = make(chan string, 1024)
    player.Out = make(chan string, 1024)
    player.account.Name = "Unknown"
    return player
}

func (player *Player) Run() {
    buf := make([]byte, 1024)

    // Setup the player here.
    player.Write([]byte("Hello, welcome to Ataxia. What is your account name?\n"))
    if _, err := player.Read(buf); err != nil {
        return
    }
    player.Write([]byte(fmt.Sprintf("Hello %s.\n", string(buf))))
    player.account.Name = string(buf)

    // Create an anonymous goroutine for reading
    go func() {
        for {
            if player.conn.socket == nil {
                break
            }

            data := make([]byte, 1024)
            var n int
            var err os.Error

            if n, err = player.Read(data); err != nil {
                log.Println(n)
                log.Println(err)
                player.Close()
                break
            }

            if n > 0 {
                player.conn.server.SendToAll(fmt.Sprintf("<%s> %s", player.account.Name, string(data)))
            }
        }
    }()

    // Create an anonymous goroutine for writing
    go func() {
        for {
            if player.conn.socket == nil {
                break
            }
            buf := <-player.In
            if _, err := player.Write([]byte(buf)); err != nil {
                log.Println(err)
                player.Close()
                break
            }
        }
    }()
}

func (player *Player) Close() {
    if (player.conn.socket != nil) {
        player.conn.socket.Close()
        player.conn.socket = nil
        player.conn.buffer = nil
        player.conn.server.RemovePlayer(player)
        log.Println("Player disconnected:", player.account.Name)
    }
}

func (player *Player) Write(buf []byte) (n int, err os.Error) {
    if player.conn.socket == nil || player.conn.buffer == nil {
        return
    }

    if n, err = player.conn.buffer.Write(buf); err != nil {
        if err == os.EOF {
            log.Println("EOF on write, disconnecting player")
            player.Close()
            return 0, nil
        }
        return 0, err
    }
    player.conn.buffer.Flush()
    return
}

func (player *Player) Read(buf []byte) (n int, err os.Error) {
    if player.conn.socket == nil || player.conn.buffer == nil {
        return
    }

    tp := textproto.NewReader(player.conn.buffer.Reader)

    var data []byte
    if data, err = tp.ReadLineBytes(); err != nil {
        if err == os.EOF {
            log.Println("Read EOF, disconnecting player")
            player.Close()
            return 0, nil
        }
        return 0, err
    }

    copy(buf, data)
    return len(buf), err
}
<|endoftext|>


<commit_before>package control

import (
	"fmt"
	"log"

	"github.com/codegangsta/cli"
	"github.com/rancherio/os/config"
	"github.com/rancherio/os/util"
)

func serviceSubCommands() []cli.Command {
	return []cli.Command{
		{
			Name:   "enable",
			Usage:  "turn on an service",
			Action: enable,
		},
		{
			Name:   "disable",
			Usage:  "turn off an service",
			Action: disable,
		},
		{
			Name:   "list",
			Usage:  "list services and state",
			Action: list,
		},
	}
}

func disable(c *cli.Context) {
	changed := false
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatal(err)
	}

	for _, service := range c.Args() {
		filtered := make([]string, 0, len(c.Args()))
		for _, existing := range cfg.EnabledServices {
			if existing != service {
				filtered = append(filtered, existing)
			}
		}

		if len(filtered) != len(c.Args()) {
			cfg.EnabledServices = filtered
			changed = true
		}
	}

	if changed {
		if err = cfg.Set("enabled_services", cfg.EnabledServices); err != nil {
			log.Fatal(err)
		}
	}
}

func enable(c *cli.Context) {
	changed := false
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatal(err)
	}

	for _, service := range c.Args() {
		if _, ok := cfg.Services[service]; ok && !util.Contains(cfg.EnabledServices, service) {
			cfg.EnabledServices = append(cfg.EnabledServices, service)
			changed = true
		}
	}

	if changed {
		if err = cfg.Set("enabled_services", cfg.EnabledServices); err != nil {
			log.Fatal(err)
		}
	}
}

func list(c *cli.Context) {
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatal(err)
	}

	enabled := map[string]bool{}

	for _, service := range cfg.EnabledServices {
		enabled[service] = true
	}

	for service, _ := range cfg.Services {
		if _, ok := enabled[service]; ok {
			fmt.Printf("%s enabled\n", service)
		} else {
			fmt.Printf("%s disabled\n", service)
		}
	}
}
<commit_msg>Fix rancherctl service enable/list<commit_after>package control

import (
	"fmt"
	"log"

	"github.com/codegangsta/cli"
	"github.com/rancherio/os/config"
	"github.com/rancherio/os/util"
)

func serviceSubCommands() []cli.Command {
	return []cli.Command{
		{
			Name:   "enable",
			Usage:  "turn on an service",
			Action: enable,
		},
		{
			Name:   "disable",
			Usage:  "turn off an service",
			Action: disable,
		},
		{
			Name:   "list",
			Usage:  "list services and state",
			Action: list,
		},
	}
}

func disable(c *cli.Context) {
	changed := false
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatal(err)
	}

	for _, service := range c.Args() {
		filtered := make([]string, 0, len(c.Args()))
		for _, existing := range cfg.EnabledServices {
			if existing != service {
				filtered = append(filtered, existing)
			}
		}

		if len(filtered) != len(c.Args()) {
			cfg.EnabledServices = filtered
			changed = true
		}
	}

	if changed {
		if err = cfg.Set("enabled_services", cfg.EnabledServices); err != nil {
			log.Fatal(err)
		}
	}
}

func enable(c *cli.Context) {
	changed := false
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatal(err)
	}

	for _, service := range c.Args() {
		if !util.Contains(cfg.EnabledServices, service) {
			cfg.EnabledServices = append(cfg.EnabledServices, service)
			changed = true
		}
	}

	if changed {
		if err = cfg.Set("enabled_services", cfg.EnabledServices); err != nil {
			log.Fatal(err)
		}
	}
}

func list(c *cli.Context) {
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatal(err)
	}

	enabled := map[string]bool{}

	for _, service := range cfg.EnabledServices {
		enabled[service] = true
	}

	for service, _ := range cfg.Services {
		if _, ok := enabled[service]; ok {
			delete(enabled, service)
			fmt.Printf("enabled  %s\n", service)
		} else {
			fmt.Printf("disabled %s\n", service)
		}
	}

	for service, _ := range enabled {
		fmt.Printf("enabled  %s\n", service)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"../libs/twodee"
)

type DirectionsHistoryEntry struct {
	prev, next *DirectionsHistoryEntry
	dir        MoveDirection
}

type DirectionsHistory struct {
	tail       *DirectionsHistoryEntry
	directions []*DirectionsHistoryEntry
}

// Adds a MoveDirection to the history if it's not already present.
func (dh *DirectionsHistory) Add(d MoveDirection) {
	fmt.Printf("Adding move direction: %v\n", d)
	entry := dh.directions[d]
	if entry.prev == nil && entry.next == nil && entry != dh.tail {
		if dh.tail != nil {
			dh.tail.next = entry
			entry.prev = dh.tail
		}
		dh.tail = entry
	}
}

// Removes a particular MoveDirection from the history chain; sets its prev and next fields to nil. Resets dh's tail if necessary.
func (dh *DirectionsHistory) Remove(d MoveDirection) {
	fmt.Printf("Trying to remove direction: %v\n", d)
	entry := dh.directions[d]
	if entry.prev != nil {
		entry.prev.next = entry.next
	}
	if entry.next != nil {
		entry.next.prev = entry.prev
	}
	if entry == dh.tail {
		dh.tail = entry.prev
	}
	entry.prev = nil
	entry.next = nil
}

func (dh *DirectionsHistory) LatestDirection() (d MoveDirection) {
	if dh.tail != nil {
		return dh.tail.dir
	}
	return None
}

func NewDirectionsHistory() (dh *DirectionsHistory) {
	// Ugh, West+1 because we're generating a sparse array indexed on the
	// int value of directions; West should be the last one enumerated.
	dh = &DirectionsHistory{
		tail:       nil,
		directions: make([]*DirectionsHistoryEntry, West+1),
	}
	dirs := []MoveDirection{North, East, South, West}
	for _, d := range dirs {
		dh.directions[d] = &DirectionsHistoryEntry{
			prev: nil,
			next: nil,
			dir:  d,
		}
	}
	return
}

type Player struct {
	*twodee.AnimatingEntity
	Health            float32
	Speed             float32
	Velocity          twodee.Point
	DirectionsHistory *DirectionsHistory
	DesiredMove       MoveDirection
	Inventory         []*Item
	State             EntityState
}

type EntityState int32

const (
	_                    = iota
	Standing EntityState = 1 << iota
	Walking
	Left
	Right
	Up
	Down
)

var PlayerAnimations = map[EntityState][]int{
	Standing | Up:    []int{24},
	Standing | Down:  []int{8},
	Standing | Left:  []int{16},
	Standing | Right: []int{16},
	Walking | Up:     []int{25, 26, 27, 28, 29, 30},
	Walking | Down:   []int{9, 10, 11, 12, 13, 14},
	Walking | Left:   []int{17, 18, 19, 20, 21, 22},
	Walking | Right:  []int{17, 18, 19, 20, 21, 22},
}

func NewPlayer(x, y float32) (player *Player) {
	var (
		inv = make([]*Item, 0, NumberOfItemTypes)
	)
	player = &Player{
		AnimatingEntity: twodee.NewAnimatingEntity(
			x, y,
			1, 1,
			0,
			twodee.Step10Hz,
			[]int{8},
		),
		Health:            100.0,
		Speed:             0.2,
		Velocity:          twodee.Pt(0, 0),
		DirectionsHistory: NewDirectionsHistory(),
		DesiredMove:       None,
		Inventory:         inv,
	}
	return
}

func (p *Player) RemState(state EntityState) {
	p.SetState(p.State & ^state)
}

func (p *Player) AddState(state EntityState) {
	p.SetState(p.State | state)
}

func (p *Player) SwapState(rem, add EntityState) {
	p.SetState(p.State & ^rem | add)
}

func (p *Player) SetState(state EntityState) {
	if state != p.State {
		p.State = state
		if frames, ok := PlayerAnimations[p.State]; ok {
			p.SetFrames(frames)
		}
	}
}

func (p *Player) FlippedX() bool {
	return p.State&Left > 0
}

// Updates the Player's desired movement direction as well as the affiliated data
// structures. If `invert`, then the movement key has been released and we should
// remove it from the affiliated data structures and perhaps pick a new movement
// direction from the tail of OrderedDirections.
func (p *Player) UpdateDesiredMove(d MoveDirection, invert bool) {
	if invert {
		p.DirectionsHistory.Remove(d)
		p.DesiredMove = p.DirectionsHistory.LatestDirection()
		return
	}
	// If the player is already moving in this direction, do nothing.
	if p.DesiredMove == d {
		return
	}
	p.DirectionsHistory.Add(d)
	p.DesiredMove = p.DirectionsHistory.LatestDirection()
}

const Fudge = 0.01

func (p *Player) AttemptMove(l *Level) {
	var (
		a, b, trunc twodee.Point
		bounds      = p.Bounds()
		pos         = p.Pos()
	)
	switch p.DesiredMove {
	case None:
		p.SwapState(Walking, Standing)
		return
	case North:
		a = twodee.Pt(bounds.Min.X+Fudge, bounds.Max.Y+p.Speed)
		b = twodee.Pt(bounds.Max.X-Fudge, bounds.Max.Y+p.Speed)
		pos.Y += p.Speed
		trunc = l.GridAlignedY(l.Active, pos)
		p.SetState(Walking | Up)
	case South:
		a = twodee.Pt(bounds.Min.X+Fudge, bounds.Min.Y-p.Speed)
		b = twodee.Pt(bounds.Max.X-Fudge, bounds.Min.Y-p.Speed)
		pos.Y -= p.Speed
		trunc = l.GridAlignedY(l.Active, pos)
		p.SetState(Walking | Down)
	case East:
		a = twodee.Pt(bounds.Max.X+p.Speed, bounds.Min.Y+Fudge)
		b = twodee.Pt(bounds.Max.X+p.Speed, bounds.Max.Y-Fudge)
		pos.X += p.Speed
		trunc = l.GridAlignedX(l.Active, pos)
		p.SetState(Walking | Right)
	case West:
		a = twodee.Pt(bounds.Min.X-p.Speed, bounds.Min.Y+Fudge)
		b = twodee.Pt(bounds.Min.X-p.Speed, bounds.Max.Y-Fudge)
		pos.X -= p.Speed
		trunc = l.GridAlignedX(l.Active, pos)
		p.SetState(Walking | Left)
	}
	if l.FrontierCollides(l.Active, a, b) {
		p.MoveTo(trunc)
	} else {
		p.MoveTo(pos)
	}
}

func (p *Player) AddToInventory(item *Item) {
	p.Inventory = append(p.Inventory, item)
	switch item.getType() {
	case Item1:
		p.Health = p.Health + 10
	case Item2:
		p.Health = p.Health + 20
	case Item3:
		p.Health = p.Health + 30
	case Item4:
		p.Speed = p.Speed + 10
	case Item5:
		p.Speed = p.Speed + 20
	}
}
<commit_msg>Remove my logging statements.<commit_after>package main

import "../libs/twodee"

type DirectionsHistoryEntry struct {
	prev, next *DirectionsHistoryEntry
	dir        MoveDirection
}

type DirectionsHistory struct {
	tail       *DirectionsHistoryEntry
	directions []*DirectionsHistoryEntry
}

// Adds a MoveDirection to the history if it's not already present.
func (dh *DirectionsHistory) Add(d MoveDirection) {
	entry := dh.directions[d]
	if entry.prev == nil && entry.next == nil && entry != dh.tail {
		if dh.tail != nil {
			dh.tail.next = entry
			entry.prev = dh.tail
		}
		dh.tail = entry
	}
}

// Removes a particular MoveDirection from the history chain; sets its prev and next fields to nil. Resets dh's tail if necessary.
func (dh *DirectionsHistory) Remove(d MoveDirection) {
	entry := dh.directions[d]
	if entry.prev != nil {
		entry.prev.next = entry.next
	}
	if entry.next != nil {
		entry.next.prev = entry.prev
	}
	if entry == dh.tail {
		dh.tail = entry.prev
	}
	entry.prev = nil
	entry.next = nil
}

func (dh *DirectionsHistory) LatestDirection() (d MoveDirection) {
	if dh.tail != nil {
		return dh.tail.dir
	}
	return None
}

func NewDirectionsHistory() (dh *DirectionsHistory) {
	// Ugh, West+1 because we're generating a sparse array indexed on the
	// int value of directions; West should be the last one enumerated.
	dh = &DirectionsHistory{
		tail:       nil,
		directions: make([]*DirectionsHistoryEntry, West+1),
	}
	dirs := []MoveDirection{North, East, South, West}
	for _, d := range dirs {
		dh.directions[d] = &DirectionsHistoryEntry{
			prev: nil,
			next: nil,
			dir:  d,
		}
	}
	return
}

type Player struct {
	*twodee.AnimatingEntity
	Health            float32
	Speed             float32
	Velocity          twodee.Point
	DirectionsHistory *DirectionsHistory
	DesiredMove       MoveDirection
	Inventory         []*Item
	State             EntityState
}

type EntityState int32

const (
	_                    = iota
	Standing EntityState = 1 << iota
	Walking
	Left
	Right
	Up
	Down
)

var PlayerAnimations = map[EntityState][]int{
	Standing | Up:    []int{24},
	Standing | Down:  []int{8},
	Standing | Left:  []int{16},
	Standing | Right: []int{16},
	Walking | Up:     []int{25, 26, 27, 28, 29, 30},
	Walking | Down:   []int{9, 10, 11, 12, 13, 14},
	Walking | Left:   []int{17, 18, 19, 20, 21, 22},
	Walking | Right:  []int{17, 18, 19, 20, 21, 22},
}

func NewPlayer(x, y float32) (player *Player) {
	var (
		inv = make([]*Item, 0, NumberOfItemTypes)
	)
	player = &Player{
		AnimatingEntity: twodee.NewAnimatingEntity(
			x, y,
			1, 1,
			0,
			twodee.Step10Hz,
			[]int{8},
		),
		Health:            100.0,
		Speed:             0.2,
		Velocity:          twodee.Pt(0, 0),
		DirectionsHistory: NewDirectionsHistory(),
		DesiredMove:       None,
		Inventory:         inv,
	}
	return
}

func (p *Player) RemState(state EntityState) {
	p.SetState(p.State & ^state)
}

func (p *Player) AddState(state EntityState) {
	p.SetState(p.State | state)
}

func (p *Player) SwapState(rem, add EntityState) {
	p.SetState(p.State & ^rem | add)
}

func (p *Player) SetState(state EntityState) {
	if state != p.State {
		p.State = state
		if frames, ok := PlayerAnimations[p.State]; ok {
			p.SetFrames(frames)
		}
	}
}

func (p *Player) FlippedX() bool {
	return p.State&Left > 0
}

// Updates the Player's desired movement direction as well as the affiliated data
// structures. If `invert`, then the movement key has been released and we should
// remove it from the affiliated data structures and perhaps pick a new movement
// direction from the tail of OrderedDirections.
func (p *Player) UpdateDesiredMove(d MoveDirection, invert bool) {
	if invert {
		p.DirectionsHistory.Remove(d)
		p.DesiredMove = p.DirectionsHistory.LatestDirection()
		return
	}
	// If the player is already moving in this direction, do nothing.
	if p.DesiredMove == d {
		return
	}
	p.DirectionsHistory.Add(d)
	p.DesiredMove = p.DirectionsHistory.LatestDirection()
}

const Fudge = 0.01

func (p *Player) AttemptMove(l *Level) {
	var (
		a, b, trunc twodee.Point
		bounds      = p.Bounds()
		pos         = p.Pos()
	)
	switch p.DesiredMove {
	case None:
		p.SwapState(Walking, Standing)
		return
	case North:
		a = twodee.Pt(bounds.Min.X+Fudge, bounds.Max.Y+p.Speed)
		b = twodee.Pt(bounds.Max.X-Fudge, bounds.Max.Y+p.Speed)
		pos.Y += p.Speed
		trunc = l.GridAlignedY(l.Active, pos)
		p.SetState(Walking | Up)
	case South:
		a = twodee.Pt(bounds.Min.X+Fudge, bounds.Min.Y-p.Speed)
		b = twodee.Pt(bounds.Max.X-Fudge, bounds.Min.Y-p.Speed)
		pos.Y -= p.Speed
		trunc = l.GridAlignedY(l.Active, pos)
		p.SetState(Walking | Down)
	case East:
		a = twodee.Pt(bounds.Max.X+p.Speed, bounds.Min.Y+Fudge)
		b = twodee.Pt(bounds.Max.X+p.Speed, bounds.Max.Y-Fudge)
		pos.X += p.Speed
		trunc = l.GridAlignedX(l.Active, pos)
		p.SetState(Walking | Right)
	case West:
		a = twodee.Pt(bounds.Min.X-p.Speed, bounds.Min.Y+Fudge)
		b = twodee.Pt(bounds.Min.X-p.Speed, bounds.Max.Y-Fudge)
		pos.X -= p.Speed
		trunc = l.GridAlignedX(l.Active, pos)
		p.SetState(Walking | Left)
	}
	if l.FrontierCollides(l.Active, a, b) {
		p.MoveTo(trunc)
	} else {
		p.MoveTo(pos)
	}
}

func (p *Player) AddToInventory(item *Item) {
	p.Inventory = append(p.Inventory, item)
	switch item.getType() {
	case Item1:
		p.Health = p.Health + 10
	case Item2:
		p.Health = p.Health + 20
	case Item3:
		p.Health = p.Health + 30
	case Item4:
		p.Speed = p.Speed + 10
	case Item5:
		p.Speed = p.Speed + 20
	}
}
<|endoftext|>


<commit_before>package main

import (
	".."
	"crypto/md5"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
)

func main() {
	tmp, err := gitmedia.TempFile()
	if err != nil {
		fmt.Println("Error trying to create temp file")
		panic(err)
	}

	sha1Hash := sha1.New()
	md5Hash := md5.New()
	writer := io.MultiWriter(sha1Hash, md5Hash, tmp)

	written, _ := io.Copy(writer, os.Stdin)
	sha := hex.EncodeToString(sha1Hash.Sum(nil))

	meta := &gitmedia.LargeAsset{
		MediaType: "application/vnd.github.large-asset",
		Size:      written,
		MD5:       hex.EncodeToString(md5Hash.Sum(nil)),
		SHA1:      sha,
	}

	mediafile := gitmedia.LocalMediaPath(sha)
	metafile := mediafile + ".json"

	if err := os.Rename(tmp.Name(), mediafile); err != nil {
		fmt.Printf("Unable to move %s to %s\n", tmp.Name(), mediafile)
		panic(err)
	}

	file, err := os.Create(metafile)
	if err != nil {
		fmt.Printf("Unable to create meta data file: %s\n", metafile)
		panic(err)
	}

	defer file.Close()
	defer func() { os.Remove(tmp.Name()) }()

	writer = io.MultiWriter(os.Stdout, file)
	writer.Write([]byte(fmt.Sprintf("# %d\n", len(gitmedia.MediaWarning))))
	writer.Write(gitmedia.MediaWarning)
	enc := json.NewEncoder(writer)
	enc.Encode(meta)
}
<commit_msg>ンンンン ン<commit_after>package main

import (
	".."
	"crypto/md5"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
)

func main() {
	tmp, err := gitmedia.TempFile()
	if err != nil {
		fmt.Println("Error trying to create temp file")
		panic(err)
	}

	sha1Hash := sha1.New()
	md5Hash := md5.New()
	writer := io.MultiWriter(sha1Hash, md5Hash, tmp)

	written, _ := io.Copy(writer, os.Stdin)
	sha := hex.EncodeToString(sha1Hash.Sum(nil))

	asset := &gitmedia.LargeAsset{
		MediaType: "application/vnd.github.large-asset",
		Size:      written,
		MD5:       hex.EncodeToString(md5Hash.Sum(nil)),
		SHA1:      sha,
	}

	output := ChooseWriter(asset, tmp)
	defer output.Close()

	output.Write([]byte(fmt.Sprintf("# %d\n", len(gitmedia.MediaWarning))))
	output.Write(gitmedia.MediaWarning)
	enc := json.NewEncoder(output)
	enc.Encode(asset)
}

func ChooseWriter(asset *gitmedia.LargeAsset, tmp *os.File) io.WriteCloser {
	mediafile := gitmedia.LocalMediaPath(asset.SHA1)
	if stat, _ := os.Stat(mediafile); stat == nil {
		return NewGitMediaCleanWriter(asset, tmp)
	} else {
		return NewGitMediaExistingWriter(asset, tmp)
	}
}

type GitMediaExistingWriter struct {
	tempfile *os.File
	writer   io.Writer
}

func NewGitMediaExistingWriter(asset *gitmedia.LargeAsset, tmp *os.File) io.WriteCloser {
	return &GitMediaExistingWriter{tmp, os.Stdout}
}

func (w *GitMediaExistingWriter) Write(buf []byte) (int, error) {
	return w.writer.Write(buf)
}

func (w *GitMediaExistingWriter) Close() error {
	return os.Remove(w.tempfile.Name())
}

type GitMediaCleanWriter struct {
	metafile *os.File
	*GitMediaExistingWriter
}

func NewGitMediaCleanWriter(asset *gitmedia.LargeAsset, tmp *os.File) io.WriteCloser {
	mediafile := gitmedia.LocalMediaPath(asset.SHA1)
	metafile := mediafile + ".json"

	if err := os.Rename(tmp.Name(), mediafile); err != nil {
		fmt.Printf("Unable to move %s to %s\n", tmp.Name(), mediafile)
		panic(err)
	}

	file, err := os.Create(metafile)
	if err != nil {
		fmt.Printf("Unable to create meta data file: %s\n", metafile)
		panic(err)
	}

	return &GitMediaCleanWriter{
		metafile:               file,
		GitMediaExistingWriter: &GitMediaExistingWriter{tmp, io.MultiWriter(os.Stdout, file)},
	}
}

func (w *GitMediaCleanWriter) Write(buf []byte) (int, error) {
	return w.writer.Write(buf)
}

func (w *GitMediaCleanWriter) Close() error {
	w.GitMediaExistingWriter.Close()
	return w.metafile.Close()
}
<|endoftext|>


<commit_before>package command

import (
	"encoding/json"
	"flag"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/funkygao/columnize"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
)

type Dbus struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Dbus) Run(args []string) (exitCode int) {
	var (
		zone    string
		topMode bool
	)
	cmdFlags := flag.NewFlagSet("dbus", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.BoolVar(&topMode, "top", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	for {
		this.checkMyslave(zkzone)
		if !topMode {
			break
		}

		time.Sleep(time.Second * 3)
		refreshScreen()
	}

	return
}

type binlogCheckpoint struct {
	File   string `json:"file"`
	Offset int64  `json:"offset"`
	Owner  string `json:"owner"`
}

func (this *Dbus) checkMyslave(zkzone *zk.ZkZone) {
	lines := []string{"Mysql|File|Offset|dbus|role|uptime|pid"}
	root := "/dbus/myslave"
	dbs, _, err := zkzone.Conn().Children(root)
	swallow(err)

	sort.Strings(dbs)

	for _, db := range dbs {
		dbRoot := fmt.Sprintf("%s/%s", root, db)
		data, _, err := zkzone.Conn().Get(dbRoot)
		var v binlogCheckpoint
		swallow(json.Unmarshal(data, &v))

		ownerPath := fmt.Sprintf("%s/owner", dbRoot)
		owner, _, err := zkzone.Conn().Get(ownerPath)
		swallow(err)

		idsPath := fmt.Sprintf("%s/ids", dbRoot)
		runners, _, err := zkzone.Conn().Children(idsPath)
		swallow(err)

		for _, r := range runners {
			_, stat, err := zkzone.Conn().Get(fmt.Sprintf("%s/%s", idsPath, r))
			swallow(err)

			i := strings.LastIndex(r, "-")
			instance := r[:i]
			pid := r[i+1:]
			uptime := zk.ZkTimestamp(stat.Ctime).Time()

			if r == string(owner) {
				lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|master|%s|%s",
					db, v.File, gofmt.Comma(v.Offset), instance, gofmt.PrettySince(uptime), pid))
			} else {
				lines = append(lines, fmt.Sprintf("%s| | |%s|slave|%s|%s",
					db, instance, gofmt.PrettySince(uptime), pid))
			}
		}

	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (*Dbus) Synopsis() string {
	return "Monitor database system runtime"
}

func (this *Dbus) Help() string {
	help := fmt.Sprintf(`
Usage: %s dbus [options]

    %s

Options:

    -z zone

    -top
      Run in top mode.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<commit_msg>print how long have dbus become master<commit_after>package command

import (
	"encoding/json"
	"flag"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/funkygao/columnize"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
)

type Dbus struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Dbus) Run(args []string) (exitCode int) {
	var (
		zone    string
		topMode bool
	)
	cmdFlags := flag.NewFlagSet("dbus", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.BoolVar(&topMode, "top", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	for {
		this.checkMyslave(zkzone)
		if !topMode {
			break
		}

		time.Sleep(time.Second * 3)
		refreshScreen()
	}

	return
}

type binlogCheckpoint struct {
	File   string `json:"file"`
	Offset int64  `json:"offset"`
	Owner  string `json:"owner"`
}

func (this *Dbus) checkMyslave(zkzone *zk.ZkZone) {
	lines := []string{"Mysql|File|Offset|dbus|role|uptime|pid|msince"}
	root := "/dbus/myslave"
	dbs, _, err := zkzone.Conn().Children(root)
	swallow(err)

	sort.Strings(dbs)

	for _, db := range dbs {
		dbRoot := fmt.Sprintf("%s/%s", root, db)
		data, _, err := zkzone.Conn().Get(dbRoot)
		var v binlogCheckpoint
		swallow(json.Unmarshal(data, &v))

		ownerPath := fmt.Sprintf("%s/owner", dbRoot)
		owner, ownerStat, err := zkzone.Conn().Get(ownerPath)
		swallow(err)

		idsPath := fmt.Sprintf("%s/ids", dbRoot)
		runners, _, err := zkzone.Conn().Children(idsPath)
		swallow(err)

		for _, r := range runners {
			_, stat, err := zkzone.Conn().Get(fmt.Sprintf("%s/%s", idsPath, r))
			swallow(err)

			i := strings.LastIndex(r, "-")
			instance := r[:i]
			pid := r[i+1:]
			uptime := zk.ZkTimestamp(stat.Ctime).Time()

			if r == string(owner) {
				ownerTime := zk.ZkTimestamp(ownerStat.Mtime).Time()
				lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|master|%s|%s|%s",
					db, v.File, gofmt.Comma(v.Offset), instance,
					gofmt.PrettySince(uptime), pid, gofmt.PrettySince(ownerTime)))
			} else {
				lines = append(lines, fmt.Sprintf("%s| | |%s|slave|%s|%s| ",
					db, instance, gofmt.PrettySince(uptime), pid))
			}
		}

	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (*Dbus) Synopsis() string {
	return "Monitor database system runtime"
}

func (this *Dbus) Help() string {
	help := fmt.Sprintf(`
Usage: %s dbus [options]

    %s

Options:

    -z zone

    -top
      Run in top mode.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package command

import (
	"encoding/json"
	"flag"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/funkygao/columnize"
	"github.com/funkygao/dbus/pkg/myslave"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-helix"
	hzk "github.com/funkygao/go-helix/store/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	zklib "github.com/samuel/go-zookeeper/zk"
)

type Dbus struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Dbus) Run(args []string) (exitCode int) {
	var (
		zone         string
		topMode      bool
		helixCluster string
		initCluster  string
		addResource  string
	)
	cmdFlags := flag.NewFlagSet("dbus", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&helixCluster, "c", "", "")
	cmdFlags.BoolVar(&topMode, "top", false, "")
	cmdFlags.StringVar(&initCluster, "init", "", "")
	cmdFlags.StringVar(&addResource, "add", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		requireAdminRights("-init", "-add").
		invalid(args) {
		return 2
	}

	if len(initCluster) != 0 {
		this.initHelixCluster(initCluster, ctx.ZoneZkAddrs(zone))
		return
	}

	if len(addResource) != 0 {
		if len(helixCluster) == 0 {
			this.Ui.Error("-c cluster is required")
			this.Ui.Output(this.Help())
			return
		}

		this.addResource(addResource, helixCluster, ctx.ZoneZkAddrs(zone))
		return
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	for {
		this.checkMyslave(zkzone)
		if !topMode {
			break
		}

		time.Sleep(time.Second * 3)
		refreshScreen()
	}

	return
}

func (this *Dbus) initHelixCluster(cluster string, zkSvr string) {
	admin := hzk.NewZkHelixAdmin(zkSvr)
	swallow(admin.Connect())

	if ok, err := admin.IsClusterSetup(cluster); err != nil {
		this.Ui.Error(err.Error())
		return
	} else if ok {
		this.Ui.Warnf("cluster %s already exists, skipped")
		return
	}

	swallow(admin.AddCluster(cluster))
	swallow(admin.AllowParticipantAutoJoin(cluster, true))
	this.Ui.Infof("cluster %s created", cluster)
}

func (this *Dbus) addResource(resource, cluster, zkSvr string) {
	admin := hzk.NewZkHelixAdmin(zkSvr)
	swallow(admin.Connect())

	if ok, err := admin.IsClusterSetup(cluster); !ok || err != nil {
		this.Ui.Errorf("cluster %s not setup", cluster)
		return
	}

	partitions := 1
	resourceOption := helix.DefaultAddResourceOption(partitions, helix.StateModelOnlineOffline)
	resourceOption.RebalancerMode = helix.RebalancerModeFullAuto
	swallow(admin.AddResource(cluster, resource, resourceOption))
	swallow(admin.Rebalance(cluster, resource, 1))
	this.Ui.Infof("%s for cluster %s added and rebalanced", resource, cluster)
}

func (this *Dbus) checkMyslave(zkzone *zk.ZkZone) {
	lines := []string{"Name|Mysql|File|Offset|dbus|ver|pid|uptime|msince"}
	root := "/dbus/myslave"
	dbs, _, err := zkzone.Conn().Children(root)
	swallow(err)

	sort.Strings(dbs)

	for _, db := range dbs {
		dbRoot := fmt.Sprintf("%s/%s", root, db)
		data, _, err := zkzone.Conn().Get(dbRoot)
		var v myslave.PositionerZk
		if err = json.Unmarshal(data, &v); err != nil {
			v.File = "?"
		}

		ownerPath := fmt.Sprintf("%s/owner", dbRoot)
		owner, ownerStat, err := zkzone.Conn().Get(ownerPath)
		if err != nil {
			if err != zklib.ErrNoNode {
				panic(err)
			}

			// an orphan binlog stream: no dbus consuming it
			lines = append(lines, fmt.Sprintf("-|%s|%s|%s|-|-|-|-|-", db,
				v.File, gofmt.Comma(int64(v.Offset))))
			continue
		}

		idsPath := fmt.Sprintf("%s/ids", dbRoot)
		runners, _, err := zkzone.Conn().Children(idsPath)
		swallow(err)

		for _, r := range runners {
			ver, stat, err := zkzone.Conn().Get(fmt.Sprintf("%s/%s", idsPath, r))
			swallow(err)

			i := strings.LastIndex(r, "-")
			instance := r[:i]
			pid := r[i+1:]
			uptime := zk.ZkTimestamp(stat.Ctime).Time()

			if r == string(owner) {
				ownerTime := zk.ZkTimestamp(ownerStat.Mtime).Time()
				lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s|%s|%s|%s|%s",
					v.Name, db, v.File, gofmt.Comma(int64(v.Offset)), instance, string(ver), pid,
					gofmt.PrettySince(uptime), gofmt.PrettySince(ownerTime)))
			} else {
				lines = append(lines, fmt.Sprintf("%s|%s| | |%s|%s|%s|%s| ",
					v.Name, db, instance, string(ver), pid, gofmt.PrettySince(uptime)))
			}
		}

	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (*Dbus) Synopsis() string {
	return "Monitor database system runtime"
}

func (this *Dbus) Help() string {
	help := fmt.Sprintf(`
Usage: %s dbus [options]

    %s

Options:

    -z zone

    -init cluster
      Initialize helix cluster.

    -add mysql_host:mysql_port
      Add a mysql instance as resource to helix.
      Work with -c helix_cluster

    -top
      Run in top mode.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<commit_msg>discard comma fmt: easy to copy to clipboard<commit_after>package command

import (
	"encoding/json"
	"flag"
	"fmt"
	"sort"
	"strings"
	"time"

	"github.com/funkygao/columnize"
	"github.com/funkygao/dbus/pkg/myslave"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/go-helix"
	hzk "github.com/funkygao/go-helix/store/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	zklib "github.com/samuel/go-zookeeper/zk"
)

type Dbus struct {
	Ui  cli.Ui
	Cmd string
}

func (this *Dbus) Run(args []string) (exitCode int) {
	var (
		zone         string
		topMode      bool
		helixCluster string
		initCluster  string
		addResource  string
	)
	cmdFlags := flag.NewFlagSet("dbus", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&helixCluster, "c", "", "")
	cmdFlags.BoolVar(&topMode, "top", false, "")
	cmdFlags.StringVar(&initCluster, "init", "", "")
	cmdFlags.StringVar(&addResource, "add", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		requireAdminRights("-init", "-add").
		invalid(args) {
		return 2
	}

	if len(initCluster) != 0 {
		this.initHelixCluster(initCluster, ctx.ZoneZkAddrs(zone))
		return
	}

	if len(addResource) != 0 {
		if len(helixCluster) == 0 {
			this.Ui.Error("-c cluster is required")
			this.Ui.Output(this.Help())
			return
		}

		this.addResource(addResource, helixCluster, ctx.ZoneZkAddrs(zone))
		return
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	for {
		this.checkMyslave(zkzone)
		if !topMode {
			break
		}

		time.Sleep(time.Second * 3)
		refreshScreen()
	}

	return
}

func (this *Dbus) initHelixCluster(cluster string, zkSvr string) {
	admin := hzk.NewZkHelixAdmin(zkSvr)
	swallow(admin.Connect())

	if ok, err := admin.IsClusterSetup(cluster); err != nil {
		this.Ui.Error(err.Error())
		return
	} else if ok {
		this.Ui.Warnf("cluster %s already exists, skipped")
		return
	}

	swallow(admin.AddCluster(cluster))
	swallow(admin.AllowParticipantAutoJoin(cluster, true))
	this.Ui.Infof("cluster %s created", cluster)
}

func (this *Dbus) addResource(resource, cluster, zkSvr string) {
	admin := hzk.NewZkHelixAdmin(zkSvr)
	swallow(admin.Connect())

	if ok, err := admin.IsClusterSetup(cluster); !ok || err != nil {
		this.Ui.Errorf("cluster %s not setup", cluster)
		return
	}

	partitions := 1
	resourceOption := helix.DefaultAddResourceOption(partitions, helix.StateModelOnlineOffline)
	resourceOption.RebalancerMode = helix.RebalancerModeFullAuto
	swallow(admin.AddResource(cluster, resource, resourceOption))
	swallow(admin.Rebalance(cluster, resource, 1))
	this.Ui.Infof("%s for cluster %s added and rebalanced", resource, cluster)
}

func (this *Dbus) checkMyslave(zkzone *zk.ZkZone) {
	lines := []string{"Name|Mysql|File|Offset|dbus|ver|pid|uptime|msince"}
	root := "/dbus/myslave"
	dbs, _, err := zkzone.Conn().Children(root)
	swallow(err)

	sort.Strings(dbs)

	for _, db := range dbs {
		dbRoot := fmt.Sprintf("%s/%s", root, db)
		data, _, err := zkzone.Conn().Get(dbRoot)
		var v myslave.PositionerZk
		if err = json.Unmarshal(data, &v); err != nil {
			v.File = "?"
		}

		ownerPath := fmt.Sprintf("%s/owner", dbRoot)
		owner, ownerStat, err := zkzone.Conn().Get(ownerPath)
		if err != nil {
			if err != zklib.ErrNoNode {
				panic(err)
			}

			// an orphan binlog stream: no dbus consuming it
			lines = append(lines, fmt.Sprintf("-|%s|%s|%d|-|-|-|-|-", db,
				v.File, v.Offset))
			continue
		}

		idsPath := fmt.Sprintf("%s/ids", dbRoot)
		runners, _, err := zkzone.Conn().Children(idsPath)
		swallow(err)

		for _, r := range runners {
			ver, stat, err := zkzone.Conn().Get(fmt.Sprintf("%s/%s", idsPath, r))
			swallow(err)

			i := strings.LastIndex(r, "-")
			instance := r[:i]
			pid := r[i+1:]
			uptime := zk.ZkTimestamp(stat.Ctime).Time()

			if r == string(owner) {
				ownerTime := zk.ZkTimestamp(ownerStat.Mtime).Time()
				lines = append(lines, fmt.Sprintf("%s|%s|%s|%d|%s|%s|%s|%s|%s",
					v.Name, db, v.File, v.Offset, instance, string(ver), pid,
					gofmt.PrettySince(uptime), gofmt.PrettySince(ownerTime)))
			} else {
				lines = append(lines, fmt.Sprintf("%s|%s| | |%s|%s|%s|%s| ",
					v.Name, db, instance, string(ver), pid, gofmt.PrettySince(uptime)))
			}
		}

	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (*Dbus) Synopsis() string {
	return "Monitor database system runtime"
}

func (this *Dbus) Help() string {
	help := fmt.Sprintf(`
Usage: %s dbus [options]

    %s

Options:

    -z zone

    -init cluster
      Initialize helix cluster.

    -add mysql_host:mysql_port
      Add a mysql instance as resource to helix.
      Work with -c helix_cluster

    -top
      Run in top mode.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<|endoftext|>