<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"runtime"
	"strings"
	"time"

	"devroom.ru/lomik/carbon/cache"
	"devroom.ru/lomik/carbon/logging"
	"devroom.ru/lomik/carbon/persister"
	"devroom.ru/lomik/carbon/receiver"

	"github.com/BurntSushi/toml"
	"github.com/Sirupsen/logrus"
)

import _ "net/http/pprof"

// Duration wrapper time.Duration for TOML
type Duration struct {
	time.Duration
}

var _ toml.TextMarshaler = &Duration{}

// UnmarshalText from TOML
func (d *Duration) UnmarshalText(text []byte) error {
	var err error
	d.Duration, err = time.ParseDuration(string(text))
	return err
}

// MarshalText encode text with TOML format
func (d *Duration) MarshalText() ([]byte, error) {
	return []byte(d.Duration.String()), nil
}

// Value return time.Duration value
func (d *Duration) Value() time.Duration {
	return d.Duration
}

type commonConfig struct {
	Logfile     string `toml:"logfile"`
	GraphPrefix string `toml:"graph-prefix"`
	MaxCPU      int    `toml:"max-cpu"`
}

type whisperConfig struct {
	DataDir string `toml:"data-dir"`
	Schemas string `toml:"schemas-file"`
	Workers int    `toml:"workers"`
	Enabled bool   `toml:"enabled"`
}

type cacheConfig struct {
	MaxSize     int `toml:"max-size"`
	InputBuffer int `toml:"input-buffer"`
}

type udpConfig struct {
	Listen  string `toml:"listen"`
	Enabled bool   `toml:"enabled"`
}

type tcpConfig struct {
	Listen  string `toml:"listen"`
	Enabled bool   `toml:"enabled"`
}

type carbonlinkConfig struct {
	Listen       string    `toml:"listen"`
	Enabled      bool      `toml:"enabled"`
	ReadTimeout  *Duration `toml:"read-timeout"`
	QueryTimeout *Duration `toml:"query-timeout"`
}

type pprofConfig struct {
	Listen  string `toml:"listen"`
	Enabled bool   `toml:"enabled"`
}

// Config ...
type Config struct {
	Common     commonConfig     `toml:"common"`
	Whisper    whisperConfig    `toml:"whisper"`
	Cache      cacheConfig      `toml:"cache"`
	Udp        udpConfig        `toml:"udp"`
	Tcp        tcpConfig        `toml:"tcp"`
	Carbonlink carbonlinkConfig `toml:"carbonlink"`
	Pprof      pprofConfig      `toml:"pprof"`
}

func newConfig() *Config {
	cfg := &Config{
		Common: commonConfig{
			Logfile:     "",
			GraphPrefix: "carbon.agents.{host}.",
			MaxCPU:      1,
		},
		Whisper: whisperConfig{
			DataDir: "/data/graphite/whisper/",
			Schemas: "/data/graphite/schemas",
			Enabled: true,
			Workers: 1,
		},
		Cache: cacheConfig{
			MaxSize:     1000000,
			InputBuffer: 51200,
		},
		Udp: udpConfig{
			Listen:  ":2003",
			Enabled: true,
		},
		Tcp: tcpConfig{
			Listen:  ":2003",
			Enabled: true,
		},
		Carbonlink: carbonlinkConfig{
			Listen:  "127.0.0.1:7002",
			Enabled: true,
			ReadTimeout: &Duration{
				Duration: 30 * time.Second,
			},
			QueryTimeout: &Duration{
				Duration: 100 * time.Millisecond,
			},
		},
		Pprof: pprofConfig{
			Listen:  "localhost:7007",
			Enabled: false,
		},
	}

	return cfg
}

// PrintConfig ...
func PrintConfig(cfg interface{}) error {
	buf := new(bytes.Buffer)

	encoder := toml.NewEncoder(buf)
	encoder.Indent = ""

	if err := encoder.Encode(cfg); err != nil {
		return err
	}

	fmt.Print(buf.String())
	return nil
}

// ParseConfig ...
func ParseConfig(filename string, cfg interface{}) error {
	if filename != "" {
		if _, err := toml.DecodeFile(filename, cfg); err != nil {
			return err
		}
	}
	return nil
}

func main() {

	/* CONFIG start */
	configFile := flag.String("config", "", "Filename of config")
	printDefaultConfig := flag.Bool("config-print-default", false, "Print default config")

	flag.Parse()

	cfg := newConfig()

	if *printDefaultConfig {
		if err := PrintConfig(cfg); err != nil {
			log.Fatal(err)
		}
		os.Exit(0)
	}

	if err := ParseConfig(*configFile, cfg); err != nil {
		log.Fatal(err)
	}

	logging.SetFile(cfg.Common.Logfile)
	logrus.SetLevel(logrus.DebugLevel)

	runtime.GOMAXPROCS(cfg.Common.MaxCPU)

	// pp.Println(cfg)
	/* CONFIG end */

	// pprof
	if cfg.Pprof.Enabled {
		go func() {
			logrus.Fatal(http.ListenAndServe(cfg.Pprof.Listen, nil))
		}()
	}

	// carbon-cache prefix
	if hostname, err := os.Hostname(); err == nil {
		hostname = strings.Replace(hostname, ".", "_", -1)
		cfg.Common.GraphPrefix = strings.Replace(cfg.Common.GraphPrefix, "{host}", hostname, -1)
	} else {
		cfg.Common.GraphPrefix = strings.Replace(cfg.Common.GraphPrefix, "{host}", "localhost", -1)
	}

	core := cache.New()
	core.SetGraphPrefix(cfg.Common.GraphPrefix)
	core.SetMaxSize(cfg.Cache.MaxSize)
	core.SetInputCapacity(cfg.Cache.InputBuffer)
	core.Start()
	defer core.Stop()

	/* UDP start */
	udpCfg := cfg.Udp
	if udpCfg.Enabled {
		udpAddr, err := net.ResolveUDPAddr("udp", udpCfg.Listen)
		if err != nil {
			log.Fatal(err)
		}

		udpListener := receiver.NewUDP(core.In())
		udpListener.SetGraphPrefix(cfg.Common.GraphPrefix)

		defer udpListener.Stop()
		if err = udpListener.Listen(udpAddr); err != nil {
			log.Fatal(err)
		}
	}
	/* UDP end */

	/* TCP start */
	tcpCfg := cfg.Tcp

	if tcpCfg.Enabled {
		tcpAddr, err := net.ResolveTCPAddr("tcp", tcpCfg.Listen)
		if err != nil {
			log.Fatal(err)
		}

		tcpListener := receiver.NewTCP(core.In())
		tcpListener.SetGraphPrefix(cfg.Common.GraphPrefix)

		defer tcpListener.Stop()
		if err = tcpListener.Listen(tcpAddr); err != nil {
			log.Fatal(err)
		}
	}
	/* TCP end */

	/* WHISPER start */
	if cfg.Whisper.Enabled {
		whisperSchemas, err := persister.ReadWhisperSchemas(cfg.Whisper.Schemas)
		if err != nil {
			log.Fatal(err)
		}

		whisperPersister := persister.NewWhisper(cfg.Whisper.DataDir, whisperSchemas, core.Out())
		whisperPersister.SetGraphPrefix(cfg.Common.GraphPrefix)
		whisperPersister.SetWorkers(cfg.Whisper.Workers)

		whisperPersister.Start()
		defer whisperPersister.Stop()
	}
	/* WHISPER end */

	/* CARBONLINK start */
	if cfg.Carbonlink.Enabled {
		linkAddr, err := net.ResolveTCPAddr("tcp", cfg.Carbonlink.Listen)
		if err != nil {
			log.Fatal(err)
		}

		carbonlink := cache.NewCarbonlinkListener(core.Query())
		carbonlink.SetReadTimeout(cfg.Carbonlink.ReadTimeout.Value())
		carbonlink.SetQueryTimeout(cfg.Carbonlink.QueryTimeout.Value())

		defer carbonlink.Stop()
		if err = carbonlink.Listen(linkAddr); err != nil {
			log.Fatal(err)
		}

	}
	/* CARBONLINK end */

	logrus.Info("started")
	select {}
}
<commit_msg>Enable pickle listener in main<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"runtime"
	"strings"
	"time"

	"devroom.ru/lomik/carbon/cache"
	"devroom.ru/lomik/carbon/logging"
	"devroom.ru/lomik/carbon/persister"
	"devroom.ru/lomik/carbon/receiver"

	"github.com/BurntSushi/toml"
	"github.com/Sirupsen/logrus"
)

import _ "net/http/pprof"

// Duration wrapper time.Duration for TOML
type Duration struct {
	time.Duration
}

var _ toml.TextMarshaler = &Duration{}

// UnmarshalText from TOML
func (d *Duration) UnmarshalText(text []byte) error {
	var err error
	d.Duration, err = time.ParseDuration(string(text))
	return err
}

// MarshalText encode text with TOML format
func (d *Duration) MarshalText() ([]byte, error) {
	return []byte(d.Duration.String()), nil
}

// Value return time.Duration value
func (d *Duration) Value() time.Duration {
	return d.Duration
}

type commonConfig struct {
	Logfile     string `toml:"logfile"`
	GraphPrefix string `toml:"graph-prefix"`
	MaxCPU      int    `toml:"max-cpu"`
}

type whisperConfig struct {
	DataDir string `toml:"data-dir"`
	Schemas string `toml:"schemas-file"`
	Workers int    `toml:"workers"`
	Enabled bool   `toml:"enabled"`
}

type cacheConfig struct {
	MaxSize     int `toml:"max-size"`
	InputBuffer int `toml:"input-buffer"`
}

type udpConfig struct {
	Listen  string `toml:"listen"`
	Enabled bool   `toml:"enabled"`
}

type tcpConfig struct {
	Listen  string `toml:"listen"`
	Enabled bool   `toml:"enabled"`
}

type carbonlinkConfig struct {
	Listen       string    `toml:"listen"`
	Enabled      bool      `toml:"enabled"`
	ReadTimeout  *Duration `toml:"read-timeout"`
	QueryTimeout *Duration `toml:"query-timeout"`
}

type pprofConfig struct {
	Listen  string `toml:"listen"`
	Enabled bool   `toml:"enabled"`
}

// Config ...
type Config struct {
	Common     commonConfig     `toml:"common"`
	Whisper    whisperConfig    `toml:"whisper"`
	Cache      cacheConfig      `toml:"cache"`
	Udp        udpConfig        `toml:"udp"`
	Tcp        tcpConfig        `toml:"tcp"`
	Pickle     tcpConfig        `toml:"pickle"`
	Carbonlink carbonlinkConfig `toml:"carbonlink"`
	Pprof      pprofConfig      `toml:"pprof"`
}

func newConfig() *Config {
	cfg := &Config{
		Common: commonConfig{
			Logfile:     "",
			GraphPrefix: "carbon.agents.{host}.",
			MaxCPU:      1,
		},
		Whisper: whisperConfig{
			DataDir: "/data/graphite/whisper/",
			Schemas: "/data/graphite/schemas",
			Enabled: true,
			Workers: 1,
		},
		Cache: cacheConfig{
			MaxSize:     1000000,
			InputBuffer: 51200,
		},
		Udp: udpConfig{
			Listen:  ":2003",
			Enabled: true,
		},
		Tcp: tcpConfig{
			Listen:  ":2003",
			Enabled: true,
		},
		Pickle: tcpConfig{
			Listen:  ":2004",
			Enabled: true,
		},
		Carbonlink: carbonlinkConfig{
			Listen:  "127.0.0.1:7002",
			Enabled: true,
			ReadTimeout: &Duration{
				Duration: 30 * time.Second,
			},
			QueryTimeout: &Duration{
				Duration: 100 * time.Millisecond,
			},
		},
		Pprof: pprofConfig{
			Listen:  "localhost:7007",
			Enabled: false,
		},
	}

	return cfg
}

// PrintConfig ...
func PrintConfig(cfg interface{}) error {
	buf := new(bytes.Buffer)

	encoder := toml.NewEncoder(buf)
	encoder.Indent = ""

	if err := encoder.Encode(cfg); err != nil {
		return err
	}

	fmt.Print(buf.String())
	return nil
}

// ParseConfig ...
func ParseConfig(filename string, cfg interface{}) error {
	if filename != "" {
		if _, err := toml.DecodeFile(filename, cfg); err != nil {
			return err
		}
	}
	return nil
}

