<commit_before>package server
import (
	"strings"
	"strconv"
	"fmt"
	"crypto/sha1"
	"encoding/hex"
	"time"
	"net"
)

/*
 Authorization: user ts sha1(user + key + ts + method + uri)

 */
func (self *Session) auth() (err error) {
	defer func() {
		// deley 1s on fail to prevent attact
		if err != nil {
			time.Sleep(1 * time.Second)
		}
	}()
	if !self.config.Auth.Enabled {
		return nil
	}

	authStr := self.req.Header.Get("Authorization")
	reqUser, reqHash, ts, err := parseAuthHeader(authStr)
	if err != nil {
		return err
	}

	user, ok := self.config.Users[reqUser]
	if !ok {
		return fmt.Errorf("user %s not found", reqUser)
	}
	if len(user.Hosts) > 0 {
		ok = false
		for _, host := range (user.Hosts) {
			_, allowedNet, err := net.ParseCIDR(host)
			if err != nil {
				continue
			}
			if allowedNet.Contains(net.ParseIP(self.req.RemoteAddr)) {
				ok = true
				break
			}
		}
		if !ok {
			return fmt.Errorf("remote host is denied")
		}
	}
	if user.Key != "" {
		nowTs := time.Now().Unix()
		maxDelta := self.config.Auth.MaxTimeDelta
		if nowTs - ts > int64(maxDelta) || ts - nowTs > int64(maxDelta) {
			return fmt.Errorf("timestamp delta too large")
		}
		strToHash := reqUser + user.Key + strconv.FormatInt(ts, 10) + self.req.Method + self.req.RequestURI
		sha1Sum := sha1.Sum([]byte(strToHash))
		realHash := hex.EncodeToString(sha1Sum[:])
		if reqHash != realHash {
			return fmt.Errorf("auth failed")
		}
	}
	return nil
}

func parseAuthHeader(authStr string) (user, hash string, ts int64, err error){
	segs := strings.SplitN(authStr, " ", 3)
	user = segs[0]
	if len(segs) < 3 {
		hash = ""
		ts = 0
		return
	}
	tsStr := segs[1]
	hash = segs[2]
	ts, err = strconv.ParseInt(tsStr, 10, 64)
	return
}<commit_msg>extract check<commit_after>package server
import (
	"strings"
	"strconv"
	"fmt"
	"crypto/sha1"
	"encoding/hex"
	"time"
	"net"
)

/*
 Authorization: user ts sha1(user + key + ts + method + uri)

 */
func (self *Session) auth() (err error) {
	defer func() {
		// deley 1s on fail to prevent attact
		if err != nil {
			time.Sleep(1 * time.Second)
		}
	}()
	if !self.config.Auth.Enabled {
		return nil
	}
	authStr := self.req.Header.Get("Authorization")
	reqUser, reqHash, ts, err := parseAuthHeader(authStr)
	if err != nil {
		return err
	}
	user, ok := self.config.Users[reqUser]
	if !ok {
		return fmt.Errorf("user %s not found", reqUser)
	}
	if err = checkHosts(self.req.RemoteAddr, user.Hosts); err != nil {
		return err
	}
	if user.Key != "" {
		nowTs := time.Now().Unix()
		maxDelta := self.config.Auth.MaxTimeDelta
		if nowTs - ts > int64(maxDelta) || ts - nowTs > int64(maxDelta) {
			return fmt.Errorf("timestamp delta too large")
		}
		strToHash := reqUser + user.Key + strconv.FormatInt(ts, 10) + self.req.Method + self.req.RequestURI
		sha1Sum := sha1.Sum([]byte(strToHash))
		realHash := hex.EncodeToString(sha1Sum[:])
		if reqHash != realHash {
			return fmt.Errorf("auth failed")
		}
	}
	return nil
}

func parseAuthHeader(authStr string) (user, hash string, ts int64, err error){
	segs := strings.SplitN(authStr, " ", 3)
	user = segs[0]
	if len(segs) < 3 {
		hash = ""
		ts = 0
		return
	}
	tsStr := segs[1]
	hash = segs[2]
	ts, err = strconv.ParseInt(tsStr, 10, 64)
	return
}

func checkHosts(remoteAddr string, hosts []string) error {
	if len(hosts) <= 0 {
		return nil
	}
	ok := false
	for _, host := range (hosts) {
		_, allowedNet, err := net.ParseCIDR(host)
		if err != nil {
			continue
		}
		if allowedNet.Contains(net.ParseIP(remoteAddr)) {
			ok = true
			break
		}
	}
	if !ok {
		return fmt.Errorf("remote host is denied")
	}
	return nil
}
<|endoftext|>


<commit_before>package routes

import (
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/qor/qor"
	"github.com/qor/qor-example/app/controllers"
	"github.com/qor/qor-example/config"
	"github.com/qor/qor-example/config/auth"
	"github.com/qor/qor-example/db"
	"github.com/qor/qor/utils"
)

func Router() *http.ServeMux {
	router := gin.Default()
	router.Use(func(ctx *gin.Context) {
		if locale := utils.GetLocale(&qor.Context{Request: ctx.Request, Writer: ctx.Writer}); locale != "" {
			newDB := db.DB.Set("l10n:locale", locale).Set("l10n:mode", "locale")
			ctx.Set("DB", newDB)
		}
	})
	gin.SetMode(gin.DebugMode)

	router.GET("/", controllers.HomeIndex)
	router.GET("/products/:code", controllers.ProductShow)
	router.GET("/switch_locale", controllers.SwitchLocale)

	var mux = http.NewServeMux()
	mux.Handle("/", router)
	mux.Handle("/auth/", auth.Auth.NewRouter())
	publicDir := http.Dir(strings.Join([]string{config.Root, "public"}, "/"))
	mux.Handle("/dist/", http.FileServer(publicDir))
	mux.Handle("/vendors/", http.FileServer(publicDir))
	mux.Handle("/images/", http.FileServer(publicDir))
	mux.Handle("/fonts/", http.FileServer(publicDir))
	return mux
}
<commit_msg>Using default locale mode<commit_after>package routes

import (
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/qor/qor"
	"github.com/qor/qor-example/app/controllers"
	"github.com/qor/qor-example/config"
	"github.com/qor/qor-example/config/auth"
	"github.com/qor/qor-example/db"
	"github.com/qor/qor/utils"
)

func Router() *http.ServeMux {
	router := gin.Default()
	router.Use(func(ctx *gin.Context) {
		if locale := utils.GetLocale(&qor.Context{Request: ctx.Request, Writer: ctx.Writer}); locale != "" {
			ctx.Set("DB", db.DB.Set("l10n:locale", locale))
		}
	})
	gin.SetMode(gin.DebugMode)

	router.GET("/", controllers.HomeIndex)
	router.GET("/products/:code", controllers.ProductShow)
	router.GET("/switch_locale", controllers.SwitchLocale)

	var mux = http.NewServeMux()
	mux.Handle("/", router)
	mux.Handle("/auth/", auth.Auth.NewRouter())
	publicDir := http.Dir(strings.Join([]string{config.Root, "public"}, "/"))
	mux.Handle("/dist/", http.FileServer(publicDir))
	mux.Handle("/vendors/", http.FileServer(publicDir))
	mux.Handle("/images/", http.FileServer(publicDir))
	mux.Handle("/fonts/", http.FileServer(publicDir))
	return mux
}
<|endoftext|>


<commit_before>package config

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"os"
)

// Routes configuration file path
var routesConfigFile = "config/routes.json"

// Routes is a variable used for storing all the routes that the api will have
var Routes []Route

// InitRoutes initializes the routes based on a configuration file
func InitRoutes(routesConfigPath string) {
	if len(routesConfigPath) != 0 {
		routesConfigFile = routesConfigPath
	}

	routesString, err := ioutil.ReadFile(routesConfigFile)

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

	deserializeRoutes(routesString)
}

// SaveRoutesConfiguration saves all the active routes (Routes slice) in json format
// into the configuration file
func SaveRoutesConfiguration() error {
	if len(Routes) == 0 {
		return errors.New("There are no routes configured in order to be saved")
	}

	data, err := json.MarshalIndent(Routes, "", "  ")

	if err != nil {
		return errors.New("Encoding routes slice to json failed!")
	}

	err = ioutil.WriteFile(routesConfigFile, data, os.ModeDevice)
	if err != nil {
		return err
	}

	return nil
}

// AddRoute adds a new route and makes it active
func AddRoute(route *Route, saveChangesToConfigFile bool) error {
	initialLength := len(Routes)

	for _, r := range Routes {
		if r.ID == route.ID {
			return errors.New("Route already exists!")
		}
	}

	Routes = append(Routes, *route)

	err := checkCollectionModification(route, initialLength)

	return saveChanges(err, saveChangesToConfigFile, SaveRoutesConfiguration)
}

// RemoveRoute disables and removes a certain route
func RemoveRoute(routeID string, saveChangesToConfigFile bool) error {
	initialLength := len(Routes)
	index := -1

	for ind, route := range Routes {
		if route.ID == routeID {
			index = ind
			break
		}
	}

	if index == -1 {
		return errors.New("Route was not found for deletion!")
	}

	removedRoute := Routes[index]
	Routes = append(Routes[:index], Routes[index+1:]...)

	err := checkCollectionModification(&removedRoute, initialLength)

	return saveChanges(err, saveChangesToConfigFile, SaveRoutesConfiguration)
}

// ModifyRoute modifies the state and information of a certain route
func ModifyRoute(routeID string, newRouteData Route, saveChangesToConfigFile bool) error {
	for i := 0; i < len(Routes); i++ {
		if Routes[i].ID == routeID {
			Routes[i] = newRouteData

			return saveChanges(nil, saveChangesToConfigFile, SaveRoutesConfiguration)
		}
	}

	return errors.New("Route was not found for modification!")
}

// GetRoute fetches a Route entity from the active routes list, base on its ID
func GetRoute(endpoint string) *Route {
	for _, route := range Routes {
		if route.Endpoint == endpoint {
			return &route
		}
	}

	return nil
}

func deserializeRoutes(routesString []byte) {
	err := json.Unmarshal(routesString, &Routes)

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

func checkCollectionModification(route *Route, initialLength int) error {
	if initialLength == len(Routes) {
		return errors.New("The route couldn't be processed for the collection!")
	}

	return nil
}

func saveChanges(err error, saveChangesToConfigFile bool, saverFunction func() error) error {
	if err != nil {
		return err
	}

	if saveChangesToConfigFile {
		return saverFunction()
	}

	return nil
}
<commit_msg>Changed the routes config package to support multiple routes adding at the same time<commit_after>package config

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"log"
	"os"
)

// Routes configuration file path
var routesConfigFile = "config/routes.json"

// Routes is a variable used for storing all the routes that the api will have
var Routes []Route

// InitRoutes initializes the routes based on a configuration file
func InitRoutes(routesConfigPath string) {
	if len(routesConfigPath) != 0 {
		routesConfigFile = routesConfigPath
	}

	routesData, err := ioutil.ReadFile(routesConfigFile)

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

	deserializeRoutes(routesData)
}

// SaveRoutesConfiguration saves all the active routes (Routes slice) in json format
// into the configuration file
func SaveRoutesConfiguration() error {
	if len(Routes) == 0 {
		return errors.New("There are no routes configured in order to be saved")
	}

	data, err := json.MarshalIndent(Routes, "", "  ")

	if err != nil {
		return errors.New("Encoding routes slice to json failed!")
	}

	err = ioutil.WriteFile(routesConfigFile, data, os.ModeDevice)
	if err != nil {
		return err
	}

	return nil
}

// AddRoutes adds a new route and makes it active
func AddRoutes(saveChangesToConfigFile bool, newRoutes ...Route) error {
	initialLength := len(Routes)

	for _, route := range newRoutes {
		existingRoute := GetRoute(route.Endpoint)
		if existingRoute != nil {
			return errors.New("Route already exists!")
		}
	}

	Routes = append(Routes, newRoutes...)

	err := checkCollectionModification(initialLength)

	return saveChanges(err, saveChangesToConfigFile, SaveRoutesConfiguration)
}

// RemoveRoute disables and removes a certain route
func RemoveRoute(routeID string, saveChangesToConfigFile bool) error {
	initialLength := len(Routes)
	index := -1

	for ind, route := range Routes {
		if route.ID == routeID {
			index = ind
			break
		}
	}

	if index == -1 {
		return errors.New("Route was not found for deletion!")
	}

	Routes = append(Routes[:index], Routes[index+1:]...)

	err := checkCollectionModification(initialLength)

	return saveChanges(err, saveChangesToConfigFile, SaveRoutesConfiguration)
}

