<commit_before>// Sample configuration file:
//
//   directory:  /home/casa/tmp
//   email:
//     smtp_host: smtp.orange.fr:25
//     recipient: casa@sweetohm.net
//     sender:    casa@sweetohm.net
//     success:   true
//
//   modules:
//     - name:    continuum
//       url:     https://github.com/c4s4/continuum.git
//       command: |
//         set -e
//         export PATH=/opt/python/current/bin:$PATH
//         virtualenv env --no-site-packages
//         . env/bin/activate
//         pip install -r etc/requirements.txt
//         bee test

package main

import (
	"fmt"
	"gopkg.in/yaml.v1"
	"io/ioutil"
	"net/smtp"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

type Config struct {
	Directory string
	Email     struct {
		SmtpHost  string "smtp_host"
		Recipient string
		Sender    string
		Success   bool
	}
	Modules []struct {
		Name    string
		Url     string
		Command string
	}
}

type Build struct {
	Module  string
	Success bool
	Output  string
}

func (build Build) String() string {
	if build.Success {
		return fmt.Sprintf("%s: OK", build.Module)
	} else {
		return fmt.Sprintf("%s: ERROR", build.Module)
	}
}

type Builds []Build

func (builds Builds) Success() bool {
	for _, module := range builds {
		if !module.Success {
			return false
		}
	}
	return true
}

func (builds Builds) String() string {
	if builds.Success() {
		return "SUCCESS"
	} else {
		return "FAILURE"
	}
}

func loadConfig(file string) Config {
	config := Config{}
	text, err := ioutil.ReadFile(file)
	if err != nil {
		panic(err.Error())
	}
	err = yaml.Unmarshal(text, &config)
	if err != nil {
		panic(err.Error())
	}
	// make directory absolute path
	absdir, err := filepath.Abs(config.Directory)
	if err != nil {
		panic(err.Error())
	}
	config.Directory = absdir
	return config
}

func buildModule(index int, config Config) Build {
	fmt.Printf("Building '%s'... ", config.Modules[index].Name)
	moduleDir := path.Join(config.Directory, config.Modules[index].Name)
	// go in build directory
	err := os.Chdir(config.Directory)
	if err != nil {
		return Build{
			Module:  config.Modules[index].Name,
			Success: false,
			Output:  err.Error(),
		}
	}
	// delete module directory if it already exists
	if _, err := os.Stat(moduleDir); err == nil {
		os.RemoveAll(moduleDir)
	}
	// git clone the module repository
	cmd := exec.Command("git", "clone", config.Modules[index].Url)
	output, err := cmd.CombinedOutput()
	defer os.RemoveAll(moduleDir)
	if err != nil {
		fmt.Println("ERROR")
		return Build{
			Module:  config.Modules[index].Name,
			Success: false,
			Output:  string(output),
		}
	} else {
		os.Chdir(moduleDir)
		// run the build command
		cmd := exec.Command("bash", "-c", config.Modules[index].Command)
		output, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Println("ERROR")
			return Build{
				Module:  config.Modules[index].Name,
				Success: false,
				Output:  strings.TrimSpace(string(output)),
			}
		} else {
			fmt.Println("OK")
			return Build{
				Module:  config.Modules[index].Name,
				Success: true,
				Output:  string(output),
			}
		}
	}
}

func buildModules(config Config) Builds {
	builds := make(Builds, len(config.Modules))
	for index, _ := range config.Modules {
		builds[index] = buildModule(index, config)
	}
	return builds
}

const timeFormat = "2006-01-02 15:04"

func sendReport(builds Builds, start time.Time, duration time.Duration, config Config) {
	if !builds.Success() || (builds.Success() && config.Email.Success) {
		subject := fmt.Sprintf("Build on %s was a %s", start.Format(timeFormat), builds)
		message := fmt.Sprintf("From: %s\n", config.Email.Sender)
		message += fmt.Sprintf("To: %s\n", config.Email.Recipient)
		message += fmt.Sprintf("Subject: %s\n\n", subject)
		message += fmt.Sprintf("Build on %s:\n\n", start.Format(timeFormat))
		for module := range builds {
			message += fmt.Sprintf("  %s\n", builds[module].String())
		}
		message += fmt.Sprintf("\nDone in %s\n", duration)
		message += builds.String()
		for _, build := range builds {
			if !build.Success {
				message += fmt.Sprintf("\n\n===================================\n")
				message += fmt.Sprintf(build.Module)
				message += fmt.Sprintf("\n-----------------------------------\n")
				message += fmt.Sprintf(build.Output)
				message += fmt.Sprintf("\n-----------------------------------\n")
			}
		}
		message += "\n--\ngontinuum"
		err := smtp.SendMail(config.Email.SmtpHost, nil, config.Email.Sender,
			[]string{config.Email.Recipient}, []byte(message))
		if err != nil {
			panic(err)
		}
	}
}

func main() {
	for i := 1; i < len(os.Args); i++ {
		start := time.Now()
		config := loadConfig(os.Args[i])
		builds := buildModules(config)
		duration := time.Since(start)
		fmt.Println("Done in", duration)
		fmt.Println(builds.String())
		sendReport(builds, start, duration, config)
	}
}
<commit_msg>Refactoring<commit_after>// Sample configuration file:
//
//   directory:  /home/casa/tmp
//   email:
//     smtp_host: smtp.orange.fr:25
//     recipient: casa@sweetohm.net
//     sender:    casa@sweetohm.net
//     success:   true
//
//   modules:
//     - name:    continuum
//       url:     https://github.com/c4s4/continuum.git
//       command: |
//         set -e
//         export PATH=/opt/python/current/bin:$PATH
//         virtualenv env --no-site-packages
//         . env/bin/activate
//         pip install -r etc/requirements.txt
//         bee test

package main

import (
	"fmt"
	"gopkg.in/yaml.v1"
	"io/ioutil"
	"net/smtp"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

///////////////////////////////////////////////////////////////////////////////
//                              CONFIG STUFF                                 //
///////////////////////////////////////////////////////////////////////////////

type EmailConfig struct {
	SmtpHost  string "smtp_host"
	Recipient string
	Sender    string
	Success   bool
}

type ModuleConfig struct {
	Name    string
	Url     string
	Command string
}

type Config struct {
	Directory string
	Email     EmailConfig
	Modules   []ModuleConfig
}

func loadConfig(file string) Config {
	config := Config{}
	text, err := ioutil.ReadFile(file)
	if err != nil {
		panic(err.Error())
	}
	err = yaml.Unmarshal(text, &config)
	if err != nil {
		panic(err.Error())
	}
	// make directory absolute path
	absdir, err := filepath.Abs(config.Directory)
	if err != nil {
		panic(err.Error())
	}
	config.Directory = absdir
	return config
}

///////////////////////////////////////////////////////////////////////////////
//                                BUILD STUFF                                //
///////////////////////////////////////////////////////////////////////////////

type Build struct {
	Module  ModuleConfig
	Success bool
	Output  string
}

type Builds []Build

func (build Build) String() string {
	if build.Success {
		return fmt.Sprintf("%s: OK", build.Module.Name)
	} else {
		return fmt.Sprintf("%s: ERROR", build.Module.Name)
	}
}

func (builds Builds) Success() bool {
	for _, build := range builds {
		if !build.Success {
			return false
		}
	}
	return true
}

func (builds Builds) String() string {
	if builds.Success() {
		return "SUCCESS"
	} else {
		return "FAILURE"
	}
}

func buildModule(module ModuleConfig, directory string) Build {
	fmt.Printf("Building '%s'... ", module.Name)
	moduleDir := path.Join(directory, module.Name)
	// go in build directory
	err := os.Chdir(directory)
	if err != nil {
		return Build{
			Module:  module,
			Success: false,
			Output:  err.Error(),
		}
	}
	// delete module directory if it already exists
	if _, err := os.Stat(moduleDir); err == nil {
		os.RemoveAll(moduleDir)
	}
	// git clone the module repository
	cmd := exec.Command("git", "clone", module.Url)
	output, err := cmd.CombinedOutput()
	defer os.RemoveAll(moduleDir)
	if err != nil {
		fmt.Println("ERROR")
		return Build{
			Module:  module,
			Success: false,
			Output:  string(output),
		}
	} else {
		os.Chdir(moduleDir)
		// run the build command
		cmd := exec.Command("bash", "-c", module.Command)
		output, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Println("ERROR")
			return Build{
				Module:  module,
				Success: false,
				Output:  strings.TrimSpace(string(output)),
			}
		} else {
			fmt.Println("OK")
			return Build{
				Module:  module,
				Success: true,
				Output:  string(output),
			}
		}
	}
}

func buildModules(config Config) Builds {
	builds := make(Builds, len(config.Modules))
	for index, module := range config.Modules {
		builds[index] = buildModule(module, config.Directory)
	}
	return builds
}

const timeFormat = "2006-01-02 15:04"

func sendEmail(builds Builds, start time.Time, duration time.Duration, config Config) {
	if !builds.Success() || (builds.Success() && config.Email.Success) {
		subject := fmt.Sprintf("Build on %s was a %s", start.Format(timeFormat), builds)
		message := fmt.Sprintf("From: %s\n", config.Email.Sender)
		message += fmt.Sprintf("To: %s\n", config.Email.Recipient)
		message += fmt.Sprintf("Subject: %s\n\n", subject)
		message += fmt.Sprintf("Build on %s:\n\n", start.Format(timeFormat))
		for _, build := range builds {
			message += fmt.Sprintf("  %s\n", build.String())
		}
		message += fmt.Sprintf("\nDone in %s\n", duration)
		message += builds.String()
		for _, build := range builds {
			if !build.Success {
				message += fmt.Sprintf("\n\n===================================\n")
				message += fmt.Sprintf(build.Module.Name)
				message += fmt.Sprintf("\n-----------------------------------\n")
				message += fmt.Sprintf(build.Output)
				message += fmt.Sprintf("\n-----------------------------------\n")
			}
		}
		message += "\n--\ngontinuum"
		err := smtp.SendMail(config.Email.SmtpHost, nil, config.Email.Sender,
			[]string{config.Email.Recipient}, []byte(message))
		if err != nil {
			panic(err)
		}
	}
}

func main() {
	for i := 1; i < len(os.Args); i++ {
		start := time.Now()
		config := loadConfig(os.Args[i])
		builds := buildModules(config)
		duration := time.Since(start)
		fmt.Println("Done in", duration)
		fmt.Println(builds.String())
		sendEmail(builds, start, duration, config)
	}
}
<|endoftext|>


<commit_before>/*
Package detour provides a net.Conn interface which detects blockage
of a site automatically and access it through alternative connection.

Basically, if a site is not whitelisted, following steps will be taken:
1. Dial proxied connection (detour) a small delay after dialed directly
2. Return to caller when any connection is established
3. Read/write through all open connections in parallel
4. Check for blockage on direct connection and closes it if it happens
5. If possible, replay operations on detour connection. [1]
6. After sucessfully read from a connection, stick with it and close others.
7. Add those sites failed on direct connection but succeeded on detour ones
   to proxied list, so above steps can be skipped next time. The list can be
   exported and persisted if required.

Blockage can happen at several stages of a connection, what detour can detect are:
1. Connection attempt is blocked (IP blocking / DNS hijack).
   Symptoms can be connection time out / TCP RST / connection refused.
2. Connection made but real data get blocked (DPI).
3. Successfully exchanged a few packets, while follow up packets are blocked. [2]
4. Connection made but get fake response or HTTP redirect to a fixed URL.

[1] Detour will not replay nonidempotent plain HTTP requests, but will add it to
    proxied list to be detoured next time.
[2] Detour can only handle exact 1 successful read followed by failed read,
    which covers most cases in reality.
*/
package detour

import (
	"bytes"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/getlantern/golog"
)

// If no any connection made after this period, stop dialing and fail
var TimeoutToConnect = 30 * time.Second

// To avoid unnecessarily proxy not-blocked url, detour will dial proxy connection
// after this small delay
var DelayBeforeDetour = 1 * time.Second

// If DirectAddrCh is set, when a direct connection is closed without any error,
// the connection's remote address (in host:port format) will be send to it
var DirectAddrCh chan string

var (
	log = golog.LoggerFor("detour")
)