func main() {

	/* CONFIG start */
	configFile := flag.String("config", "", "Filename of config")
	printDefaultConfig := flag.Bool("config-print-default", false, "Print default config")

	flag.Parse()

	cfg := newConfig()

	if *printDefaultConfig {
		if err := PrintConfig(cfg); err != nil {
			log.Fatal(err)
		}
		os.Exit(0)
	}

	if err := ParseConfig(*configFile, cfg); err != nil {
		log.Fatal(err)
	}

	logging.SetFile(cfg.Common.Logfile)
	logrus.SetLevel(logrus.DebugLevel)

	runtime.GOMAXPROCS(cfg.Common.MaxCPU)

	// pp.Println(cfg)
	/* CONFIG end */

	// pprof
	if cfg.Pprof.Enabled {
		go func() {
			logrus.Fatal(http.ListenAndServe(cfg.Pprof.Listen, nil))
		}()
	}

	// carbon-cache prefix
	if hostname, err := os.Hostname(); err == nil {
		hostname = strings.Replace(hostname, ".", "_", -1)
		cfg.Common.GraphPrefix = strings.Replace(cfg.Common.GraphPrefix, "{host}", hostname, -1)
	} else {
		cfg.Common.GraphPrefix = strings.Replace(cfg.Common.GraphPrefix, "{host}", "localhost", -1)
	}

	core := cache.New()
	core.SetGraphPrefix(cfg.Common.GraphPrefix)
	core.SetMaxSize(cfg.Cache.MaxSize)
	core.SetInputCapacity(cfg.Cache.InputBuffer)
	core.Start()
	defer core.Stop()

	/* UDP start */
	udpCfg := cfg.Udp
	if udpCfg.Enabled {
		udpAddr, err := net.ResolveUDPAddr("udp", udpCfg.Listen)
		if err != nil {
			log.Fatal(err)
		}

		udpListener := receiver.NewUDP(core.In())
		udpListener.SetGraphPrefix(cfg.Common.GraphPrefix)

		defer udpListener.Stop()
		if err = udpListener.Listen(udpAddr); err != nil {
			log.Fatal(err)
		}
	}
	/* UDP end */

	/* TCP start */
	tcpCfg := cfg.Tcp

	if tcpCfg.Enabled {
		tcpAddr, err := net.ResolveTCPAddr("tcp", tcpCfg.Listen)
		if err != nil {
			log.Fatal(err)
		}

		tcpListener := receiver.NewTCP(core.In())
		tcpListener.SetGraphPrefix(cfg.Common.GraphPrefix)

		defer tcpListener.Stop()
		if err = tcpListener.Listen(tcpAddr); err != nil {
			log.Fatal(err)
		}
	}
	/* TCP end */

	/* PICKLE start */
	pickleCfg := cfg.Pickle

	if pickleCfg.Enabled {
		pickleAddr, err := net.ResolveTCPAddr("tcp", pickleCfg.Listen)
		if err != nil {
			log.Fatal(err)
		}

		pickleListener := receiver.NewPickle(core.In())
		pickleListener.SetGraphPrefix(cfg.Common.GraphPrefix)

		defer pickleListener.Stop()
		if err = pickleListener.Listen(pickleAddr); err != nil {
			log.Fatal(err)
		}
	}
	/* PICKLE end */

	/* WHISPER start */
	if cfg.Whisper.Enabled {
		whisperSchemas, err := persister.ReadWhisperSchemas(cfg.Whisper.Schemas)
		if err != nil {
			log.Fatal(err)
		}

		whisperPersister := persister.NewWhisper(cfg.Whisper.DataDir, whisperSchemas, core.Out())
		whisperPersister.SetGraphPrefix(cfg.Common.GraphPrefix)
		whisperPersister.SetWorkers(cfg.Whisper.Workers)

		whisperPersister.Start()
		defer whisperPersister.Stop()
	}
	/* WHISPER end */

	/* CARBONLINK start */
	if cfg.Carbonlink.Enabled {
		linkAddr, err := net.ResolveTCPAddr("tcp", cfg.Carbonlink.Listen)
		if err != nil {
			log.Fatal(err)
		}

		carbonlink := cache.NewCarbonlinkListener(core.Query())
		carbonlink.SetReadTimeout(cfg.Carbonlink.ReadTimeout.Value())
		carbonlink.SetQueryTimeout(cfg.Carbonlink.QueryTimeout.Value())

		defer carbonlink.Stop()
		if err = carbonlink.Listen(linkAddr); err != nil {
			log.Fatal(err)
		}

	}
	/* CARBONLINK end */

	logrus.Info("started")
	select {}
}
<|endoftext|>


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

package api

import (
	log "github.com/Sirupsen/logrus"
	"github.com/fukata/golang-stats-api-handler"
	"github.com/gorilla/mux"
	"net/http"
	"strconv"
)

const (
	_ = iota
	REQ_NEIGHBOR
	REQ_NEIGHBORS
	REQ_ADJ_RIB_IN
	REQ_ADJ_RIB_OUT
	REQ_LOCAL_RIB
)

const (
	BASE_VERSION       = "/v1"
	NEIGHBOR           = "/bgp/neighbor"
	NEIGHBORS          = "/bgp/neighbors"
	ADJ_RIB_IN         = "/bgp/adj-rib-in"
	ADJ_RIB_OUT        = "/bgp/adj-rib-out"
	ADJ_RIB_LOCAL      = "/bgp/adj-rib-local"
	ADJ_RIB_LOCAL_BEST = "/bgp/adj-rib-local/best"

	PARAM_REMOTE_PEER_ADDR = "remotePeerAddr"
	STATS                  = "/stats"
)

const REST_PORT = 8080

// trigger struct for exchanging information in the rest and peer.
// rest and peer operated at different thread.

type RestRequest struct {
	RequestType int
	RemoteAddr  string
	ResponseCh  chan *RestResponse
	Err         error
}

func NewRestRequest(reqType int, remoteAddr string) *RestRequest {
	r := &RestRequest{
		RequestType: reqType,
		RemoteAddr:  remoteAddr,
		ResponseCh:  make(chan *RestResponse),
	}
	return r
}

type RestResponse struct {
	ResponseErr error
	Data        []byte
}

func (r *RestResponse) Err() error {
	return r.ResponseErr
}

type RestServer struct {
	port        int
	bgpServerCh chan *RestRequest
}

func NewRestServer(port int, bgpServerCh chan *RestRequest) *RestServer {
	rs := &RestServer{
		port:        port,
		bgpServerCh: bgpServerCh}
	return rs
}

// Main thread of rest service.
// URL than can receive.
//   get state of neighbors.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbors
//   get state of neighbor.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbor/<remote address of target neighbor>
//   get adj-rib-in of each neighbor.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbor/<remote address of target neighbor>/adj-rib-in
//   get adj-rib-out of each neighbor.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbor/<remote address of target neighbor>/adj-rib-out
//   get local-rib of each neighbor.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbor/<remote address of target neighbor>/local-rib
func (rs *RestServer) Serve() {
	neighbor := BASE_VERSION + NEIGHBOR
	neighbors := BASE_VERSION + NEIGHBORS
	// adjRibIn := BASE_VERSION + ADJ_RIB_IN
	// adjRibOut := BASE_VERSION + ADJ_RIB_OUT
	// adjRibLocal := BASE_VERSION + ADJ_RIB_LOCAL
	// adjRibLocalBest := BASE_VERSION + ADJ_RIB_LOCAL_BEST

	r := mux.NewRouter()
	// set URLs
	r.HandleFunc(neighbors, rs.Neighbors).Methods("GET")
	r.HandleFunc(neighbor+"/{"+PARAM_REMOTE_PEER_ADDR+"}", rs.Neighbor).Methods("GET")
	// r.HandleFunc(adjRibIn+"/{"+PARAM_REMOTE_PEER_ADDR+"}", rs.AdjRibIn).Methods("GET")
	// r.HandleFunc(adjRibOut+"/{"+PARAM_REMOTE_PEER_ADDR+"}", rs.AdjRibOut).Methods("GET")
	r.HandleFunc(neighbor+"/{"+PARAM_REMOTE_PEER_ADDR+"}/"+"local-rib", rs.NeighborLocalRib).Methods("GET")
	// stats
	r.HandleFunc(STATS, stats_api.Handler)

	// Handler when not found url
	r.NotFoundHandler = http.HandlerFunc(NotFoundHandler)
	http.Handle("/", r)

	http.ListenAndServe(":"+strconv.Itoa(rs.port), nil)

}