// ModifyRoute modifies the state and information of a certain route
func ModifyRoute(routeID string, newRouteData Route, saveChangesToConfigFile bool) error {
	for i := 0; i < len(Routes); i++ {
		if Routes[i].ID == routeID {
			Routes[i] = newRouteData

			return saveChanges(nil, saveChangesToConfigFile, SaveRoutesConfiguration)
		}
	}

	return errors.New("Route was not found for modification!")
}

// GetRoute fetches a Route entity from the active routes list, base on its ID
func GetRoute(endpoint string) *Route {
	for _, route := range Routes {
		if route.Endpoint == endpoint {
			return &route
		}
	}

	return nil
}

func deserializeRoutes(routesData []byte) {
	err := json.Unmarshal(routesData, &Routes)

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

func checkCollectionModification(initialLength int) error {
	if initialLength == len(Routes) {
		return errors.New("The route couldn't be processed for the collection!")
	}

	return nil
}

func saveChanges(err error, saveChangesToConfigFile bool, saverFunction func() error) error {
	if err != nil {
		return err
	}

	if saveChangesToConfigFile {
		return saverFunction()
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Hari haran. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/gorilla/websocket"
)

func speculoHandler(w http.ResponseWriter, r *http.Request) {
	scanner := bufio.NewScanner(os.Stdin)
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		log.Println(err)
		return
	}
	go func() {
		for {
			_, r, err := ws.ReadMessage() // messageType is ignored
			if err != nil {
				log.Fatal(err)
			}
			fmt.Println(string(r))
		}
	}()
	for scanner.Scan() {
		inputcommand := []byte(scanner.Text())
		if err := ws.WriteMessage(1, inputcommand); err != nil {
			log.Println(err)
		}
	}
}

func main() {
	path, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}

	http.Handle("/", http.FileServer(http.Dir(path)))
	http.HandleFunc("/repl", speculoHandler)

	fmt.Println("Open http://localhost:8080 and type js commands in the terminal!")
	http.ListenAndServe(":8080", nil)
}
<commit_msg>Make a proper terminal interface for the REPL<commit_after>// Copyright 2014 Hari haran. All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"os"

	"github.com/gorilla/websocket"
	"golang.org/x/crypto/ssh/terminal"
)

func speculoHandler(w http.ResponseWriter, r *http.Request) {
	ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
	if _, ok := err.(websocket.HandshakeError); ok {
		http.Error(w, "Not a websocket handshake", 400)
		return
	} else if err != nil {
		log.Fatal(err)
		return
	}

	oldState, err := terminal.MakeRaw(0)
	if err != nil {
		panic(err)
	}
	defer terminal.Restore(0, oldState)
	term := terminal.NewTerminal(os.Stdin, "> ")

	term.AutoCompleteCallback = func(line string, pos int, key rune) (newLine string, newPos int, ok bool) {
		if key == '\x03' {
			fmt.Println()
			if oldState != nil {
				terminal.Restore(0, oldState)
			}
			os.Exit(0)
		}
		return "", 0, false
	}

	go func() {
		for {
			_, r, err := ws.ReadMessage() // messageType is ignored
			if err != nil {
				log.Fatal(err)
			}
			term.Write(r)
			term.Write([]byte("\n"))
		}
	}()

	for {
		text, err := term.ReadLine()
		if err != nil {
			// Ctrl-D is EOF. So, quit without printing stacktrace.
			if err == io.EOF {
				fmt.Println()
				if oldState != nil {
					terminal.Restore(0, oldState)
				}
				os.Exit(0)
			}
			panic(err)
		}
		inputcommand := []byte(text)
		if err := ws.WriteMessage(1, inputcommand); err != nil {
			log.Println(err)
		}
	}
}

func main() {
	path, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}

	http.Handle("/", http.FileServer(http.Dir(path)))
	http.HandleFunc("/repl", speculoHandler)

	fmt.Println("Open http://localhost:8080 and type js commands in the terminal!")
	http.ListenAndServe(":8080", nil)
}
<|endoftext|>


<commit_before>// 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 spinner is a simple package to add a spinner / progress indicator to any terminal application.
package spinner

import (
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/fatih/color"
)

// errInvalidColor is returned when attempting to set an invalid color
var errInvalidColor = errors.New("invalid color")

// validColors holds an array of the only colors allowed
var validColors = map[string]bool{
	// default colors for backwards compatibility
	"black":   true,
	"red":     true,
	"green":   true,
	"yellow":  true,
	"blue":    true,
	"magenta": true,
	"cyan":    true,
	"white":   true,

	// attributes
	"reset":        true,
	"bold":         true,
	"faint":        true,
	"italic":       true,
	"underline":    true,
	"blinkslow":    true,
	"blinkrapid":   true,
	"reversevideo": true,
	"concealed":    true,
	"crossedout":   true,

	// foreground text
	"fgBlack":   true,
	"fgRed":     true,
	"fgGreen":   true,
	"fgYellow":  true,
	"fgBlue":    true,
	"fgMagenta": true,
	"fgCyan":    true,
	"fgWhite":   true,

	// foreground Hi-Intensity text
	"fgHiBlack":   true,
	"fgHiRed":     true,
	"fgHiGreen":   true,
	"fgHiYellow":  true,
	"fgHiBlue":    true,
	"fgHiMagenta": true,
	"fgHiCyan":    true,
	"fgHiWhite":   true,

	// background text
	"bgBlack":   true,
	"bgRed":     true,
	"bgGreen":   true,
	"bgYellow":  true,
	"bgBlue":    true,
	"bgMagenta": true,
	"bgCyan":    true,
	"bgWhite":   true,

	// background Hi-Intensity text
	"bgHiBlack":   true,
	"bgHiRed":     true,
	"bgHiGreen":   true,
	"bgHiYellow":  true,
	"bgHiBlue":    true,
	"bgHiMagenta": true,
	"bgHiCyan":    true,
	"bgHiWhite":   true,
}

// returns a valid color's foreground text color attribute
var colorAttributeMap = map[string]color.Attribute{
	// default colors for backwards compatibility
	"black":   color.FgBlack,
	"red":     color.FgRed,
	"green":   color.FgGreen,
	"yellow":  color.FgYellow,
	"blue":    color.FgBlue,
	"magenta": color.FgMagenta,
	"cyan":    color.FgCyan,
	"white":   color.FgWhite,

	// attributes
	"reset":        color.Reset,
	"bold":         color.Bold,
	"faint":        color.Faint,
	"italic":       color.Italic,
	"underline":    color.Underline,
	"blinkslow":    color.BlinkSlow,
	"blinkrapid":   color.BlinkRapid,
	"reversevideo": color.ReverseVideo,
	"concealed":    color.Concealed,
	"crossedout":   color.CrossedOut,

	// foreground text colors
	"fgBlack":   color.FgBlack,
	"fgRed":     color.FgRed,
	"fgGreen":   color.FgGreen,
	"fgYellow":  color.FgYellow,
	"fgBlue":    color.FgBlue,
	"fgMagenta": color.FgMagenta,
	"fgCyan":    color.FgCyan,
	"fgWhite":   color.FgWhite,

	// foreground Hi-Intensity text colors
	"fgHiBlack":   color.FgHiBlack,
	"fgHiRed":     color.FgHiRed,
	"fgHiGreen":   color.FgHiGreen,
	"fgHiYellow":  color.FgHiYellow,
	"fgHiBlue":    color.FgHiBlue,
	"fgHiMagenta": color.FgHiMagenta,
	"fgHiCyan":    color.FgHiCyan,
	"fgHiWhite":   color.FgHiWhite,

	// background text colors
	"bgBlack":   color.BgBlack,
	"bgRed":     color.BgRed,
	"bgGreen":   color.BgGreen,
	"bgYellow":  color.BgYellow,
	"bgBlue":    color.BgBlue,
	"bgMagenta": color.BgMagenta,
	"bgCyan":    color.BgCyan,
	"bgWhite":   color.BgWhite,

	// background Hi-Intensity text colors
	"bgHiBlack":   color.BgHiBlack,
	"bgHiRed":     color.BgHiRed,
	"bgHiGreen":   color.BgHiGreen,
	"bgHiYellow":  color.BgHiYellow,
	"bgHiBlue":    color.BgHiBlue,
	"bgHiMagenta": color.BgHiMagenta,
	"bgHiCyan":    color.BgHiCyan,
	"bgHiWhite":   color.BgHiWhite,
}

// validColor will make sure the given color is actually allowed.
func validColor(c string) bool {
	if validColors[c] {
		return true
	}
	return false
}

// Spinner struct to hold the provided options.
type Spinner struct {
	mu         *sync.RWMutex                 //
	Delay      time.Duration                 // Delay is the speed of the indicator
	chars      []string                      // chars holds the chosen character set
	Prefix     string                        // Prefix is the text preppended to the indicator
	Suffix     string                        // Suffix is the text appended to the indicator
	FinalMSG   string                        // string displayed after Stop() is called
	lastOutput string                        // last character(set) written
	color      func(a ...interface{}) string // default color is white
	Writer     io.Writer                     // to make testing better, exported so users have access. Use `WithWriter` to update after initialization.
	active     bool                          // active holds the state of the spinner
	stopChan   chan struct{}                 // stopChan is a channel used to stop the indicator
	HideCursor bool                          // hideCursor determines if the cursor is visible
	PreUpdate  func(s *Spinner)              // will be triggered before every spinner update
	PostUpdate func(s *Spinner)              // will be triggered after every spinner update
}

// New provides a pointer to an instance of Spinner with the supplied options.
func New(cs []string, d time.Duration, options ...Option) *Spinner {
	s := &Spinner{
		Delay:    d,
		chars:    cs,
		color:    color.New(color.FgWhite).SprintFunc(),
		mu:       &sync.RWMutex{},
		Writer:   color.Output,
		active:   false,
		stopChan: make(chan struct{}, 1),
	}

	for _, option := range options {
		option(s)
	}
	return s
}

// Option is a function that takes a spinner and applies
// a given configuration.
type Option func(*Spinner)

// Options contains fields to configure the spinner.
type Options struct {
	Color      string
	Suffix     string
	FinalMSG   string
	HideCursor bool
}

// WithColor adds the given color to the spinner.
func WithColor(color string) Option {
	return func(s *Spinner) {
		s.Color(color)
	}
}

// WithSuffix adds the given string to the spinner
// as the suffix.
func WithSuffix(suffix string) Option {
	return func(s *Spinner) {
		s.Suffix = suffix
	}
}

// WithFinalMSG adds the given string ot the spinner
// as the final message to be written.
func WithFinalMSG(finalMsg string) Option {
	return func(s *Spinner) {
		s.FinalMSG = finalMsg
	}
}

// WithHiddenCursor hides the cursor
// if hideCursor = true given.
func WithHiddenCursor(hideCursor bool) Option {
	return func(s *Spinner) {
		s.HideCursor = hideCursor
	}
}

// WithWriter adds the given writer to the spinner. This
// function should be favored over directly assigning to
// the struct value.
func WithWriter(w io.Writer) Option {
	return func(s *Spinner) {
		s.mu.Lock()
		s.Writer = w
		s.mu.Unlock()
	}
}

// Active will return whether or not the spinner is currently active.
func (s *Spinner) Active() bool {
	return s.active
}

// Start will start the indicator.
func (s *Spinner) Start() {
	s.mu.Lock()
	if s.active {
		s.mu.Unlock()
		return
	}
	if s.HideCursor && runtime.GOOS != "windows" {
		// hides the cursor
		fmt.Print("\033[?25l")
	}
	s.active = true
	s.mu.Unlock()

	go func() {
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					if !s.active {
						return
					}
					s.mu.Lock()
					s.erase()

					if s.PreUpdate != nil {
						s.PreUpdate(s)
					}

					var outColor string
					if runtime.GOOS == "windows" {
						if s.Writer == os.Stderr {
							outColor = fmt.Sprintf("\r%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
						} else {
							outColor = fmt.Sprintf("\r%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
						}
					} else {
						outColor = fmt.Sprintf("%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
					}
					outPlain := fmt.Sprintf("%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
					fmt.Fprint(s.Writer, outColor)
					s.lastOutput = outPlain
					delay := s.Delay

					if s.PostUpdate != nil {
						s.PostUpdate(s)
					}

					s.mu.Unlock()
					time.Sleep(delay)
				}
			}
		}
	}()
}

// Stop stops the indicator.
func (s *Spinner) Stop() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.active {
		s.active = false
		if s.HideCursor && runtime.GOOS != "windows" {
			// makes the cursor visible
			fmt.Print("\033[?25h")
		}
		s.erase()
		if s.FinalMSG != "" {
			fmt.Fprintf(s.Writer, s.FinalMSG)
		}
		s.stopChan <- struct{}{}
	}
}