// Conn implements an net.Conn interface by utilizing underlie direct and
// detour connections.
type Conn struct {
	// The underlie connections, uses buffered channel as ring queue to avoid
	// locking. We have at most 2 connetions so a length of 2 is enough.
	conns chan conn

	// The chan to notify dialer to dial detour immediately
	chDialDetourNow chan bool
	// The channel to notify read/write that a detour connection is available
	chDetourConn chan conn

	// The chan to receive result of any read operation
	chRead chan ioResult
	// The chan to receive result of any write operation
	chWrite chan ioResult

	// Keeps track of the total bytes read from this connection, atomic
	readBytes uint64

	addr string

	muWriteBuffer sync.RWMutex
	// Keeps written bytes through direct connection to replay it if required.
	writeBuffer *bytes.Buffer
	// Is it a plain HTTP request or not, atomic
	nonIdempotentHTTPRequest uint32
}

// The data structure to pass result of io operation back from underlie connection
type ioResult struct {
	// Number of bytes read/wrote
	n int
	// IO error, if any
	err error
	// The underlie connection itself
	conn conn
}

type connType int

const (
	connTypeDirect connType = iota
	connTypeDetour connType = iota
)

type conn interface {
	ConnType() connType
	FirstRead(b []byte, ch chan ioResult)
	FollowupRead(b []byte, ch chan ioResult)
	Write(b []byte, ch chan ioResult)
	Close() error
	Closed() bool
}

func typeOf(c conn) string {
	var connTypeDesc = []string{"direct", "detour"}
	return connTypeDesc[c.ConnType()]
}

type dialFunc func(network, addr string) (net.Conn, error)

// Dialer returns a function with same signature of net.Dialer.Dial().
func Dialer(detourDialer dialFunc) func(network, addr string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		dc := &Conn{
			addr:            addr,
			writeBuffer:     new(bytes.Buffer),
			conns:           make(chan conn, 2),
			chDetourConn:    make(chan conn),
			chRead:          make(chan ioResult),
			chWrite:         make(chan ioResult),
			chDialDetourNow: make(chan bool),
		}
		// use buffered channel as we may send twice to it but only receive once
		chAnyConn := make(chan bool, 1)
		ch := make(chan conn)
		// dialing sequence
		if whitelisted(addr) {
			dialDetour(network, addr, detourDialer, ch)
		} else {
			go func() {
				dialDirect(network, addr, ch)
				dt := time.NewTimer(DelayBeforeDetour)
				select {
				case <-dt.C:
				case <-dc.chDialDetourNow:
				}
				if dc.anyDataReceived() {
					ch <- nil
					return
				}
				dialDetour(network, addr, detourDialer, ch)
			}()
		}
		// handle dialing result
		go func() {
			t := time.NewTimer(TimeoutToConnect)
			defer t.Stop()
			for i := 0; i < 2; i++ {
				log.Tracef("Waiting for connection to %s, round %d", dc.addr, i)
				select {
				case c := <-ch:
					if c == nil {
						log.Tracef("No new connection to %s remaining, return", dc.addr)
						return
					}
					// first connection made, pass it back to caller
					if i == 0 {
						dc.conns <- c
						chAnyConn <- true
					} else {
						if c.ConnType() == connTypeDirect {
							// Could happen if direct route is much slower.
							log.Debugf("Direct connection to %s established too late, close it", dc.addr)
							_ = c.Close()
							return
						}
						log.Tracef("Feed detour connection to %s to read/write op", dc.addr)
						dc.chDetourConn <- c
					}
				case <-t.C:
					// still no connection made
					chAnyConn <- false
					return
				}
			}
		}()
		// return to caller if any connection available
		if anyConn := <-chAnyConn; anyConn {
			return dc, nil
		}
		return nil, fmt.Errorf("Timeout dialing any connection to %s", addr)
	}
}

func (dc *Conn) anyDataReceived() bool {
	return atomic.LoadUint64(&dc.readBytes) > 0
}

func (dc *Conn) incReadBytes(n int) {
	atomic.AddUint64(&dc.readBytes, uint64(n))
}