// TODO: merge the above function
func (rs *RestServer) neighbor(w http.ResponseWriter, r *http.Request, reqType int) {
	params := mux.Vars(r)
	remoteAddr, found := params[PARAM_REMOTE_PEER_ADDR]
	if !found {
		errStr := "neighbor address is not specified"
		log.Debug(errStr)
		http.Error(w, errStr, http.StatusInternalServerError)
		return
	}

	log.Debugf("Look up neighbor with the remote address : %v", remoteAddr)

	//Send channel of request parameter.
	req := NewRestRequest(reqType, remoteAddr)
	rs.bgpServerCh <- req

	//Wait response
	res := <-req.ResponseCh
	if e := res.Err(); e != nil {
		log.Debug(e.Error())
		http.Error(w, e.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Write(res.Data)
}

func (rs *RestServer) Neighbor(w http.ResponseWriter, r *http.Request) {
	rs.neighbor(w, r, REQ_NEIGHBOR)
}

func (rs *RestServer) NeighborLocalRib(w http.ResponseWriter, r *http.Request) {
	rs.neighbor(w, r, REQ_LOCAL_RIB)
}

func (rs *RestServer) Neighbors(w http.ResponseWriter, r *http.Request) {
	//Send channel of request parameter.
	req := NewRestRequest(REQ_NEIGHBORS, "")
	rs.bgpServerCh <- req

	//Wait response
	res := <-req.ResponseCh
	if e := res.Err(); e != nil {
		log.Debug(e.Error())
		http.Error(w, e.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Write(res.Data)
}

func NotFoundHandler(w http.ResponseWriter, r *http.Request) {
	http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}
<commit_msg>rest: specified HTTP method<commit_after>// Copyright (C) 2014 Nippon Telegraph and Telephone Corporation.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package api

import (
	log "github.com/Sirupsen/logrus"
	"github.com/fukata/golang-stats-api-handler"
	"github.com/gorilla/mux"
	"net/http"
	"strconv"
)

const (
	_ = iota
	REQ_NEIGHBOR
	REQ_NEIGHBORS
	REQ_ADJ_RIB_IN
	REQ_ADJ_RIB_OUT
	REQ_LOCAL_RIB
)

const (
	BASE_VERSION       = "/v1"
	NEIGHBOR           = "/bgp/neighbor"
	NEIGHBORS          = "/bgp/neighbors"
	ADJ_RIB_IN         = "/bgp/adj-rib-in"
	ADJ_RIB_OUT        = "/bgp/adj-rib-out"
	ADJ_RIB_LOCAL      = "/bgp/adj-rib-local"
	ADJ_RIB_LOCAL_BEST = "/bgp/adj-rib-local/best"

	PARAM_REMOTE_PEER_ADDR = "remotePeerAddr"
	STATS                  = "/stats"
)

const REST_PORT = 8080

// trigger struct for exchanging information in the rest and peer.
// rest and peer operated at different thread.

type RestRequest struct {
	RequestType int
	RemoteAddr  string
	ResponseCh  chan *RestResponse
	Err         error
}

func NewRestRequest(reqType int, remoteAddr string) *RestRequest {
	r := &RestRequest{
		RequestType: reqType,
		RemoteAddr:  remoteAddr,
		ResponseCh:  make(chan *RestResponse),
	}
	return r
}

type RestResponse struct {
	ResponseErr error
	Data        []byte
}

func (r *RestResponse) Err() error {
	return r.ResponseErr
}

type RestServer struct {
	port        int
	bgpServerCh chan *RestRequest
}

func NewRestServer(port int, bgpServerCh chan *RestRequest) *RestServer {
	rs := &RestServer{
		port:        port,
		bgpServerCh: bgpServerCh}
	return rs
}

// Main thread of rest service.
// URL than can receive.
//   get state of neighbors.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbors
//   get state of neighbor.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbor/<remote address of target neighbor>
//   get adj-rib-in of each neighbor.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbor/<remote address of target neighbor>/adj-rib-in
//   get adj-rib-out of each neighbor.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbor/<remote address of target neighbor>/adj-rib-out
//   get local-rib of each neighbor.
//     -- curl -i -X GET http://<ownIP>:8080/v1/bgp/neighbor/<remote address of target neighbor>/local-rib
func (rs *RestServer) Serve() {
	neighbor := BASE_VERSION + NEIGHBOR
	neighbors := BASE_VERSION + NEIGHBORS
	// adjRibIn := BASE_VERSION + ADJ_RIB_IN
	// adjRibOut := BASE_VERSION + ADJ_RIB_OUT
	// adjRibLocal := BASE_VERSION + ADJ_RIB_LOCAL
	// adjRibLocalBest := BASE_VERSION + ADJ_RIB_LOCAL_BEST

	r := mux.NewRouter()
	// set URLs
	r.HandleFunc(neighbors, rs.Neighbors).Methods("GET")
	r.HandleFunc(neighbor+"/{"+PARAM_REMOTE_PEER_ADDR+"}", rs.Neighbor).Methods("GET")
	// r.HandleFunc(adjRibIn+"/{"+PARAM_REMOTE_PEER_ADDR+"}", rs.AdjRibIn).Methods("GET")
	// r.HandleFunc(adjRibOut+"/{"+PARAM_REMOTE_PEER_ADDR+"}", rs.AdjRibOut).Methods("GET")
	r.HandleFunc(neighbor+"/{"+PARAM_REMOTE_PEER_ADDR+"}/"+"local-rib", rs.NeighborLocalRib).Methods("GET")
	// stats
	r.HandleFunc(STATS, stats_api.Handler).Methods("GET")

	// Handler when not found url
	r.NotFoundHandler = http.HandlerFunc(NotFoundHandler)
	http.Handle("/", r)

	http.ListenAndServe(":"+strconv.Itoa(rs.port), nil)

}

// TODO: merge the above function
func (rs *RestServer) neighbor(w http.ResponseWriter, r *http.Request, reqType int) {
	params := mux.Vars(r)
	remoteAddr, found := params[PARAM_REMOTE_PEER_ADDR]
	if !found {
		errStr := "neighbor address is not specified"
		log.Debug(errStr)
		http.Error(w, errStr, http.StatusInternalServerError)
		return
	}

	log.Debugf("Look up neighbor with the remote address : %v", remoteAddr)

	//Send channel of request parameter.
	req := NewRestRequest(reqType, remoteAddr)
	rs.bgpServerCh <- req

	//Wait response
	res := <-req.ResponseCh
	if e := res.Err(); e != nil {
		log.Debug(e.Error())
		http.Error(w, e.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Write(res.Data)
}

func (rs *RestServer) Neighbor(w http.ResponseWriter, r *http.Request) {
	rs.neighbor(w, r, REQ_NEIGHBOR)
}

func (rs *RestServer) NeighborLocalRib(w http.ResponseWriter, r *http.Request) {
	rs.neighbor(w, r, REQ_LOCAL_RIB)
}

func (rs *RestServer) Neighbors(w http.ResponseWriter, r *http.Request) {
	//Send channel of request parameter.
	req := NewRestRequest(REQ_NEIGHBORS, "")
	rs.bgpServerCh <- req

	//Wait response
	res := <-req.ResponseCh
	if e := res.Err(); e != nil {
		log.Debug(e.Error())
		http.Error(w, e.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Write(res.Data)
}

func NotFoundHandler(w http.ResponseWriter, r *http.Request) {
	http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
}
<|endoftext|>


<commit_before>package caspercloud

import (
	"bufio"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

type CasperCmd struct {
	proxyServer   string
	id            string
	tmpl          string
	userName      string
	passWord      string
	message       chan *Output
	input         chan map[string]string
	isKill        bool
	isFinish      bool
	args          map[string]string
	privateKey    *rsa.PrivateKey
	mailProcessor *MailProcessor
}

func NewCasperCmd(id, tmpl, proxyServer string) *CasperCmd {
	ret := &CasperCmd{
		proxyServer:   proxyServer,
		id:            id,
		tmpl:          tmpl,
		userName:      "",
		passWord:      "",
		message:       make(chan *Output, 5),
		input:         make(chan map[string]string, 5),
		args:          make(map[string]string),
		isKill:        false,
		isFinish:      false,
		mailProcessor: NewMailProcessor("server_list.json"),
	}
	var err error
	ret.privateKey, err = generateRSAKey()
	if err != nil {
		log.Fatalln("fail to generate rsa key", err)
	}
	go ret.run()
	return ret
}

func (self *CasperCmd) GetId() string {
	return self.id
}

func (self *CasperCmd) SetInputArgs(input map[string]string) {
	if self.Finished() {
		log.Println("start another casperjs")
		go self.run()
	}
	log.Println("insert input:", input)
	self.input <- input
}

func (self *CasperCmd) GetMessage() *Output {
	return <-self.message
}

func (self *CasperCmd) readInputArgs(key string) string {
	args := <-self.input
	log.Println("read args:", args)
	for k, v := range args {
		if k == "username" {
			self.userName = v
		}

		if k == "password" {
			self.passWord = v
		}

		self.args[k] = v
	}
	if val, ok := self.args[key]; ok {
		return val
	}

	message := &Output{
		Id:        self.GetArgsValue("id"),
		NeedParam: key,
		Status:    NEED_PARAM,
	}
	self.message <- message
	return ""
}

func (self *CasperCmd) GetArgsValue(key string) string {
	if val, ok := self.args[key]; ok {
		log.Println("successfully get args value", val)
		return val
	}
	for {
		val := self.readInputArgs(key)
		if len(val) != 0 {
			log.Println("successfully get args value", val)
			return val
		}
	}

}

func (self *CasperCmd) getArgsList(args string) []string {
	segs := strings.Split(args, "/")
	if len(segs) < 2 {
		return nil
	}
	return segs[1:]
}

func (self *CasperCmd) getMetaInfo() *ParseRequest {
	ret := &ParseRequest{}
	ret.PrivateKey = string(privateKeyString(self.privateKey))
	ret.PublicKey = string(publicKeyString(&self.privateKey.PublicKey))
	ret.Tmpl = self.tmpl
	ret.UserName = self.userName
	ret.Secret = self.passWord
	ret.RowKey = self.tmpl + "|" + self.userName
	return ret
}

func (self *CasperCmd) Finished() bool {
	return self.isKill || self.isFinish
}

func (self *CasperCmd) DecodePassword(p string) string {
	bp, err := hex.DecodeString(p)
	if err != nil {
		log.Println("decode password hex error:", err)
		return ""
	}
	out, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, self.privateKey,
		bp, []byte(""))
	if err != nil {
		log.Println("decode password error:", err)
		return ""
	}
	log.Println("decode password:", string(out))
	return string(out)
}

func (self *CasperCmd) run() {
	log.Println("begin run cmd", self.tmpl)
	self.isFinish = false
	self.isKill = false

	path := "./" + self.tmpl + "/" + self.id
	os.RemoveAll(path)
	if err := os.MkdirAll(path, 0755); err != nil {
		log.Fatalln("can not create", path, err)
	}

	cookieFile, err := os.Create(path + "/cookie.txt")
	defer cookieFile.Close()
	var cmd *exec.Cmd
	if len(self.proxyServer) == 0 {
		cmd = exec.Command("casperjs", self.tmpl+".js",
			"--web-security=no",
			"--cookies-file="+path+"/cookie.txt",
			"--context="+path)
	} else {
		cmd = exec.Command("casperjs", self.tmpl+".js",
			"--web-security=no",
			"--cookies-file="+path+"/cookie.txt",
			"--proxy="+self.proxyServer, "--proxy-type=http",
			"--context="+path)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Panicln("can not get stdout pipe:", err)
	}
	bufout := bufio.NewReader(stdout)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		log.Panicln("can not get stdin pipe:", err)
	}
	bufin := bufio.NewWriter(stdin)

	if err := cmd.Start(); err != nil {
		log.Panicln("can not start cmd:", err)
	}

	go func() {
		timer := time.NewTimer(5 * time.Minute)
		<-timer.C
		cmd.Process.Kill()
		self.isKill = true
	}()

	log.Println("begin read line from capser")
	for {
		line, err := bufout.ReadString('\n')
		line = strings.Trim(line, "\n")
		if err != nil {
			log.Println(err)
			cmd.Process.Wait()
			cmd.Process.Kill()
			break
		}
		log.Println(line)

		if strings.HasPrefix(line, "CMD INFO STARTED") {
			message := &Output{
				Id:     self.GetArgsValue("id"),
				Status: OUTPUT_PUBLICKEY,
				Data:   string(publicKeyString(&self.privateKey.PublicKey)),
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD GET ARGS") {
			for _, v := range self.getArgsList(line) {
				key := strings.TrimRight(v, "\n")
				val := self.GetArgsValue(key)
				if key == "password" {
					val = self.DecodePassword(val)
				}
				bufin.WriteString(val)
				delete(self.args, key)
				bufin.WriteRune('\n')
				bufin.Flush()
			}
			continue
		}

		if strings.HasPrefix(line, "CMD INFO LOGIN SUCCESS") {
			message := &Output{
				Id:     self.GetArgsValue("id"),
				Status: LOGIN_SUCCESS,
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD INFO RANDCODE") {
			result := strings.TrimPrefix(line, "CMD INFO RANDCODE")
			result = strings.Trim(result, " \n")
			result = UploadImage("./site/" + result)
			log.Println("success upload captcha image to", result)
			message := &Output{
				Id:        self.GetArgsValue("id"),
				Status:    OUTPUT_VERIFYCODE,
				Data:      result,
				NeedParam: PARAM_VERIFY_CODE,
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD INFO CONTENT") {
			message := &Output{
				Status: strings.TrimSpace(strings.TrimPrefix(line,
					"CMD INFO CONTENT")),
				Id: self.GetArgsValue("id"),
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD INFO FETCHED MAIL") {
			result := strings.TrimSpace(strings.TrimPrefix(line,
				"CMD INFO FETCHED MAIL"))
			var out CasperOutput
			json.Unmarshal([]byte(result), &out)
			if self.mailProcessor != nil {
				go self.mailProcessor.Process(self.getMetaInfo(), out.Downloads)
			}
			message := &Output{
				Status: FINISH_FETCH_DATA,
				Id:     self.GetArgsValue("id"),
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD FAIL") {
			message := &Output{
				Status: FAIL,
				Id:     self.GetArgsValue("id"),
			}
			self.message <- message
			cmd.Process.Wait()
			cmd.Process.Kill()
			break
		}
	}
	self.isFinish = true
}
<commit_msg>support need param<commit_after>package caspercloud

import (
	"bufio"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

type CasperCmd struct {
	proxyServer   string
	id            string
	tmpl          string
	userName      string
	passWord      string
	message       chan *Output
	input         chan map[string]string
	isKill        bool
	isFinish      bool
	args          map[string]string
	privateKey    *rsa.PrivateKey
	mailProcessor *MailProcessor
}

func NewCasperCmd(id, tmpl, proxyServer string) *CasperCmd {
	ret := &CasperCmd{
		proxyServer:   proxyServer,
		id:            id,
		tmpl:          tmpl,
		userName:      "",
		passWord:      "",
		message:       make(chan *Output, 5),
		input:         make(chan map[string]string, 5),
		args:          make(map[string]string),
		isKill:        false,
		isFinish:      false,
		mailProcessor: NewMailProcessor("server_list.json"),
	}
	var err error
	ret.privateKey, err = generateRSAKey()
	if err != nil {
		log.Fatalln("fail to generate rsa key", err)
	}
	go ret.run()
	return ret
}

func (self *CasperCmd) GetId() string {
	return self.id
}

func (self *CasperCmd) SetInputArgs(input map[string]string) {
	if self.Finished() {
		log.Println("start another casperjs")
		go self.run()
	}
	log.Println("insert input:", input)
	self.input <- input
}

func (self *CasperCmd) GetMessage() *Output {
	return <-self.message
}

func (self *CasperCmd) readInputArgs(key string) string {
	args := <-self.input
	for k, v := range args {
		if k == "username" {
			self.userName = v
		}

		if k == "password" {
			self.passWord = v
		}

		self.args[k] = v
	}
	if val, ok := self.args[key]; ok {
		log.Println("find param", key, val)
		return val
	}

	message := &Output{
		Id:        self.GetArgsValue("id"),
		NeedParam: key,
		Status:    NEED_PARAM,
	}
	log.Println("need param", key)
	self.message <- message
	return ""
}

func (self *CasperCmd) GetArgsValue(key string) string {
	if val, ok := self.args[key]; ok {
		log.Println("successfully get args value", val)
		return val
	}
	for {
		val := self.readInputArgs(key)
		if len(val) != 0 {
			log.Println("successfully get args value", val)
			return val
		}
	}

}

func (self *CasperCmd) getArgsList(args string) []string {
	segs := strings.Split(args, "/")
	if len(segs) < 2 {
		return nil
	}
	return segs[1:]
}

func (self *CasperCmd) getMetaInfo() *ParseRequest {
	ret := &ParseRequest{}
	ret.PrivateKey = string(privateKeyString(self.privateKey))
	ret.PublicKey = string(publicKeyString(&self.privateKey.PublicKey))
	ret.Tmpl = self.tmpl
	ret.UserName = self.userName
	ret.Secret = self.passWord
	ret.RowKey = self.tmpl + "|" + self.userName
	return ret
}

func (self *CasperCmd) Finished() bool {
	return self.isKill || self.isFinish
}

func (self *CasperCmd) DecodePassword(p string) string {
	bp, err := hex.DecodeString(p)
	if err != nil {
		log.Println("decode password hex error:", err)
		return ""
	}
	out, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, self.privateKey,
		bp, []byte(""))
	if err != nil {
		log.Println("decode password error:", err)
		return ""
	}
	log.Println("decode password:", string(out))
	return string(out)
}

func (self *CasperCmd) run() {
	log.Println("begin run cmd", self.tmpl)
	self.isFinish = false
	self.isKill = false

	path := "./" + self.tmpl + "/" + self.id
	os.RemoveAll(path)
	if err := os.MkdirAll(path, 0755); err != nil {
		log.Fatalln("can not create", path, err)
	}

	cookieFile, err := os.Create(path + "/cookie.txt")
	defer cookieFile.Close()
	var cmd *exec.Cmd
	if len(self.proxyServer) == 0 {
		cmd = exec.Command("casperjs", self.tmpl+".js",
			"--web-security=no",
			"--cookies-file="+path+"/cookie.txt",
			"--context="+path)
	} else {
		cmd = exec.Command("casperjs", self.tmpl+".js",
			"--web-security=no",
			"--cookies-file="+path+"/cookie.txt",
			"--proxy="+self.proxyServer, "--proxy-type=http",
			"--context="+path)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Panicln("can not get stdout pipe:", err)
	}
	bufout := bufio.NewReader(stdout)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		log.Panicln("can not get stdin pipe:", err)
	}
	bufin := bufio.NewWriter(stdin)

	if err := cmd.Start(); err != nil {
		log.Panicln("can not start cmd:", err)
	}

	go func() {
		timer := time.NewTimer(5 * time.Minute)
		<-timer.C
		cmd.Process.Kill()
		self.isKill = true
	}()

	log.Println("begin read line from capser")
	for {
		line, err := bufout.ReadString('\n')
		line = strings.Trim(line, "\n")
		if err != nil {
			log.Println(err)
			cmd.Process.Wait()
			cmd.Process.Kill()
			break
		}
		log.Println(line)

		if strings.HasPrefix(line, "CMD INFO STARTED") {
			message := &Output{
				Id:     self.GetArgsValue("id"),
				Status: OUTPUT_PUBLICKEY,
				Data:   string(publicKeyString(&self.privateKey.PublicKey)),
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD GET ARGS") {
			for _, v := range self.getArgsList(line) {
				key := strings.TrimRight(v, "\n")
				val := self.GetArgsValue(key)
				if key == "password" {
					val = self.DecodePassword(val)
				}
				bufin.WriteString(val)
				delete(self.args, key)
				bufin.WriteRune('\n')
				bufin.Flush()
			}
			continue
		}

		if strings.HasPrefix(line, "CMD INFO LOGIN SUCCESS") {
			message := &Output{
				Id:     self.GetArgsValue("id"),
				Status: LOGIN_SUCCESS,
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD NEED") {
			result := strings.TrimPrefix(line, "CMD NEED")
			result = strings.TrimSpace(result)
			message := &Output{
				Id:        self.GetArgsValue("id"),
				Status:    NEED_PARAM,
				NeedParam: result,
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD INFO RANDCODE") {
			result := strings.TrimPrefix(line, "CMD INFO RANDCODE")
			result = strings.Trim(result, " \n")
			result = UploadImage("./site/" + result)
			log.Println("success upload captcha image to", result)
			message := &Output{
				Id:        self.GetArgsValue("id"),
				Status:    OUTPUT_VERIFYCODE,
				Data:      result,
				NeedParam: PARAM_VERIFY_CODE,
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD INFO CONTENT") {
			message := &Output{
				Status: strings.TrimSpace(strings.TrimPrefix(line,
					"CMD INFO CONTENT")),
				Id: self.GetArgsValue("id"),
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD INFO FETCHED MAIL") {
			result := strings.TrimSpace(strings.TrimPrefix(line,
				"CMD INFO FETCHED MAIL"))
			var out CasperOutput
			json.Unmarshal([]byte(result), &out)
			if self.mailProcessor != nil {
				go self.mailProcessor.Process(self.getMetaInfo(), out.Downloads)
			}
			message := &Output{
				Status: FINISH_FETCH_DATA,
				Id:     self.GetArgsValue("id"),
			}
			self.message <- message
			continue
		}

		if strings.HasPrefix(line, "CMD FAIL") {
			message := &Output{
				Status: FAIL,
				Id:     self.GetArgsValue("id"),
			}
			self.message <- message
			cmd.Process.Wait()
			cmd.Process.Kill()
			break
		}
	}
	self.isFinish = true
}
<|endoftext|>


<commit_before>package lxc

import (
	"errors"
	"fmt"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"log"
	"strings"
	"time"
)

type StepWaitInit struct {
	WaitTimeout time.Duration
}

func (s *StepWaitInit) Run(state multistep.StateBag) multistep.StepAction {
	ui := state.Get("ui").(packer.Ui)

	var err error

	cancel := make(chan struct{})
	waitDone := make(chan bool, 1)
	go func() {
		ui.Say("Waiting for container to finish init...")
		err = s.waitForInit(state, cancel)
		waitDone <- true
	}()

	log.Printf("Waiting for container to finish init, up to timeout: %s", s.WaitTimeout)
	timeout := time.After(s.WaitTimeout)
WaitLoop:
	for {
		select {
		case <-waitDone:
			if err != nil {
				ui.Error(fmt.Sprintf("Error waiting for container to finish init: %s", err))
				return multistep.ActionHalt
			}

			ui.Say("Container finished init!")
			break WaitLoop
		case <-timeout:
			err := fmt.Errorf("Timeout waiting for container to finish init.")
			state.Put("error", err)
			ui.Error(err.Error())
			close(cancel)
			return multistep.ActionHalt
		case <-time.After(1 * time.Second):
			if _, ok := state.GetOk(multistep.StateCancelled); ok {
				close(cancel)
				log.Println("Interrupt detected, quitting waiting for container to finish init.")
				return multistep.ActionHalt
			}
		}
	}

	return multistep.ActionContinue
}

func (s *StepWaitInit) Cleanup(multistep.StateBag) {
}

func (s *StepWaitInit) waitForInit(state multistep.StateBag, cancel <-chan struct{}) error {
	config := state.Get("config").(*Config)
	mountPath := state.Get("mount_path").(string)
	wrappedCommand := state.Get("wrappedCommand").(CommandWrapper)

	for {
		select {
		case <-cancel:
			log.Println("Cancelled. Exiting loop.")
			return errors.New("Wait cancelled")
		case <-time.After(1 * time.Second):
		}

		comm := &LxcAttachCommunicator{
			ContainerName: config.ContainerName,
			RootFs:        mountPath,
			CmdWrapper:    wrappedCommand,
		}

		runlevel, _ := comm.CheckInit()
		currentRunlevel := "unknown"
		if arr := strings.Split(runlevel, " "); len(arr) >= 2 {
			currentRunlevel = arr[1]
		}

		log.Printf("Current runlevel in container: '%s'", runlevel)

		targetRunlevel := fmt.Sprintf("%d", config.TargetRunlevel)
		if currentRunlevel == targetRunlevel {
			log.Printf("Container finished init.")
			break
		}

		/*log.Println("Attempting SSH connection...")
		comm, err = ssh.New(config)
		if err != nil {
			log.Printf("SSH handshake err: %s", err)

			// Only count this as an attempt if we were able to attempt
			// to authenticate. Note this is very brittle since it depends
			// on the string of the error... but I don't see any other way.
			if strings.Contains(err.Error(), "authenticate") {
				log.Printf("Detected authentication error. Increasing handshake attempts.")
				handshakeAttempts += 1
			}

			if handshakeAttempts < 10 {
				// Try to connect via SSH a handful of times
				continue
			}

			return nil, err
		}

		break
		*/
	}

	return nil
}
<commit_msg>remove commented out code<commit_after>package lxc

import (
	"errors"
	"fmt"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"log"
	"strings"
	"time"
)

type StepWaitInit struct {
	WaitTimeout time.Duration
}

func (s *StepWaitInit) Run(state multistep.StateBag) multistep.StepAction {
	ui := state.Get("ui").(packer.Ui)

	var err error

	cancel := make(chan struct{})
	waitDone := make(chan bool, 1)
	go func() {
		ui.Say("Waiting for container to finish init...")
		err = s.waitForInit(state, cancel)
		waitDone <- true
	}()

	log.Printf("Waiting for container to finish init, up to timeout: %s", s.WaitTimeout)
	timeout := time.After(s.WaitTimeout)
WaitLoop:
	for {
		select {
		case <-waitDone:
			if err != nil {
				ui.Error(fmt.Sprintf("Error waiting for container to finish init: %s", err))
				return multistep.ActionHalt
			}

			ui.Say("Container finished init!")
			break WaitLoop
		case <-timeout:
			err := fmt.Errorf("Timeout waiting for container to finish init.")
			state.Put("error", err)
			ui.Error(err.Error())
			close(cancel)
			return multistep.ActionHalt
		case <-time.After(1 * time.Second):
			if _, ok := state.GetOk(multistep.StateCancelled); ok {
				close(cancel)
				log.Println("Interrupt detected, quitting waiting for container to finish init.")
				return multistep.ActionHalt
			}
		}
	}

	return multistep.ActionContinue
}

func (s *StepWaitInit) Cleanup(multistep.StateBag) {
}

func (s *StepWaitInit) waitForInit(state multistep.StateBag, cancel <-chan struct{}) error {
	config := state.Get("config").(*Config)
	mountPath := state.Get("mount_path").(string)
	wrappedCommand := state.Get("wrappedCommand").(CommandWrapper)

	for {
		select {
		case <-cancel:
			log.Println("Cancelled. Exiting loop.")
			return errors.New("Wait cancelled")
		case <-time.After(1 * time.Second):
		}

		comm := &LxcAttachCommunicator{
			ContainerName: config.ContainerName,
			RootFs:        mountPath,
			CmdWrapper:    wrappedCommand,
		}

		runlevel, _ := comm.CheckInit()
		currentRunlevel := "unknown"
		if arr := strings.Split(runlevel, " "); len(arr) >= 2 {
			currentRunlevel = arr[1]
		}

		log.Printf("Current runlevel in container: '%s'", runlevel)

		targetRunlevel := fmt.Sprintf("%d", config.TargetRunlevel)
		if currentRunlevel == targetRunlevel {
			log.Printf("Container finished init.")
			break
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package api

import (
	"fmt"
	"net/url"

	"github.com/google/go-querystring/query"
)

// UserDomain holds user domains
type UserDomain struct {
	ID   int    `json:"id"`
	Name string `json:"name"`
}

// UserOrganization holds user organizations
type UserOrganization struct {
	ID   int    `json:"id"`
	Name string `json:"name"`
}

// UserAddress addresses
type UserAddress struct {
}

// User holds users
type User struct {
	ID            int                `json:"id,omitempty"`
	Username      string             `json:"username"`
	Firstname     string             `json:"firstname"`
	Lastname      string             `json:"lastname"`
	Password1     string             `json:"password1,omitempty"`
	Password2     string             `json:"password2,omitempty"`
	Email         string             `json:"email"`
	Timezone      string             `json:"timezone"`
	AccountType   int                `json:"account_type"`
	Active        bool               `json:"active"`
	SendReport    bool               `json:"send_report"`
	SpamChecks    bool               `json:"spam_checks"`
	LowScore      float64            `json:"low_score"`
	HighScore     float64            `json:"high_score"`
	BlockMacros   bool               `json:"block_macros"`
	CreatedOn     MyTime             `json:"created_on"`
	LastLogin     MyTime             `json:"last_login"`
	Domains       []UserDomain       `json:"domains,omitempty"`
	Organizations []UserOrganization `json:"organizations,omitempty"`
}

// UserList holds users
type UserList struct {
	Items []User `json:"items"`
	Links Links  `json:"links"`
	Meta  Meta   `json:"meta"`
}

// GetUsers returns a UserList object
// This contains a paginated list of user accounts and links
// to the neigbouring pages.
// https://www.baruwa.com/docs/api/#list-all-accounts
func (c *Client) GetUsers(opts *ListOptions) (l *UserList, err error) {
	err = c.get("users", nil, l)

	return
}

// GetUser returns a user account
// https://www.baruwa.com/docs/api/#retrieve-an-existing-account
func (c *Client) GetUser(userID int) (user *User, err error) {
	if userID <= 0 {
		err = fmt.Errorf(userIDError)
		return
	}

	user = &User{}

	err = c.get(fmt.Sprintf("users/%d", userID), nil, user)

	return
}

// CreateUser creates a user account
// https://www.baruwa.com/docs/api/#create-a-new-account
func (c *Client) CreateUser(user *User) (err error) {
	var v url.Values

	if user == nil {
		err = fmt.Errorf("The user param cannot be nil")
		return
	}

	if v, err = query.Values(user); err != nil {
		return
	}

	err = c.post("users", v, user)

	return
}

// UpdateUser updates a user account
// https://www.baruwa.com/docs/api/#update-an-account
func (c *Client) UpdateUser(user *User) (err error) {
	var v url.Values

	if user == nil {
		err = fmt.Errorf("The user param cannot be nil")
		return
	}

	if user.ID <= 0 {
		err = fmt.Errorf("The user.ID param should be > 0")
		return
	}

	if v, err = query.Values(user); err != nil {
		return
	}

	err = c.put(fmt.Sprintf("users/%d", user.ID), v, user)

	return
}

// DeleteUser deletes a user account
// https://www.baruwa.com/docs/api/#delete-an-account
func (c *Client) DeleteUser(userID int) (err error) {
	if userID <= 0 {
		err = fmt.Errorf("The userID param should be > 0")
		return
	}

	err = c.delete(fmt.Sprintf("users/%d", userID), nil)

	return
}
<commit_msg>FIX: Improvements to users code<commit_after>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

package api

import (
	"fmt"
	"net/url"

	"github.com/google/go-querystring/query"
)

// UserDomain holds user domains
type UserDomain struct {
	ID   int    `json:"id"`
	Name string `json:"name"`
}

// UserOrganization holds user organizations
type UserOrganization struct {
	ID   int    `json:"id"`
	Name string `json:"name"`
}

// UserAddress addresses
type UserAddress struct {
}

// User holds users
type User struct {
	ID            int                `json:"id,omitempty"`
	Username      string             `json:"username"`
	Firstname     string             `json:"firstname"`
	Lastname      string             `json:"lastname"`
	Password1     string             `json:"password1,omitempty"`
	Password2     string             `json:"password2,omitempty"`
	Email         string             `json:"email"`
	Timezone      string             `json:"timezone"`
	AccountType   int                `json:"account_type"`
	Active        bool               `json:"active"`
	SendReport    bool               `json:"send_report"`
	SpamChecks    bool               `json:"spam_checks"`
	LowScore      float64            `json:"low_score"`
	HighScore     float64            `json:"high_score"`
	BlockMacros   bool               `json:"block_macros"`
	CreatedOn     MyTime             `json:"created_on"`
	LastLogin     MyTime             `json:"last_login"`
	Domains       []UserDomain       `json:"domains,omitempty"`
	Organizations []UserOrganization `json:"organizations,omitempty"`
}

// UserList holds users
type UserList struct {
	Items []User `json:"items"`
	Links Links  `json:"links"`
	Meta  Meta   `json:"meta"`
}

// GetUsers returns a UserList object
// This contains a paginated list of user accounts and links
// to the neigbouring pages.
// https://www.baruwa.com/docs/api/#list-all-accounts
func (c *Client) GetUsers(opts *ListOptions) (l *UserList, err error) {
	l = &UserList{}

	err = c.get("users", nil, l)

	return
}

// GetUser returns a user account
// https://www.baruwa.com/docs/api/#retrieve-an-existing-account
func (c *Client) GetUser(userID int) (user *User, err error) {
	if userID <= 0 {
		err = fmt.Errorf(userIDError)
		return
	}

	user = &User{}

	err = c.get(fmt.Sprintf("users/%d", userID), nil, user)

	return
}

// CreateUser creates a user account
// https://www.baruwa.com/docs/api/#create-a-new-account
func (c *Client) CreateUser(user *User) (err error) {
	var v url.Values

	if user == nil {
		err = fmt.Errorf(userParamError)
		return
	}

	if v, err = query.Values(user); err != nil {
		return
	}

	err = c.post("users", v, nil)

	return
}

// UpdateUser updates a user account
// https://www.baruwa.com/docs/api/#update-an-account
func (c *Client) UpdateUser(user *User) (err error) {
	var v url.Values

	if user == nil {
		err = fmt.Errorf(userParamError)
		return
	}

	if user.ID <= 0 {
		err = fmt.Errorf(userIDError)
		return
	}

	if v, err = query.Values(user); err != nil {
		return
	}

	err = c.put(fmt.Sprintf("users/%d", user.ID), v, nil)

	return
}

// DeleteUser deletes a user account
// https://www.baruwa.com/docs/api/#delete-an-account
func (c *Client) DeleteUser(userID int) (err error) {
	if userID <= 0 {
		err = fmt.Errorf(userIDError)
		return
	}

	err = c.delete(fmt.Sprintf("users/%d", userID), nil)

	return
}
<|endoftext|>


<commit_before>package api

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/sirupsen/logrus"
	"github.com/netlify/gotrue/models"
)

// UserUpdateParams parameters for updating a user
type UserUpdateParams struct {
	Email            string                 `json:"email"`
	Password         string                 `json:"password"`
	EmailChangeToken string                 `json:"email_change_token"`
	Data             map[string]interface{} `json:"data"`
	AppData          map[string]interface{} `json:"app_metadata,omitempty"`
}

// UserGet returns a user
func (a *API) UserGet(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	token := getToken(ctx)

	id, ok := token.Claims["id"].(string)
	if !ok {
		BadRequestError(w, "Could not read User ID claim")
		return
	}

	tokenAud, ok := token.Claims["aud"].(string)
	if !ok {
		BadRequestError(w, "Could not read User Aud claim")
		return
	}

	aud := a.requestAud(ctx, r)
	if aud != tokenAud {
		BadRequestError(w, "Token audience doesn't match request audience")
		return
	}

	user, err := a.db.FindUserByID(id)
	if err != nil {
		if models.IsNotFoundError(err) {
			NotFoundError(w, err.Error())
		} else {
			InternalServerError(w, err.Error())
		}
		return
	}

	sendJSON(w, 200, user)
}

// UserUpdate updates fields on a user
func (a *API) UserUpdate(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	token := getToken(ctx)

	params := &UserUpdateParams{}
	jsonDecoder := json.NewDecoder(r.Body)
	err := jsonDecoder.Decode(params)
	if err != nil {
		BadRequestError(w, fmt.Sprintf("Could not read User Update params: %v", err))
		return
	}

	id, ok := token.Claims["id"].(string)
	if !ok {
		BadRequestError(w, "Could not read User ID claim")
		return
	}

	user, err := a.db.FindUserByID(id)
	if err != nil {
		if models.IsNotFoundError(err) {
			NotFoundError(w, err.Error())
		} else {
			InternalServerError(w, err.Error())
		}
		return
	}

	var sendChangeEmailVerification bool
	if err = a.mailer.ValidateEmail(params.Email); err == nil {
		exists, err := a.db.IsDuplicatedEmail(params.Email, user.Aud)
		if err != nil {
			InternalServerError(w, err.Error())
			return
		}

		if exists {
			UnprocessableEntity(w, "Email address already registered by another user")
			return
		}

		user.GenerateEmailChange(params.Email)
		sendChangeEmailVerification = true
	} else {
		UnprocessableEntity(w, "Unable to verify new email address: "+err.Error())
		return
	}

	logrus.Debugf("Checking params for token %v", params)

	if params.EmailChangeToken != "" {
		logrus.Debugf("Got change token %v", params.EmailChangeToken)

		if params.EmailChangeToken != user.EmailChangeToken {
			UnauthorizedError(w, "Email Change Token didn't match token on file")
			return
		}

		user.ConfirmEmailChange()
	}

	if params.Password != "" {
		if err = user.EncryptPassword(params.Password); err != nil {
			InternalServerError(w, fmt.Sprintf("Error during password encryption: %v", err))
			return
		}
	}

	if params.Data != nil {
		user.UpdateUserMetaData(params.Data)
	}

	if params.AppData != nil {
		if a.isAdmin(user, a.config.JWT.Aud) {
			UnauthorizedError(w, "Updating app_metadata requires admin privileges")
			return
		}

		user.UpdateAppMetaData(params.AppData)
	}

	if err := a.db.UpdateUser(user); err != nil {
		InternalServerError(w, err.Error())
		return
	}

	if sendChangeEmailVerification {
		a.mailer.EmailChangeMail(user)
	}

	sendJSON(w, 200, user)
}
<commit_msg>Check if email is sent before validating<commit_after>package api

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/netlify/gotrue/models"
	"github.com/sirupsen/logrus"
)

// UserUpdateParams parameters for updating a user
type UserUpdateParams struct {
	Email            string                 `json:"email"`
	Password         string                 `json:"password"`
	EmailChangeToken string                 `json:"email_change_token"`
	Data             map[string]interface{} `json:"data"`
	AppData          map[string]interface{} `json:"app_metadata,omitempty"`
}

// UserGet returns a user
func (a *API) UserGet(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	token := getToken(ctx)

	id, ok := token.Claims["id"].(string)
	if !ok {
		BadRequestError(w, "Could not read User ID claim")
		return
	}

	tokenAud, ok := token.Claims["aud"].(string)
	if !ok {
		BadRequestError(w, "Could not read User Aud claim")
		return
	}

	aud := a.requestAud(ctx, r)
	if aud != tokenAud {
		BadRequestError(w, "Token audience doesn't match request audience")
		return
	}

	user, err := a.db.FindUserByID(id)
	if err != nil {
		if models.IsNotFoundError(err) {
			NotFoundError(w, err.Error())
		} else {
			InternalServerError(w, err.Error())
		}
		return
	}

	sendJSON(w, 200, user)
}

// UserUpdate updates fields on a user
func (a *API) UserUpdate(ctx context.Context, w http.ResponseWriter, r *http.Request) {
	token := getToken(ctx)

	params := &UserUpdateParams{}
	jsonDecoder := json.NewDecoder(r.Body)
	err := jsonDecoder.Decode(params)
	if err != nil {
		BadRequestError(w, fmt.Sprintf("Could not read User Update params: %v", err))
		return
	}

	id, ok := token.Claims["id"].(string)
	if !ok {
		BadRequestError(w, "Could not read User ID claim")
		return
	}

	user, err := a.db.FindUserByID(id)
	if err != nil {
		if models.IsNotFoundError(err) {
			NotFoundError(w, err.Error())
		} else {
			InternalServerError(w, err.Error())
		}
		return
	}

	var sendChangeEmailVerification bool
	if params.Email != "" {
		if err = a.mailer.ValidateEmail(params.Email); err == nil {
			exists, err := a.db.IsDuplicatedEmail(params.Email, user.Aud)
			if err != nil {
				InternalServerError(w, err.Error())
				return
			}

			if exists {
				UnprocessableEntity(w, "Email address already registered by another user")
				return
			}

			user.GenerateEmailChange(params.Email)
			sendChangeEmailVerification = true
		} else {
			UnprocessableEntity(w, "Unable to verify new email address: "+err.Error())
			return
		}
	}

	logrus.Debugf("Checking params for token %v", params)

	if params.EmailChangeToken != "" {
		logrus.Debugf("Got change token %v", params.EmailChangeToken)

		if params.EmailChangeToken != user.EmailChangeToken {
			UnauthorizedError(w, "Email Change Token didn't match token on file")
			return
		}

		user.ConfirmEmailChange()
	}

	if params.Password != "" {
		if err = user.EncryptPassword(params.Password); err != nil {
			InternalServerError(w, fmt.Sprintf("Error during password encryption: %v", err))
			return
		}
	}

	if params.Data != nil {
		user.UpdateUserMetaData(params.Data)
	}

	if params.AppData != nil {
		if a.isAdmin(user, a.config.JWT.Aud) {
			UnauthorizedError(w, "Updating app_metadata requires admin privileges")
			return
		}

		user.UpdateAppMetaData(params.AppData)
	}

	if err := a.db.UpdateUser(user); err != nil {
		InternalServerError(w, err.Error())
		return
	}

	if sendChangeEmailVerification {
		a.mailer.EmailChangeMail(user)
	}

	sendJSON(w, 200, user)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strings"
	"text/template"
)

const (
	startPrefix = "=== RUN "
	passPrefix  = "--- PASS: "
	failPrefix  = "--- FAIL: "

	version = "0.1.2"
)

// "end of test" regexp for name and time, examples:
// --- PASS: TestSub (0.00 seconds)
// --- FAIL: TestSubFail (0.00 seconds)
var endRegexp *regexp.Regexp = regexp.MustCompile(`([^ ]+) \((\d+\.\d+)`)

type Test struct {
	Name, Time, Message string
	Failed              bool
}

type TestResults struct {
	Tests  []*Test
	Count  int
	Failed int
	Bamboo bool
}

// parseEnd parses "end of test" line and returns (name, time, error)
func parseEnd(prefix, line string) (string, string, error) {
	matches := endRegexp.FindStringSubmatch(line[len(prefix):])

	if len(matches) == 0 {
		return "", "", fmt.Errorf("can't parse %s", line)
	}

	return matches[1], matches[2], nil
}

// parseOutput parses output of "go test -v", returns a list of tests
func parseOutput(rd io.Reader) ([]*Test, error) {
	tests := []*Test{}
	var test *Test = nil

	nextTest := func() {
		// We are switching to the next test, store the current one.
		if test == nil {
			return
		}

		tests = append(tests, test)
		test = nil
	}

	reader := bufio.NewReader(rd)
	for {
		buf, _, err := reader.ReadLine()

		switch err {
		case io.EOF:
			nextTest()
			return tests, nil
		case nil:
			// nil is OK

		default: // Error other than io.EOF
			return nil, err
		}

		line := string(buf)

		switch {
		case strings.HasPrefix(line, startPrefix):
		case strings.HasPrefix(line, failPrefix):
			nextTest()

			// Extract the test name and the duration:
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}

			test = &Test{
				Name:   name,
				Time:   time,
				Failed: true,
			}

		case strings.HasPrefix(line, passPrefix):
			nextTest()
			// Extract the test name and the duration:
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}

			// Create the test structure and store it.
			tests = append(tests, &Test{
				Name:   name,
				Time:   time,
				Failed: false,
			})
			test = nil
		case line == "FAIL":
			nextTest()
		default:
			if test != nil { // test != nil marks we're in the middle of a test
				test.Message += line + "\n"
			}
		}
	}

	// If we're here, it's an error
	return nil, fmt.Errorf("Error parsing")
}

// numFailures count how man tests failed
func numFailures(tests []*Test) int {
	count := 0
	for _, test := range tests {
		if test.Failed {
			count++
		}
	}

	return count
}

var xmlTemplate string = `<?xml version="1.0" encoding="utf-8"?>
{{if .Bamboo}}<testsuites>{{end}}
  <testsuite name="go2xunit" tests="{{.Count}}" errors="0" failures="{{.Failed}}" skip="0">
{{range $test := .Tests}}    <testcase classname="go2xunit" name="{{$test.Name}}" time="{{$test.Time}}">
{{if $test.Failed }}      <failure type="go.error" message="error">
        <![CDATA[{{$test.Message}}]]></failure>
{{end}}    </testcase>
{{end}}  </testsuite>
{{if .Bamboo}}</testsuites>{{end}}
	`

// writeXML exits xunit XML of tests to out
func writeXML(tests []*Test, out io.Writer, bamboo bool) {
	count := len(tests)
	failed := numFailures(tests)

	testsResult := TestResults{Tests: tests, Count: count, Failed: failed, Bamboo: bamboo}
	t := template.New("test template")
	t, err := t.Parse(xmlTemplate)
	if err != nil {
		fmt.Println("Error en parse %v", err)
		return
	}
	err = t.Execute(out, testsResult)
	if err != nil {
		fmt.Println("Error en execute %v", err)
		return
	}
}

// getInput return input io.Reader from file name, if file name is - it will
// return os.Stdin
func getInput(filename string) (io.Reader, error) {
	if filename == "-" || filename == "" {
		return os.Stdin, nil
	}

	return os.Open(filename)
}

// getInput return output io.Writer from file name, if file name is - it will
// return os.Stdout
func getOutput(filename string) (io.Writer, error) {
	if filename == "-" || filename == "" {
		return os.Stdout, nil
	}

	return os.Create(filename)
}

// getIO returns input and output streams from file names
func getIO(inputFile, outputFile string) (io.Reader, io.Writer, error) {
	input, err := getInput(inputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for reading: %s", inputFile, err)
	}

	output, err := getOutput(outputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for writing: %s", outputFile, err)
	}

	return input, output, nil
}

func main() {
	inputFile := flag.String("input", "", "input file (default to stdin)")
	outputFile := flag.String("output", "", "output file (default to stdout)")
	fail := flag.Bool("fail", false, "fail (non zero exit) if any test failed")
	showVersion := flag.Bool("version", false, "print version and exit")
	bamboo := flag.Bool("bamboo", false, "xml compatible with Atlassian's Bamboo")
	flag.Parse()

	if *showVersion {
		fmt.Println(version)
		os.Exit(0)
	}

	// No time ... prefix for error messages
	log.SetFlags(0)

	if flag.NArg() > 0 {
		log.Fatalf("error: %s does not take parameters (did you mean -input?)", os.Args[0])
	}

	input, output, err := getIO(*inputFile, *outputFile)
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	tests, err := parseOutput(input)
	if err != nil {
		log.Fatalf("error: %s", err)
	}
	if len(tests) == 0 {
		log.Fatalf("error: no tests found")
		os.Exit(1)
	}

	writeXML(tests, output, *bamboo)
	if *fail && numFailures(tests) > 0 {
		os.Exit(1)
	}
}
<commit_msg>CT-42 quick fix for making classname distint so test with the same name will not cloack each other<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strings"
	"text/template"
)

const (
	startPrefix = "=== RUN "
	passPrefix  = "--- PASS: "
	failPrefix  = "--- FAIL: "

	version = "0.1.2"
)

// "end of test" regexp for name and time, examples:
// --- PASS: TestSub (0.00 seconds)
// --- FAIL: TestSubFail (0.00 seconds)
var endRegexp *regexp.Regexp = regexp.MustCompile(`([^ ]+) \((\d+\.\d+)`)

type Test struct {
	Name, Time, Message string
	Failed              bool
}

type TestResults struct {
	Tests  []*Test
	Count  int
	Failed int
	Bamboo bool
}

// parseEnd parses "end of test" line and returns (name, time, error)
func parseEnd(prefix, line string) (string, string, error) {
	matches := endRegexp.FindStringSubmatch(line[len(prefix):])

	if len(matches) == 0 {
		return "", "", fmt.Errorf("can't parse %s", line)
	}

	return matches[1], matches[2], nil
}

// parseOutput parses output of "go test -v", returns a list of tests
func parseOutput(rd io.Reader) ([]*Test, error) {
	tests := []*Test{}
	var test *Test = nil

	nextTest := func() {
		// We are switching to the next test, store the current one.
		if test == nil {
			return
		}

		tests = append(tests, test)
		test = nil
	}

	reader := bufio.NewReader(rd)
	for {
		buf, _, err := reader.ReadLine()

		switch err {
		case io.EOF:
			nextTest()
			return tests, nil
		case nil:
			// nil is OK

		default: // Error other than io.EOF
			return nil, err
		}

		line := string(buf)

		switch {
		case strings.HasPrefix(line, startPrefix):
		case strings.HasPrefix(line, failPrefix):
			nextTest()

			// Extract the test name and the duration:
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}

			test = &Test{
				Name:   name,
				Time:   time,
				Failed: true,
			}

		case strings.HasPrefix(line, passPrefix):
			nextTest()
			// Extract the test name and the duration:
			name, time, err := parseEnd(passPrefix, line)
			if err != nil {
				return nil, err
			}

			// Create the test structure and store it.
			tests = append(tests, &Test{
				Name:   name,
				Time:   time,
				Failed: false,
			})
			test = nil
		case line == "FAIL":
			nextTest()
		default:
			if test != nil { // test != nil marks we're in the middle of a test
				test.Message += line + "\n"
			}
		}
	}

	// If we're here, it's an error
	return nil, fmt.Errorf("Error parsing")
}

// numFailures count how man tests failed
func numFailures(tests []*Test) int {
	count := 0
	for _, test := range tests {
		if test.Failed {
			count++
		}
	}

	return count
}

var xmlTemplate string = `<?xml version="1.0" encoding="utf-8"?>
{{if .Bamboo}}<testsuites>{{end}}
  <testsuite name="go2xunit" tests="{{.Count}}" errors="0" failures="{{.Failed}}" skip="0">
{{range $index, $test := .Tests}}    <testcase classname="go2xunit-{{$index}}" name="{{$test.Name}}" time="{{$test.Time}}">
{{if $test.Failed }}      <failure type="go.error" message="error">
        <![CDATA[{{$test.Message}}]]></failure>
{{end}}    </testcase>
{{end}}  </testsuite>
{{if .Bamboo}}</testsuites>{{end}}
	`

// writeXML exits xunit XML of tests to out
func writeXML(tests []*Test, out io.Writer, bamboo bool) {
	count := len(tests)
	failed := numFailures(tests)

	testsResult := TestResults{Tests: tests, Count: count, Failed: failed, Bamboo: bamboo}
	t := template.New("test template")
	t, err := t.Parse(xmlTemplate)
	if err != nil {
		fmt.Println("Error en parse %v", err)
		return
	}
	err = t.Execute(out, testsResult)
	if err != nil {
		fmt.Println("Error en execute %v", err)
		return
	}
}

// getInput return input io.Reader from file name, if file name is - it will
// return os.Stdin
func getInput(filename string) (io.Reader, error) {
	if filename == "-" || filename == "" {
		return os.Stdin, nil
	}

	return os.Open(filename)
}

// getInput return output io.Writer from file name, if file name is - it will
// return os.Stdout
func getOutput(filename string) (io.Writer, error) {
	if filename == "-" || filename == "" {
		return os.Stdout, nil
	}

	return os.Create(filename)
}

// getIO returns input and output streams from file names
func getIO(inputFile, outputFile string) (io.Reader, io.Writer, error) {
	input, err := getInput(inputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for reading: %s", inputFile, err)
	}

	output, err := getOutput(outputFile)
	if err != nil {
		return nil, nil, fmt.Errorf("can't open %s for writing: %s", outputFile, err)
	}

	return input, output, nil
}

func main() {
	inputFile := flag.String("input", "", "input file (default to stdin)")
	outputFile := flag.String("output", "", "output file (default to stdout)")
	fail := flag.Bool("fail", false, "fail (non zero exit) if any test failed")
	showVersion := flag.Bool("version", false, "print version and exit")
	bamboo := flag.Bool("bamboo", false, "xml compatible with Atlassian's Bamboo")
	flag.Parse()

	if *showVersion {
		fmt.Println(version)
		os.Exit(0)
	}

	// No time ... prefix for error messages
	log.SetFlags(0)

	if flag.NArg() > 0 {
		log.Fatalf("error: %s does not take parameters (did you mean -input?)", os.Args[0])
	}

	input, output, err := getIO(*inputFile, *outputFile)
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	tests, err := parseOutput(input)
	if err != nil {
		log.Fatalf("error: %s", err)
	}
	if len(tests) == 0 {
		log.Fatalf("error: no tests found")
		os.Exit(1)
	}

	writeXML(tests, output, *bamboo)
	if *fail && numFailures(tests) > 0 {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package goheroku

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"github.com/Shopify/sarama"
	cluster "github.com/bsm/sarama-cluster"
	"net/url"
	"os"
	"strings"
)

//  NewClusterConsumer creates a bsm sarama-cluster consumer
//  Provide the topic, consumer group and a cluster config
func NewClusterConsumer(topic string, consumerGroup string, cfg *cluster.Config) (*cluster.Consumer, error) {
	tc, err := createTLSConfig()
	if err != nil {
		return nil, err
	}
	cfg.Net.TLS.Config = tc
	cfg.Net.TLS.Enable = true

	topic = AppendPrefixTo(topic)
	group := AppendPrefixTo(consumerGroup)
	brokers, err := brokerAddresses()
	if err != nil {
		return nil, err
	}

	consumer, err := cluster.NewConsumer(brokers, group, []string{topic}, cfg)
	if err != nil {
		return nil, err
	}

	return consumer, nil
}

// A quick and easy cluster consumer built with default configs
func NewQuickClusterConsumer(topic string, groupID string) (*cluster.Consumer, error) {
	cfg := cluster.NewConfig()
	consumer, err := NewClusterConsumer(topic, groupID, cfg)
	if err != nil {
		return nil, err
	}
	return consumer, nil
}

// TODO: Investigate use of/need for topic
func NewConsumer(cfg *sarama.Config) (sarama.Consumer, error) {
	tc, err := createTLSConfig()
	if err != nil {
		return nil, err
	}
	cfg.Net.TLS.Config = tc
	cfg.Net.TLS.Enable = true

	brokers, err := brokerAddresses()
	if err != nil {
		return nil, err
	}
	consumer, err := sarama.NewConsumer(brokers, cfg)
	if err != nil {
		return nil, err
	}

	return consumer, nil
}

// A quick and easy sarama consumer built with default configs
func NewQuickConsumer() (sarama.Consumer, error) {
	cfg := sarama.NewConfig()
	consumer, err := NewConsumer(cfg)
	if err != nil {
		return nil, err
	}
	return consumer, nil
}

//  NewAsyncProducer creates a sarama Async Producer
//  Provide a sarama config
//  For more information about the difference between Async and Sync producers
//  see the sarama documentation
func NewAsyncProducer(cfg *sarama.Config) (sarama.AsyncProducer, error) {
	tc, err := createTLSConfig()
	if err != nil {
		return nil, err
	}
	cfg.Net.TLS.Config = tc
	cfg.Net.TLS.Enable = true

	brokers, err := brokerAddresses()
	if err != nil {
		return nil, err
	}
	producer, err := sarama.NewAsyncProducer(brokers, cfg)
	if err != nil {
		return nil, err
	}

	return producer, nil
}

// A quick and easy async producer built with default configs
func NewQuickAsyncProducer() (sarama.AsyncProducer, error) {
	cfg := sarama.NewConfig()
	producer, err := NewAsyncProducer(cfg)
	if err != nil {
		return nil, err
	}
	return producer, nil
}

//  NewSyncProducer creates a sarama Sync Producer
//  Provide a sarama config
//  For more information about the difference between Async and Sync producers
//  see the sarama documentation
func NewSyncProducer(cfg *sarama.Config) (sarama.SyncProducer, error) {
	tc, err := createTLSConfig()
	if err != nil {
		return nil, err
	}
	cfg.Net.TLS.Config = tc
	cfg.Net.TLS.Enable = true

	brokers, err := brokerAddresses()
	if err != nil {
		return nil, err
	}
	producer, err := sarama.NewSyncProducer(brokers, cfg)
	if err != nil {
		return nil, err
	}

	return producer, nil
}

// A quick and easy sync producer built with default configs
func NewQuickSyncProducer() (sarama.SyncProducer, error) {
	cfg := sarama.NewConfig()
	cfg.Producer.Return.Successes = true
	producer, err := NewSyncProducer(cfg)
	if err != nil {
		return nil, err
	}
	return producer, nil
}

// Endangered function. TODO: Investigate usefulness of client
func NewClient(cfg *sarama.Config) (sarama.Client, error) {
	addrs, err := brokerAddresses()
	if err != nil {
		return nil, err
	}
	client, err := sarama.NewClient(addrs, cfg)
	if err != nil {
		return nil, err
	}
	return client, nil
}

// To specify the topic or consumer group, the Kafka prefix needs
// to appended to it. This function makes it possible without having
// access to the app config.
func AppendPrefixTo(str string) string {
	prefix := os.Getenv("KAFKA_PREFIX")
	if prefix != "" {
		str = strings.Join([]string{prefix, str}, "")
	}
	return str
}

// Create the TLS context, using the key and certificates provided.
func createTLSConfig() (*tls.Config, error) {
	trustedCert, ok := os.LookupEnv("KAFKA_TRUSTED_CERT")
	if !ok {
		return nil, errors.New("Kafka Trusted Certificate not found!")
	}

	clientCertKey, ok := os.LookupEnv("KAFKA_CLIENT_CERT_KEY")
	if !ok {
		return nil, errors.New("Kafka Client Certificate Key not found!")
	}

	clientCert, ok := os.LookupEnv("KAFKA_CLIENT_CERT")
	if !ok {
		return nil, errors.New("Kafka Client Certificate not found!")
	}

	roots := x509.NewCertPool()
	ok = roots.AppendCertsFromPEM([]byte(trustedCert))
	if !ok {
		return nil, errors.New("Unable to parse Root Cert. Please check your Heroku environment.")
	}
	// Setup certs for Sarama
	cert, err := tls.X509KeyPair([]byte(clientCert), []byte(clientCertKey))
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		Certificates:       []tls.Certificate{cert},
		InsecureSkipVerify: true,
		RootCAs:            roots,
	}, nil
}

// Extract the host:port pairs from the Kafka URL(s)
func brokerAddresses() ([]string, error) {
	URL, ok := os.LookupEnv("KAFKA_URL")
	if !ok {
		return nil, errors.New("Kafka URL not found!")
	}
	urls := strings.Split(URL, ",")
	addrs := make([]string, len(urls))
	for i, v := range urls {
		u, err := url.Parse(v)
		if err != nil {
			return nil, err
		}
		addrs[i] = u.Host
	}
	return addrs, nil
}
<commit_msg>Replace Lookupenv with Getenv<commit_after>package goheroku

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"github.com/Shopify/sarama"
	cluster "github.com/bsm/sarama-cluster"
	"net/url"
	"os"
	"strings"
)

//  NewClusterConsumer creates a bsm sarama-cluster consumer
//  Provide the topic, consumer group and a cluster config
func NewClusterConsumer(topic string, consumerGroup string, cfg *cluster.Config) (*cluster.Consumer, error) {
	tc, err := createTLSConfig()
	if err != nil {
		return nil, err
	}
	cfg.Net.TLS.Config = tc
	cfg.Net.TLS.Enable = true

	topic = AppendPrefixTo(topic)
	group := AppendPrefixTo(consumerGroup)
	brokers, err := brokerAddresses()
	if err != nil {
		return nil, err
	}

	consumer, err := cluster.NewConsumer(brokers, group, []string{topic}, cfg)
	if err != nil {
		return nil, err
	}

	return consumer, nil
}

// A quick and easy cluster consumer built with default configs
func NewQuickClusterConsumer(topic string, groupID string) (*cluster.Consumer, error) {
	cfg := cluster.NewConfig()
	consumer, err := NewClusterConsumer(topic, groupID, cfg)
	if err != nil {
		return nil, err
	}
	return consumer, nil
}

// TODO: Investigate use of/need for topic
func NewConsumer(cfg *sarama.Config) (sarama.Consumer, error) {
	tc, err := createTLSConfig()
	if err != nil {
		return nil, err
	}
	cfg.Net.TLS.Config = tc
	cfg.Net.TLS.Enable = true

	brokers, err := brokerAddresses()
	if err != nil {
		return nil, err
	}
	consumer, err := sarama.NewConsumer(brokers, cfg)
	if err != nil {
		return nil, err
	}

	return consumer, nil
}

// A quick and easy sarama consumer built with default configs
func NewQuickConsumer() (sarama.Consumer, error) {
	cfg := sarama.NewConfig()
	consumer, err := NewConsumer(cfg)
	if err != nil {
		return nil, err
	}
	return consumer, nil
}

//  NewAsyncProducer creates a sarama Async Producer
//  Provide a sarama config
//  For more information about the difference between Async and Sync producers
//  see the sarama documentation
func NewAsyncProducer(cfg *sarama.Config) (sarama.AsyncProducer, error) {
	tc, err := createTLSConfig()
	if err != nil {
		return nil, err
	}
	cfg.Net.TLS.Config = tc
	cfg.Net.TLS.Enable = true

	brokers, err := brokerAddresses()
	if err != nil {
		return nil, err
	}
	producer, err := sarama.NewAsyncProducer(brokers, cfg)
	if err != nil {
		return nil, err
	}

	return producer, nil
}

// A quick and easy async producer built with default configs
func NewQuickAsyncProducer() (sarama.AsyncProducer, error) {
	cfg := sarama.NewConfig()
	producer, err := NewAsyncProducer(cfg)
	if err != nil {
		return nil, err
	}
	return producer, nil
}

//  NewSyncProducer creates a sarama Sync Producer
//  Provide a sarama config
//  For more information about the difference between Async and Sync producers
//  see the sarama documentation
func NewSyncProducer(cfg *sarama.Config) (sarama.SyncProducer, error) {
	tc, err := createTLSConfig()
	if err != nil {
		return nil, err
	}
	cfg.Net.TLS.Config = tc
	cfg.Net.TLS.Enable = true

	brokers, err := brokerAddresses()
	if err != nil {
		return nil, err
	}
	producer, err := sarama.NewSyncProducer(brokers, cfg)
	if err != nil {
		return nil, err
	}

	return producer, nil
}

// A quick and easy sync producer built with default configs
func NewQuickSyncProducer() (sarama.SyncProducer, error) {
	cfg := sarama.NewConfig()
	cfg.Producer.Return.Successes = true
	producer, err := NewSyncProducer(cfg)
	if err != nil {
		return nil, err
	}
	return producer, nil
}

// Endangered function. TODO: Investigate usefulness of client
func NewClient(cfg *sarama.Config) (sarama.Client, error) {
	addrs, err := brokerAddresses()
	if err != nil {
		return nil, err
	}
	client, err := sarama.NewClient(addrs, cfg)
	if err != nil {
		return nil, err
	}
	return client, nil
}

// To specify the topic or consumer group, the Kafka prefix needs
// to appended to it. This function makes it possible without having
// access to the app config.
func AppendPrefixTo(str string) string {
	prefix := os.Getenv("KAFKA_PREFIX")
	if prefix != "" {
		str = strings.Join([]string{prefix, str}, "")
	}
	return str
}

// Create the TLS context, using the key and certificates provided.
func createTLSConfig() (*tls.Config, error) {
	trustedCert := os.Getenv("KAFKA_TRUSTED_CERT")
	if trustedCert == "" {
		return nil, errors.New("Kafka Trusted Certificate not found!")
	}

	clientCertKey := os.Getenv("KAFKA_CLIENT_CERT_KEY")
	if clientCertKey == "" {
		return nil, errors.New("Kafka Client Certificate Key not found!")
	}

	clientCert := os.Getenv("KAFKA_CLIENT_CERT")
	if clientCert == "" {
		return nil, errors.New("Kafka Client Certificate not found!")
	}

	roots := x509.NewCertPool()
	ok := roots.AppendCertsFromPEM([]byte(trustedCert))
	if !ok {
		return nil, errors.New("Unable to parse Root Cert. Please check your Heroku environment.")
	}
	// Setup certs for Sarama
	cert, err := tls.X509KeyPair([]byte(clientCert), []byte(clientCertKey))
	if err != nil {
		return nil, err
	}
	return &tls.Config{
		Certificates:       []tls.Certificate{cert},
		InsecureSkipVerify: true,
		RootCAs:            roots,
	}, nil
}

// Extract the host:port pairs from the Kafka URL(s)
func brokerAddresses() ([]string, error) {
	URL := os.Getenv("KAFKA_URL")
	if URL == "" {
		return nil, errors.New("Kafka URL not found!")
	}
	urls := strings.Split(URL, ",")
	addrs := make([]string, len(urls))
	for i, v := range urls {
		u, err := url.Parse(v)
		if err != nil {
			return nil, err
		}
		addrs[i] = u.Host
	}
	return addrs, nil
}
<|endoftext|>


<commit_before>/*
Package Goria provides the functionality of an LRU Cache with an eye to JSR 107
*/
package goria

import (
	"container/list"
	"errors"
)

type EvictionCallback func(key interface{}, value interface{})

type GoriaLRU struct {
	Name         string
	Size         int
	items        map[interface{}]*list.Element
	evictionList *list.List
	onEvict      EvictionCallback
	statsEnabled bool
	stats        CacheStats
}

type CacheStats struct {
	Items     int64
	Gets      int64
	Hits      int64
	Evictions int64
	Miss      int64
}

type entry struct {
	key   interface{}
	value interface{}
}

func newGoriaLRU(name string, size int, evictionC EvictionCallback, statsEnabled bool) (*GoriaLRU, error) {
	if size <= 0 {
		return nil, errors.New("The Goria Cache need a positive value as size")
	}
	c := &GoriaLRU{
		Name:         name,
		Size:         size,
		evictionList: list.New(),
		items:        make(map[interface{}]*list.Element),
		onEvict:      evictionC,
		statsEnabled: statsEnabled,
		stats: CacheStats{
			Items:     0,
			Evictions: 0,
			Gets:      0,
			Hits:      0,
			Miss:      0,
		},
	}
	return c, nil
}

func (c *GoriaLRU) Put(key, value interface{}) {
	if item, ok := c.items[key]; ok {
		c.evictionList.MoveToFront(item)
		item.Value.(*entry).value = value
		return
	}

	item := &entry{key, value}
	element := c.evictionList.PushFront(item)
	c.items[key] = element

	if c.evictionList.Len() > c.Size {
		c.removeFromTail()
	}

	if c.IsStatsEnabled() {
		c.stats.Items++
	}
}

func (c *GoriaLRU) PutAll(m map[interface{}]interface{}) {
	for key, value := range m {
		c.Put(key, value)
	}
}

func (c *GoriaLRU) PutIfAbsent(key, value interface{}) bool {
	var element, exists = c.items[key]
	if !exists && element == nil {
		item := &entry{key, value}
		element := c.evictionList.PushFront(item)
		c.items[key] = element

		if c.evictionList.Len() > c.Size {
			c.removeFromTail()
		}
		if c.IsStatsEnabled() {
			c.stats.Items++
		}
		return true
	}
	return false
}

func (c *GoriaLRU) Get(key interface{}) (value interface{}, exists bool) {
	if c.IsStatsEnabled() {
		c.stats.Gets++
	}
	if item, exists := c.items[key]; exists {
		c.evictionList.MoveToFront(item)
		if c.IsStatsEnabled() {
			c.stats.Hits++
		}
		return item.Value.(*entry).value, true
	}
	if c.IsStatsEnabled() {
		c.stats.Miss++
	}
	return
}

func (c *GoriaLRU) GetAll(m map[interface{}]interface{}) map[interface{}]interface{} {
	returnedMap := make(map[interface{}]interface{})

	for k := range m {
		value, exists := c.Get(k)
		if exists {
			returnedMap[k] = value
		}
	}
	return returnedMap
}

func (c *GoriaLRU) Replace(key, oldValue interface{}, newValue interface{}) bool {
	var element, exists = c.items[key]
	if exists && element.Value.(*entry).value == oldValue {
		c.evictionList.MoveToFront(element)
		element.Value.(*entry).value = newValue
		return true
	}
	return false
}

func (c *GoriaLRU) ReplaceWithKeyOnly(key, newValue interface{}) bool {
	var element, exists = c.items[key]
	if exists && element != nil {
		c.evictionList.MoveToFront(element)
		element.Value.(*entry).value = newValue
		return true
	}
	return false
}

func (c *GoriaLRU) GetAndReplace(key interface{}, newValue interface{}) interface{} {
	v, ok := c.Get(key)
	if ok {
		c.ReplaceWithKeyOnly(key, newValue)
		return v
	}
	return nil
}

func (c *GoriaLRU) RemoveWithKeyOnly(key interface{}) bool {
	if element, exists := c.items[key]; exists {
		c.removeElement(element)
		return true
	}
	return false
}

func (c *GoriaLRU) Remove(key interface{}, oldValue interface{}) bool {
	var element, exists = c.items[key]
	if exists && element.Value.(*entry).value == oldValue {
		c.removeElement(element)
		return true
	}
	return false
}

func (c *GoriaLRU) RemoveAll(m map[interface{}]interface{}) {
	for key, value := range m {
		c.Remove(key, value)
	}
}

func (c *GoriaLRU) RemoveAllWithoutParameters() {
	var keys = c.Keys()
	for i := 0; i < len(keys); i++ {
		c.RemoveWithKeyOnly(keys[i])
	}
}

func (c *GoriaLRU) GetAndRemove(key interface{}) interface{} {
	v, ok := c.Get(key)
	if ok {
		c.RemoveWithKeyOnly(key)
		return v
	}
	return nil
}

func (c *GoriaLRU) Keys() []interface{} {
	keys := make([]interface{}, len(c.items))
	i := 0
	for ent := c.evictionList.Back(); ent != nil; ent = ent.Prev() {
		keys[i] = ent.Value.(*entry).key
		i++
	}
	return keys
}

func (c *GoriaLRU) ContainsKey(key interface{}) bool {
	for ent := c.evictionList.Back(); ent != nil; ent = ent.Prev() {
		if ent.Value.(*entry).key == key {
			return true
		}
	}
	return false
}

func (c *GoriaLRU) Len() int {
	return c.evictionList.Len()
}

func (c *GoriaLRU) GetName() string {
	return c.Name
}

func (c *GoriaLRU) IsStatsEnabled() bool {
	return c.statsEnabled
}

func (c *GoriaLRU) GetStats() CacheStats {
	return c.stats
}

func (c *GoriaLRU) removeFromTail() {
	element := c.evictionList.Back()

	if element != nil {
		c.removeElement(element)
	}
}

func (c *GoriaLRU) removeElement(el *list.Element) {
	c.evictionList.Remove(el)
	entry := el.Value.(*entry)
	delete(c.items, entry.key)

	if c.onEvict != nil {
		c.onEvict(entry.key, entry.value)
	}

	if c.IsStatsEnabled() {
		c.stats.Evictions++
		c.stats.Items--
	}
}
<commit_msg>Add an RWMutex to GoriaLRU struct and add Lock/Unlock to Put Operation<commit_after>/*
Package Goria provides the functionality of an LRU Cache with an eye to JSR 107
*/
package goria

import (
	"container/list"
	"errors"
	"sync"
)

type EvictionCallback func(key interface{}, value interface{})

type GoriaLRU struct {
	Name         string
	Size         int
	items        map[interface{}]*list.Element
	evictionList *list.List
	onEvict      EvictionCallback
	statsEnabled bool
	stats        CacheStats
	rwMutex      sync.RWMutex
}

type CacheStats struct {
	Items     int64
	Gets      int64
	Hits      int64
	Evictions int64
	Miss      int64
}

type entry struct {
	key   interface{}
	value interface{}
}

func newGoriaLRU(name string, size int, evictionC EvictionCallback, statsEnabled bool) (*GoriaLRU, error) {
	if size <= 0 {
		return nil, errors.New("The Goria Cache need a positive value as size")
	}
	c := &GoriaLRU{
		Name:         name,
		Size:         size,
		evictionList: list.New(),
		items:        make(map[interface{}]*list.Element),
		onEvict:      evictionC,
		statsEnabled: statsEnabled,
		stats: CacheStats{
			Items:     0,
			Evictions: 0,
			Gets:      0,
			Hits:      0,
			Miss:      0,
		},
	}
	return c, nil
}

func (c *GoriaLRU) Put(key, value interface{}) {
	if item, ok := c.items[key]; ok {
		c.rwMutex.Lock()
		defer c.rwMutex.Unlock()
		c.evictionList.MoveToFront(item)
		item.Value.(*entry).value = value
		return
	}

	c.rwMutex.Lock()
	defer c.rwMutex.Unlock()
	item := &entry{key, value}
	element := c.evictionList.PushFront(item)
	c.items[key] = element

	if c.evictionList.Len() > c.Size {
		c.removeFromTail()
	}

	if c.IsStatsEnabled() {
		c.stats.Items++
	}
}

func (c *GoriaLRU) PutAll(m map[interface{}]interface{}) {
	for key, value := range m {
		c.Put(key, value)
	}
}

func (c *GoriaLRU) PutIfAbsent(key, value interface{}) bool {
	var element, exists = c.items[key]
	if !exists && element == nil {
		item := &entry{key, value}
		element := c.evictionList.PushFront(item)
		c.items[key] = element

		if c.evictionList.Len() > c.Size {
			c.removeFromTail()
		}
		if c.IsStatsEnabled() {
			c.stats.Items++
		}
		return true
	}
	return false
}

func (c *GoriaLRU) Get(key interface{}) (value interface{}, exists bool) {
	if c.IsStatsEnabled() {
		c.stats.Gets++
	}
	if item, exists := c.items[key]; exists {
		c.evictionList.MoveToFront(item)
		if c.IsStatsEnabled() {
			c.stats.Hits++
		}
		return item.Value.(*entry).value, true
	}
	if c.IsStatsEnabled() {
		c.stats.Miss++
	}
	return
}

func (c *GoriaLRU) GetAll(m map[interface{}]interface{}) map[interface{}]interface{} {
	returnedMap := make(map[interface{}]interface{})

	for k := range m {
		value, exists := c.Get(k)
		if exists {
			returnedMap[k] = value
		}
	}
	return returnedMap
}

func (c *GoriaLRU) Replace(key, oldValue interface{}, newValue interface{}) bool {
	var element, exists = c.items[key]
	if exists && element.Value.(*entry).value == oldValue {
		c.evictionList.MoveToFront(element)
		element.Value.(*entry).value = newValue
		return true
	}
	return false
}

func (c *GoriaLRU) ReplaceWithKeyOnly(key, newValue interface{}) bool {
	var element, exists = c.items[key]
	if exists && element != nil {
		c.evictionList.MoveToFront(element)
		element.Value.(*entry).value = newValue
		return true
	}
	return false
}

func (c *GoriaLRU) GetAndReplace(key interface{}, newValue interface{}) interface{} {
	v, ok := c.Get(key)
	if ok {
		c.ReplaceWithKeyOnly(key, newValue)
		return v
	}
	return nil
}

func (c *GoriaLRU) RemoveWithKeyOnly(key interface{}) bool {
	if element, exists := c.items[key]; exists {
		c.removeElement(element)
		return true
	}
	return false
}

func (c *GoriaLRU) Remove(key interface{}, oldValue interface{}) bool {
	var element, exists = c.items[key]
	if exists && element.Value.(*entry).value == oldValue {
		c.removeElement(element)
		return true
	}
	return false
}

func (c *GoriaLRU) RemoveAll(m map[interface{}]interface{}) {
	for key, value := range m {
		c.Remove(key, value)
	}
}

func (c *GoriaLRU) RemoveAllWithoutParameters() {
	var keys = c.Keys()
	for i := 0; i < len(keys); i++ {
		c.RemoveWithKeyOnly(keys[i])
	}
}

func (c *GoriaLRU) GetAndRemove(key interface{}) interface{} {
	v, ok := c.Get(key)
	if ok {
		c.RemoveWithKeyOnly(key)
		return v
	}
	return nil
}

func (c *GoriaLRU) Keys() []interface{} {
	keys := make([]interface{}, len(c.items))
	i := 0
	for ent := c.evictionList.Back(); ent != nil; ent = ent.Prev() {
		keys[i] = ent.Value.(*entry).key
		i++
	}
	return keys
}

func (c *GoriaLRU) ContainsKey(key interface{}) bool {
	for ent := c.evictionList.Back(); ent != nil; ent = ent.Prev() {
		if ent.Value.(*entry).key == key {
			return true
		}
	}
	return false
}

func (c *GoriaLRU) Len() int {
	return c.evictionList.Len()
}

func (c *GoriaLRU) GetName() string {
	return c.Name
}

func (c *GoriaLRU) IsStatsEnabled() bool {
	return c.statsEnabled
}

func (c *GoriaLRU) GetStats() CacheStats {
	return c.stats
}

func (c *GoriaLRU) removeFromTail() {
	element := c.evictionList.Back()

	if element != nil {
		c.removeElement(element)
	}
}

func (c *GoriaLRU) removeElement(el *list.Element) {
	c.evictionList.Remove(el)
	entry := el.Value.(*entry)
	delete(c.items, entry.key)

	if c.onEvict != nil {
		c.onEvict(entry.key, entry.value)
	}

	if c.IsStatsEnabled() {
		c.stats.Evictions++
		c.stats.Items--
	}
}
<|endoftext|>


<commit_before>package connector

import (
	"bytes"
	"io"

	"io/ioutil"

	"time"

	logging "github.com/op/go-logging"
)

var logStreamLogger = logging.MustGetLogger("stream_logger")

// LogStreamer defines a structure for collecting
// data into a writer and writing to a logger.
type LogStreamer struct {
	writer   io.Writer
	logger   *logging.Logger
	stopRead bool
}

// NewLogStreamer reads from a stream and logs to stdout
func NewLogStreamer() *LogStreamer {
	return &LogStreamer{
		writer: bytes.NewBuffer(nil),
		logger: logStreamLogger,
	}
}

// GetWriter returns a writer
func (ls *LogStreamer) GetWriter() io.Writer {
	return ls.writer
}

// SetLogger sets the logger
func (ls *LogStreamer) SetLogger(logger *logging.Logger) {
	ls.logger = logger
}

// Stop sets the stop read flag to true which effectively stops the
// the writer from being read and logged.
func (ls *LogStreamer) Stop() {
	ls.stopRead = true
}

// Start starts reading the writer and logging its data
func (ls *LogStreamer) Start() error {
	for ls.stopRead == false {
		bs, err := ioutil.ReadAll(ls.writer.(*bytes.Buffer))
		if err != nil {
			return err
		}
		if len(bs) > 0 {
			ls.logger.Info(string(bs))
		}
		time.Sleep(100 * time.Millisecond)
	}
	return nil
}
<commit_msg>reset writer<commit_after>package connector

import (
	"bytes"
	"io"

	"io/ioutil"

	"time"

	"sync"

	logging "github.com/op/go-logging"
)

var logStreamLogger = logging.MustGetLogger("stream_logger")

// LogStreamer defines a structure for collecting
// data into a writer and writing to a logger.
type LogStreamer struct {
	writer   io.Writer
	logger   *logging.Logger
	stopRead bool
	mutex    sync.Mutex
}

// NewLogStreamer reads from a stream and logs to stdout
func NewLogStreamer() *LogStreamer {
	return &LogStreamer{
		writer: bytes.NewBuffer(nil),
		logger: logStreamLogger,
		mutex:  sync.Mutex{},
	}
}

// GetWriter returns a writer
func (ls *LogStreamer) GetWriter() io.Writer {
	return ls.writer
}

// SetLogger sets the logger
func (ls *LogStreamer) SetLogger(logger *logging.Logger) {
	ls.logger = logger
}

// Stop sets the stop read flag to true which effectively stops the
// the writer from being read and logged.
func (ls *LogStreamer) Stop() {
	ls.stopRead = true
}

// Start starts reading the writer and logging its data
func (ls *LogStreamer) Start() error {
	for ls.stopRead == false {
		ls.mutex.Lock()
		bs, err := ioutil.ReadAll(ls.writer.(*bytes.Buffer))
		if err != nil {
			ls.mutex.Unlock()
			return err
		}

		// reset the writer
		ls.writer.(*bytes.Buffer).Reset()
		ls.mutex.Unlock()

		if len(bs) > 0 {
			ls.logger.Info(string(bs))
		}
		time.Sleep(100 * time.Millisecond)
	}
	return nil
}
<|endoftext|>


<commit_before>package goscrape

import (
	"encoding/hex"
	"time"
)

/**
 * Stores multiple sessions and keeps track
 * of when to refresh the connection ID.
 */
type Bulk struct {
	Sess   []Session
	Expire time.Time
}

/**
 * Scrapes 1 or more trackers for 1 or more
 * torrents. It doesn't save the connection.
 */
func Single(urls []string, btihs []string) []Result {
	bulk := NewBulk(urls)
	return bulk.ScrapeBulk(btihs)
}

/**
 * Creates a new bulk object for running multiple
 * scrapes on the same set of trackers without recreating
 * a new connection each time.
 */
func NewBulk(trackers []string) Bulk {
	size := len(trackers)
	var sessions []Session = make([]Session, size)
	var channels = make([]chan Session, size)

	for i := 0; i < size; i++ {
		channels[i] = make(chan Session)
		go asyncSession(trackers[i], channels[i])
	}

	for i := 0; i < size; i++ {
		sessions[i] = <-channels[i]
	}

	return Bulk{Sess: sessions, Expire: time.Now().Add(1 * time.Minute)}
}

/**
 * Scrapes a set of info hashes from the
 * connections the bulk was initialized with.
 */
func (bulk *Bulk) ScrapeBulk(btihs []string) []Result {
	// Refresh sessions if it's been over a minute
	if time.Now().After(bulk.Expire) {
		bulk.refreshSessions()
	}

	// Validate the btihs and get size
	var cleanBtihs []string = make([]string, 0)
	for _, btih := range btihs {
		// Take the BTIH and convert it into bytes
		infohash, err := hex.DecodeString(btih)
		// Check errors
		if err == nil {
			if len(infohash) == 20 {
				cleanBtihs = append(cleanBtihs, btih)
			}
		}
	}

	// Make a result variable
	var results []Result = make([]Result, len(cleanBtihs))
	for i := 0; i < len(results); i++ {
		results[i] = Result{cleanBtihs[i], 0, 0, 0}
	}

	// Loop through the sessions
	for _, sess := range bulk.Sess {
		// Perform a multi scrape with all btihs on the single session
		scrape, err := sess.scrape(cleanBtihs)
		if err == nil {
			// Merge result array into results
			for i, result := range scrape {
				if result.Seeders > results[i].Seeders {
					results[i].Seeders = result.Seeders
				}
				if result.Leechers > results[i].Leechers {
					results[i].Leechers = result.Leechers
				}
				if result.Completed > results[i].Completed {
					results[i].Completed = result.Completed
				}
			}
		}
	}

	return results
}

func asyncSession(url string, output chan Session) {
	output <- newConn(url)
}

func (bulk *Bulk) refreshSessions() {
	// Get the size of the sessions
	size := len(bulk.Sess)
	// Make channels
	var channels = make([]chan Session, size)

	// Make channels and make new sessions asynchronously
	for i := 0; i < size; i++ {
		channels[i] = make(chan Session)
		go asyncSession(bulk.Sess[i].URL, channels[i])
	}

	// Replace old sessions with new ones
	for i := 0; i < size; i++ {
		bulk.Sess[i] = <-channels[i]
	}

	// Update the expire time.
	bulk.Expire = time.Now().Add(1 * time.Minute)
}
<commit_msg>Fix comments duh.<commit_after>package goscrape

import (
	"encoding/hex"
	"time"
)

/*
Stores multiple sessions and keeps track
of when to refresh the connection ID.
*/
type Bulk struct {
	Sess   []Session
	Expire time.Time
}

/*
Scrapes 1 or more trackers for 1 or more
torrents. It doesn't save the connection.
*/
func Single(urls []string, btihs []string) []Result {
	bulk := NewBulk(urls)
	return bulk.ScrapeBulk(btihs)
}

/*
Creates a new bulk object for running multiple
scrapes on the same set of trackers without recreating
a new connection each time.
*/
func NewBulk(trackers []string) Bulk {
	size := len(trackers)
	var sessions []Session = make([]Session, size)
	var channels = make([]chan Session, size)

	for i := 0; i < size; i++ {
		channels[i] = make(chan Session)
		go asyncSession(trackers[i], channels[i])
	}

	for i := 0; i < size; i++ {
		sessions[i] = <-channels[i]
	}

	return Bulk{Sess: sessions, Expire: time.Now().Add(1 * time.Minute)}
}

/*
Scrapes a set of info hashes from the
connections the bulk was initialized with.
*/
func (bulk *Bulk) ScrapeBulk(btihs []string) []Result {
	// Refresh sessions if it's been over a minute
	if time.Now().After(bulk.Expire) {
		bulk.refreshSessions()
	}

	// Validate the btihs and get size
	var cleanBtihs []string = make([]string, 0)
	for _, btih := range btihs {
		// Take the BTIH and convert it into bytes
		infohash, err := hex.DecodeString(btih)
		// Check errors
		if err == nil {
			if len(infohash) == 20 {
				cleanBtihs = append(cleanBtihs, btih)
			}
		}
	}

	// Make a result variable
	var results []Result = make([]Result, len(cleanBtihs))
	for i := 0; i < len(results); i++ {
		results[i] = Result{cleanBtihs[i], 0, 0, 0}
	}

	// Loop through the sessions
	for _, sess := range bulk.Sess {
		// Perform a multi scrape with all btihs on the single session
		scrape, err := sess.scrape(cleanBtihs)
		if err == nil {
			// Merge result array into results
			for i, result := range scrape {
				if result.Seeders > results[i].Seeders {
					results[i].Seeders = result.Seeders
				}
				if result.Leechers > results[i].Leechers {
					results[i].Leechers = result.Leechers
				}
				if result.Completed > results[i].Completed {
					results[i].Completed = result.Completed
				}
			}
		}
	}

	return results
}

func asyncSession(url string, output chan Session) {
	output <- newConn(url)
}

func (bulk *Bulk) refreshSessions() {
	// Get the size of the sessions
	size := len(bulk.Sess)
	// Make channels
	var channels = make([]chan Session, size)

	// Make channels and make new sessions asynchronously
	for i := 0; i < size; i++ {
		channels[i] = make(chan Session)
		go asyncSession(bulk.Sess[i].URL, channels[i])
	}

	// Replace old sessions with new ones
	for i := 0; i < size; i++ {
		bulk.Sess[i] = <-channels[i]
	}

	// Update the expire time.
	bulk.Expire = time.Now().Add(1 * time.Minute)
}
<|endoftext|>