// Restart will stop and start the indicator.
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to the indicator.
func (s *Spinner) Reverse() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
}

// Color will set the struct field for the given color to be used.
func (s *Spinner) Color(colors ...string) error {
	colorAttributes := make([]color.Attribute, len(colors))

	// Verify colours are valid and place the appropriate attribute in the array
	for index, c := range colors {
		if !validColor(c) {
			return errInvalidColor
		}
		colorAttributes[index] = colorAttributeMap[c]
	}

	s.mu.Lock()
	s.color = color.New(colorAttributes...).SprintFunc()
	s.mu.Unlock()
	s.Restart()
	return nil
}

// UpdateSpeed will set the indicator delay to the given value.
func (s *Spinner) UpdateSpeed(d time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Delay = d
}

// UpdateCharSet will change the current character set to the given one.
func (s *Spinner) UpdateCharSet(cs []string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.chars = cs
}

// erase deletes written characters.
// Caller must already hold s.lock.
func (s *Spinner) erase() {
	n := utf8.RuneCountInString(s.lastOutput)
	if runtime.GOOS == "windows" {
		clearString := "\r"
		for i := 0; i < n; i++ {
			clearString += " "
		}
		clearString += "\r"
		fmt.Fprintf(s.Writer, clearString)
		s.lastOutput = ""
		return
	}
	del, _ := hex.DecodeString("7f")
	for _, c := range []string{"\b", string(del), "\b", "\033[K"} { // "\033[K" for macOS Terminal
		for i := 0; i < n; i++ {
			fmt.Fprintf(s.Writer, c)
		}
	}
	s.lastOutput = ""
}

// Lock allows for manual control to lock the spinner.
func (s *Spinner) Lock() {
	s.mu.Lock()
}

// Unlock allows for manual control to unlock the spinner.
func (s *Spinner) Unlock() {
	s.mu.Unlock()
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string.
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, length)
	for i := 0; i < length; i++ {
		numSeq[i] = strconv.Itoa(i)
	}
	return numSeq
}
<commit_msg>fix issue 55 - creep to upper left<commit_after>// 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 spinner is a simple package to add a spinner / progress indicator to any terminal application.
package spinner

import (
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/fatih/color"
)

// errInvalidColor is returned when attempting to set an invalid color
var errInvalidColor = errors.New("invalid color")

// validColors holds an array of the only colors allowed
var validColors = map[string]bool{
	// default colors for backwards compatibility
	"black":   true,
	"red":     true,
	"green":   true,
	"yellow":  true,
	"blue":    true,
	"magenta": true,
	"cyan":    true,
	"white":   true,

	// attributes
	"reset":        true,
	"bold":         true,
	"faint":        true,
	"italic":       true,
	"underline":    true,
	"blinkslow":    true,
	"blinkrapid":   true,
	"reversevideo": true,
	"concealed":    true,
	"crossedout":   true,

	// foreground text
	"fgBlack":   true,
	"fgRed":     true,
	"fgGreen":   true,
	"fgYellow":  true,
	"fgBlue":    true,
	"fgMagenta": true,
	"fgCyan":    true,
	"fgWhite":   true,

	// foreground Hi-Intensity text
	"fgHiBlack":   true,
	"fgHiRed":     true,
	"fgHiGreen":   true,
	"fgHiYellow":  true,
	"fgHiBlue":    true,
	"fgHiMagenta": true,
	"fgHiCyan":    true,
	"fgHiWhite":   true,

	// background text
	"bgBlack":   true,
	"bgRed":     true,
	"bgGreen":   true,
	"bgYellow":  true,
	"bgBlue":    true,
	"bgMagenta": true,
	"bgCyan":    true,
	"bgWhite":   true,

	// background Hi-Intensity text
	"bgHiBlack":   true,
	"bgHiRed":     true,
	"bgHiGreen":   true,
	"bgHiYellow":  true,
	"bgHiBlue":    true,
	"bgHiMagenta": true,
	"bgHiCyan":    true,
	"bgHiWhite":   true,
}

// returns a valid color's foreground text color attribute
var colorAttributeMap = map[string]color.Attribute{
	// default colors for backwards compatibility
	"black":   color.FgBlack,
	"red":     color.FgRed,
	"green":   color.FgGreen,
	"yellow":  color.FgYellow,
	"blue":    color.FgBlue,
	"magenta": color.FgMagenta,
	"cyan":    color.FgCyan,
	"white":   color.FgWhite,

	// attributes
	"reset":        color.Reset,
	"bold":         color.Bold,
	"faint":        color.Faint,
	"italic":       color.Italic,
	"underline":    color.Underline,
	"blinkslow":    color.BlinkSlow,
	"blinkrapid":   color.BlinkRapid,
	"reversevideo": color.ReverseVideo,
	"concealed":    color.Concealed,
	"crossedout":   color.CrossedOut,

	// foreground text colors
	"fgBlack":   color.FgBlack,
	"fgRed":     color.FgRed,
	"fgGreen":   color.FgGreen,
	"fgYellow":  color.FgYellow,
	"fgBlue":    color.FgBlue,
	"fgMagenta": color.FgMagenta,
	"fgCyan":    color.FgCyan,
	"fgWhite":   color.FgWhite,

	// foreground Hi-Intensity text colors
	"fgHiBlack":   color.FgHiBlack,
	"fgHiRed":     color.FgHiRed,
	"fgHiGreen":   color.FgHiGreen,
	"fgHiYellow":  color.FgHiYellow,
	"fgHiBlue":    color.FgHiBlue,
	"fgHiMagenta": color.FgHiMagenta,
	"fgHiCyan":    color.FgHiCyan,
	"fgHiWhite":   color.FgHiWhite,

	// background text colors
	"bgBlack":   color.BgBlack,
	"bgRed":     color.BgRed,
	"bgGreen":   color.BgGreen,
	"bgYellow":  color.BgYellow,
	"bgBlue":    color.BgBlue,
	"bgMagenta": color.BgMagenta,
	"bgCyan":    color.BgCyan,
	"bgWhite":   color.BgWhite,

	// background Hi-Intensity text colors
	"bgHiBlack":   color.BgHiBlack,
	"bgHiRed":     color.BgHiRed,
	"bgHiGreen":   color.BgHiGreen,
	"bgHiYellow":  color.BgHiYellow,
	"bgHiBlue":    color.BgHiBlue,
	"bgHiMagenta": color.BgHiMagenta,
	"bgHiCyan":    color.BgHiCyan,
	"bgHiWhite":   color.BgHiWhite,
}

// validColor will make sure the given color is actually allowed.
func validColor(c string) bool {
	if validColors[c] {
		return true
	}
	return false
}

// Spinner struct to hold the provided options.
type Spinner struct {
	mu         *sync.RWMutex                 //
	Delay      time.Duration                 // Delay is the speed of the indicator
	chars      []string                      // chars holds the chosen character set
	Prefix     string                        // Prefix is the text preppended to the indicator
	Suffix     string                        // Suffix is the text appended to the indicator
	FinalMSG   string                        // string displayed after Stop() is called
	lastOutput string                        // last character(set) written
	color      func(a ...interface{}) string // default color is white
	Writer     io.Writer                     // to make testing better, exported so users have access. Use `WithWriter` to update after initialization.
	active     bool                          // active holds the state of the spinner
	stopChan   chan struct{}                 // stopChan is a channel used to stop the indicator
	HideCursor bool                          // hideCursor determines if the cursor is visible
	PreUpdate  func(s *Spinner)              // will be triggered before every spinner update
	PostUpdate func(s *Spinner)              // will be triggered after every spinner update
}

// New provides a pointer to an instance of Spinner with the supplied options.
func New(cs []string, d time.Duration, options ...Option) *Spinner {
	s := &Spinner{
		Delay:    d,
		chars:    cs,
		color:    color.New(color.FgWhite).SprintFunc(),
		mu:       &sync.RWMutex{},
		Writer:   color.Output,
		active:   false,
		stopChan: make(chan struct{}, 1),
	}

	for _, option := range options {
		option(s)
	}
	return s
}

// Option is a function that takes a spinner and applies
// a given configuration.
type Option func(*Spinner)

// Options contains fields to configure the spinner.
type Options struct {
	Color      string
	Suffix     string
	FinalMSG   string
	HideCursor bool
}

// WithColor adds the given color to the spinner.
func WithColor(color string) Option {
	return func(s *Spinner) {
		s.Color(color)
	}
}

// WithSuffix adds the given string to the spinner
// as the suffix.
func WithSuffix(suffix string) Option {
	return func(s *Spinner) {
		s.Suffix = suffix
	}
}

// WithFinalMSG adds the given string ot the spinner
// as the final message to be written.
func WithFinalMSG(finalMsg string) Option {
	return func(s *Spinner) {
		s.FinalMSG = finalMsg
	}
}

// WithHiddenCursor hides the cursor
// if hideCursor = true given.
func WithHiddenCursor(hideCursor bool) Option {
	return func(s *Spinner) {
		s.HideCursor = hideCursor
	}
}

// WithWriter adds the given writer to the spinner. This
// function should be favored over directly assigning to
// the struct value.
func WithWriter(w io.Writer) Option {
	return func(s *Spinner) {
		s.mu.Lock()
		s.Writer = w
		s.mu.Unlock()
	}
}

// Active will return whether or not the spinner is currently active.
func (s *Spinner) Active() bool {
	return s.active
}

// Start will start the indicator.
func (s *Spinner) Start() {
	s.mu.Lock()
	if s.active {
		s.mu.Unlock()
		return
	}
	if s.HideCursor && runtime.GOOS != "windows" {
		// hides the cursor
		fmt.Print("\033[?25l")
	}
	s.active = true
	s.mu.Unlock()

	go func() {
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					if !s.active {
						return
					}
					s.mu.Lock()
					s.erase()

					if s.PreUpdate != nil {
						s.PreUpdate(s)
					}

					var outColor string
					if runtime.GOOS == "windows" {
						if s.Writer == os.Stderr {
							outColor = fmt.Sprintf("\r%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
						} else {
							outColor = fmt.Sprintf("\r%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
						}
					} else {
						outColor = fmt.Sprintf("%s%s%s ", s.Prefix, s.color(s.chars[i]), s.Suffix)
					}
					outPlain := fmt.Sprintf("%s%s%s ", s.Prefix, s.chars[i], s.Suffix)
					fmt.Fprint(s.Writer, outColor)
					s.lastOutput = outPlain
					delay := s.Delay

					if s.PostUpdate != nil {
						s.PostUpdate(s)
					}

					s.mu.Unlock()
					time.Sleep(delay)
				}
			}
		}
	}()
}

// Stop stops the indicator.
func (s *Spinner) Stop() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.active {
		s.active = false
		if s.HideCursor && runtime.GOOS != "windows" {
			// makes the cursor visible
			fmt.Print("\033[?25h")
		}
		s.erase()
		if s.FinalMSG != "" {
			fmt.Fprintf(s.Writer, s.FinalMSG)
		}
		s.stopChan <- struct{}{}
	}
}

// Restart will stop and start the indicator.
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to the indicator.
func (s *Spinner) Reverse() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
}

// Color will set the struct field for the given color to be used.
func (s *Spinner) Color(colors ...string) error {
	colorAttributes := make([]color.Attribute, len(colors))

	// Verify colours are valid and place the appropriate attribute in the array
	for index, c := range colors {
		if !validColor(c) {
			return errInvalidColor
		}
		colorAttributes[index] = colorAttributeMap[c]
	}

	s.mu.Lock()
	s.color = color.New(colorAttributes...).SprintFunc()
	s.mu.Unlock()
	s.Restart()
	return nil
}

// UpdateSpeed will set the indicator delay to the given value.
func (s *Spinner) UpdateSpeed(d time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Delay = d
}

// UpdateCharSet will change the current character set to the given one.
func (s *Spinner) UpdateCharSet(cs []string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.chars = cs
}

// erase deletes written characters.
// Caller must already hold s.lock.
func (s *Spinner) erase() {
	n := utf8.RuneCountInString(s.lastOutput)
	if runtime.GOOS == "windows" {
		clearString := "\r"
		for i := 0; i < n; i++ {
			clearString += " "
		}
		clearString += "\r"
		fmt.Fprintf(s.Writer, clearString)
		s.lastOutput = ""
		return
	}
	del, _ := hex.DecodeString("7f")
	for _, c := range []string{"\b", string(del)} {
		for i := 0; i < n; i++ {
			fmt.Fprintf(s.Writer, c)
		}
	}
	fmt.Fprintf(s.Writer, "\033[K") // erases to end of line
	s.lastOutput = ""
}

// Lock allows for manual control to lock the spinner.
func (s *Spinner) Lock() {
	s.mu.Lock()
}