// Read() implements the function from net.Conn
func (dc *Conn) Read(b []byte) (n int, err error) {
	log.Tracef("Initiate a read request to %s", dc.addr)
	if dc.anyDataReceived() {
		return dc.followupRead(b)
	}
	// At initial stage, we only have one connection,
	// but detour connection can be available at anytime.
	if !dc.withValidConn(func(c conn) { c.FirstRead(b, dc.chRead) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	for count := 1; count > 0; count-- {
		select {
		case newConn := <-dc.chDetourConn:
			if atomic.LoadUint32(&dc.nonIdempotentHTTPRequest) == 1 {
				log.Tracef("Not replay nonideompotent request to %s, only add to whitelist", dc.addr)
				AddToWl(dc.addr, false)
				_ = newConn.Close()
				return
			}
			log.Tracef("Got detour connection to %s, replay previous op on it", dc.addr)
			dc.muWriteBuffer.RLock()
			sentBytes := dc.writeBuffer.Bytes()
			dc.muWriteBuffer.RUnlock()
			newConn.Write(sentBytes, dc.chWrite)
			newConn.FirstRead(b, dc.chRead)
			count++
			// add new connection to connections
			dc.conns <- newConn
		case result := <-dc.chRead:
			conn, n, err := result.conn, result.n, result.err
			if err != nil {
				log.Tracef("Read from %s connection to %s failed, closing: %s", typeOf(conn), dc.addr, err)
				_ = conn.Close()
				// skip failed connection as we have more
				if count > 1 {
					continue
				}
				switch conn.ConnType() {
				case connTypeDirect:
					// if we haven't dial detour yet, do so now
					select {
					case dc.chDialDetourNow <- true:
						count++
					default:
					}
					continue
				case connTypeDetour:
					log.Tracef("Detour connection to %s failed, removing from whitelist", dc.addr)
					RemoveFromWl(dc.addr)
					// no more connections, return directly to avoid dead lock
					return n, err
				}
			}
			log.Tracef("Read %d bytes from %s connection to %s", n, typeOf(conn), dc.addr)
			dc.incReadBytes(n)
			return n, err
		}
	}
	return
}

// followUpRead is called by Read() if a connection's state already settled
func (dc *Conn) followupRead(b []byte) (n int, err error) {
	if !dc.withValidConn(func(c conn) { c.FollowupRead(b, dc.chRead) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	result := <-dc.chRead
	dc.incReadBytes(result.n)
	return result.n, result.err
}

// Write() implements the function from net.Conn
func (dc *Conn) Write(b []byte) (n int, err error) {
	log.Tracef("Initiate a write request to %s", dc.addr)
	if dc.anyDataReceived() {
		return dc.followupWrite(b)
	}
	if isNonIdempotentHTTPRequest(b) {
		atomic.StoreUint32(&dc.nonIdempotentHTTPRequest, 1)
	} else {
		dc.muWriteBuffer.Lock()
		_, _ = dc.writeBuffer.Write(b)
		dc.muWriteBuffer.Unlock()
	}
	if !dc.withValidConn(func(c conn) { c.Write(b, dc.chWrite) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}

	result := <-dc.chWrite
	if n, err = result.n, result.err; err != nil {
		log.Tracef("Error writing %s connection to %s: %s", typeOf(result.conn), dc.addr, err)
		_ = result.conn.Close()
		return
	}
	log.Tracef("Wrote %d bytes to %s connection to %s", n, typeOf(result.conn), dc.addr)
	return
}

// followupWrite is called by Write() if a connection's state already settled
func (dc *Conn) followupWrite(b []byte) (n int, err error) {
	if !dc.withValidConn(func(c conn) { c.Write(b, dc.chWrite) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	result := <-dc.chWrite
	return result.n, result.err
}

// Close implements the function from net.Conn
func (dc *Conn) Close() error {
	log.Tracef("Closing connection to %s", dc.addr)
	for len(dc.conns) > 0 {
		conn := <-dc.conns
		_ = conn.Close()
	}
	return nil
}

// LocalAddr implements the function from net.Conn
func (dc *Conn) LocalAddr() net.Addr {
	return nil
}

// RemoteAddr implements the function from net.Conn
func (dc *Conn) RemoteAddr() net.Addr {
	return nil
}

// SetDeadline implements the function from net.Conn
func (dc *Conn) SetDeadline(t time.Time) error {
	if err := dc.SetReadDeadline(t); err != nil {
		log.Debugf("Unable to set read deadline: %v", err)
	}
	if err := dc.SetWriteDeadline(t); err != nil {
		log.Debugf("Unable to set write deadline: %v", err)
	}
	return nil
}

// SetReadDeadline implements the function from net.Conn
func (dc *Conn) SetReadDeadline(t time.Time) error {
	return nil
}

// SetWriteDeadline implements the function from net.Conn
func (dc *Conn) SetWriteDeadline(t time.Time) error {
	return nil
}

func (dc *Conn) withValidConn(f func(conn)) bool {
	for i := 0; i < len(dc.conns); i++ {
		select {
		case c := <-dc.conns:
			if c.Closed() {
				log.Tracef("Drain closed %s connection to %s", typeOf(c), dc.addr)
				continue
			}
			f(c)
			dc.conns <- c
			return true
		default:
			break
		}
	}
	return false
}

var nonIdempotentMethods = [][]byte{
	[]byte("POST "),
	[]byte("PATCH "),
}

// Ref section 9.1.2 of https://www.ietf.org/rfc/rfc2616.txt.
// We consider the https handshake phase to be idemponent.
func isNonIdempotentHTTPRequest(b []byte) bool {
	if len(b) > 4 {
		for _, m := range nonIdempotentMethods {
			if bytes.HasPrefix(b, m) {
				return true
			}
		}
	}
	return false
}
<commit_msg>add trace for not implemented detour methods<commit_after>/*
Package detour provides a net.Conn interface which detects blockage
of a site automatically and access it through alternative connection.

Basically, if a site is not whitelisted, following steps will be taken:
1. Dial proxied connection (detour) a small delay after dialed directly
2. Return to caller when any connection is established
3. Read/write through all open connections in parallel
4. Check for blockage on direct connection and closes it if it happens
5. If possible, replay operations on detour connection. [1]
6. After sucessfully read from a connection, stick with it and close others.
7. Add those sites failed on direct connection but succeeded on detour ones
   to proxied list, so above steps can be skipped next time. The list can be
   exported and persisted if required.

Blockage can happen at several stages of a connection, what detour can detect are:
1. Connection attempt is blocked (IP blocking / DNS hijack).
   Symptoms can be connection time out / TCP RST / connection refused.
2. Connection made but real data get blocked (DPI).
3. Successfully exchanged a few packets, while follow up packets are blocked. [2]
4. Connection made but get fake response or HTTP redirect to a fixed URL.

[1] Detour will not replay nonidempotent plain HTTP requests, but will add it to
    proxied list to be detoured next time.
[2] Detour can only handle exact 1 successful read followed by failed read,
    which covers most cases in reality.
*/
package detour

import (
	"bytes"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/getlantern/golog"
)

// If no any connection made after this period, stop dialing and fail
var TimeoutToConnect = 30 * time.Second

// To avoid unnecessarily proxy not-blocked url, detour will dial proxy connection
// after this small delay
var DelayBeforeDetour = 1 * time.Second

// If DirectAddrCh is set, when a direct connection is closed without any error,
// the connection's remote address (in host:port format) will be send to it
var DirectAddrCh chan string

var (
	log = golog.LoggerFor("detour")
)

// Conn implements an net.Conn interface by utilizing underlie direct and
// detour connections.
type Conn struct {
	// The underlie connections, uses buffered channel as ring queue to avoid
	// locking. We have at most 2 connetions so a length of 2 is enough.
	conns chan conn

	// The chan to notify dialer to dial detour immediately
	chDialDetourNow chan bool
	// The channel to notify read/write that a detour connection is available
	chDetourConn chan conn

	// The chan to receive result of any read operation
	chRead chan ioResult
	// The chan to receive result of any write operation
	chWrite chan ioResult

	// Keeps track of the total bytes read from this connection, atomic
	readBytes uint64

	addr string

	muWriteBuffer sync.RWMutex
	// Keeps written bytes through direct connection to replay it if required.
	writeBuffer *bytes.Buffer
	// Is it a plain HTTP request or not, atomic
	nonIdempotentHTTPRequest uint32
}

// The data structure to pass result of io operation back from underlie connection
type ioResult struct {
	// Number of bytes read/wrote
	n int
	// IO error, if any
	err error
	// The underlie connection itself
	conn conn
}

type connType int

const (
	connTypeDirect connType = iota
	connTypeDetour connType = iota
)

type conn interface {
	ConnType() connType
	FirstRead(b []byte, ch chan ioResult)
	FollowupRead(b []byte, ch chan ioResult)
	Write(b []byte, ch chan ioResult)
	Close() error
	Closed() bool
}

func typeOf(c conn) string {
	var connTypeDesc = []string{"direct", "detour"}
	return connTypeDesc[c.ConnType()]
}

type dialFunc func(network, addr string) (net.Conn, error)

// Dialer returns a function with same signature of net.Dialer.Dial().
func Dialer(detourDialer dialFunc) func(network, addr string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		dc := &Conn{
			addr:            addr,
			writeBuffer:     new(bytes.Buffer),
			conns:           make(chan conn, 2),
			chDetourConn:    make(chan conn),
			chRead:          make(chan ioResult),
			chWrite:         make(chan ioResult),
			chDialDetourNow: make(chan bool),
		}
		// use buffered channel as we may send twice to it but only receive once
		chAnyConn := make(chan bool, 1)
		ch := make(chan conn)
		// dialing sequence
		if whitelisted(addr) {
			dialDetour(network, addr, detourDialer, ch)
		} else {
			go func() {
				dialDirect(network, addr, ch)
				dt := time.NewTimer(DelayBeforeDetour)
				select {
				case <-dt.C:
				case <-dc.chDialDetourNow:
				}
				if dc.anyDataReceived() {
					ch <- nil
					return
				}
				dialDetour(network, addr, detourDialer, ch)
			}()
		}
		// handle dialing result
		go func() {
			t := time.NewTimer(TimeoutToConnect)
			defer t.Stop()
			for i := 0; i < 2; i++ {
				log.Tracef("Waiting for connection to %s, round %d", dc.addr, i)
				select {
				case c := <-ch:
					if c == nil {
						log.Tracef("No new connection to %s remaining, return", dc.addr)
						return
					}
					// first connection made, pass it back to caller
					if i == 0 {
						dc.conns <- c
						chAnyConn <- true
					} else {
						if c.ConnType() == connTypeDirect {
							// Could happen if direct route is much slower.
							log.Debugf("Direct connection to %s established too late, close it", dc.addr)
							_ = c.Close()
							return
						}
						log.Tracef("Feed detour connection to %s to read/write op", dc.addr)
						dc.chDetourConn <- c
					}
				case <-t.C:
					// still no connection made
					chAnyConn <- false
					return
				}
			}
		}()
		// return to caller if any connection available
		if anyConn := <-chAnyConn; anyConn {
			return dc, nil
		}
		return nil, fmt.Errorf("Timeout dialing any connection to %s", addr)
	}
}

func (dc *Conn) anyDataReceived() bool {
	return atomic.LoadUint64(&dc.readBytes) > 0
}

func (dc *Conn) incReadBytes(n int) {
	atomic.AddUint64(&dc.readBytes, uint64(n))
}

// Read() implements the function from net.Conn
func (dc *Conn) Read(b []byte) (n int, err error) {
	log.Tracef("Initiate a read request to %s", dc.addr)
	if dc.anyDataReceived() {
		return dc.followupRead(b)
	}
	// At initial stage, we only have one connection,
	// but detour connection can be available at anytime.
	if !dc.withValidConn(func(c conn) { c.FirstRead(b, dc.chRead) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	for count := 1; count > 0; count-- {
		select {
		case newConn := <-dc.chDetourConn:
			if atomic.LoadUint32(&dc.nonIdempotentHTTPRequest) == 1 {
				log.Tracef("Not replay nonideompotent request to %s, only add to whitelist", dc.addr)
				AddToWl(dc.addr, false)
				_ = newConn.Close()
				return
			}
			log.Tracef("Got detour connection to %s, replay previous op on it", dc.addr)
			dc.muWriteBuffer.RLock()
			sentBytes := dc.writeBuffer.Bytes()
			dc.muWriteBuffer.RUnlock()
			newConn.Write(sentBytes, dc.chWrite)
			newConn.FirstRead(b, dc.chRead)
			count++
			// add new connection to connections
			dc.conns <- newConn
		case result := <-dc.chRead:
			conn, n, err := result.conn, result.n, result.err
			if err != nil {
				log.Tracef("Read from %s connection to %s failed, closing: %s", typeOf(conn), dc.addr, err)
				_ = conn.Close()
				// skip failed connection as we have more
				if count > 1 {
					continue
				}
				switch conn.ConnType() {
				case connTypeDirect:
					// if we haven't dial detour yet, do so now
					select {
					case dc.chDialDetourNow <- true:
						count++
					default:
					}
					continue
				case connTypeDetour:
					log.Tracef("Detour connection to %s failed, removing from whitelist", dc.addr)
					RemoveFromWl(dc.addr)
					// no more connections, return directly to avoid dead lock
					return n, err
				}
			}
			log.Tracef("Read %d bytes from %s connection to %s", n, typeOf(conn), dc.addr)
			dc.incReadBytes(n)
			return n, err
		}
	}
	return
}

// followUpRead is called by Read() if a connection's state already settled
func (dc *Conn) followupRead(b []byte) (n int, err error) {
	if !dc.withValidConn(func(c conn) { c.FollowupRead(b, dc.chRead) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	result := <-dc.chRead
	dc.incReadBytes(result.n)
	return result.n, result.err
}

// Write() implements the function from net.Conn
func (dc *Conn) Write(b []byte) (n int, err error) {
	log.Tracef("Initiate a write request to %s", dc.addr)
	if dc.anyDataReceived() {
		return dc.followupWrite(b)
	}
	if isNonIdempotentHTTPRequest(b) {
		atomic.StoreUint32(&dc.nonIdempotentHTTPRequest, 1)
	} else {
		dc.muWriteBuffer.Lock()
		_, _ = dc.writeBuffer.Write(b)
		dc.muWriteBuffer.Unlock()
	}
	if !dc.withValidConn(func(c conn) { c.Write(b, dc.chWrite) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}

	result := <-dc.chWrite
	if n, err = result.n, result.err; err != nil {
		log.Tracef("Error writing %s connection to %s: %s", typeOf(result.conn), dc.addr, err)
		_ = result.conn.Close()
		return
	}
	log.Tracef("Wrote %d bytes to %s connection to %s", n, typeOf(result.conn), dc.addr)
	return
}

// followupWrite is called by Write() if a connection's state already settled
func (dc *Conn) followupWrite(b []byte) (n int, err error) {
	if !dc.withValidConn(func(c conn) { c.Write(b, dc.chWrite) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	result := <-dc.chWrite
	return result.n, result.err
}

// Close implements the function from net.Conn
func (dc *Conn) Close() error {
	log.Tracef("Closing connection to %s", dc.addr)
	for len(dc.conns) > 0 {
		conn := <-dc.conns
		_ = conn.Close()
	}
	return nil
}

// LocalAddr implements the function from net.Conn
func (dc *Conn) LocalAddr() net.Addr {
	log.Trace("LocalAddr not implemented")
	return nil
}

// RemoteAddr implements the function from net.Conn
func (dc *Conn) RemoteAddr() net.Addr {
	log.Trace("RemoteAddr not implemented")
	return nil
}

// SetDeadline implements the function from net.Conn
func (dc *Conn) SetDeadline(t time.Time) error {
	log.Trace("SetDeadline not implemented")
	return nil
}

// SetReadDeadline implements the function from net.Conn
func (dc *Conn) SetReadDeadline(t time.Time) error {
	log.Trace("SetReadDeadline not implemented")
	return nil
}

// SetWriteDeadline implements the function from net.Conn
func (dc *Conn) SetWriteDeadline(t time.Time) error {
	log.Trace("SetWriteDeadline not implemented")
	return nil
}

func (dc *Conn) withValidConn(f func(conn)) bool {
	for i := 0; i < len(dc.conns); i++ {
		select {
		case c := <-dc.conns:
			if c.Closed() {
				log.Tracef("Drain closed %s connection to %s", typeOf(c), dc.addr)
				continue
			}
			f(c)
			dc.conns <- c
			return true
		default:
			break
		}
	}
	return false
}

var nonIdempotentMethods = [][]byte{
	[]byte("POST "),
	[]byte("PATCH "),
}

// Ref section 9.1.2 of https://www.ietf.org/rfc/rfc2616.txt.
// We consider the https handshake phase to be idemponent.
func isNonIdempotentHTTPRequest(b []byte) bool {
	if len(b) > 4 {
		for _, m := range nonIdempotentMethods {
			if bytes.HasPrefix(b, m) {
				return true
			}
		}
	}
	return false
}
<|endoftext|>


<commit_before>// +build linux

package cpu

import (
	"errors"
	"fmt"
	"os/exec"
	"strconv"
	"strings"

	"github.com/shirou/gopsutil/internal/common"
)

var cpu_tick = float64(100)

func init() {
	getconf, err := exec.LookPath("/usr/bin/getconf")
	if err != nil {
		return
	}
	out, err := invoke.Command(getconf, "CLK_TCK")
	// ignore errors
	if err == nil {
		i, err := strconv.ParseFloat(strings.TrimSpace(string(out)), 64)
		if err == nil {
			cpu_tick = float64(i)
		}
	}
}

func Times(percpu bool) ([]TimesStat, error) {
	filename := common.HostProc("stat")
	var lines = []string{}
	if percpu {
		var startIdx uint = 1
		for {
			linen, _ := common.ReadLinesOffsetN(filename, startIdx, 1)
			line := linen[0]
			if !strings.HasPrefix(line, "cpu") {
				break
			}
			lines = append(lines, line)
			startIdx++
		}
	} else {
		lines, _ = common.ReadLinesOffsetN(filename, 0, 1)
	}

	ret := make([]TimesStat, 0, len(lines))

	for _, line := range lines {
		ct, err := parseStatLine(line)
		if err != nil {
			continue
		}
		ret = append(ret, *ct)

	}
	return ret, nil
}

func sysCPUPath(cpu int32, relPath string) string {
	return common.HostSys(fmt.Sprintf("devices/system/cpu/cpu%d", cpu), relPath)
}

func finishCPUInfo(c *InfoStat) error {
	var lines []string
	var err error
	var value float64

	if len(c.CoreID) == 0 {
		lines, err = common.ReadLines(sysCPUPath(c.CPU, "topology/core_id"))
		if err == nil {
			c.CoreID = lines[0]
		}
	}

	// override the value of c.Mhz with cpufreq/cpuinfo_max_freq regardless
	// of the value from /proc/cpuinfo because we want to report the maximum
	// clock-speed of the CPU for c.Mhz, matching the behaviour of Windows
	lines, err = common.ReadLines(sysCPUPath(c.CPU, "cpufreq/cpuinfo_max_freq"))
	// if we encounter errors below but has a value from parsing /proc/cpuinfo
	// then we ignore the error
	if err != nil {
		if c.Mhz == 0 {
			return err
		} else {
			return nil
		}
	}
	value, err = strconv.ParseFloat(lines[0], 64)
	if err != nil {
		if c.Mhz == 0 {
			return err
		} else {
			return nil
		}
	}
	c.Mhz = value/1000.0  // value is in kHz
	return nil
}

// CPUInfo on linux will return 1 item per physical thread.
//
// CPUs have three levels of counting: sockets, cores, threads.
// Cores with HyperThreading count as having 2 threads per core.
// Sockets often come with many physical CPU cores.
// For example a single socket board with two cores each with HT will
// return 4 CPUInfoStat structs on Linux and the "Cores" field set to 1.
func Info() ([]InfoStat, error) {
	filename := common.HostProc("cpuinfo")
	lines, _ := common.ReadLines(filename)

	var ret []InfoStat

	c := InfoStat{CPU: -1, Cores: 1}
	for _, line := range lines {
		fields := strings.Split(line, ":")
		if len(fields) < 2 {
			continue
		}
		key := strings.TrimSpace(fields[0])
		value := strings.TrimSpace(fields[1])

		switch key {
		case "processor":
			if c.CPU >= 0 {
				err := finishCPUInfo(&c)
				if err != nil {
					return ret, err
				}
				ret = append(ret, c)
			}
			c = InfoStat{Cores: 1}
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.CPU = int32(t)
		case "vendorId", "vendor_id":
			c.VendorID = value
		case "cpu family":
			c.Family = value
		case "model":
			c.Model = value
		case "model name":
			c.ModelName = value
		case "stepping":
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.Stepping = int32(t)
		case "cpu MHz":
			// treat this as the fallback value, thus we ignore error
			if t, err := strconv.ParseFloat(value, 64); err == nil {
				c.Mhz = t
			}
		case "cache size":
			t, err := strconv.ParseInt(strings.Replace(value, " KB", "", 1), 10, 64)
			if err != nil {
				return ret, err
			}
			c.CacheSize = int32(t)
		case "physical id":
			c.PhysicalID = value
		case "core id":
			c.CoreID = value
		case "flags", "Features":
			c.Flags = strings.FieldsFunc(value, func(r rune) bool {
				return r == ',' || r == ' '
			})
		}
	}
	if c.CPU >= 0 {
		err := finishCPUInfo(&c)
		if err != nil {
			return ret, err
		}
		ret = append(ret, c)
	}
	return ret, nil
}

func parseStatLine(line string) (*TimesStat, error) {
	fields := strings.Fields(line)

	if strings.HasPrefix(fields[0], "cpu") == false {
		//		return CPUTimesStat{}, e
		return nil, errors.New("not contain cpu")
	}

	cpu := fields[0]
	if cpu == "cpu" {
		cpu = "cpu-total"
	}
	user, err := strconv.ParseFloat(fields[1], 64)
	if err != nil {
		return nil, err
	}
	nice, err := strconv.ParseFloat(fields[2], 64)
	if err != nil {
		return nil, err
	}
	system, err := strconv.ParseFloat(fields[3], 64)
	if err != nil {
		return nil, err
	}
	idle, err := strconv.ParseFloat(fields[4], 64)
	if err != nil {
		return nil, err
	}
	iowait, err := strconv.ParseFloat(fields[5], 64)
	if err != nil {
		return nil, err
	}
	irq, err := strconv.ParseFloat(fields[6], 64)
	if err != nil {
		return nil, err
	}
	softirq, err := strconv.ParseFloat(fields[7], 64)
	if err != nil {
		return nil, err
	}

	ct := &TimesStat{
		CPU:     cpu,
		User:    float64(user) / cpu_tick,
		Nice:    float64(nice) / cpu_tick,
		System:  float64(system) / cpu_tick,
		Idle:    float64(idle) / cpu_tick,
		Iowait:  float64(iowait) / cpu_tick,
		Irq:     float64(irq) / cpu_tick,
		Softirq: float64(softirq) / cpu_tick,
	}
	if len(fields) > 8 { // Linux >= 2.6.11
		steal, err := strconv.ParseFloat(fields[8], 64)
		if err != nil {
			return nil, err
		}
		ct.Steal = float64(steal) / cpu_tick
	}
	if len(fields) > 9 { // Linux >= 2.6.24
		guest, err := strconv.ParseFloat(fields[9], 64)
		if err != nil {
			return nil, err
		}
		ct.Guest = float64(guest) / cpu_tick
	}
	if len(fields) > 10 { // Linux >= 3.2.0
		guestNice, err := strconv.ParseFloat(fields[10], 64)
		if err != nil {
			return nil, err
		}
		ct.GuestNice = float64(guestNice) / cpu_tick
	}

	return ct, nil
}
<commit_msg>Allow clock field in /proc/cpuinfo as cpu MHz fallback value<commit_after>// +build linux

package cpu

import (
	"errors"
	"fmt"
	"os/exec"
	"strconv"
	"strings"

	"github.com/shirou/gopsutil/internal/common"
)

var cpu_tick = float64(100)

func init() {
	getconf, err := exec.LookPath("/usr/bin/getconf")
	if err != nil {
		return
	}
	out, err := invoke.Command(getconf, "CLK_TCK")
	// ignore errors
	if err == nil {
		i, err := strconv.ParseFloat(strings.TrimSpace(string(out)), 64)
		if err == nil {
			cpu_tick = float64(i)
		}
	}
}

func Times(percpu bool) ([]TimesStat, error) {
	filename := common.HostProc("stat")
	var lines = []string{}
	if percpu {
		var startIdx uint = 1
		for {
			linen, _ := common.ReadLinesOffsetN(filename, startIdx, 1)
			line := linen[0]
			if !strings.HasPrefix(line, "cpu") {
				break
			}
			lines = append(lines, line)
			startIdx++
		}
	} else {
		lines, _ = common.ReadLinesOffsetN(filename, 0, 1)
	}

	ret := make([]TimesStat, 0, len(lines))

	for _, line := range lines {
		ct, err := parseStatLine(line)
		if err != nil {
			continue
		}
		ret = append(ret, *ct)

	}
	return ret, nil
}

func sysCPUPath(cpu int32, relPath string) string {
	return common.HostSys(fmt.Sprintf("devices/system/cpu/cpu%d", cpu), relPath)
}

func finishCPUInfo(c *InfoStat) error {
	var lines []string
	var err error
	var value float64

	if len(c.CoreID) == 0 {
		lines, err = common.ReadLines(sysCPUPath(c.CPU, "topology/core_id"))
		if err == nil {
			c.CoreID = lines[0]
		}
	}

	// override the value of c.Mhz with cpufreq/cpuinfo_max_freq regardless
	// of the value from /proc/cpuinfo because we want to report the maximum
	// clock-speed of the CPU for c.Mhz, matching the behaviour of Windows
	lines, err = common.ReadLines(sysCPUPath(c.CPU, "cpufreq/cpuinfo_max_freq"))
	// if we encounter errors below but has a value from parsing /proc/cpuinfo
	// then we ignore the error
	if err != nil {
		if c.Mhz == 0 {
			return err
		} else {
			return nil
		}
	}
	value, err = strconv.ParseFloat(lines[0], 64)
	if err != nil {
		if c.Mhz == 0 {
			return err
		} else {
			return nil
		}
	}
	c.Mhz = value/1000.0  // value is in kHz
	return nil
}

// CPUInfo on linux will return 1 item per physical thread.
//
// CPUs have three levels of counting: sockets, cores, threads.
// Cores with HyperThreading count as having 2 threads per core.
// Sockets often come with many physical CPU cores.
// For example a single socket board with two cores each with HT will
// return 4 CPUInfoStat structs on Linux and the "Cores" field set to 1.
func Info() ([]InfoStat, error) {
	filename := common.HostProc("cpuinfo")
	lines, _ := common.ReadLines(filename)

	var ret []InfoStat

	c := InfoStat{CPU: -1, Cores: 1}
	for _, line := range lines {
		fields := strings.Split(line, ":")
		if len(fields) < 2 {
			continue
		}
		key := strings.TrimSpace(fields[0])
		value := strings.TrimSpace(fields[1])

		switch key {
		case "processor":
			if c.CPU >= 0 {
				err := finishCPUInfo(&c)
				if err != nil {
					return ret, err
				}
				ret = append(ret, c)
			}
			c = InfoStat{Cores: 1}
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.CPU = int32(t)
		case "vendorId", "vendor_id":
			c.VendorID = value
		case "cpu family":
			c.Family = value
		case "model":
			c.Model = value
		case "model name":
			c.ModelName = value
		case "stepping":
			t, err := strconv.ParseInt(value, 10, 64)
			if err != nil {
				return ret, err
			}
			c.Stepping = int32(t)
		case "cpu MHz", "clock":
			// treat this as the fallback value, thus we ignore error
			if t, err := strconv.ParseFloat(strings.Replace(value, "MHz", "", 1), 64); err == nil {
				c.Mhz = t
			}
		case "cache size":
			t, err := strconv.ParseInt(strings.Replace(value, " KB", "", 1), 10, 64)
			if err != nil {
				return ret, err
			}
			c.CacheSize = int32(t)
		case "physical id":
			c.PhysicalID = value
		case "core id":
			c.CoreID = value
		case "flags", "Features":
			c.Flags = strings.FieldsFunc(value, func(r rune) bool {
				return r == ',' || r == ' '
			})
		}
	}
	if c.CPU >= 0 {
		err := finishCPUInfo(&c)
		if err != nil {
			return ret, err
		}
		ret = append(ret, c)
	}
	return ret, nil
}

func parseStatLine(line string) (*TimesStat, error) {
	fields := strings.Fields(line)

	if strings.HasPrefix(fields[0], "cpu") == false {
		//		return CPUTimesStat{}, e
		return nil, errors.New("not contain cpu")
	}

	cpu := fields[0]
	if cpu == "cpu" {
		cpu = "cpu-total"
	}
	user, err := strconv.ParseFloat(fields[1], 64)
	if err != nil {
		return nil, err
	}
	nice, err := strconv.ParseFloat(fields[2], 64)
	if err != nil {
		return nil, err
	}
	system, err := strconv.ParseFloat(fields[3], 64)
	if err != nil {
		return nil, err
	}
	idle, err := strconv.ParseFloat(fields[4], 64)
	if err != nil {
		return nil, err
	}
	iowait, err := strconv.ParseFloat(fields[5], 64)
	if err != nil {
		return nil, err
	}
	irq, err := strconv.ParseFloat(fields[6], 64)
	if err != nil {
		return nil, err
	}
	softirq, err := strconv.ParseFloat(fields[7], 64)
	if err != nil {
		return nil, err
	}

	ct := &TimesStat{
		CPU:     cpu,
		User:    float64(user) / cpu_tick,
		Nice:    float64(nice) / cpu_tick,
		System:  float64(system) / cpu_tick,
		Idle:    float64(idle) / cpu_tick,
		Iowait:  float64(iowait) / cpu_tick,
		Irq:     float64(irq) / cpu_tick,
		Softirq: float64(softirq) / cpu_tick,
	}
	if len(fields) > 8 { // Linux >= 2.6.11
		steal, err := strconv.ParseFloat(fields[8], 64)
		if err != nil {
			return nil, err
		}
		ct.Steal = float64(steal) / cpu_tick
	}
	if len(fields) > 9 { // Linux >= 2.6.24
		guest, err := strconv.ParseFloat(fields[9], 64)
		if err != nil {
			return nil, err
		}
		ct.Guest = float64(guest) / cpu_tick
	}
	if len(fields) > 10 { // Linux >= 3.2.0
		guestNice, err := strconv.ParseFloat(fields[10], 64)
		if err != nil {
			return nil, err
		}
		ct.GuestNice = float64(guestNice) / cpu_tick
	}

	return ct, nil
}
<|endoftext|>


<commit_before>package logging

import (
	"errors"
	"fmt"
	"github.com/qb0C80aE/clay/logging"
)

var utility = &Utility{}

// Utility handles logging operation
type Utility struct {
}

// GetUtility returns the instance of utility
func GetUtility() *Utility {
	return utility
}

// Debug outputs debug message
func (receiver *Utility) Debug(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Debug(message)
	return nil
}

// Debugf outputs formatted debug message
func (receiver *Utility) Debugf(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Debugf(message, parameters...)
	return nil
}

// Info outputs information message
func (receiver *Utility) Info(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Info(message)
	return nil
}

// Infof outputs formatted information message
func (receiver *Utility) Infof(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Infof(message, parameters...)
	return nil
}

// Warn outputs warning message
func (receiver *Utility) Warn(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Warn(message)
	return nil
}

// Warnf outputs formatted warning message
func (receiver *Utility) Warnf(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Warnf(message, parameters...)
	return nil
}

// Critical outputs critical message
func (receiver *Utility) Critical(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Critical(message)
	return nil
}

// Criticalf outputs formatted critial message
func (receiver *Utility) Criticalf(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Criticalf(message, parameters...)
	return nil
}

// Panic throws panic with message
func (receiver *Utility) Panic(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	panic(message)
	return nil
}

// Panicf throws panic with formatted message
func (receiver *Utility) Panicf(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	message = fmt.Sprintf(message, parameters...)
	panic(message)
	return nil
}
<commit_msg>Fix go vet warnings<commit_after>package logging

import (
	"errors"
	"fmt"
	"github.com/qb0C80aE/clay/logging"
)

var utility = &Utility{}

// Utility handles logging operation
type Utility struct {
}

// GetUtility returns the instance of utility
func GetUtility() *Utility {
	return utility
}

// Debug outputs debug message
func (receiver *Utility) Debug(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Debug(message)
	return nil
}

// Debugf outputs formatted debug message
func (receiver *Utility) Debugf(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Debugf(message, parameters...)
	return nil
}

// Info outputs information message
func (receiver *Utility) Info(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Info(message)
	return nil
}

// Infof outputs formatted information message
func (receiver *Utility) Infof(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Infof(message, parameters...)
	return nil
}

// Warn outputs warning message
func (receiver *Utility) Warn(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Warn(message)
	return nil
}

// Warnf outputs formatted warning message
func (receiver *Utility) Warnf(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Warnf(message, parameters...)
	return nil
}

// Critical outputs critical message
func (receiver *Utility) Critical(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Critical(message)
	return nil
}

// Criticalf outputs formatted critial message
func (receiver *Utility) Criticalf(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	logging.Logger().Criticalf(message, parameters...)
	return nil
}

// Panic throws panic with message
func (receiver *Utility) Panic(log interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	panic(message)
}

// Panicf throws panic with formatted message
func (receiver *Utility) Panicf(log interface{}, parameters []interface{}) error {
	message, ok := log.(string)
	if !ok {
		return errors.New("log is not a string")
	}

	message = fmt.Sprintf(message, parameters...)
	panic(message)
}
<|endoftext|>


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

package rocserv

import (
	"flag"
	"fmt"
	"github.com/shawnfeng/roc/util/service/sla"
	"github.com/shawnfeng/sutil/trace"
	"reflect"

	"github.com/julienschmidt/httprouter"

	"git.apache.org/thrift.git/lib/go/thrift"

	"github.com/shawnfeng/sutil/slog"
)

const (
	PROCESSOR_HTTP   = "http"
	PROCESSOR_THRIFT = "thrift"
	PROCESSOR_GRPC   = "gprc"
)

type Service struct {
}

var service Service

//func NewService

type cmdArgs struct {
	servLoc string
	logDir  string
	sessKey string
	group   string
}

func (m *Service) parseFlag() (*cmdArgs, error) {
	var serv, logDir, skey string
	flag.StringVar(&serv, "serv", "", "servic name")
	flag.StringVar(&logDir, "logdir", "", "serice log dir")
	flag.StringVar(&skey, "skey", "", "service session key")

	flag.Parse()

	if len(serv) == 0 {
		return nil, fmt.Errorf("serv args need!")
	}

	if len(skey) == 0 {
		return nil, fmt.Errorf("skey args need!")
	}

	return &cmdArgs{
		servLoc: serv,
		logDir:  logDir,
		sessKey: skey,
	}, nil

}

func (m *Service) loadDriver(sb ServBase, procs map[string]Processor) (map[string]*ServInfo, error) {
	fun := "Service.loadDriver -->"

	infos := make(map[string]*ServInfo)

	// load processor's driver
	for n, p := range procs {
		addr, driver := p.Driver()
		if driver == nil {
			slog.Infof("%s processor:%s no driver", fun, n)
			continue
		}

		slog.Infof("%s processor:%s type:%s addr:%s", fun, n, reflect.TypeOf(driver), addr)

		switch d := driver.(type) {
		case *httprouter.Router:
			sa, err := powerHttp(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_HTTP,
				Addr: sa,
			}

		case thrift.TProcessor:
			sa, err := powerThrift(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_THRIFT,
				Addr: sa,
			}
		case *GrpcServer:
			sa, err := powerGrpc(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_GRPC,
				Addr: sa,
			}
		default:
			return nil, fmt.Errorf("processor:%s driver not recognition", n)

		}

	}

	return infos, nil

}

func (m *Service) Serve(confEtcd configEtcd, initfn func(ServBase) error, procs map[string]Processor) error {
	fun := "Service.Serve -->"

	args, err := m.parseFlag()
	if err != nil {
		slog.Panicf("%s parse arg err:%s", fun, err)
		return err
	}

	return m.Init(confEtcd, args.servLoc, args.sessKey, args.logDir, args.group, initfn, procs)
}

func (m *Service) Init(confEtcd configEtcd, servLoc, sessKey, logDir, group string, initfn func(ServBase) error, procs map[string]Processor) error {
	fun := "Service.Init -->"

	sb, err := NewServBaseV2(confEtcd, servLoc, sessKey)
	if err != nil {
		slog.Panicf("%s init servbase loc:%s key:%s err:%s", fun, servLoc, sessKey, err)
		return err
	}

	var logConfig struct {
		Log struct {
			Level string
			Dir   string
		}
	}
	// default use level INFO
	logConfig.Log.Level = "INFO"

	err = sb.ServConfig(&logConfig)
	if err != nil {
		slog.Panicf("%s serv config err:%s", fun, err)
		fmt.Sprintf("%s serv config err:%s", fun, err)
		return err
	}

	var logdir string
	if len(logConfig.Log.Dir) > 0 {
		logdir = fmt.Sprintf("%s/%s", logConfig.Log.Dir, sb.Copyname())
	}

	if len(logDir) > 0 {
		logdir = fmt.Sprintf("%s/%s", logDir, sb.Copyname())
	}

	if logDir == "console" {
		logdir = ""
	}

	fmt.Sprintf("%s init log dir:%s name:%s level:%s", fun, logdir, servLoc, logConfig.Log.Level)
	slog.Infof("%s init log dir:%s name:%s level:%s", fun, logdir, servLoc, logConfig.Log.Level)

	slog.Init(logdir, "serv.log", logConfig.Log.Level)
	defer slog.Sync()

	// init tracer
	err = trace.InitDefaultTracer(servLoc)
	if err != nil {
		slog.Warnf("%s init tracer fail:%v", err)
	}

	// sla metric埋点 ==================
	//init metric
	//user defualt metric opts
	metrics := rocserv.NewMetricsprocessor()
	if err != nil {
		slog.Warnf("init metrics fail:%v", err)
	}
	err = metrics.Init()
	if err != nil {
		slog.Warnf("%s init metrics err:%s", fun, err)
	}

	minfos, err := m.loadDriver(sb, map[string]Processor{"_PROC_METRICS": metrics})
	if err == nil {
		err = sb.RegisterMetrics(minfos)
		if err != nil {
			slog.Warnf("%s regist backdoor err:%s", fun, err)
		}

	} else {
		slog.Warnf("%s load metrics driver err:%s", fun, err)
	}
	//==============================

	// init callback
	err = initfn(sb)
	if err != nil {
		slog.Panicf("%s serv init err:%s", fun, err)
		return err
	}

	// init processor
	for n, p := range procs {
		if len(n) == 0 {
			slog.Panicf("%s processor name empty", fun)
			return err
		}

		if n[0] == '_' {
			slog.Panicf("%s processor name can not prefix '_'", fun)
			return err
		}

		if p == nil {
			slog.Panicf("%s processor:%s is nil", fun, n)
			return fmt.Errorf("processor:%s is nil", n)
		} else {
			err := p.Init()
			if err != nil {
				slog.Panicf("%s processor:%s init err:%s", fun, err)
				return err
			}
		}
	}

	infos, err := m.loadDriver(sb, procs)
	if err != nil {
		slog.Panicf("%s load driver err:%s", fun, err)
		return err
	}

	err = sb.RegisterService(infos)
	if err != nil {
		slog.Panicf("%s regist service err:%s", fun, err)
		return err
	}

	// 后门接口 ==================
	backdoor := &backDoorHttp{}
	err = backdoor.Init()
	if err != nil {
		slog.Panicf("%s init backdoor err:%s", fun, err)
		return err
	}

	binfos, err := m.loadDriver(sb, map[string]Processor{"_PROC_BACKDOOR": backdoor})
	if err == nil {
		err = sb.RegisterBackDoor(binfos)
		if err != nil {
			slog.Errorf("%s regist backdoor err:%s", fun, err)
		}

	} else {
		slog.Warnf("%s load backdoor driver err:%s", fun, err)
	}
	//==============================

	sb.SetGroup(group)

	// pause here
	var pause chan bool
	pause <- true

	return nil

}
func (m *Service) getMetricOps(sb *ServBaseV2) *rocserv.MetricsOpts {
	fun := "Service.getMetricOps -->"
	var metricConfig struct {
		metric *rocserv.MetricsOpts
	}
	err := sb.ServConfig(&metricConfig)
	if err != nil {
		slog.Panicf("%s serv config err:%s", fun, err)
		fmt.Sprintf("%s serv config err:%s", fun, err)
		return nil
	}
	return metricConfig.metric
}

func Serve(etcds []string, baseLoc string, initfn func(ServBase) error, procs map[string]Processor) error {
	return service.Serve(configEtcd{etcds, baseLoc}, initfn, procs)
}

func Init(etcds []string, baseLoc string, servLoc, servKey string, initfn func(ServBase) error, procs map[string]Processor) error {
	return service.Init(configEtcd{etcds, baseLoc}, servLoc, servKey, "", "", initfn, procs)
}

func Test(etcds []string, baseLoc string, initfn func(ServBase) error) error {
	return service.Init(configEtcd{etcds, baseLoc}, "test/test", "test", "console", "", initfn, nil)
}
<commit_msg>MOD:group 参数<commit_after>// Copyright 2014 The roc Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rocserv

import (
	"flag"
	"fmt"
	"github.com/shawnfeng/roc/util/service/sla"
	"github.com/shawnfeng/sutil/trace"
	"reflect"

	"github.com/julienschmidt/httprouter"

	"git.apache.org/thrift.git/lib/go/thrift"

	"github.com/shawnfeng/sutil/slog"
)

const (
	PROCESSOR_HTTP   = "http"
	PROCESSOR_THRIFT = "thrift"
	PROCESSOR_GRPC   = "gprc"
)

type Service struct {
}

var service Service

//func NewService

type cmdArgs struct {
	servLoc string
	logDir  string
	sessKey string
	group   string
}

func (m *Service) parseFlag() (*cmdArgs, error) {
	var serv, logDir, skey string
	flag.StringVar(&serv, "serv", "", "servic name")
	flag.StringVar(&logDir, "logdir", "", "serice log dir")
	flag.StringVar(&skey, "skey", "", "service session key")

	flag.Parse()

	if len(serv) == 0 {
		return nil, fmt.Errorf("serv args need!")
	}

	if len(skey) == 0 {
		return nil, fmt.Errorf("skey args need!")
	}

	return &cmdArgs{
		servLoc: serv,
		logDir:  logDir,
		sessKey: skey,
	}, nil

}

func (m *Service) loadDriver(sb ServBase, procs map[string]Processor) (map[string]*ServInfo, error) {
	fun := "Service.loadDriver -->"

	infos := make(map[string]*ServInfo)

	// load processor's driver
	for n, p := range procs {
		addr, driver := p.Driver()
		if driver == nil {
			slog.Infof("%s processor:%s no driver", fun, n)
			continue
		}

		slog.Infof("%s processor:%s type:%s addr:%s", fun, n, reflect.TypeOf(driver), addr)

		switch d := driver.(type) {
		case *httprouter.Router:
			sa, err := powerHttp(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_HTTP,
				Addr: sa,
			}

		case thrift.TProcessor:
			sa, err := powerThrift(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_THRIFT,
				Addr: sa,
			}
		case *GrpcServer:
			sa, err := powerGrpc(addr, d)
			if err != nil {
				return nil, err
			}

			slog.Infof("%s load ok processor:%s serv addr:%s", fun, n, sa)
			infos[n] = &ServInfo{
				Type: PROCESSOR_GRPC,
				Addr: sa,
			}
		default:
			return nil, fmt.Errorf("processor:%s driver not recognition", n)

		}

	}

	return infos, nil

}

func (m *Service) Serve(confEtcd configEtcd, initfn func(ServBase) error, procs map[string]Processor) error {
	fun := "Service.Serve -->"

	args, err := m.parseFlag()
	if err != nil {
		slog.Panicf("%s parse arg err:%s", fun, err)
		return err
	}

	return m.Init(confEtcd, args.servLoc, args.sessKey, args.logDir, args.group, initfn, procs)
}

func (m *Service) Init(confEtcd configEtcd, servLoc, sessKey, logDir, group string, initfn func(ServBase) error, procs map[string]Processor) error {
	fun := "Service.Init -->"

	sb, err := NewServBaseV2(confEtcd, servLoc, sessKey)
	if err != nil {
		slog.Panicf("%s init servbase loc:%s key:%s err:%s", fun, servLoc, sessKey, err)
		return err
	}

	var logConfig struct {
		Log struct {
			Level string
			Dir   string
		}
	}
	// default use level INFO
	logConfig.Log.Level = "INFO"

	err = sb.ServConfig(&logConfig)
	if err != nil {
		slog.Panicf("%s serv config err:%s", fun, err)
		fmt.Sprintf("%s serv config err:%s", fun, err)
		return err
	}

	var logdir string
	if len(logConfig.Log.Dir) > 0 {
		logdir = fmt.Sprintf("%s/%s", logConfig.Log.Dir, sb.Copyname())
	}

	if len(logDir) > 0 {
		logdir = fmt.Sprintf("%s/%s", logDir, sb.Copyname())
	}

	if logDir == "console" {
		logdir = ""
	}

	fmt.Sprintf("%s init log dir:%s name:%s level:%s", fun, logdir, servLoc, logConfig.Log.Level)
	slog.Infof("%s init log dir:%s name:%s level:%s", fun, logdir, servLoc, logConfig.Log.Level)

	slog.Init(logdir, "serv.log", logConfig.Log.Level)
	defer slog.Sync()

	// init tracer
	err = trace.InitDefaultTracer(servLoc)
	if err != nil {
		slog.Warnf("%s init tracer fail:%v", err)
	}

	// sla metric埋点 ==================
	//init metric
	//user defualt metric opts
	metrics := rocserv.NewMetricsprocessor()
	if err != nil {
		slog.Warnf("init metrics fail:%v", err)
	}
	err = metrics.Init()
	if err != nil {
		slog.Warnf("%s init metrics err:%s", fun, err)
	}

	minfos, err := m.loadDriver(sb, map[string]Processor{"_PROC_METRICS": metrics})
	if err == nil {
		err = sb.RegisterMetrics(minfos)
		if err != nil {
			slog.Warnf("%s regist backdoor err:%s", fun, err)
		}

	} else {
		slog.Warnf("%s load metrics driver err:%s", fun, err)
	}
	//==============================

	// init callback
	err = initfn(sb)
	if err != nil {
		slog.Panicf("%s serv init err:%s", fun, err)
		return err
	}

	// init processor
	for n, p := range procs {
		if len(n) == 0 {
			slog.Panicf("%s processor name empty", fun)
			return err
		}

		if n[0] == '_' {
			slog.Panicf("%s processor name can not prefix '_'", fun)
			return err
		}

		if p == nil {
			slog.Panicf("%s processor:%s is nil", fun, n)
			return fmt.Errorf("processor:%s is nil", n)
		} else {
			err := p.Init()
			if err != nil {
				slog.Panicf("%s processor:%s init err:%s", fun, err)
				return err
			}
		}
	}

	infos, err := m.loadDriver(sb, procs)
	if err != nil {
		slog.Panicf("%s load driver err:%s", fun, err)
		return err
	}

	err = sb.RegisterService(infos)
	if err != nil {
		slog.Panicf("%s regist service err:%s", fun, err)
		return err
	}

	sb.SetGroup(group)

	// 后门接口 ==================
	backdoor := &backDoorHttp{}
	err = backdoor.Init()
	if err != nil {
		slog.Panicf("%s init backdoor err:%s", fun, err)
		return err
	}

	binfos, err := m.loadDriver(sb, map[string]Processor{"_PROC_BACKDOOR": backdoor})
	if err == nil {
		err = sb.RegisterBackDoor(binfos)
		if err != nil {
			slog.Errorf("%s regist backdoor err:%s", fun, err)
		}

	} else {
		slog.Warnf("%s load backdoor driver err:%s", fun, err)
	}
	//==============================

	sb.SetGroup(group)

	// pause here
	var pause chan bool
	pause <- true

	return nil

}
func (m *Service) getMetricOps(sb *ServBaseV2) *rocserv.MetricsOpts {
	fun := "Service.getMetricOps -->"
	var metricConfig struct {
		metric *rocserv.MetricsOpts
	}
	err := sb.ServConfig(&metricConfig)
	if err != nil {
		slog.Panicf("%s serv config err:%s", fun, err)
		fmt.Sprintf("%s serv config err:%s", fun, err)
		return nil
	}
	return metricConfig.metric
}

func Serve(etcds []string, baseLoc string, initfn func(ServBase) error, procs map[string]Processor) error {
	return service.Serve(configEtcd{etcds, baseLoc}, initfn, procs)
}

func Init(etcds []string, baseLoc string, servLoc, servKey string, initfn func(ServBase) error, procs map[string]Processor) error {
	return service.Init(configEtcd{etcds, baseLoc}, servLoc, servKey, "", "", initfn, procs)
}

func Test(etcds []string, baseLoc string, initfn func(ServBase) error) error {
	return service.Init(configEtcd{etcds, baseLoc}, "test/test", "test", "console", "", initfn, nil)
}
<|endoftext|>


<commit_before>package infrastructure

import (
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/goadapp/goad/version"
	"github.com/satori/go.uuid"
)

// Infrastructure manages the resource creation and updates necessary to use
// Goad.
type Infrastructure struct {
	config   *aws.Config
	queueURL string
	regions  []string
}

// New creates the required infrastructure to run the load tests in Lambda
// functions.
func New(regions []string, config *aws.Config) (*Infrastructure, error) {
	infra := &Infrastructure{config: config, regions: regions}
	if err := infra.setup(); err != nil {
		return nil, err
	}
	return infra, nil
}

// QueueURL returns the URL of the SQS queue to use for the load test session.
func (infra *Infrastructure) QueueURL() string {
	return infra.queueURL
}

// Clean removes any AWS resources that cannot be reused for a subsequent test.
func (infra *Infrastructure) Clean() {
	infra.removeSQSQueue()
}

func (infra *Infrastructure) setup() error {
	roleArn, err := infra.createIAMLambdaRole("goad-lambda-role")
	if err != nil {
		return err
	}
	zip, err := Asset("data/lambda.zip")
	if err != nil {
		return err
	}
	for _, region := range infra.regions {
		err = infra.createOrUpdateLambdaFunction(region, roleArn, zip)
		if err != nil {
			return err
		}
	}
	queueURL, err := infra.createSQSQueue()
	if err != nil {
		return err
	}
	infra.queueURL = queueURL
	return nil
}

func (infra *Infrastructure) createOrUpdateLambdaFunction(region, roleArn string, payload []byte) error {
	config := aws.NewConfig().WithRegion(region)
	svc := lambda.New(session.New(), config)

	exists, err := lambdaExists(svc)

	if err != nil {
		return err
	}

	if exists {
		aliasExists, err := lambdaAliasExists(svc)
		if err != nil || aliasExists {
			return err
		}
		return infra.updateLambdaFunction(svc, roleArn, payload)
	}

	return infra.createLambdaFunction(svc, roleArn, payload)
}

func (infra *Infrastructure) createLambdaFunction(svc *lambda.Lambda, roleArn string, payload []byte) error {
	function, err := svc.CreateFunction(&lambda.CreateFunctionInput{
		Code: &lambda.FunctionCode{
			ZipFile: payload,
		},
		FunctionName: aws.String("goad"),
		Handler:      aws.String("index.handler"),
		Role:         aws.String(roleArn),
		Runtime:      aws.String("nodejs"),
		MemorySize:   aws.Int64(1536),
		Publish:      aws.Bool(true),
		Timeout:      aws.Int64(300),
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			// Calling this function too soon after creating the role might
			// fail, so we should retry after a little while.
			// TODO: limit the number of retries.
			if awsErr.Code() == "InvalidParameterValueException" {
				time.Sleep(time.Second)
				return infra.createLambdaFunction(svc, roleArn, payload)
			}
		}
		return err
	}
	return createLambdaAlias(svc, function.Version)
}

func (infra *Infrastructure) updateLambdaFunction(svc *lambda.Lambda, roleArn string, payload []byte) error {
	function, err := svc.UpdateFunctionCode(&lambda.UpdateFunctionCodeInput{
		ZipFile:      payload,
		FunctionName: aws.String("goad"),
		Publish:      aws.Bool(true),
	})
	if err != nil {
		return err
	}
	return createLambdaAlias(svc, function.Version)
}

func lambdaExists(svc *lambda.Lambda) (bool, error) {
	_, err := svc.GetFunction(&lambda.GetFunctionInput{
		FunctionName: aws.String("goad"),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				return false, nil
			}
		}
		return false, err
	}

	return true, nil
}

func createLambdaAlias(svc *lambda.Lambda, functionVersion *string) error {
	_, err := svc.CreateAlias(&lambda.CreateAliasInput{
		FunctionName:    aws.String("goad"),
		FunctionVersion: functionVersion,
		Name:            aws.String(version.LambdaVersion()),
	})
	return err
}

func lambdaAliasExists(svc *lambda.Lambda) (bool, error) {
	_, err := svc.GetAlias(&lambda.GetAliasInput{
		FunctionName: aws.String("goad"),
		Name:         aws.String(version.LambdaVersion()),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				return false, nil
			}
		}
		return false, err
	}

	return true, nil
}

func (infra *Infrastructure) createIAMLambdaRole(roleName string) (arn string, err error) {
	svc := iam.New(session.New(), infra.config)

	resp, err := svc.GetRole(&iam.GetRoleInput{
		RoleName: aws.String(roleName),
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "NoSuchEntity" {
				res, err := svc.CreateRole(&iam.CreateRoleInput{
					AssumeRolePolicyDocument: aws.String(`{
        	          "Version": "2012-10-17",
        	          "Statement": {
        	            "Effect": "Allow",
        	            "Principal": {"Service": "lambda.amazonaws.com"},
        	            "Action": "sts:AssumeRole"
        	          }
            	    }`),
					RoleName: aws.String(roleName),
					Path:     aws.String("/"),
				})
				if err != nil {
					return "", err
				}
				if err := infra.createIAMLambdaRolePolicy(*res.Role.RoleName); err != nil {
					return "", err
				}
				return *res.Role.Arn, nil
			}
		}
		return "", err
	}

	return *resp.Role.Arn, nil
}

func (infra *Infrastructure) createIAMLambdaRolePolicy(roleName string) error {
	svc := iam.New(session.New(), infra.config)

	_, err := svc.PutRolePolicy(&iam.PutRolePolicyInput{
		PolicyDocument: aws.String(`{
          "Version": "2012-10-17",
          "Statement": [
            {
              "Action": [
                "sqs:SendMessage"
              ],
              "Effect": "Allow",
              "Resource": "arn:aws:sqs:*:*:goad-*"
		  	},
			{
              "Action": [
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents"
              ],
              "Effect": "Allow",
              "Resource": "arn:aws:logs:*:*:*"
	        }
          ]
        }`),
		PolicyName: aws.String("goad-lambda-role-policy"),
		RoleName:   aws.String(roleName),
	})
	return err
}

func (infra *Infrastructure) createSQSQueue() (url string, err error) {
	svc := sqs.New(session.New(), infra.config)

	resp, err := svc.CreateQueue(&sqs.CreateQueueInput{
		QueueName: aws.String("goad-" + uuid.NewV4().String()),
	})

	if err != nil {
		return "", err
	}

	return *resp.QueueUrl, nil
}

func (infra *Infrastructure) removeSQSQueue() {
	svc := sqs.New(session.New(), infra.config)

	svc.DeleteQueue(&sqs.DeleteQueueInput{
		QueueUrl: aws.String(infra.queueURL),
	})
}
<commit_msg>this is preventing to see the error on Lambda function creation<commit_after>package infrastructure

import (
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/goadapp/goad/version"
	"github.com/satori/go.uuid"
)

// Infrastructure manages the resource creation and updates necessary to use
// Goad.
type Infrastructure struct {
	config   *aws.Config
	queueURL string
	regions  []string
}

// New creates the required infrastructure to run the load tests in Lambda
// functions.
func New(regions []string, config *aws.Config) (*Infrastructure, error) {
	infra := &Infrastructure{config: config, regions: regions}
	if err := infra.setup(); err != nil {
		return nil, err
	}
	return infra, nil
}

// QueueURL returns the URL of the SQS queue to use for the load test session.
func (infra *Infrastructure) QueueURL() string {
	return infra.queueURL
}

// Clean removes any AWS resources that cannot be reused for a subsequent test.
func (infra *Infrastructure) Clean() {
	infra.removeSQSQueue()
}

func (infra *Infrastructure) setup() error {
	roleArn, err := infra.createIAMLambdaRole("goad-lambda-role")
	if err != nil {
		return err
	}
	zip, err := Asset("data/lambda.zip")
	if err != nil {
		return err
	}

	for _, region := range infra.regions {
		err = infra.createOrUpdateLambdaFunction(region, roleArn, zip)
		if err != nil {
			return err
		}
	}
	queueURL, err := infra.createSQSQueue()
	if err != nil {
		return err
	}
	infra.queueURL = queueURL
	return nil
}

func (infra *Infrastructure) createOrUpdateLambdaFunction(region, roleArn string, payload []byte) error {
	config := aws.NewConfig().WithRegion(region)
	svc := lambda.New(session.New(), config)

	exists, err := lambdaExists(svc)

	if err != nil {
		return err
	}

	if exists {
		aliasExists, err := lambdaAliasExists(svc)
		if err != nil || aliasExists {
			return err
		}
		return infra.updateLambdaFunction(svc, roleArn, payload)
	}

	return infra.createLambdaFunction(svc, roleArn, payload)
}

func (infra *Infrastructure) createLambdaFunction(svc *lambda.Lambda, roleArn string, payload []byte) error {
	function, err := svc.CreateFunction(&lambda.CreateFunctionInput{
		Code: &lambda.FunctionCode{
			ZipFile: payload,
		},
		FunctionName: aws.String("goad"),
		Handler:      aws.String("index.handler"),
		Role:         aws.String(roleArn),
		Runtime:      aws.String("nodejs"),
		MemorySize:   aws.Int64(1536),
		Publish:      aws.Bool(true),
		Timeout:      aws.Int64(300),
	})
	if err != nil {
		return err
	}
	return createLambdaAlias(svc, function.Version)
}

func (infra *Infrastructure) updateLambdaFunction(svc *lambda.Lambda, roleArn string, payload []byte) error {
	function, err := svc.UpdateFunctionCode(&lambda.UpdateFunctionCodeInput{
		ZipFile:      payload,
		FunctionName: aws.String("goad"),
		Publish:      aws.Bool(true),
	})
	if err != nil {
		return err
	}
	return createLambdaAlias(svc, function.Version)
}

func lambdaExists(svc *lambda.Lambda) (bool, error) {
	_, err := svc.GetFunction(&lambda.GetFunctionInput{
		FunctionName: aws.String("goad"),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				return false, nil
			}
		}
		return false, err
	}

	return true, nil
}

func createLambdaAlias(svc *lambda.Lambda, functionVersion *string) error {
	_, err := svc.CreateAlias(&lambda.CreateAliasInput{
		FunctionName:    aws.String("goad"),
		FunctionVersion: functionVersion,
		Name:            aws.String(version.LambdaVersion()),
	})
	return err
}

func lambdaAliasExists(svc *lambda.Lambda) (bool, error) {
	_, err := svc.GetAlias(&lambda.GetAliasInput{
		FunctionName: aws.String("goad"),
		Name:         aws.String(version.LambdaVersion()),
	})

	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "ResourceNotFoundException" {
				return false, nil
			}
		}
		return false, err
	}

	return true, nil
}

func (infra *Infrastructure) createIAMLambdaRole(roleName string) (arn string, err error) {
	svc := iam.New(session.New(), infra.config)

	resp, err := svc.GetRole(&iam.GetRoleInput{
		RoleName: aws.String(roleName),
	})
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok {
			if awsErr.Code() == "NoSuchEntity" {
				res, err := svc.CreateRole(&iam.CreateRoleInput{
					AssumeRolePolicyDocument: aws.String(`{
        	          "Version": "2012-10-17",
        	          "Statement": {
        	            "Effect": "Allow",
        	            "Principal": {"Service": "lambda.amazonaws.com"},
        	            "Action": "sts:AssumeRole"
        	          }
            	    }`),
					RoleName: aws.String(roleName),
					Path:     aws.String("/"),
				})
				if err != nil {
					return "", err
				}
				if err := infra.createIAMLambdaRolePolicy(*res.Role.RoleName); err != nil {
					return "", err
				}
				return *res.Role.Arn, nil
			}
		}
		return "", err
	}

	return *resp.Role.Arn, nil
}

func (infra *Infrastructure) createIAMLambdaRolePolicy(roleName string) error {
	svc := iam.New(session.New(), infra.config)

	_, err := svc.PutRolePolicy(&iam.PutRolePolicyInput{
		PolicyDocument: aws.String(`{
          "Version": "2012-10-17",
          "Statement": [
            {
              "Action": [
                "sqs:SendMessage"
              ],
              "Effect": "Allow",
              "Resource": "arn:aws:sqs:*:*:goad-*"
		  	},
			{
              "Action": [
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents"
              ],
              "Effect": "Allow",
              "Resource": "arn:aws:logs:*:*:*"
	        }
          ]
        }`),
		PolicyName: aws.String("goad-lambda-role-policy"),
		RoleName:   aws.String(roleName),
	})
	return err
}

func (infra *Infrastructure) createSQSQueue() (url string, err error) {
	svc := sqs.New(session.New(), infra.config)

	resp, err := svc.CreateQueue(&sqs.CreateQueueInput{
		QueueName: aws.String("goad-" + uuid.NewV4().String()),
	})

	if err != nil {
		return "", err
	}

	return *resp.QueueUrl, nil
}

func (infra *Infrastructure) removeSQSQueue() {
	svc := sqs.New(session.New(), infra.config)

	svc.DeleteQueue(&sqs.DeleteQueueInput{
		QueueUrl: aws.String(infra.queueURL),
	})
}
<|endoftext|>


<commit_before>package layout

import (
	"github.com/ian-kent/go-log/levels"
	"github.com/stretchr/testify/assert"
	"testing"
	"os"
)

func TestPattern(t *testing.T) {
	p := Pattern("")
	assert.NotNil(t, p)

	assert.Equal(t, p.Format(levels.DEBUG, "Test message %s", "test"), "")
	assert.Equal(t, p.Format(levels.DEBUG, "%c %C %d"), "")

	p = Pattern("%c")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "layout")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "layout")

	p = Pattern("%C")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "layout")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "layout")

	p = Pattern("%d")
	assert.NotNil(t, p)
	// FIXME
	//assert.Equal(t, p.Format(levels.DEBUG, ""), time.Now().String())
	//assert.Equal(t, p.Format(levels.DEBUG, "foo"), time.Now().String())

	p = Pattern("%F")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), os.Getenv("GOPATH") + "/src/github.com/ian-kent/go-log/layout/pattern_test.go")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), os.Getenv("GOPATH") + "/src/github.com/ian-kent/go-log/layout/pattern_test.go")

	p = Pattern("%l")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "layout/pattern_test.go:40")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "layout/pattern_test.go:41")

	p = Pattern("%L")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "45")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "46")

	p = Pattern("%m")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, "test message"), "test message")
	assert.Equal(t, p.Format(levels.DEBUG, "test message %s", "test"), "test message test")
	assert.Equal(t, p.Format(levels.DEBUG, "test message %d", 2), "test message 2")

	p = Pattern("%n")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "\n")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "\n")

	p = Pattern("%p")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "DEBUG")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "DEBUG")

	p = Pattern("%r")
	assert.NotNil(t, p)
	// FIXME
	//assert.Equal(t, p.Format(levels.DEBUG, ""), "FIXME")
	//assert.Equal(t, p.Format(levels.DEBUG, "foo"), "FIXME")

	p = Pattern("%x")
	assert.NotNil(t, p)
	// FIXME
	//assert.Equal(t, p.Format(levels.DEBUG, ""), "FIXME")
	//assert.Equal(t, p.Format(levels.DEBUG, "foo"), "FIXME")

	p = Pattern("%X")
	assert.NotNil(t, p)
	// FIXME
	//assert.Equal(t, p.Format(levels.DEBUG, ""), "FIXME")
	//assert.Equal(t, p.Format(levels.DEBUG, "foo"), "FIXME")

	p = Pattern("%%")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "%")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "%")

	p = Pattern("%c %C %F %l %L %p %m%n")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "layout layout " + os.Getenv("GOPATH") + "/src/github.com/ian-kent/go-log/layout/pattern_test.go layout/pattern_test.go:89 89 DEBUG \n")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "layout layout " + os.Getenv("GOPATH") + "/src/github.com/ian-kent/go-log/layout/pattern_test.go layout/pattern_test.go:90 90 DEBUG foo\n")
	assert.Equal(t, p.Format(levels.DEBUG, "foo=%s", "bar"), "layout layout " + os.Getenv("GOPATH") + "/src/github.com/ian-kent/go-log/layout/pattern_test.go layout/pattern_test.go:91 91 DEBUG foo=bar\n")
}
<commit_msg>Remove filename dependent tests<commit_after>package layout