// Unlock allows for manual control to unlock the spinner.
func (s *Spinner) Unlock() {
	s.mu.Unlock()
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string.
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, length)
	for i := 0; i < length; i++ {
		numSeq[i] = strconv.Itoa(i)
	}
	return numSeq
}
<|endoftext|>


<commit_before>// +build windows,unit

// Copyright 2014-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 task

import (
	"encoding/json"
	"fmt"
	"runtime"
	"testing"

	"github.com/aws/amazon-ecs-agent/agent/acs/model/ecsacs"
	apicontainer "github.com/aws/amazon-ecs-agent/agent/api/container"
	apicontainerstatus "github.com/aws/amazon-ecs-agent/agent/api/container/status"
	apitaskstatus "github.com/aws/amazon-ecs-agent/agent/api/task/status"
	"github.com/aws/amazon-ecs-agent/agent/config"
	taskresourcevolume "github.com/aws/amazon-ecs-agent/agent/taskresource/volume"

	"github.com/aws/amazon-ecs-agent/agent/dockerclient"
	dockercontainer "github.com/docker/docker/api/types/container"
	"github.com/stretchr/testify/assert"
)

const (
	minDockerClientAPIVersion = dockerclient.Version_1_24

	nonZeroMemoryReservationValue  = 1
	expectedMemoryReservationValue = 0
)

func TestPostUnmarshalWindowsCanonicalPaths(t *testing.T) {
	// Testing type conversions, bleh. At least the type conversion itself
	// doesn't look this messy.
	taskFromAcs := ecsacs.Task{
		Arn:           strptr("myArn"),
		DesiredStatus: strptr("RUNNING"),
		Family:        strptr("myFamily"),
		Version:       strptr("1"),
		Containers: []*ecsacs.Container{
			{
				Name: strptr("myName"),
				MountPoints: []*ecsacs.MountPoint{
					{
						ContainerPath: strptr(`C:/Container/Path`),
						SourceVolume:  strptr("sourceVolume"),
					},
				},
			},
		},
		Volumes: []*ecsacs.Volume{
			{
				Name: strptr("sourceVolume"),
				Type: strptr("host"),
				Host: &ecsacs.HostVolumeProperties{
					SourcePath: strptr(`C:/Host/path`),
				},
			},
		},
	}
	expectedTask := &Task{
		Arn:                 "myArn",
		DesiredStatusUnsafe: apitaskstatus.TaskRunning,
		Family:              "myFamily",
		Version:             "1",
		Containers: []*apicontainer.Container{
			{
				Name: "myName",
				MountPoints: []apicontainer.MountPoint{
					{
						ContainerPath: `c:\container\path`,
						SourceVolume:  "sourceVolume",
					},
				},
				TransitionDependenciesMap: make(map[apicontainerstatus.ContainerStatus]apicontainer.TransitionDependencySet),
			},
		},
		Volumes: []TaskVolume{
			{
				Name: "sourceVolume",
				Type: "host",
				Volume: &taskresourcevolume.FSHostVolume{
					FSSourcePath: `c:\host\path`,
				},
			},
		},
		StartSequenceNumber: 42,
	}

	seqNum := int64(42)
	task, err := TaskFromACS(&taskFromAcs, &ecsacs.PayloadMessage{SeqNum: &seqNum})
	assert.Nil(t, err, "Should be able to handle acs task")
	cfg := config.Config{TaskCPUMemLimit: config.ExplicitlyDisabled}
	task.PostUnmarshalTask(&cfg, nil, nil, nil, nil)

	for _, container := range task.Containers { // remove v3 endpoint from each container because it's randomly generated
		removeV3EndpointConfig(container)
	}
	assert.Equal(t, expectedTask.Containers, task.Containers, "Containers should be equal")
	assert.Equal(t, expectedTask.Volumes, task.Volumes, "Volumes should be equal")
}

// removeV3EndpointConfig removes the v3 endpoint id and the injected env for a container
// so that checking all other fields can be easier
func removeV3EndpointConfig(container *apicontainer.Container) {
	container.SetV3EndpointID("")
	if container.Environment != nil {
		delete(container.Environment, apicontainer.MetadataURIEnvironmentVariableName)
	}
	if len(container.Environment) == 0 {
		container.Environment = nil
	}
}

func TestWindowsPlatformHostConfigOverride(t *testing.T) {
	// Testing Windows platform override for HostConfig.
	// Expects MemorySwappiness option to be set to -1

	task := &Task{}

	hostConfig := &dockercontainer.HostConfig{Resources: dockercontainer.Resources{CPUShares: int64(1 * cpuSharesPerCore)}}

	task.platformHostConfigOverride(hostConfig)
	assert.Equal(t, int64(1*cpuSharesPerCore*percentageFactor)/int64(cpuShareScaleFactor), hostConfig.CPUPercent)
	assert.Equal(t, int64(0), hostConfig.CPUShares)

	hostConfig = &dockercontainer.HostConfig{Resources: dockercontainer.Resources{CPUShares: 10}}
	task.platformHostConfigOverride(hostConfig)
	assert.Equal(t, int64(minimumCPUPercent), hostConfig.CPUPercent)
	assert.Empty(t, hostConfig.CPUShares)
}

func TestDockerHostConfigRawConfigMerging(t *testing.T) {
	// Use a struct that will marshal to the actual message we expect; not
	// dockercontainer.HostConfig which will include a lot of zero values.
	rawHostConfigInput := struct {
		Privileged  bool     `json:"Privileged,omitempty" yaml:"Privileged,omitempty"`
		SecurityOpt []string `json:"SecurityOpt,omitempty" yaml:"SecurityOpt,omitempty"`
	}{
		Privileged:  true,
		SecurityOpt: []string{"foo", "bar"},
	}

	rawHostConfig, err := json.Marshal(&rawHostConfigInput)
	if err != nil {
		t.Fatal(err)
	}

	testTask := &Task{
		Arn:     "arn:aws:ecs:us-east-1:012345678910:task/c09f0188-7f87-4b0f-bfc3-16296622b6fe",
		Family:  "myFamily",
		Version: "1",
		Containers: []*apicontainer.Container{
			{
				Name:        "c1",
				Image:       "image",
				CPU:         10,
				Memory:      100,
				VolumesFrom: []apicontainer.VolumeFrom{{SourceContainer: "c2"}},
				DockerConfig: apicontainer.DockerConfig{
					HostConfig: strptr(string(rawHostConfig)),
				},
			},
			{
				Name: "c2",
			},
		},
	}

	hostConfig, configErr := testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask), minDockerClientAPIVersion)
	assert.Nil(t, configErr)

	expected := dockercontainer.HostConfig{
		Resources: dockercontainer.Resources{
			// Convert MB to B and set Memory
			Memory:     apicontainer.DockerContainerMinimumMemoryInBytes,
			CPUPercent: minimumCPUPercent,
		},
		Privileged:  true,
		SecurityOpt: []string{"foo", "bar"},
		VolumesFrom: []string{"dockername-c2"},
	}

	assert.Nil(t, expected.MemorySwappiness, "Expected default memorySwappiness to be nil")
	assertSetStructFieldsEqual(t, expected, *hostConfig)
}

func TestCPUPercentBasedOnUnboundedEnabled(t *testing.T) {
	cpuShareScaleFactor := runtime.NumCPU() * cpuSharesPerCore
	testcases := []struct {
		cpu          int64
		cpuUnbounded bool
		cpuPercent   int64
	}{
		{
			cpu:          0,
			cpuUnbounded: true,
			cpuPercent:   0,
		},
		{
			cpu:          1,
			cpuUnbounded: true,
			cpuPercent:   1,
		},
		{
			cpu:          0,
			cpuUnbounded: false,
			cpuPercent:   1,
		},
		{
			cpu:          1,
			cpuUnbounded: false,
			cpuPercent:   1,
		},
		{
			cpu:          100,
			cpuUnbounded: true,
			cpuPercent:   100 * percentageFactor / int64(cpuShareScaleFactor),
		},
		{
			cpu:          100,
			cpuUnbounded: false,
			cpuPercent:   100 * percentageFactor / int64(cpuShareScaleFactor),
		},
	}
	for _, tc := range testcases {
		t.Run(fmt.Sprintf("container cpu-%d,cpu unbounded tasks enabled- %t,expected cpu percent-%d",
			tc.cpu, tc.cpuUnbounded, tc.cpuPercent), func(t *testing.T) {
			testTask := &Task{
				Containers: []*apicontainer.Container{
					{
						Name: "c1",
						CPU:  uint(tc.cpu),
					},
				},
				PlatformFields: PlatformFields{
					CpuUnbounded: tc.cpuUnbounded,
				},
			}

			hostconfig, err := testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask), minDockerClientAPIVersion)
			assert.Nil(t, err)
			assert.Empty(t, hostconfig.CPUShares)
			assert.Equal(t, tc.cpuPercent, hostconfig.CPUPercent)
		})
	}
}

func TestWindowsMemoryReservationOption(t *testing.T) {
	// Testing sending a task to windows overriding MemoryReservation value
	rawHostConfigInput := docker.HostConfig{
		MemoryReservation: nonZeroMemoryReservationValue,
	}

	rawHostConfig, err := json.Marshal(&rawHostConfigInput)
	if err != nil {
		t.Fatal(err)
	}

	testTask := &Task{
		Arn:     "arn:aws:ecs:us-east-1:012345678910:task/c09f0188-7f87-4b0f-bfc3-16296622b6fe",
		Family:  "myFamily",
		Version: "1",
		Containers: []*Container{
			{
				Name: "c1",
				DockerConfig: DockerConfig{
					HostConfig: strptr(string(rawHostConfig)),
				},
			},
		},
		PlatformFields: PlatformFields{
			RamUnbounded: false,
		},
	}

	// With RamUnbounded set to false, MemoryReservation is not overridden
	config, configErr := testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask), defaultDockerClientAPIVersion)
	if configErr != nil {
		t.Fatal(configErr)
	}
	assert.EqualValues(t, nonZeroMemoryReservationValue, config.MemoryReservation)

	// With RamUnbounded set to true, tasks with no memory hard limit will have their memory reservation set to zero
	testTask.PlatformFields.RamUnbounded = true
	config, configErr = testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask), defaultDockerClientAPIVersion)
	if configErr != nil {
		t.Fatal(configErr)
	}
	assert.EqualValues(t, expectedMemoryReservationValue, config.MemoryReservation)
}

func TestGetCanonicalPath(t *testing.T) {
	testcases := []struct {
		name           string
		path           string
		expectedResult string
	}{
		{
			name:           "folderPath",
			path:           `C:\myFile`,
			expectedResult: `c:\myfile`,
		},
		{
			name:           "drivePath",
			path:           `D:`,
			expectedResult: `d:`,
		},
		{
			name:           "pipePath",
			path:           `\\.\pipe\docker_engine`,
			expectedResult: `\\.\pipe\docker_engine`,
		},
	}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			result := getCanonicalPath(tc.path)
			assert.Equal(t, result, tc.expectedResult)
		})
	}
}
<commit_msg>Fix failing windows tests<commit_after>// +build windows,unit

// Copyright 2014-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 task

import (
	"encoding/json"
	"fmt"
	"runtime"
	"testing"

	"github.com/aws/amazon-ecs-agent/agent/acs/model/ecsacs"
	apicontainer "github.com/aws/amazon-ecs-agent/agent/api/container"
	apicontainerstatus "github.com/aws/amazon-ecs-agent/agent/api/container/status"
	apitaskstatus "github.com/aws/amazon-ecs-agent/agent/api/task/status"
	"github.com/aws/amazon-ecs-agent/agent/config"
	taskresourcevolume "github.com/aws/amazon-ecs-agent/agent/taskresource/volume"

	"github.com/aws/amazon-ecs-agent/agent/dockerclient"
	dockercontainer "github.com/docker/docker/api/types/container"
	"github.com/stretchr/testify/assert"
)

const (
	minDockerClientAPIVersion = dockerclient.Version_1_24

	nonZeroMemoryReservationValue  = 1
	expectedMemoryReservationValue = 0
)