import (
	"github.com/ian-kent/go-log/levels"
	"github.com/stretchr/testify/assert"
	"testing"
	"os"
)

func TestPattern(t *testing.T) {
	p := Pattern("")
	assert.NotNil(t, p)

	assert.Equal(t, p.Format(levels.DEBUG, "Test message %s", "test"), "")
	assert.Equal(t, p.Format(levels.DEBUG, "%c %C %d"), "")

	p = Pattern("%c")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "layout")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "layout")

	p = Pattern("%C")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "layout")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "layout")

	p = Pattern("%d")
	assert.NotNil(t, p)
	// FIXME
	//assert.Equal(t, p.Format(levels.DEBUG, ""), time.Now().String())
	//assert.Equal(t, p.Format(levels.DEBUG, "foo"), time.Now().String())

	p = Pattern("%F")
	assert.NotNil(t, p)
	//assert.Equal(t, p.Format(levels.DEBUG, ""), os.Getenv("GOPATH") + "/src/github.com/ian-kent/go-log/layout/pattern_test.go")
	//assert.Equal(t, p.Format(levels.DEBUG, "foo"), os.Getenv("GOPATH") + "/src/github.com/ian-kent/go-log/layout/pattern_test.go")

	p = Pattern("%l")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "layout/pattern_test.go:40")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "layout/pattern_test.go:41")

	p = Pattern("%L")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "45")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "46")

	p = Pattern("%m")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, "test message"), "test message")
	assert.Equal(t, p.Format(levels.DEBUG, "test message %s", "test"), "test message test")
	assert.Equal(t, p.Format(levels.DEBUG, "test message %d", 2), "test message 2")

	p = Pattern("%n")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "\n")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "\n")

	p = Pattern("%p")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "DEBUG")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "DEBUG")

	p = Pattern("%r")
	assert.NotNil(t, p)
	// FIXME
	//assert.Equal(t, p.Format(levels.DEBUG, ""), "FIXME")
	//assert.Equal(t, p.Format(levels.DEBUG, "foo"), "FIXME")

	p = Pattern("%x")
	assert.NotNil(t, p)
	// FIXME
	//assert.Equal(t, p.Format(levels.DEBUG, ""), "FIXME")
	//assert.Equal(t, p.Format(levels.DEBUG, "foo"), "FIXME")

	p = Pattern("%X")
	assert.NotNil(t, p)
	// FIXME
	//assert.Equal(t, p.Format(levels.DEBUG, ""), "FIXME")
	//assert.Equal(t, p.Format(levels.DEBUG, "foo"), "FIXME")

	p = Pattern("%%")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "%")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "%")

	p = Pattern("%c %C %l %L %p %m%n")
	assert.NotNil(t, p)
	assert.Equal(t, p.Format(levels.DEBUG, ""), "layout layout layout/pattern_test.go:89 89 DEBUG \n")
	assert.Equal(t, p.Format(levels.DEBUG, "foo"), "layout layout layout/pattern_test.go:90 90 DEBUG foo\n")
	assert.Equal(t, p.Format(levels.DEBUG, "foo=%s", "bar"), "layout layout layout/pattern_test.go:91 91 DEBUG foo=bar\n")
}
<|endoftext|>