func TestPostUnmarshalWindowsCanonicalPaths(t *testing.T) {
	// Testing type conversions, bleh. At least the type conversion itself
	// doesn't look this messy.
	taskFromAcs := ecsacs.Task{
		Arn:           strptr("myArn"),
		DesiredStatus: strptr("RUNNING"),
		Family:        strptr("myFamily"),
		Version:       strptr("1"),
		Containers: []*ecsacs.Container{
			{
				Name: strptr("myName"),
				MountPoints: []*ecsacs.MountPoint{
					{
						ContainerPath: strptr(`C:/Container/Path`),
						SourceVolume:  strptr("sourceVolume"),
					},
				},
			},
		},
		Volumes: []*ecsacs.Volume{
			{
				Name: strptr("sourceVolume"),
				Type: strptr("host"),
				Host: &ecsacs.HostVolumeProperties{
					SourcePath: strptr(`C:/Host/path`),
				},
			},
		},
	}
	expectedTask := &Task{
		Arn:                 "myArn",
		DesiredStatusUnsafe: apitaskstatus.TaskRunning,
		Family:              "myFamily",
		Version:             "1",
		Containers: []*apicontainer.Container{
			{
				Name: "myName",
				MountPoints: []apicontainer.MountPoint{
					{
						ContainerPath: `c:\container\path`,
						SourceVolume:  "sourceVolume",
					},
				},
				TransitionDependenciesMap: make(map[apicontainerstatus.ContainerStatus]apicontainer.TransitionDependencySet),
			},
		},
		Volumes: []TaskVolume{
			{
				Name: "sourceVolume",
				Type: "host",
				Volume: &taskresourcevolume.FSHostVolume{
					FSSourcePath: `c:\host\path`,
				},
			},
		},
		StartSequenceNumber: 42,
	}

	seqNum := int64(42)
	task, err := TaskFromACS(&taskFromAcs, &ecsacs.PayloadMessage{SeqNum: &seqNum})
	assert.Nil(t, err, "Should be able to handle acs task")
	cfg := config.Config{TaskCPUMemLimit: config.ExplicitlyDisabled}
	task.PostUnmarshalTask(&cfg, nil, nil, nil, nil)

	for _, container := range task.Containers { // remove v3 endpoint from each container because it's randomly generated
		removeV3EndpointConfig(container)
	}
	assert.Equal(t, expectedTask.Containers, task.Containers, "Containers should be equal")
	assert.Equal(t, expectedTask.Volumes, task.Volumes, "Volumes should be equal")
}

// removeV3EndpointConfig removes the v3 endpoint id and the injected env for a container
// so that checking all other fields can be easier
func removeV3EndpointConfig(container *apicontainer.Container) {
	container.SetV3EndpointID("")
	if container.Environment != nil {
		delete(container.Environment, apicontainer.MetadataURIEnvironmentVariableName)
	}
	if len(container.Environment) == 0 {
		container.Environment = nil
	}
}

func TestWindowsPlatformHostConfigOverride(t *testing.T) {
	// Testing Windows platform override for HostConfig.
	// Expects MemorySwappiness option to be set to -1

	task := &Task{}

	hostConfig := &dockercontainer.HostConfig{Resources: dockercontainer.Resources{CPUShares: int64(1 * cpuSharesPerCore)}}

	task.platformHostConfigOverride(hostConfig)
	assert.Equal(t, int64(1*cpuSharesPerCore*percentageFactor)/int64(cpuShareScaleFactor), hostConfig.CPUPercent)
	assert.Equal(t, int64(0), hostConfig.CPUShares)

	hostConfig = &dockercontainer.HostConfig{Resources: dockercontainer.Resources{CPUShares: 10}}
	task.platformHostConfigOverride(hostConfig)
	assert.Equal(t, int64(minimumCPUPercent), hostConfig.CPUPercent)
	assert.Empty(t, hostConfig.CPUShares)
}

func TestDockerHostConfigRawConfigMerging(t *testing.T) {
	// Use a struct that will marshal to the actual message we expect; not
	// dockercontainer.HostConfig which will include a lot of zero values.
	rawHostConfigInput := struct {
		Privileged  bool     `json:"Privileged,omitempty" yaml:"Privileged,omitempty"`
		SecurityOpt []string `json:"SecurityOpt,omitempty" yaml:"SecurityOpt,omitempty"`
	}{
		Privileged:  true,
		SecurityOpt: []string{"foo", "bar"},
	}

	rawHostConfig, err := json.Marshal(&rawHostConfigInput)
	if err != nil {
		t.Fatal(err)
	}

	testTask := &Task{
		Arn:     "arn:aws:ecs:us-east-1:012345678910:task/c09f0188-7f87-4b0f-bfc3-16296622b6fe",
		Family:  "myFamily",
		Version: "1",
		Containers: []*apicontainer.Container{
			{
				Name:        "c1",
				Image:       "image",
				CPU:         10,
				Memory:      100,
				VolumesFrom: []apicontainer.VolumeFrom{{SourceContainer: "c2"}},
				DockerConfig: apicontainer.DockerConfig{
					HostConfig: strptr(string(rawHostConfig)),
				},
			},
			{
				Name: "c2",
			},
		},
	}

	hostConfig, configErr := testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask), minDockerClientAPIVersion)
	assert.Nil(t, configErr)

	expected := dockercontainer.HostConfig{
		Resources: dockercontainer.Resources{
			// Convert MB to B and set Memory
			Memory:     apicontainer.DockerContainerMinimumMemoryInBytes,
			CPUPercent: minimumCPUPercent,
		},
		Privileged:  true,
		SecurityOpt: []string{"foo", "bar"},
		VolumesFrom: []string{"dockername-c2"},
	}

	assert.Nil(t, expected.MemorySwappiness, "Expected default memorySwappiness to be nil")
	assertSetStructFieldsEqual(t, expected, *hostConfig)
}

func TestCPUPercentBasedOnUnboundedEnabled(t *testing.T) {
	cpuShareScaleFactor := runtime.NumCPU() * cpuSharesPerCore
	testcases := []struct {
		cpu          int64
		cpuUnbounded bool
		cpuPercent   int64
	}{
		{
			cpu:          0,
			cpuUnbounded: true,
			cpuPercent:   0,
		},
		{
			cpu:          1,
			cpuUnbounded: true,
			cpuPercent:   1,
		},
		{
			cpu:          0,
			cpuUnbounded: false,
			cpuPercent:   1,
		},
		{
			cpu:          1,
			cpuUnbounded: false,
			cpuPercent:   1,
		},
		{
			cpu:          100,
			cpuUnbounded: true,
			cpuPercent:   100 * percentageFactor / int64(cpuShareScaleFactor),
		},
		{
			cpu:          100,
			cpuUnbounded: false,
			cpuPercent:   100 * percentageFactor / int64(cpuShareScaleFactor),
		},
	}
	for _, tc := range testcases {
		t.Run(fmt.Sprintf("container cpu-%d,cpu unbounded tasks enabled- %t,expected cpu percent-%d",
			tc.cpu, tc.cpuUnbounded, tc.cpuPercent), func(t *testing.T) {
			testTask := &Task{
				Containers: []*apicontainer.Container{
					{
						Name: "c1",
						CPU:  uint(tc.cpu),
					},
				},
				PlatformFields: PlatformFields{
					CpuUnbounded: tc.cpuUnbounded,
				},
			}

			hostconfig, err := testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask), minDockerClientAPIVersion)
			assert.Nil(t, err)
			assert.Empty(t, hostconfig.CPUShares)
			assert.Equal(t, tc.cpuPercent, hostconfig.CPUPercent)
		})
	}
}

func TestWindowsMemoryReservationOption(t *testing.T) {
	// Testing sending a task to windows overriding MemoryReservation value
	rawHostConfigInput := dockercontainer.HostConfig{
		Resources: dockercontainer.Resources{
			MemoryReservation: nonZeroMemoryReservationValue,
		},
	}

	rawHostConfig, err := json.Marshal(&rawHostConfigInput)
	if err != nil {
		t.Fatal(err)
	}

	testTask := &Task{
		Arn:     "arn:aws:ecs:us-east-1:012345678910:task/c09f0188-7f87-4b0f-bfc3-16296622b6fe",
		Family:  "myFamily",
		Version: "1",
		Containers: []*apicontainer.Container{
			{
				Name: "c1",
				DockerConfig: apicontainer.DockerConfig{
					HostConfig: strptr(string(rawHostConfig)),
				},
			},
		},
		PlatformFields: PlatformFields{
			RamUnbounded: false,
		},
	}

	// With RamUnbounded set to false, MemoryReservation is not overridden
	config, configErr := testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask), defaultDockerClientAPIVersion)
	if configErr != nil {
		t.Fatal(configErr)
	}
	assert.EqualValues(t, nonZeroMemoryReservationValue, config.MemoryReservation)

	// With RamUnbounded set to true, tasks with no memory hard limit will have their memory reservation set to zero
	testTask.PlatformFields.RamUnbounded = true
	config, configErr = testTask.DockerHostConfig(testTask.Containers[0], dockerMap(testTask), defaultDockerClientAPIVersion)
	if configErr != nil {
		t.Fatal(configErr)
	}
	assert.EqualValues(t, expectedMemoryReservationValue, config.MemoryReservation)
}

func TestGetCanonicalPath(t *testing.T) {
	testcases := []struct {
		name           string
		path           string
		expectedResult string
	}{
		{
			name:           "folderPath",
			path:           `C:\myFile`,
			expectedResult: `c:\myfile`,
		},
		{
			name:           "drivePath",
			path:           `D:`,
			expectedResult: `d:`,
		},
		{
			name:           "pipePath",
			path:           `\\.\pipe\docker_engine`,
			expectedResult: `\\.\pipe\docker_engine`,
		},
	}

	for _, tc := range testcases {
		t.Run(tc.name, func(t *testing.T) {
			result := getCanonicalPath(tc.path)
			assert.Equal(t, result, tc.expectedResult)
		})
	}
}
<|endoftext|>


<commit_before>package containerd