<commit_before>// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

// arpscan implements ARP scanning of all interfaces' local networks using
// gopacket and its subpackages.  This example shows, among other things:
//   * Generating and sending packet data
//   * Reading in packet data and interpreting it
//   * Use of the 'pcap' subpackage for reading/writing
package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"log"
	"net"
	"sync"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

func main() {
	// Get a list of all interfaces.
	ifaces, err := net.Interfaces()
	if err != nil {
		panic(err)
	}

	var wg sync.WaitGroup
	for _, iface := range ifaces {
		wg.Add(1)
		// Start up a scan on each interface.
		go func(iface net.Interface) {
			defer wg.Done()
			if err := scan(&iface); err != nil {
				log.Printf("interface %v: %v", iface.Name, err)
			}
		}(iface)
	}
	// Wait for all interfaces' scans to complete.  They'll try to run
	// forever, but will stop on an error, so if we get past this Wait
	// it means all attempts to write have failed.
	wg.Wait()
}

// scan scans an individual interface's local network for machines using ARP requests/replies.
//
// scan loops forever, sending packets out regularly.  It returns an error if
// it's ever unable to write a packet.
func scan(iface *net.Interface) error {
	// We just look for IPv4 addresses, so try to find if the interface has one.
	var addr *net.IPNet
	if addrs, err := iface.Addrs(); err != nil {
		return err
	} else {
		for _, a := range addrs {
			if ipnet, ok := a.(*net.IPNet); ok {
				if ip4 := ipnet.IP.To4(); ip4 != nil {
					addr = &net.IPNet{
						IP:   ip4,
						Mask: ipnet.Mask[len(ipnet.Mask)-4:],
					}
					break
				}
			}
		}
	}
	// Sanity-check that the interface has a good address.
	if addr == nil {
		return errors.New("no good IP network found")
	} else if addr.IP[0] == 127 {
		return errors.New("skipping localhost")
	} else if addr.Mask[0] != 0xff || addr.Mask[1] != 0xff {
		return errors.New("mask means network is too large")
	}
	log.Printf("Using network range %v for interface %v", addr, iface.Name)

	// Open up a pcap handle for packet reads/writes.
	handle, err := pcap.OpenLive(iface.Name, 65536, true, pcap.BlockForever)
	if err != nil {
		return err
	}
	defer handle.Close()

	// Start up a goroutine to read in packet data.
	stop := make(chan struct{})
	go readARP(handle, iface, stop)
	defer close(stop)
	for {
		// Write our scan packets out to the handle.
		if err := writeARP(handle, iface, addr); err != nil {
			log.Printf("error writing packets on %v: %v", iface.Name, err)
			return err
		}
		// We don't know exactly how long it'll take for packets to be
		// sent back to us, but 10 seconds should be more than enough
		// time ;)
		time.Sleep(10 * time.Second)
	}
}

// readARP watches a handle for incoming ARP responses we might care about, and prints them.
//
// readARP loops until 'stop' is closed.
func readARP(handle *pcap.Handle, iface *net.Interface, stop chan struct{}) {
	src := gopacket.NewPacketSource(handle, layers.LayerTypeEthernet)
	in := src.Packets()
	for {
		var packet gopacket.Packet
		select {
		case <-stop:
			return
		case packet = <-in:
			arpLayer := packet.Layer(layers.LayerTypeARP)
			if arpLayer == nil {
				continue
			}
			arp := arpLayer.(*layers.ARP)
			if arp.Operation != layers.ARPReply || bytes.Equal([]byte(iface.HardwareAddr), arp.SourceHwAddress) {
				// This is a packet I sent.
				continue
			}
			// Note:  we might get some packets here that aren't responses to ones we've sent,
			// if for example someone else sends US an ARP request.  Doesn't much matter, though...
			// all information is good information :)
			log.Printf("IP %v is at %v", net.IP(arp.SourceProtAddress), net.HardwareAddr(arp.SourceHwAddress))
		}
	}
}

// writeARP writes an ARP request for each address on our local network to the
// pcap handle.
func writeARP(handle *pcap.Handle, iface *net.Interface, addr *net.IPNet) error {
	// Set up all the layers' fields we can.
	eth := layers.Ethernet{
		SrcMAC:       iface.HardwareAddr,
		DstMAC:       net.HardwareAddr{0xff, 0xff, 0xff, 0xff, 0xff, 0xff},
		EthernetType: layers.EthernetTypeARP,
	}
	arp := layers.ARP{
		AddrType:          layers.LinkTypeEthernet,
		Protocol:          layers.EthernetTypeIPv4,
		HwAddressSize:     6,
		ProtAddressSize:   4,
		Operation:         layers.ARPRequest,
		SourceHwAddress:   []byte(iface.HardwareAddr),
		SourceProtAddress: []byte(addr.IP),
		DstHwAddress:      []byte{0, 0, 0, 0, 0, 0},
	}
	// Set up buffer and options for serialization.
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	// Send one packet for every address.
	for _, ip := range ips(addr) {
		arp.DstProtAddress = []byte(ip)
		gopacket.SerializeLayers(buf, opts, &eth, &arp)
		if err := handle.WritePacketData(buf.Bytes()); err != nil {
			return err
		}
	}
	return nil
}