import (
	"fmt"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/api/types/task"
	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/log"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

type controller struct {
	task    *api.Task
	adapter *containerAdapter
	closed  chan struct{}
	err     error

	pulled     chan struct{} // closed after pull
	cancelPull func()        // cancels pull context if not nil
	pullErr    error         // pull error, protected by close of pulled
}

var _ exec.Controller = &controller{}

func newController(client *containerd.Client, task *api.Task, secrets exec.SecretGetter) (exec.Controller, error) {
	adapter, err := newContainerAdapter(client, task, secrets)
	if err != nil {
		return nil, err
	}

	return &controller{
		task:    task,
		adapter: adapter,
		closed:  make(chan struct{}),
	}, nil
}

// ContainerStatus returns the container-specific status for the task.
func (r *controller) ContainerStatus(ctx context.Context) (*api.ContainerStatus, error) {
	ctnr, err := r.adapter.inspect(ctx)
	if err != nil {
		if isUnknownContainer(err) {
			return nil, nil
		}

		return nil, err
	}

	status := &api.ContainerStatus{
		ContainerID: ctnr.ID,
		PID:         int32(ctnr.Pid),
	}

	switch ctnr.Status {
	case containerd.Stopped:
		err := r.adapter.shutdown(ctx)
		if ec, ok := err.(exec.ExitCoder); ok {
			status.ExitCode = int32(ec.ExitCode())
			return status, nil
		}
		return status, err
	default:
		return status, nil
	}

}

// Update takes a recent task update and applies it to the container.
func (r *controller) Update(ctx context.Context, t *api.Task) error {
	log.G(ctx).Warnf("task updates not yet supported")
	// TODO(stevvooe): While assignment of tasks is idempotent, we do allow
	// updates of metadata, such as labelling, as well as any other properties
	// that make sense.
	return nil
}

// Prepare creates a container and ensures the image is pulled.
//
// If the container has already be created, exec.ErrTaskPrepared is returned.
func (r *controller) Prepare(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	//// Make sure all the networks that the task needs are created.
	// TODO(ijc)
	//if err := r.adapter.createNetworks(ctx); err != nil {
	//	return err
	//}

	//// Make sure all the volumes that the task needs are created.
	// TODO(ijc)
	//if err := r.adapter.createVolumes(ctx); err != nil {
	//	return err
	//}

	if r.pulled == nil {
		// Launches a re-entrant pull operation associated with controller,
		// dissociating the context from the caller's context. Allows pull
		// operation to be re-entrant on calls to prepare, resuming from the
		// same point after cancellation.
		var pctx context.Context

		r.pulled = make(chan struct{})
		pctx, r.cancelPull = context.WithCancel(context.Background()) // TODO(stevvooe): Bind a context to the entire controller.

		go func() {
			defer close(r.pulled)
			r.pullErr = r.adapter.pullImage(pctx)
		}()
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-r.pulled:
		if r.pullErr != nil {
			// NOTE(stevvooe): We always try to pull the image to make sure we have
			// the most up to date version. This will return an error, but we only
			// log it. If the image truly doesn't exist, the create below will
			// error out.
			//
			// This gives us some nice behavior where we use up to date versions of
			// mutable tags, but will still run if the old image is available but a
			// registry is down.
			//
			// If you don't want this behavior, lock down your image to an
			// immutable tag or digest.
			log.G(ctx).WithError(r.pullErr).Error("pulling image failed")
		}
	}

	if err := r.adapter.prepare(ctx); err != nil {
		if isContainerCreateNameConflict(err) {
			if _, err := r.adapter.inspect(ctx); err != nil {
				return err
			}

			// container is already created. success!
			return exec.ErrTaskPrepared
		}

		return errors.Wrap(err, "create container failed")
	}

	return nil
}

// Start the container. An error will be returned if the container is already started.
func (r *controller) Start(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	ctnr, err := r.adapter.inspect(ctx)
	if err != nil {
		return err
	}

	// Detect whether the container has *ever* been started. If so, we don't
	// issue the start.
	//
	// TODO(stevvooe): This is very racy. While reading inspect, another could
	// start the process and we could end up starting it twice.
	if ctnr.Status != containerd.Created {
		return exec.ErrTaskStarted
	}

	if err := r.adapter.start(ctx); err != nil {
		return errors.Wrap(err, "starting container failed")
	}

	// TODO(ijc): Wait for HealthCheck to report OK.

	return nil
}

// Wait on the container to exit.
func (r *controller) Wait(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	// check the initial state and report that.
	ctnr, err := r.adapter.inspect(ctx)
	if err != nil {
		return errors.Wrap(err, "inspecting container failed")
	}

	// TODO(ijc) this shouldn't be needed here, figure out why
	// .shutdown/.remove are not being called otherwise.
	shutdownWithExitStatus := func() error {
		err := r.adapter.shutdown(ctx)
		if err2 := r.adapter.remove(ctx); err != nil {
			// Just log it, report the exit status
			log.G(ctx).WithError(err2).Info("remove after wait failed")
		}
		return err
	}
	switch ctnr.Status {
	case containerd.Stopped:
		return shutdownWithExitStatus()
	}

	// We do not disable FailFast for this initial call (like we
	// do on the retry below) since we are still halfway through
	// setting up the container and if containerd goes away half
	// way through we consider that a failure.
	eventq, closed, err := r.adapter.events(ctx)
	if err != nil {
		return err
	}

	for {
		select {
		case event := <-eventq:
			switch event.Type {
			case task.Event_EXIT:
				return shutdownWithExitStatus()
			case task.Event_OOM, task.Event_CREATE, task.Event_START, task.Event_EXEC_ADDED, task.Event_PAUSED:
				continue
			default:
				return errors.Errorf("Unknown event type %s\n", event.Type.String())
			}
		case <-closed:
			// restart!
			log.G(ctx).Debugf("Restarting event stream")
			// We disable FailFast for this call so that gRPC will keep
			// retrying while we wait for containerd to come back. Otherwise
			// a temporary glitch in the connection (e.g. a containerd restart)
			// will result in the task being declared dead even though it is
			// likely to be recoverable.
			eventq, closed, err = r.adapter.events(ctx, grpc.FailFast(false))
			if err != nil {
				return err
			}

			// recheck the container state, if this fails then we may have missed a
			ctnr, err := r.adapter.inspect(ctx)
			if err != nil {
				return errors.Wrap(err, "inspecting container on event restart failed")
			}
			switch ctnr.Status {
			case containerd.Stopped:
				return shutdownWithExitStatus()
			}

		case <-ctx.Done():
			return ctx.Err()
		case <-r.closed:
			return r.err
		}
	}
}

// Shutdown the container cleanly.
func (r *controller) Shutdown(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	if r.cancelPull != nil {
		r.cancelPull()
	}

	if err := r.adapter.shutdown(ctx); err != nil {
		if isUnknownContainer(err) {
			return nil
		}

		return err
	}

	return nil
}

// Terminate the container, with force.
func (r *controller) Terminate(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	if r.cancelPull != nil {
		r.cancelPull()
	}

	if err := r.adapter.terminate(ctx); err != nil {
		if isUnknownContainer(err) {
			return nil
		}

		return err
	}

	return nil
}

// Remove the container and its resources.
func (r *controller) Remove(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	if r.cancelPull != nil {
		r.cancelPull()
	}

	// It may be necessary to shut down the task before removing it.
	if err := r.Shutdown(ctx); err != nil {
		if isUnknownContainer(err) {
			return nil
		}

		// This may fail if the task was already shut down.
		log.G(ctx).WithError(err).Debug("shutdown failed on removal")
	}

	// Try removing networks referenced in this task in case this
	// task is the last one referencing it
	// TODO(ijc)
	//if err := r.adapter.removeNetworks(ctx); err != nil {
	//	if isUnknownContainer(err) {
	//		return nil
	//	}

	//	return err
	//}

	if err := r.adapter.remove(ctx); err != nil {
		if isUnknownContainer(err) {
			return nil
		}

		return err
	}

	return nil
}

// Close the controller and clean up any ephemeral resources.
func (r *controller) Close() error {
	select {
	case <-r.closed:
		return r.err
	default:
		if r.cancelPull != nil {
			r.cancelPull()
		}

		r.err = exec.ErrControllerClosed
		close(r.closed)
	}
	return nil
}

func (r *controller) checkClosed() error {
	select {
	case <-r.closed:
		return r.err
	default:
		return nil
	}
}

type exitError struct {
	code  uint32
	cause error
}

func (e *exitError) Error() string {
	if e.cause != nil {
		return fmt.Sprintf("task: non-zero exit (%v): %v", e.code, e.cause)
	}
	return fmt.Sprintf("task: non-zero exit (%v)", e.code)
}

func (e *exitError) ExitCode() int {
	return int(e.code)
}

func (e *exitError) Cause() error {
	return e.cause
}

func makeExitError(exitStatus uint32, reason string) error {
	if exitStatus != 0 {
		var cause error
		if reason != "" {
			cause = errors.New(reason)
		}

		return &exitError{
			code:  exitStatus,
			cause: cause,
		}
	}

	return nil
}
<commit_msg>containerd: Stop shutting down the container in ContainerStatus()<commit_after>package containerd

import (
	"fmt"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/api/types/task"
	"github.com/docker/swarmkit/agent/exec"
	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/log"
	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

type controller struct {
	task    *api.Task
	adapter *containerAdapter
	closed  chan struct{}
	err     error

	pulled     chan struct{} // closed after pull
	cancelPull func()        // cancels pull context if not nil
	pullErr    error         // pull error, protected by close of pulled
}

var _ exec.Controller = &controller{}

func newController(client *containerd.Client, task *api.Task, secrets exec.SecretGetter) (exec.Controller, error) {
	adapter, err := newContainerAdapter(client, task, secrets)
	if err != nil {
		return nil, err
	}

	return &controller{
		task:    task,
		adapter: adapter,
		closed:  make(chan struct{}),
	}, nil
}

// ContainerStatus returns the container-specific status for the task.
func (r *controller) ContainerStatus(ctx context.Context) (*api.ContainerStatus, error) {
	ctnr, err := r.adapter.inspect(ctx)
	if err != nil {
		if isUnknownContainer(err) {
			return nil, nil
		}

		return nil, err
	}

	status := &api.ContainerStatus{
		ContainerID: ctnr.ID,
		PID:         int32(ctnr.Pid),
	}

	if ec, ok := ctnr.ExitStatus.(exec.ExitCoder); ok {
		status.ExitCode = int32(ec.ExitCode())
	}

	return status, nil
}

// Update takes a recent task update and applies it to the container.
func (r *controller) Update(ctx context.Context, t *api.Task) error {
	log.G(ctx).Warnf("task updates not yet supported")
	// TODO(stevvooe): While assignment of tasks is idempotent, we do allow
	// updates of metadata, such as labelling, as well as any other properties
	// that make sense.
	return nil
}

// Prepare creates a container and ensures the image is pulled.
//
// If the container has already be created, exec.ErrTaskPrepared is returned.
func (r *controller) Prepare(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	//// Make sure all the networks that the task needs are created.
	// TODO(ijc)
	//if err := r.adapter.createNetworks(ctx); err != nil {
	//	return err
	//}

	//// Make sure all the volumes that the task needs are created.
	// TODO(ijc)
	//if err := r.adapter.createVolumes(ctx); err != nil {
	//	return err
	//}

	if r.pulled == nil {
		// Launches a re-entrant pull operation associated with controller,
		// dissociating the context from the caller's context. Allows pull
		// operation to be re-entrant on calls to prepare, resuming from the
		// same point after cancellation.
		var pctx context.Context

		r.pulled = make(chan struct{})
		pctx, r.cancelPull = context.WithCancel(context.Background()) // TODO(stevvooe): Bind a context to the entire controller.

		go func() {
			defer close(r.pulled)
			r.pullErr = r.adapter.pullImage(pctx)
		}()
	}

	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-r.pulled:
		if r.pullErr != nil {
			// NOTE(stevvooe): We always try to pull the image to make sure we have
			// the most up to date version. This will return an error, but we only
			// log it. If the image truly doesn't exist, the create below will
			// error out.
			//
			// This gives us some nice behavior where we use up to date versions of
			// mutable tags, but will still run if the old image is available but a
			// registry is down.
			//
			// If you don't want this behavior, lock down your image to an
			// immutable tag or digest.
			log.G(ctx).WithError(r.pullErr).Error("pulling image failed")
		}
	}

	if err := r.adapter.prepare(ctx); err != nil {
		if isContainerCreateNameConflict(err) {
			if _, err := r.adapter.inspect(ctx); err != nil {
				return err
			}

			// container is already created. success!
			return exec.ErrTaskPrepared
		}

		return errors.Wrap(err, "create container failed")
	}

	return nil
}

// Start the container. An error will be returned if the container is already started.
func (r *controller) Start(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	ctnr, err := r.adapter.inspect(ctx)
	if err != nil {
		return err
	}

	// Detect whether the container has *ever* been started. If so, we don't
	// issue the start.
	//
	// TODO(stevvooe): This is very racy. While reading inspect, another could
	// start the process and we could end up starting it twice.
	if ctnr.Status != containerd.Created {
		return exec.ErrTaskStarted
	}

	if err := r.adapter.start(ctx); err != nil {
		return errors.Wrap(err, "starting container failed")
	}

	// TODO(ijc): Wait for HealthCheck to report OK.

	return nil
}

// Wait on the container to exit.
func (r *controller) Wait(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	// check the initial state and report that.
	ctnr, err := r.adapter.inspect(ctx)
	if err != nil {
		return errors.Wrap(err, "inspecting container failed")
	}

	// TODO(ijc) this shouldn't be needed here, figure out why
	// .shutdown/.remove are not being called otherwise.
	shutdownWithExitStatus := func() error {
		err := r.adapter.shutdown(ctx)
		if err2 := r.adapter.remove(ctx); err != nil {
			// Just log it, report the exit status
			log.G(ctx).WithError(err2).Info("remove after wait failed")
		}
		return err
	}
	switch ctnr.Status {
	case containerd.Stopped:
		return shutdownWithExitStatus()
	}

	// We do not disable FailFast for this initial call (like we
	// do on the retry below) since we are still halfway through
	// setting up the container and if containerd goes away half
	// way through we consider that a failure.
	eventq, closed, err := r.adapter.events(ctx)
	if err != nil {
		return err
	}

	for {
		select {
		case event := <-eventq:
			switch event.Type {
			case task.Event_EXIT:
				return shutdownWithExitStatus()
			case task.Event_OOM, task.Event_CREATE, task.Event_START, task.Event_EXEC_ADDED, task.Event_PAUSED:
				continue
			default:
				return errors.Errorf("Unknown event type %s\n", event.Type.String())
			}
		case <-closed:
			// restart!
			log.G(ctx).Debugf("Restarting event stream")
			// We disable FailFast for this call so that gRPC will keep
			// retrying while we wait for containerd to come back. Otherwise
			// a temporary glitch in the connection (e.g. a containerd restart)
			// will result in the task being declared dead even though it is
			// likely to be recoverable.
			eventq, closed, err = r.adapter.events(ctx, grpc.FailFast(false))
			if err != nil {
				return err
			}

			// recheck the container state, if this fails then we may have missed a
			ctnr, err := r.adapter.inspect(ctx)
			if err != nil {
				return errors.Wrap(err, "inspecting container on event restart failed")
			}
			switch ctnr.Status {
			case containerd.Stopped:
				return shutdownWithExitStatus()
			}

		case <-ctx.Done():
			return ctx.Err()
		case <-r.closed:
			return r.err
		}
	}
}

// Shutdown the container cleanly.
func (r *controller) Shutdown(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	if r.cancelPull != nil {
		r.cancelPull()
	}

	if err := r.adapter.shutdown(ctx); err != nil {
		if isUnknownContainer(err) {
			return nil
		}

		return err
	}

	return nil
}

// Terminate the container, with force.
func (r *controller) Terminate(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	if r.cancelPull != nil {
		r.cancelPull()
	}

	if err := r.adapter.terminate(ctx); err != nil {
		if isUnknownContainer(err) {
			return nil
		}

		return err
	}

	return nil
}

// Remove the container and its resources.
func (r *controller) Remove(ctx context.Context) error {
	ctx = log.WithModule(ctx, "containerd")

	if err := r.checkClosed(); err != nil {
		return err
	}

	if r.cancelPull != nil {
		r.cancelPull()
	}

	// It may be necessary to shut down the task before removing it.
	if err := r.Shutdown(ctx); err != nil {
		if isUnknownContainer(err) {
			return nil
		}

		// This may fail if the task was already shut down.
		log.G(ctx).WithError(err).Debug("shutdown failed on removal")
	}

	// Try removing networks referenced in this task in case this
	// task is the last one referencing it
	// TODO(ijc)
	//if err := r.adapter.removeNetworks(ctx); err != nil {
	//	if isUnknownContainer(err) {
	//		return nil
	//	}

	//	return err
	//}

	if err := r.adapter.remove(ctx); err != nil {
		if isUnknownContainer(err) {
			return nil
		}

		return err
	}

	return nil
}

// Close the controller and clean up any ephemeral resources.
func (r *controller) Close() error {
	select {
	case <-r.closed:
		return r.err
	default:
		if r.cancelPull != nil {
			r.cancelPull()
		}

		r.err = exec.ErrControllerClosed
		close(r.closed)
	}
	return nil
}

func (r *controller) checkClosed() error {
	select {
	case <-r.closed:
		return r.err
	default:
		return nil
	}
}

type exitError struct {
	code  uint32
	cause error
}

func (e *exitError) Error() string {
	if e.cause != nil {
		return fmt.Sprintf("task: non-zero exit (%v): %v", e.code, e.cause)
	}
	return fmt.Sprintf("task: non-zero exit (%v)", e.code)
}

func (e *exitError) ExitCode() int {
	return int(e.code)
}

func (e *exitError) Cause() error {
	return e.cause
}

func makeExitError(exitStatus uint32, reason string) error {
	if exitStatus != 0 {
		var cause error
		if reason != "" {
			cause = errors.New(reason)
		}

		return &exitError{
			code:  exitStatus,
			cause: cause,
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"io"
	"log"
	"net"
	"os/exec"
	"path"
	"strings"
	"time"

	"github.com/rneugeba/virtsock/go/vsock"
	"github.com/rneugeba/virtsock/go/hvsock"
)

func run(timeout time.Duration, w *tar.Writer, command string, args ...string) {
	log.Printf("Running %s", command)
	c := exec.Command(command, args...)
	stdoutPipe, err := c.StdoutPipe()
	if err != nil {
		log.Fatalf("Failed to create stdout pipe: %#v", err)
	}
	stderrPipe, err := c.StderrPipe()
	if err != nil {
		log.Fatalf("Failed to create stderr pipe: %#v", err)
	}
	var stdoutBuffer bytes.Buffer
	var stderrBuffer bytes.Buffer
	done := make(chan int)
	go func() {
		io.Copy(&stdoutBuffer, stdoutPipe)
		done <- 0
	}()
	go func() {
		io.Copy(&stderrBuffer, stderrPipe)
		done <- 0
	}()
	var timer *time.Timer
	timer = time.AfterFunc(timeout, func() {
		timer.Stop()
		if c.Process != nil {
			c.Process.Kill()
		}
	})
	_ = c.Run()
	<-done
	<-done
	timer.Stop()

	name := strings.Join(append([]string{path.Base(command)}, args...), " ")

	hdr := &tar.Header{
		Name: name + ".stdout",
		Mode: 0644,
		Size: int64(stdoutBuffer.Len()),
	}
	if err = w.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
	if _, err = w.Write(stdoutBuffer.Bytes()); err != nil {
		log.Fatalln(err)
	}
	hdr = &tar.Header{
		Name: name + ".stderr",
		Mode: 0644,
		Size: int64(stderrBuffer.Len()),
	}
	if err = w.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
	if _, err = w.Write(stderrBuffer.Bytes()); err != nil {
		log.Fatalln(err)
	}

}

func capture(w *tar.Writer) {
	t := 2 * time.Second

	run(t, w, "/bin/date")
	run(t, w, "/bin/uname", "-a")
	run(t, w, "/bin/ps", "uax")
	run(t, w, "/bin/netstat", "-tulpn")
	run(t, w, "/sbin/iptables-save")
	run(t, w, "/sbin/ifconfig", "-a")
	run(t, w, "/sbin/route", "-n")
	run(t, w, "/usr/sbin/brctl", "show")
	run(t, w, "/bin/dmesg")
	run(t, w, "/usr/bin/docker", "ps")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/docker.log")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/messages")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/proxy-vsockd.log")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/vsudd.log")
	run(t, w, "/bin/mount")
	run(t, w, "/bin/df")
	run(t, w, "/bin/ls", "-l", "/var")
	run(t, w, "/bin/ls", "-l", "/var/lib")
	run(t, w, "/bin/ls", "-l", "/var/lib/docker")
	run(t, w, "/usr/bin/diagnostics")
	run(t, w, "/bin/ping", "-w", "5", "8.8.8.8")
	run(t, w, "/bin/cp", "/etc/resolv.conf", ".")
	run(t, w, "/usr/bin/dig", "docker.com")
	run(t, w, "/usr/bin/wget", "-O", "-", "http://www.docker.com/")
}

func main() {
	listeners := make([]net.Listener, 0)

	ip, err := net.Listen("tcp", ":62374")
	if err != nil {
		log.Printf("Failed to bind to TCP port 62374: %#v", err)
	} else {
		listeners = append(listeners, ip)
	}
	vsock, err := vsock.Listen(uint(62374))
	if err != nil {
		log.Printf("Failed to bind to vsock port 62374: %#v", err)
	} else {
		listeners = append(listeners, vsock)
	}
	svcid, _ := hvsock.GuidFromString("445BA2CB-E69B-4912-8B42-D7F494D007EA")
	hvsock, err := hvsock.Listen(hvsock.HypervAddr{VmId: hvsock.GUID_WILDCARD, ServiceId: svcid})
	if err != nil {
		log.Printf("Failed to bind to hvsock port: %#v", err)
	} else {
		listeners = append(listeners, hvsock)
	}

	for _, l := range listeners {
		go func(l net.Listener) {
			for {
				conn, err := l.Accept()
				if err != nil {
					log.Printf("Error accepting connection: %#v", err)
					return // no more listening
				}
				go func(conn net.Conn) {
					w := tar.NewWriter(conn)
					capture(w)
					if err := w.Close(); err != nil {
						log.Println(err)
					}
					conn.Close()
				}(conn)
			}
		}(l)
	}
	forever := make(chan int)
	<-forever
}
<commit_msg>diagnostics: dump the master branch of the database<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"

	"github.com/rneugeba/virtsock/go/vsock"
	"github.com/rneugeba/virtsock/go/hvsock"
)