// ips is a simple and not very good method for getting all IPv4 addresses from a
// net.IPNet.  It returns all IPs it can over the channel it sends back, closing
// the channel when done.
func ips(n *net.IPNet) (out []net.IP) {
	num := binary.BigEndian.Uint32([]byte(n.IP))
	mask := binary.BigEndian.Uint32([]byte(n.Mask))
	net := num & mask
	broadcast := net | ^mask
	for net++; net < broadcast-1; net++ {
		var buf [4]byte
		binary.BigEndian.PutUint32(buf[:], net)
		out = append(out, net.IP(buf[:]))
	}
	return
}
<commit_msg>Fix for loop in ips<commit_after>// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

// arpscan implements ARP scanning of all interfaces' local networks using
// gopacket and its subpackages.  This example shows, among other things:
//   * Generating and sending packet data
//   * Reading in packet data and interpreting it
//   * Use of the 'pcap' subpackage for reading/writing
package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"log"
	"net"
	"sync"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

func main() {
	// Get a list of all interfaces.
	ifaces, err := net.Interfaces()
	if err != nil {
		panic(err)
	}

	var wg sync.WaitGroup
	for _, iface := range ifaces {
		wg.Add(1)
		// Start up a scan on each interface.
		go func(iface net.Interface) {
			defer wg.Done()
			if err := scan(&iface); err != nil {
				log.Printf("interface %v: %v", iface.Name, err)
			}
		}(iface)
	}
	// Wait for all interfaces' scans to complete.  They'll try to run
	// forever, but will stop on an error, so if we get past this Wait
	// it means all attempts to write have failed.
	wg.Wait()
}

// scan scans an individual interface's local network for machines using ARP requests/replies.
//
// scan loops forever, sending packets out regularly.  It returns an error if
// it's ever unable to write a packet.
func scan(iface *net.Interface) error {
	// We just look for IPv4 addresses, so try to find if the interface has one.
	var addr *net.IPNet
	if addrs, err := iface.Addrs(); err != nil {
		return err
	} else {
		for _, a := range addrs {
			if ipnet, ok := a.(*net.IPNet); ok {
				if ip4 := ipnet.IP.To4(); ip4 != nil {
					addr = &net.IPNet{
						IP:   ip4,
						Mask: ipnet.Mask[len(ipnet.Mask)-4:],
					}
					break
				}
			}
		}
	}
	// Sanity-check that the interface has a good address.
	if addr == nil {
		return errors.New("no good IP network found")
	} else if addr.IP[0] == 127 {
		return errors.New("skipping localhost")
	} else if addr.Mask[0] != 0xff || addr.Mask[1] != 0xff {
		return errors.New("mask means network is too large")
	}
	log.Printf("Using network range %v for interface %v", addr, iface.Name)

	// Open up a pcap handle for packet reads/writes.
	handle, err := pcap.OpenLive(iface.Name, 65536, true, pcap.BlockForever)
	if err != nil {
		return err
	}
	defer handle.Close()

	// Start up a goroutine to read in packet data.
	stop := make(chan struct{})
	go readARP(handle, iface, stop)
	defer close(stop)
	for {
		// Write our scan packets out to the handle.
		if err := writeARP(handle, iface, addr); err != nil {
			log.Printf("error writing packets on %v: %v", iface.Name, err)
			return err
		}
		// We don't know exactly how long it'll take for packets to be
		// sent back to us, but 10 seconds should be more than enough
		// time ;)
		time.Sleep(10 * time.Second)
	}
}

// readARP watches a handle for incoming ARP responses we might care about, and prints them.
//
// readARP loops until 'stop' is closed.
func readARP(handle *pcap.Handle, iface *net.Interface, stop chan struct{}) {
	src := gopacket.NewPacketSource(handle, layers.LayerTypeEthernet)
	in := src.Packets()
	for {
		var packet gopacket.Packet
		select {
		case <-stop:
			return
		case packet = <-in:
			arpLayer := packet.Layer(layers.LayerTypeARP)
			if arpLayer == nil {
				continue
			}
			arp := arpLayer.(*layers.ARP)
			if arp.Operation != layers.ARPReply || bytes.Equal([]byte(iface.HardwareAddr), arp.SourceHwAddress) {
				// This is a packet I sent.
				continue
			}
			// Note:  we might get some packets here that aren't responses to ones we've sent,
			// if for example someone else sends US an ARP request.  Doesn't much matter, though...
			// all information is good information :)
			log.Printf("IP %v is at %v", net.IP(arp.SourceProtAddress), net.HardwareAddr(arp.SourceHwAddress))
		}
	}
}

// writeARP writes an ARP request for each address on our local network to the
// pcap handle.
func writeARP(handle *pcap.Handle, iface *net.Interface, addr *net.IPNet) error {
	// Set up all the layers' fields we can.
	eth := layers.Ethernet{
		SrcMAC:       iface.HardwareAddr,
		DstMAC:       net.HardwareAddr{0xff, 0xff, 0xff, 0xff, 0xff, 0xff},
		EthernetType: layers.EthernetTypeARP,
	}
	arp := layers.ARP{
		AddrType:          layers.LinkTypeEthernet,
		Protocol:          layers.EthernetTypeIPv4,
		HwAddressSize:     6,
		ProtAddressSize:   4,
		Operation:         layers.ARPRequest,
		SourceHwAddress:   []byte(iface.HardwareAddr),
		SourceProtAddress: []byte(addr.IP),
		DstHwAddress:      []byte{0, 0, 0, 0, 0, 0},
	}
	// Set up buffer and options for serialization.
	buf := gopacket.NewSerializeBuffer()
	opts := gopacket.SerializeOptions{
		FixLengths:       true,
		ComputeChecksums: true,
	}
	// Send one packet for every address.
	for _, ip := range ips(addr) {
		arp.DstProtAddress = []byte(ip)
		gopacket.SerializeLayers(buf, opts, &eth, &arp)
		if err := handle.WritePacketData(buf.Bytes()); err != nil {
			return err
		}
	}
	return nil
}

// ips is a simple and not very good method for getting all IPv4 addresses from a
// net.IPNet.  It returns all IPs it can over the channel it sends back, closing
// the channel when done.
func ips(n *net.IPNet) (out []net.IP) {
	num := binary.BigEndian.Uint32([]byte(n.IP))
	mask := binary.BigEndian.Uint32([]byte(n.Mask))
	net := num & mask
	broadcast := net | ^mask
	for net++; net < broadcast; net++ {
		var buf [4]byte
		binary.BigEndian.PutUint32(buf[:], net)
		out = append(out, net.IP(buf[:]))
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/couchbaselabs/logg"
	"github.com/tleyden/ocr-worker"
)

// This assumes that there is a rabbit mq running
// To test it, fire up a webserver and send it a curl request

func init() {
	logg.LogKeys["OCR_CLIENT"] = true
	logg.LogKeys["OCR_WORKER"] = true
	logg.LogKeys["OCR_HTTP"] = true
	logg.LogKeys["OCR_TESSERACT"] = true
}

func main() {

	rabbitConfig := ocrworker.RabbitConfig{
		AmqpURI:            "amqp://guest:guest@localhost:5672/",
		Exchange:           "test-exchange",
		ExchangeType:       "direct",
		RoutingKey:         "test-key",
		CallbackRoutingKey: "callback-key",
		Reliable:           true,
		QueueName:          "test-queue",
		CallbackQueueName:  "callback-queue",
	}

	ocrWorker, err := ocrworker.NewOcrRpcWorker(rabbitConfig)
	if err != nil {
		logg.LogPanic("Could not create rpc worker")
	}
	ocrWorker.Run()

	select {}

	// blockForeverChan := make(chan bool)
	// <-blockForeverChan

}
<commit_msg>remove commented code<commit_after>package main

import (
	"github.com/couchbaselabs/logg"
	"github.com/tleyden/ocr-worker"
)

// This assumes that there is a rabbit mq running
// To test it, fire up a webserver and send it a curl request

func init() {
	logg.LogKeys["OCR_CLIENT"] = true
	logg.LogKeys["OCR_WORKER"] = true
	logg.LogKeys["OCR_HTTP"] = true
	logg.LogKeys["OCR_TESSERACT"] = true
}

func main() {

	rabbitConfig := ocrworker.RabbitConfig{
		AmqpURI:            "amqp://guest:guest@localhost:5672/",
		Exchange:           "test-exchange",
		ExchangeType:       "direct",
		RoutingKey:         "test-key",
		CallbackRoutingKey: "callback-key",
		Reliable:           true,
		QueueName:          "test-queue",
		CallbackQueueName:  "callback-queue",
	}

	ocrWorker, err := ocrworker.NewOcrRpcWorker(rabbitConfig)
	if err != nil {
		logg.LogPanic("Could not create rpc worker")
	}
	ocrWorker.Run()

	select {}

}
<|endoftext|>