func run(timeout time.Duration, w *tar.Writer, command string, args ...string) {
	log.Printf("Running %s", command)
	c := exec.Command(command, args...)
	stdoutPipe, err := c.StdoutPipe()
	if err != nil {
		log.Fatalf("Failed to create stdout pipe: %#v", err)
	}
	stderrPipe, err := c.StderrPipe()
	if err != nil {
		log.Fatalf("Failed to create stderr pipe: %#v", err)
	}
	var stdoutBuffer bytes.Buffer
	var stderrBuffer bytes.Buffer
	done := make(chan int)
	go func() {
		io.Copy(&stdoutBuffer, stdoutPipe)
		done <- 0
	}()
	go func() {
		io.Copy(&stderrBuffer, stderrPipe)
		done <- 0
	}()
	var timer *time.Timer
	timer = time.AfterFunc(timeout, func() {
		timer.Stop()
		if c.Process != nil {
			c.Process.Kill()
		}
	})
	_ = c.Run()
	<-done
	<-done
	timer.Stop()

	name := strings.Join(append([]string{path.Base(command)}, args...), " ")

	hdr := &tar.Header{
		Name: name + ".stdout",
		Mode: 0644,
		Size: int64(stdoutBuffer.Len()),
	}
	if err = w.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
	if _, err = w.Write(stdoutBuffer.Bytes()); err != nil {
		log.Fatalln(err)
	}
	hdr = &tar.Header{
		Name: name + ".stderr",
		Mode: 0644,
		Size: int64(stderrBuffer.Len()),
	}
	if err = w.WriteHeader(hdr); err != nil {
		log.Fatalln(err)
	}
	if _, err = w.Write(stderrBuffer.Bytes()); err != nil {
		log.Fatalln(err)
	}

}

func capture(w *tar.Writer) {
	t := 2 * time.Second

	run(t, w, "/bin/date")
	run(t, w, "/bin/uname", "-a")
	run(t, w, "/bin/ps", "uax")
	run(t, w, "/bin/netstat", "-tulpn")
	run(t, w, "/sbin/iptables-save")
	run(t, w, "/sbin/ifconfig", "-a")
	run(t, w, "/sbin/route", "-n")
	run(t, w, "/usr/sbin/brctl", "show")
	run(t, w, "/bin/dmesg")
	run(t, w, "/usr/bin/docker", "ps")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/docker.log")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/messages")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/proxy-vsockd.log")
	run(t, w, "/usr/bin/tail", "-100", "/var/log/vsudd.log")
	run(t, w, "/bin/mount")
	run(t, w, "/bin/df")
	run(t, w, "/bin/ls", "-l", "/var")
	run(t, w, "/bin/ls", "-l", "/var/lib")
	run(t, w, "/bin/ls", "-l", "/var/lib/docker")
	run(t, w, "/usr/bin/diagnostics")
	run(t, w, "/bin/ping", "-w", "5", "8.8.8.8")
	run(t, w, "/bin/cp", "/etc/resolv.conf", ".")
	run(t, w, "/usr/bin/dig", "docker.com")
	run(t, w, "/usr/bin/wget", "-O", "-", "http://www.docker.com/")

	// Dump the database
	dbBase := "/Database/branch/master/ro"
	filepath.Walk(dbBase, func(path string, f os.FileInfo, err error) error {
		if f.Mode().IsRegular() {
			run(t, w, "/bin/cat", path)
		}
		return nil
	})
}

func main() {
	listeners := make([]net.Listener, 0)

	ip, err := net.Listen("tcp", ":62374")
	if err != nil {
		log.Printf("Failed to bind to TCP port 62374: %#v", err)
	} else {
		listeners = append(listeners, ip)
	}
	vsock, err := vsock.Listen(uint(62374))
	if err != nil {
		log.Printf("Failed to bind to vsock port 62374: %#v", err)
	} else {
		listeners = append(listeners, vsock)
	}
	svcid, _ := hvsock.GuidFromString("445BA2CB-E69B-4912-8B42-D7F494D007EA")
	hvsock, err := hvsock.Listen(hvsock.HypervAddr{VmId: hvsock.GUID_WILDCARD, ServiceId: svcid})
	if err != nil {
		log.Printf("Failed to bind to hvsock port: %#v", err)
	} else {
		listeners = append(listeners, hvsock)
	}

	for _, l := range listeners {
		go func(l net.Listener) {
			for {
				conn, err := l.Accept()
				if err != nil {
					log.Printf("Error accepting connection: %#v", err)
					return // no more listening
				}
				go func(conn net.Conn) {
					w := tar.NewWriter(conn)
					capture(w)
					if err := w.Close(); err != nil {
						log.Println(err)
					}
					conn.Close()
				}(conn)
			}
		}(l)
	}
	forever := make(chan int)
	<-forever
}
<|endoftext|>


<commit_before>// Copyright 2015 Jacek Masiulaniec. All rights reserved.
// Use of this source code is governed by a free license that can be
// found in the LICENSE file.

// Package riemann implements a Riemann client.
package riemann

import (
	"errors"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"time"

	riemann "github.com/masiulaniec/riemann/proto"

	"github.com/golang/protobuf/proto"
)

const (
	frameMaxEvents  = 100    // how frequently to write?
	poolDefaultSize = 1      // how many parallel conns?
	poolMaxQueue    = 100000 // sender-side safety buffer
)

// errClosed is used internally to signal connection close.
var errClosed = errors.New("cclient closed")

// Event represents a Riemann event.
type Event struct {
	Host    string  // sending host
	Service string  // the service the event pertains to
	IsFloat bool    // value type
	Float   float64 // value if IsFloat
	Int     int64   // value if !IsFloat

	// Optional
	Description string   // arbitrary text
	Time        int64    // creation time (unix epoch)
	TTL         float32  // time to live (seconds)
	State       string   // associated health state
	Tags        []string // list of string labels
	Attributes  []string // list of key, value pairs
}

// ClientConfig can be used to override Client settings.
type ClientConfig struct {
	// PoolSize defines how many TCP connections to create.
	// Default: 1
	PoolSize int
}

// Client represents a Riemann client.
type Client struct {
	addr string
	pool chan *Event
}

// NewClient returns a client for Riemann server listening at
// the given address. If nil config is provided, default settings
// are used.
func NewClient(addr string, config *ClientConfig) *Client {
	if config == nil {
		config = &ClientConfig{}
	}
	if config.PoolSize == 0 {
		config.PoolSize = poolDefaultSize
	}
	client := &Client{
		addr: addr,
		pool: make(chan *Event, poolMaxQueue),
	}
	for i := 0; i < config.PoolSize; i++ {
		go client.conn()
	}
	return client
}

// Close terminates the client by closing its connection pool.
func (c *Client) Close() error {
	close(c.pool)
	return nil
}

// conn keeps alive a connection to the server.
func (c *Client) conn() {
	for {
		err := c.conn1()
		if err == errClosed {
			return
		}
		log.Printf("riemann: server %s: %v", c.addr, err)
		time.Sleep(5 * time.Second)
	}
}

func (c *Client) conn1() error {
	addr := c.addr
	if !strings.HasSuffix(addr, ":5555") {
		addr += ":5555"
	}
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return err
	}
	defer conn.Close()
	go c.readLoop(conn)
	return c.writeLoop(conn)
}

// writeLoop sends requests to the server.
func (c *Client) writeLoop(w io.Writer) error {
	frame := new(frame)
	frame.Reset()
	timer := time.NewTimer(1 * time.Second)
	defer timer.Stop()
	for {
		select {
		case event, ok := <-c.pool:
			if !ok {
				return errClosed
			}
			frame.append(event)
			if frame.Len == frameMaxEvents {
				if _, err := w.Write(frame.Bytes()); err != nil {
					return err
				}
				frame.Reset()
				timer.Reset()
			}
		case <-timer.C:
			if frame.Len > 0 {
				if _, err := w.Write(frame.Bytes()); err != nil {
					return err
				}
				frame.Reset()
			}
			timer.Reset()
		}
	}
}

// readLoop processes server responses.
func (c *Client) readLoop(r io.Reader) {
	// Just discard the acks. We know they are acks and not errors
	// because invalid messages are never transmitted. Instead,
	// bad data is detected eagerly in (*Client).Send, which helps
	// identify faulty call sites.
	io.Copy(ioutil.Discard, r)
}

// Send tries to delivier the given event. Delivery is not guaranteed.
// Send will panic if the given event is invalid.
func (c *Client) Send(event *Event) {
	if err := event.validate(); err != nil {
		log.Panic(err)
	}
	select {
	case c.pool <- event:
		// ok
	default:
		// Data loss.
	}
}

func (e *Event) validate() error {
	if e.Host == "" {
		hostname, err := os.Hostname()
		if err != nil {
			return err
		}
		e.Host = hostname
	}
	if e.Time == 0 {
		now := time.Now().Unix()
		e.Time = now
	}
	if e.Service == "" {
		return errors.New("undefined service")
	}
	if e.IsFloat && e.Int != 0 {
		return errors.New("float event with non-float value")
	}
	return nil
}

// frame represents a Riemann frame.
type frame struct {
	proto.Buffer               // holds Riemann header and protobuf body
	Len          int           // number of events buffered
	event        riemann.Event // helps avoid an allocation
	scratch      [512]byte     // helps avoid an allocation
}

// Reset prepares a new frame.
func (fr *frame) Reset() {
	fr.Buffer.Reset()
	// Reserve space for Riemann's frame header.
	fr.EncodeFixed32(0)
	fr.Len = 0
}

// Bytes returns the frame in a ready-to-write form.
func (fr *frame) Bytes() []byte {
	var (
		header = fr.Buffer.Bytes()[:4]
		body   = fr.Buffer.Bytes()[4:]
	)
	bodyLen := uint32(len(body))
	header[0] = byte(bodyLen >> 24)
	header[1] = byte(bodyLen >> 16)
	header[2] = byte(bodyLen >> 8)
	header[3] = byte(bodyLen >> 0)
	return fr.Buffer.Bytes()
}

// append encodes the given event in the frame.
func (fr *frame) append(event *Event) {
	e := &fr.event

	// Prepare the event for marshaling.
	e.Time = &event.Time
	if event.State != "" {
		e.State = &event.State
	} else {
		e.State = nil
	}
	e.Service = &event.Service
	e.Host = &event.Host
	if event.Description != "" {
		e.Description = &event.Description
	} else {
		e.Description = nil
	}
	e.Tags = event.Tags
	if event.TTL != 0 {
		e.Ttl = &event.TTL
	} else {
		e.Ttl = nil
	}
	if event.IsFloat {
		e.MetricSint64 = nil
		e.MetricD = &event.Float
	} else {
		e.MetricSint64 = &event.Int
		e.MetricD = nil
	}
	e.Attributes = e.Attributes[:0]
	for i := 0; i < len(event.Attributes); i += 2 {
		attr := &riemann.Attribute{
			Key:   &event.Attributes[i],
			Value: &event.Attributes[i+1],
		}
		e.Attributes = append(e.Attributes, attr)
	}

	// Marshal it.
	ebuf := proto.NewBuffer(fr.scratch[:0])
	if err := ebuf.Marshal(e); err != nil {
		panic(err)
	}
	if err := fr.EncodeVarint(6<<3 | proto.WireBytes); err != nil {
		panic(err)
	}
	if err := fr.EncodeRawBytes(ebuf.Bytes()); err != nil {
		panic(err)
	}
	fr.Len++
}
<commit_msg>Unbreak build<commit_after>// Copyright 2015 Jacek Masiulaniec. All rights reserved.
// Use of this source code is governed by a free license that can be
// found in the LICENSE file.

// Package riemann implements a Riemann client.
package riemann

import (
	"errors"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"time"

	riemann "github.com/masiulaniec/riemann/proto"

	"github.com/golang/protobuf/proto"
)

// Client flushes every 100 events or every 1 second, whichever occurs first
const (
	frameMaxEvents = 100
	frameMaxDelay  = 1 * time.Second
)

const (
	poolDefaultSize = 1      // how many parallel conns?
	poolMaxQueue    = 100000 // sender-side safety buffer
)

// errClosed is used internally to signal connection close.
var errClosed = errors.New("cclient closed")

// Event represents a Riemann event.
type Event struct {
	Host    string  // sending host
	Service string  // the service the event pertains to
	IsFloat bool    // value type
	Float   float64 // value if IsFloat
	Int     int64   // value if !IsFloat

	// Optional
	Description string   // arbitrary text
	Time        int64    // creation time (unix epoch)
	TTL         float32  // time to live (seconds)
	State       string   // associated health state
	Tags        []string // list of string labels
	Attributes  []string // list of key, value pairs
}

// ClientConfig can be used to override Client settings.
type ClientConfig struct {
	// PoolSize defines how many TCP connections to create.
	// Default: 1
	PoolSize int
}

// Client represents a Riemann client.
type Client struct {
	addr string
	pool chan *Event
}

// NewClient returns a client for Riemann server listening at
// the given address. If nil config is provided, default settings
// are used.
func NewClient(addr string, config *ClientConfig) *Client {
	if config == nil {
		config = &ClientConfig{}
	}
	if config.PoolSize == 0 {
		config.PoolSize = poolDefaultSize
	}
	client := &Client{
		addr: addr,
		pool: make(chan *Event, poolMaxQueue),
	}
	for i := 0; i < config.PoolSize; i++ {
		go client.conn()
	}
	return client
}

// Close terminates the client by closing its connection pool.
func (c *Client) Close() error {
	close(c.pool)
	return nil
}

// conn keeps alive a connection to the server.
func (c *Client) conn() {
	for {
		err := c.conn1()
		if err == errClosed {
			return
		}
		log.Printf("riemann: server %s: %v", c.addr, err)
		time.Sleep(5 * time.Second)
	}
}

func (c *Client) conn1() error {
	addr := c.addr
	if !strings.HasSuffix(addr, ":5555") {
		addr += ":5555"
	}
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return err
	}
	defer conn.Close()
	go c.readLoop(conn)
	return c.writeLoop(conn)
}

// writeLoop sends requests to the server.
func (c *Client) writeLoop(w io.Writer) error {
	frame := new(frame)
	frame.Reset()
	timer := time.NewTimer(frameMaxDelay)
	defer timer.Stop()
	for {
		select {
		case event, ok := <-c.pool:
			if !ok {
				return errClosed
			}
			frame.append(event)
			if frame.Len == frameMaxEvents {
				if _, err := w.Write(frame.Bytes()); err != nil {
					return err
				}
				frame.Reset()
				timer.Reset(frameMaxDelay)
			}
		case <-timer.C:
			if frame.Len > 0 {
				if _, err := w.Write(frame.Bytes()); err != nil {
					return err
				}
				frame.Reset()
			}
			timer.Reset(frameMaxDelay)
		}
	}
}

// readLoop processes server responses.
func (c *Client) readLoop(r io.Reader) {
	// Just discard the acks. We know they are acks and not errors
	// because invalid messages are never transmitted. Instead,
	// bad data is detected eagerly in (*Client).Send, which helps
	// identify faulty call sites.
	io.Copy(ioutil.Discard, r)
}

// Send tries to delivier the given event. Delivery is not guaranteed.
// Send will panic if the given event is invalid.
func (c *Client) Send(event *Event) {
	if err := event.validate(); err != nil {
		log.Panic(err)
	}
	select {
	case c.pool <- event:
		// ok
	default:
		// Data loss.
	}
}

func (e *Event) validate() error {
	if e.Host == "" {
		hostname, err := os.Hostname()
		if err != nil {
			return err
		}
		e.Host = hostname
	}
	if e.Time == 0 {
		now := time.Now().Unix()
		e.Time = now
	}
	if e.Service == "" {
		return errors.New("undefined service")
	}
	if e.IsFloat && e.Int != 0 {
		return errors.New("float event with non-float value")
	}
	return nil
}

// frame represents a Riemann frame.
type frame struct {
	proto.Buffer               // holds Riemann header and protobuf body
	Len          int           // number of events buffered
	event        riemann.Event // helps avoid an allocation
	scratch      [512]byte     // helps avoid an allocation
}

// Reset prepares a new frame.
func (fr *frame) Reset() {
	fr.Buffer.Reset()
	// Reserve space for Riemann's frame header.
	fr.EncodeFixed32(0)
	fr.Len = 0
}

// Bytes returns the frame in a ready-to-write form.
func (fr *frame) Bytes() []byte {
	var (
		header = fr.Buffer.Bytes()[:4]
		body   = fr.Buffer.Bytes()[4:]
	)
	bodyLen := uint32(len(body))
	header[0] = byte(bodyLen >> 24)
	header[1] = byte(bodyLen >> 16)
	header[2] = byte(bodyLen >> 8)
	header[3] = byte(bodyLen >> 0)
	return fr.Buffer.Bytes()
}

// append encodes the given event in the frame.
func (fr *frame) append(event *Event) {
	e := &fr.event

	// Prepare the event for marshaling.
	e.Time = &event.Time
	if event.State != "" {
		e.State = &event.State
	} else {
		e.State = nil
	}
	e.Service = &event.Service
	e.Host = &event.Host
	if event.Description != "" {
		e.Description = &event.Description
	} else {
		e.Description = nil
	}
	e.Tags = event.Tags
	if event.TTL != 0 {
		e.Ttl = &event.TTL
	} else {
		e.Ttl = nil
	}
	if event.IsFloat {
		e.MetricSint64 = nil
		e.MetricD = &event.Float
	} else {
		e.MetricSint64 = &event.Int
		e.MetricD = nil
	}
	e.Attributes = e.Attributes[:0]
	for i := 0; i < len(event.Attributes); i += 2 {
		attr := &riemann.Attribute{
			Key:   &event.Attributes[i],
			Value: &event.Attributes[i+1],
		}
		e.Attributes = append(e.Attributes, attr)
	}

	// Marshal it.
	ebuf := proto.NewBuffer(fr.scratch[:0])
	if err := ebuf.Marshal(e); err != nil {
		panic(err)
	}
	if err := fr.EncodeVarint(6<<3 | proto.WireBytes); err != nil {
		panic(err)
	}
	if err := fr.EncodeRawBytes(ebuf.Bytes()); err != nil {
		panic(err)
	}
	fr.Len++
}
<|endoftext|>