<commit_before>package main

import (
	"../libgolb"
	"github.com/docopt/docopt-go"
	"github.com/gorilla/mux"
	"io"
	"net/http"
	"os"
	"strconv"
	"strings"
)

func getServer() (server string) {
	server = libgolb.Conf.BackServers[libgolb.RoundRobin]
	libgolb.RoundRobin++
	if libgolb.RoundRobin >= libgolb.NumberBack {
		libgolb.RoundRobin = 0
	}
	return
}

func golbGet(w http.ResponseWriter, req *http.Request) {
	var secondResp *http.Response
	var errsp error

	serv := strings.Split(req.RemoteAddr, ":") // extract just IP without port
	libgolb.Log("misc", "Access From :"+serv[0])
	server, errGS := libgolb.RadixGetString(libgolb.LBClient, serv[0])
	if errGS != nil {
		server = getServer()
	}
	limit := 0
	for limit < libgolb.NumberBack {
		resp, _ := http.NewRequest(req.Method, "http://"+server+req.RequestURI, nil)
		for k, v := range req.Header {
			resp.Header[k] = v
		}
		resp.Header.Set("X-Forwarded-For", req.RemoteAddr)
		secondResp, errsp = http.DefaultClient.Do(resp)
		defer secondResp.Body.Close()
		if errsp != nil {
			libgolb.Log("error", "Connection with the HTTP file server failed: "+errsp.Error())
			server = getServer()
			limit++
		} else {
			break
		}
	}
	if limit >= libgolb.NumberBack {
		libgolb.HttpResponse(w, 500, "Internal server error\n")
		libgolb.Log("error", "No Backend Server avalaible")
		return
	}
	for k, v := range secondResp.Header {
		w.Header().Add(k, strings.Join(v, ""))
	}
	w.Header().Set("Status", "200")
	io.Copy(w, secondResp.Body)
	libgolb.RadixSet(libgolb.LBClient, serv[0], server)
	libgolb.Log("ok", "Answer From :"+serv[0])
	//TTL
	libgolb.RadixExpire(libgolb.LBClient, serv[0], strconv.Itoa(libgolb.Conf.TTL))
	libgolb.LogW3C(w, req, false)
}

func parseArgument(configuration string) {

	// Load configuration
	libgolb.ConfLoad(configuration)
	// Check Redis connection
	redis := libgolb.ConnectToRedis()
	if redis != nil {
		libgolb.Log("error", "Redis connection failed")
		os.Exit(1)
	}

	// Router
	rtr := mux.NewRouter()
	rtr.HandleFunc("/", golbGet).Methods("GET")
	http.Handle("/", rtr)

	// Listening
	libgolb.Log("ok", "Listening on "+libgolb.Conf.Server.Hostname+":"+libgolb.Conf.Server.Port)
	err := http.ListenAndServe(libgolb.Conf.Server.Hostname+":"+libgolb.Conf.Server.Port, nil)
	libgolb.ErrCatcher("ListenAndServe: ", err)
}

func main() {
	usage := `Golb.

Usage:
  golb <configuration>
  golb -h | --help
  golb --version

Options:
  -h --help     Show this screen.
  --version     Show version.`

	arguments, _ := docopt.Parse(usage, nil, true, "GoLB 0.1", false)
	parseArgument(arguments["<configuration>"].(string))
}
<commit_msg>Fix a close()<commit_after>package main

import (
	"../libgolb"
	"github.com/docopt/docopt-go"
	"github.com/gorilla/mux"
	"io"
	"net/http"
	"os"
	"strconv"
	"strings"
)

func getServer() (server string) {
	server = libgolb.Conf.BackServers[libgolb.RoundRobin]
	libgolb.RoundRobin++
	if libgolb.RoundRobin >= libgolb.NumberBack {
		libgolb.RoundRobin = 0
	}
	return
}

func golbGet(w http.ResponseWriter, req *http.Request) {
	var secondResp *http.Response
	var errsp error

	serv := strings.Split(req.RemoteAddr, ":") // extract just IP without port
	libgolb.Log("misc", "Access From :"+serv[0])
	server, errGS := libgolb.RadixGetString(libgolb.LBClient, serv[0])
	if errGS != nil {
		server = getServer()
	}
	limit := 0
	for limit < libgolb.NumberBack {
		resp, _ := http.NewRequest(req.Method, "http://"+server+req.RequestURI, nil)
		for k, v := range req.Header {
			resp.Header[k] = v
		}
		resp.Header.Set("X-Forwarded-For", req.RemoteAddr)
		secondResp, errsp = http.DefaultClient.Do(resp)
		if errsp != nil {
			libgolb.Log("error", "Connection with the HTTP file server failed: "+errsp.Error())
			server = getServer()
			limit++
		} else {
			break
		}
	}
	if limit >= libgolb.NumberBack {
		libgolb.HttpResponse(w, 500, "Internal server error\n")
		libgolb.Log("error", "No Backend Server avalaible")
		return
	}
	for k, v := range secondResp.Header {
		w.Header().Add(k, strings.Join(v, ""))
	}
	defer secondResp.Body.Close()
	w.Header().Set("Status", "200")
	io.Copy(w, secondResp.Body)
	libgolb.RadixSet(libgolb.LBClient, serv[0], server)
	libgolb.Log("ok", "Answer From :"+serv[0])
	//TTL
	libgolb.RadixExpire(libgolb.LBClient, serv[0], strconv.Itoa(libgolb.Conf.TTL))
	libgolb.LogW3C(w, req, false)
}

func parseArgument(configuration string) {

	// Load configuration
	libgolb.ConfLoad(configuration)
	// Check Redis connection
	redis := libgolb.ConnectToRedis()
	if redis != nil {
		libgolb.Log("error", "Redis connection failed")
		os.Exit(1)
	}

	// Router
	rtr := mux.NewRouter()
	rtr.HandleFunc("/", golbGet).Methods("GET")
	http.Handle("/", rtr)

	// Listening
	libgolb.Log("ok", "Listening on "+libgolb.Conf.Server.Hostname+":"+libgolb.Conf.Server.Port)
	err := http.ListenAndServe(libgolb.Conf.Server.Hostname+":"+libgolb.Conf.Server.Port, nil)
	libgolb.ErrCatcher("ListenAndServe: ", err)
}

func main() {
	usage := `Golb.

Usage:
  golb <configuration>
  golb -h | --help
  golb --version

Options:
  -h --help     Show this screen.
  --version     Show version.`

	arguments, _ := docopt.Parse(usage, nil, true, "GoLB 0.1", false)
	parseArgument(arguments["<configuration>"].(string))
}
<|endoftext|>


<commit_before>// gopp provides an easy to use (but somewhat illogical) preprocessor for Go.
package gopp

import (
	"bytes"
	"fmt"
	"go/parser"
	"go/printer"
	"go/scanner"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"strings"
)

const version = "0.3"

// Token provides a simple struct face to the `pos, tok, lit` returned by
// Go's go/scanner package. This is passed around by gopp internally in a chan.
type Token struct {
	Position token.Pos
	Token    token.Token
	String   string
}

type gopp struct {
	defines                 map[string]interface{}
	output                  chan Token
	StripComments, ignoring bool
	Prefix                  string
}

// DefineValue takes a key and a value, and defines them as a macro to use when
// preprocessing a Go file.
func (g *gopp) DefineValue(key string, value interface{}) {
	g.defines[key] = value
}

// Define takes a key, and sets it to the sane default of 1, which can then be
// used when preprocessing a Go file.
func (g *gopp) Define(key string) {
	g.DefineValue(key, 1)
}

// Undefine will remove a given macro from the macro list used when processing
// Go files.
func (g *gopp) Undefine(key string) {
	delete(g.defines, key)
}

// Parse takes an io.Reader (usually from an os.Open call), and will process the
// resulting code read from it, preprocessing and substituting in macros as
// required.
func (g *gopp) Parse(r io.Reader) error {
	src, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}

	fset := token.NewFileSet()
	file := fset.AddFile("<stdin>", fset.Base(), len(src))

	s := scanner.Scanner{}
	s.Init(file, src, nil, scanner.ScanComments)

	go func() {
		for {
			pos, tok_, str := s.Scan()
			if len(str) == 0 {
				str = tok_.String()
			}
			tok := Token{pos, tok_, str}

			if tok.Token == token.EOF {
				break
			}

			if tok.Token != token.COMMENT && !g.ignoring {
				val, ok := g.defines[tok.String]
				if ok {
					tok.String = val.(string)
				}
				g.output <- tok
				continue
			}

			if !strings.HasPrefix(str, g.Prefix) {
				if !g.StripComments && !g.ignoring {
					tok.String += "\n"
					g.output <- tok
				}
				continue
			}

			// Trim the prefix from the start.
			strTrim := strings.Replace(tok.String, g.Prefix, "", 1)
			lnr := strings.SplitN(strTrim, " ", 2)
			if len(lnr) < 1 {
				fmt.Println("Invalid gopp comment:", str)
				continue
			}

			cmd := strings.ToLower(lnr[0])

			if cmd == "ifdef" {
				if len(lnr) != 2 {
					continue
				}
				def := lnr[1]
				_, ok := g.defines[def]
				g.ignoring = !ok
			} else if cmd == "ifndef" {
				if len(lnr) != 2 {
					continue
				}
				def := lnr[1]
				_, ok := g.defines[def]
				g.ignoring = ok
			} else if cmd == "else" {
				g.ignoring = !g.ignoring
			} else if cmd == "endif" && g.ignoring {
				g.ignoring = false
			} else if cmd == "define" && !g.ignoring {
				if len(lnr) != 2 {
					continue
				}
				lnr = strings.SplitN(lnr[1], " ", 2)
				g.DefineValue(lnr[0], lnr[1])
			} else if cmd == "undef" && !g.ignoring {
				if len(lnr) != 2 {
					continue
				}
				g.Undefine(lnr[1])
			}
		}
		close(g.output)
	}()
	return nil
}

// Print takes the resulting new AST after a call to Parse, and will print it
// to the given io.Writer
func (g *gopp) Print(w io.Writer) {
	outbuf := new(bytes.Buffer)
	for tok := range g.output {
		fmt.Fprintf(outbuf, " %s", tok.String)
	}

	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, "<stdin>", outbuf, parser.ParseComments)
	if err != nil {
		panic(err)
	}

	printer.Fprint(os.Stdout, fset, file)
}

// Reset will redefine the bits of a gopp instance, which you can use (e.g.)
// each time you want to parse a new file.
func (g *gopp) Reset() {
	g.defines = make(map[string]interface{})
	g.output = make(chan Token)
	g.ignoring = false
}

// New creates a new gopp instance with sane defaults.
func New(strip bool) (g *gopp) {
	 g = &gopp{
		defines:       make(map[string]interface{}),
		output:        make(chan Token),
		StripComments: strip,
		ignoring:      false,
		Prefix:        "//gopp:",
	}
	g.DefineValue("_GOPP", version)
	return
}
<commit_msg>gopp/gopp.go: tidied up documentation, exported Gopp type, renamed internal fields.<commit_after>// gopp provides an easy to use (but somewhat illogical) preprocessor for Go.
// It allows similar functionality to cpp, such as ifdef, ifndef, define, undef
// and other odd combinations thereof.
//
// Note: There is currently no support for nested ifdef/ifndefs.
package gopp

import (
	"bytes"
	"fmt"
	"go/parser"
	"go/printer"
	"go/scanner"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"strings"
)

const version = "0.4"

// Token provides a simple struct face to the `pos, tok, lit` returned by
// Go's go/scanner package. This is passed around by gopp internally in a chan.
type Token struct {
	Position token.Pos
	Token    token.Token
	String   string
}

type Gopp struct {
	Macros                  map[string]interface{}
	output                  chan Token
	StripComments, ignoring bool
	Prefix                  string
}

// DefineValue takes a name and a value, and defines them as a macro to use when
// preprocessing a Go file.
func (g *Gopp) DefineValue(name string, value interface{}) {
	g.Macros[name] = value
}

// Define takes a name, and sets it to the sane default of 1, which can then be
// used when preprocessing a Go file.
func (g *Gopp) Define(name string) {
	g.DefineValue(name, 1)
}

// Undefine will remove a given name from the macro list used when processing
// Go files.
func (g *Gopp) Undefine(name string) {
	delete(g.Macros, name)
}

// Parse takes an io.Reader (usually from an os.Open call), and will process the
// resulting code read from it, preprocessing and substituting in macros as
// required.
func (g *Gopp) Parse(r io.Reader) error {
	src, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}

	fset := token.NewFileSet()
	file := fset.AddFile("<stdin>", fset.Base(), len(src))

	s := scanner.Scanner{}
	s.Init(file, src, nil, scanner.ScanComments)

	go func() {
		for {
			pos, tok_, str := s.Scan()
			if len(str) == 0 {
				str = tok_.String()
			}
			tok := Token{pos, tok_, str}

			if tok.Token == token.EOF {
				break
			}

			if tok.Token != token.COMMENT && !g.ignoring {
				val, ok := g.Macros[tok.String]
				if ok {
					tok.String = val.(string)
				}
				g.output <- tok
				continue
			}

			if !strings.HasPrefix(str, g.Prefix) {
				if !g.StripComments && !g.ignoring {
					tok.String += "\n"
					g.output <- tok
				}
				continue
			}

			// Trim the prefix from the start.
			strTrim := strings.Replace(tok.String, g.Prefix, "", 1)
			lnr := strings.SplitN(strTrim, " ", 2)
			if len(lnr) < 1 {
				fmt.Println("Invalid gopp comment:", str)
				continue
			}

			cmd := strings.ToLower(lnr[0])

			if cmd == "ifdef" {
				if len(lnr) != 2 {
					continue
				}
				def := lnr[1]
				_, ok := g.Macros[def]
				g.ignoring = !ok
			} else if cmd == "ifndef" {
				if len(lnr) != 2 {
					continue
				}
				def := lnr[1]
				_, ok := g.Macros[def]
				g.ignoring = ok
			} else if cmd == "else" {
				g.ignoring = !g.ignoring
			} else if cmd == "endif" && g.ignoring {
				g.ignoring = false
			} else if cmd == "define" && !g.ignoring {
				if len(lnr) != 2 {
					continue
				}
				lnr = strings.SplitN(lnr[1], " ", 2)
				g.DefineValue(lnr[0], lnr[1])
			} else if cmd == "undef" && !g.ignoring {
				if len(lnr) != 2 {
					continue
				}
				g.Undefine(lnr[1])
			}
		}
		close(g.output)
	}()
	return nil
}

// Print takes the resulting new AST after a call to Parse, and will print it
// to the given io.Writer
func (g *Gopp) Print(w io.Writer) error {
	outbuf := new(bytes.Buffer)
	for tok := range g.output {
		fmt.Fprintf(outbuf, " %s", tok.String)
	}

	fset := token.NewFileSet()
	file, err := parser.ParseFile(fset, "<stdin>", outbuf, parser.ParseComments)
	if err != nil {
		return err
	}

	printer.Fprint(os.Stdout, fset, file)
	return nil
}

// Reset will redefine the bits of a gopp instance, which you can use (e.g.)
// each time you want to parse a new file.
func (g *Gopp) Reset() {
	g.Macros = make(map[string]interface{})
	g.output = make(chan Token)
	g.ignoring = false
}

// New creates a new gopp instance with sane defaults. The macro `_GOPP` is
// automatically set to the gopp version.
func New(strip bool) (g *Gopp) {
	g = &Gopp{
		Macros:        make(map[string]interface{}),
		output:        make(chan Token),
		StripComments: strip,
		ignoring:      false,
		Prefix:        "//gopp:",
	}
	g.DefineValue("_GOPP", version)
	return
}
<|endoftext|>


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

package gospinner

import (
	"fmt"
	"time"
)

const (
	DefaultChars string        = "|/-\\"
	DefaultSpeed time.Duration = 100 * time.Millisecond
)

// Spinner is a struct that stores Spinner information.
type Spinner struct {
	StartChan chan bool
	StopChan  chan bool
	Chars     string
	Speed     time.Duration
}

// NewSpinner creates a new Spinner.
func NewSpinner() *Spinner {
	return &Spinner{
		StartChan: make(chan bool),
		StopChan:  make(chan bool),
		Chars:     DefaultChars,
		Speed:     DefaultSpeed,
	}
}

// Start starts the spinner. Start takes a function where the long-running
// function happens.
// To start the spinner, set the start channel to true.
// To Stop the spinner, set the stop channel to true.
func (s *Spinner) Start(f func(start, stop chan bool)) {
	i := 0
	spin := false
	go f(s.StartChan, s.StopChan)
	for {
		select {
		case <-s.StartChan:
			spin = true
		case <-s.StopChan:
			fmt.Print("\r")
			return
		default:
			if spin {
				i++
				i = i % len(s.Chars)
				byte := s.Chars[i]
				fmt.Printf("\r%c", byte)
				time.Sleep(s.Speed)
			}
		}
	}
}

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

package gospinner

import (
	"fmt"
	"time"
)

const (
	DefaultChars string        = "|/-\\"
	DefaultSpeed time.Duration = 100 * time.Millisecond
)

// Spinner is a struct that stores Spinner information.
type Spinner struct {
	startChan chan bool
	stopChan  chan bool
	Chars     string
	Speed     time.Duration
}

// NewSpinner creates a new Spinner.
func NewSpinner() *Spinner {
	return &Spinner{
		startChan: make(chan bool),
		stopChan:  make(chan bool),
		Chars:     DefaultChars,
		Speed:     DefaultSpeed,
	}
}

// Start starts the spinner. Start takes a function where the long-running
// function happens.
// To start the spinner, set the start channel to true.
// To Stop the spinner, set the stop channel to true.
func (s *Spinner) Start(f func(start, stop chan bool)) {
	i := 0
	spin := false
	go f(s.startChan, s.stopChan)
	for {
		select {
		case <-s.startChan:
			spin = true
		case <-s.stopChan:
			fmt.Print("\r")
			return
		default:
			if spin {
				i++
				i = i % len(s.Chars)
				byte := s.Chars[i]
				fmt.Printf("\r%c", byte)
				time.Sleep(s.Speed)
			}
		}
	}
}

// Stop stops the spinner.
func (s *Spinner) Stop() {
	s.stopChan <- true
}
<|endoftext|>


<commit_before>package gracedown

import (
	"crypto/tls"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

// Server provides a graceful equivalent of net/http.Server.
type Server struct {
	*http.Server

	KillTimeOut time.Duration

	wg                sync.WaitGroup
	mu                sync.Mutex
	originalConnState func(conn net.Conn, newState http.ConnState)
	connStateOnce     sync.Once
	closed            int32 // accessed atomically.
	idlePool          map[net.Conn]struct{}
	listeners         map[net.Listener]struct{}
}

// NewWithServer wraps an existing http.Server.
func NewWithServer(s *http.Server) *Server {
	return &Server{
		Server:      s,
		KillTimeOut: 10 * time.Second,
		idlePool:    map[net.Conn]struct{}{},
		listeners:   map[net.Listener]struct{}{},
	}
}

// ListenAndServe provides a graceful equivalent of net/http.Server.ListenAndServe
func (srv *Server) ListenAndServe() error {
	addr := srv.Server.Addr
	if addr == "" {
		addr = ":http"
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return srv.Serve(ln)
}

// ListenAndServeTLS provides a graceful equivalent of net/http.Server.ListenAndServeTLS
func (srv *Server) ListenAndServeTLS(certFile, keyFile string) error {
	// direct lift from net/http/server.go
	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}
	config := &tls.Config{}
	if srv.TLSConfig != nil {
		*config = *srv.TLSConfig
	}
	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	return srv.Serve(tls.NewListener(ln, config))
}

// Serve provides a graceful equivalent of net/http.Server.Serve
func (srv *Server) Serve(l net.Listener) error {
	// remember net.Listener
	srv.mu.Lock()
	srv.listeners[l] = struct{}{}
	srv.mu.Unlock()
	defer func() {
		srv.mu.Lock()
		delete(srv.listeners, l)
		srv.mu.Unlock()
	}()

	// replace ConnState
	srv.connStateOnce.Do(func() {
		srv.originalConnState = srv.Server.ConnState
		srv.Server.ConnState = srv.connState
	})

	err := srv.Server.Serve(l)

	go func() {
		// wait for closing keep-alive connection by sending `Connection: Close` header.
		time.Sleep(srv.KillTimeOut)

		// time out, close all idle connections
		srv.mu.Lock()
		for conn := range srv.idlePool {
			conn.Close()
		}
		srv.mu.Unlock()
	}()

	// wait all connections have done
	srv.wg.Wait()

	if atomic.LoadInt32(&srv.closed) != 0 {
		// ignore closed network error when srv.Close() is called
		return nil
	}
	return err
}

// Close shuts down the default server used by ListenAndServe, ListenAndServeTLS and
// Serve. It returns true if it's the first time Close is called.
func (srv *Server) Close() bool {
	if atomic.CompareAndSwapInt32(&srv.closed, 0, 1) {
		srv.Server.SetKeepAlivesEnabled(false)
		srv.mu.Lock()
		listeners := srv.listeners
		srv.listeners = map[net.Listener]struct{}{}
		srv.mu.Unlock()
		for l := range listeners {
			l.Close()
		}
		return true
	}
	return false
}

func (srv *Server) connState(conn net.Conn, newState http.ConnState) {
	srv.mu.Lock()
	switch newState {
	case http.StateNew:
		srv.wg.Add(1)
	case http.StateActive:
		delete(srv.idlePool, conn)
	case http.StateIdle:
		srv.idlePool[conn] = struct{}{}
	case http.StateClosed, http.StateHijacked:
		delete(srv.idlePool, conn)
		srv.wg.Done()
	}
	srv.mu.Unlock()
	if srv.originalConnState != nil {
		srv.originalConnState(conn, newState)
	}
}
<commit_msg>ignoring the unexported sync.Once of tls.Config.<commit_after>package gracedown

import (
	"crypto/tls"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

// Server provides a graceful equivalent of net/http.Server.
type Server struct {
	*http.Server

	KillTimeOut time.Duration

	wg                sync.WaitGroup
	mu                sync.Mutex
	originalConnState func(conn net.Conn, newState http.ConnState)
	connStateOnce     sync.Once
	closed            int32 // accessed atomically.
	idlePool          map[net.Conn]struct{}
	listeners         map[net.Listener]struct{}
}

// NewWithServer wraps an existing http.Server.
func NewWithServer(s *http.Server) *Server {
	return &Server{
		Server:      s,
		KillTimeOut: 10 * time.Second,
		idlePool:    map[net.Conn]struct{}{},
		listeners:   map[net.Listener]struct{}{},
	}
}

// ListenAndServe provides a graceful equivalent of net/http.Server.ListenAndServe
func (srv *Server) ListenAndServe() error {
	addr := srv.Server.Addr
	if addr == "" {
		addr = ":http"
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	return srv.Serve(ln)
}

// ListenAndServeTLS provides a graceful equivalent of net/http.Server.ListenAndServeTLS
func (srv *Server) ListenAndServeTLS(certFile, keyFile string) error {
	// direct lift from net/http/server.go
	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

	config := cloneTLSConfig(srv.TLSConfig)
	if !strSliceContains(config.NextProtos, "http/1.1") {
		config.NextProtos = append(config.NextProtos, "http/1.1")
	}

	configHasCert := len(config.Certificates) > 0 || config.GetCertificate != nil
	if !configHasCert || certFile != "" || keyFile != "" {
		var err error
		config.Certificates = make([]tls.Certificate, 1)
		config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
		if err != nil {
			return err
		}
	}

	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	return srv.Serve(tls.NewListener(ln, config))
}

// cloneTLSConfig returns a shallow clone of the exported
// fields of cfg, ignoring the unexported sync.Once, which
// contains a mutex and must not be copied.
//
// The cfg must not be in active use by tls.Server, or else
// there can still be a race with tls.Server updating SessionTicketKey
// and our copying it, and also a race with the server setting
// SessionTicketsDisabled=false on failure to set the random
// ticket key.
//
// If cfg is nil, a new zero tls.Config is returned.
//
// Direct lift from net/http/transport.go
func cloneTLSConfig(cfg *tls.Config) *tls.Config {
	if cfg == nil {
		return &tls.Config{}
	}
	return &tls.Config{
		Rand:                     cfg.Rand,
		Time:                     cfg.Time,
		Certificates:             cfg.Certificates,
		NameToCertificate:        cfg.NameToCertificate,
		GetCertificate:           cfg.GetCertificate,
		RootCAs:                  cfg.RootCAs,
		NextProtos:               cfg.NextProtos,
		ServerName:               cfg.ServerName,
		ClientAuth:               cfg.ClientAuth,
		ClientCAs:                cfg.ClientCAs,
		InsecureSkipVerify:       cfg.InsecureSkipVerify,
		CipherSuites:             cfg.CipherSuites,
		PreferServerCipherSuites: cfg.PreferServerCipherSuites,
		SessionTicketsDisabled:   cfg.SessionTicketsDisabled,
		SessionTicketKey:         cfg.SessionTicketKey,
		ClientSessionCache:       cfg.ClientSessionCache,
		MinVersion:               cfg.MinVersion,
		MaxVersion:               cfg.MaxVersion,
		CurvePreferences:         cfg.CurvePreferences,
	}
}

func strSliceContains(ss []string, s string) bool {
	for _, v := range ss {
		if v == s {
			return true
		}
	}
	return false
}

// Serve provides a graceful equivalent of net/http.Server.Serve
func (srv *Server) Serve(l net.Listener) error {
	// remember net.Listener
	srv.mu.Lock()
	srv.listeners[l] = struct{}{}
	srv.mu.Unlock()
	defer func() {
		srv.mu.Lock()
		delete(srv.listeners, l)
		srv.mu.Unlock()
	}()

	// replace ConnState
	srv.connStateOnce.Do(func() {
		srv.originalConnState = srv.Server.ConnState
		srv.Server.ConnState = srv.connState
	})

	err := srv.Server.Serve(l)

	go func() {
		// wait for closing keep-alive connection by sending `Connection: Close` header.
		time.Sleep(srv.KillTimeOut)

		// time out, close all idle connections
		srv.mu.Lock()
		for conn := range srv.idlePool {
			conn.Close()
		}
		srv.mu.Unlock()
	}()

	// wait all connections have done
	srv.wg.Wait()

	if atomic.LoadInt32(&srv.closed) != 0 {
		// ignore closed network error when srv.Close() is called
		return nil
	}
	return err
}

// Close shuts down the default server used by ListenAndServe, ListenAndServeTLS and
// Serve. It returns true if it's the first time Close is called.
func (srv *Server) Close() bool {
	if atomic.CompareAndSwapInt32(&srv.closed, 0, 1) {
		srv.Server.SetKeepAlivesEnabled(false)
		srv.mu.Lock()
		listeners := srv.listeners
		srv.listeners = map[net.Listener]struct{}{}
		srv.mu.Unlock()
		for l := range listeners {
			l.Close()
		}
		return true
	}
	return false
}

func (srv *Server) connState(conn net.Conn, newState http.ConnState) {
	srv.mu.Lock()
	switch newState {
	case http.StateNew:
		srv.wg.Add(1)
	case http.StateActive:
		delete(srv.idlePool, conn)
	case http.StateIdle:
		srv.idlePool[conn] = struct{}{}
	case http.StateClosed, http.StateHijacked:
		delete(srv.idlePool, conn)
		srv.wg.Done()
	}
	srv.mu.Unlock()
	if srv.originalConnState != nil {
		srv.originalConnState(conn, newState)
	}
}
<|endoftext|>


<commit_before>package gracegrpc

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/facebookgo/grace/gracenet"
	"google.golang.org/grpc"
)

var (
	verbose    = flag.Bool("gracelog", true, "Enable logging.")
	didInherit = os.Getenv("LISTEN_FDS") != ""
	ppid       = os.Getppid()
)

type graceGrpc struct {
	server   *grpc.Server
	net      *gracenet.Net
	listener net.Listener
	errors   chan error
}

func NewGraceGrpc(s *grpc.Server, net, addr string) *graceGrpc {
	gr := &graceGrpc{
		server: s,
		net:    &gracenet.Net{},

		//for  StartProcess error.
		errors: make(chan error),
	}
	l, err := gr.net.Listen(net, addr)
	if err != nil {
		panic(err)
	}
	gr.listener = l
	return gr
}

func (gr *graceGrpc) serve() {
	go gr.server.Serve(gr.listener)
}

func (gr *graceGrpc) wait() {
	var wg sync.WaitGroup
	wg.Add(1)
	go gr.signalHandler(&wg)
	wg.Wait()
}

func (gr *graceGrpc) signalHandler(wg *sync.WaitGroup) {
	ch := make(chan os.Signal, 10)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR2)
	for {
		sig := <-ch
		switch sig {
		case syscall.SIGINT, syscall.SIGTERM:
			signal.Stop(ch)
			gr.server.GracefulStop()
			wg.Done()
			return
		case syscall.SIGUSR2:
			if _, err := gr.net.StartProcess(); err != nil {
				gr.errors <- err
			}
		}
	}
}

func (gr *graceGrpc) Serve() error {

	if *verbose {
		if didInherit {
			if ppid == 1 {
				log.Printf("Listening on init activated %s\n", pprintAddr(gr.listener))
			} else {
				const msg = "Graceful handoff of %s with new pid %d replace old pid %d"
				log.Printf(msg, pprintAddr(gr.listener), os.Getpid(), ppid)
			}
		} else {
			const msg = "Serving %s with pid %d\n"
			log.Printf(msg, pprintAddr(gr.listener), os.Getpid())
		}
	}

	gr.serve()

	if didInherit && ppid != 1 {
		if err := syscall.Kill(ppid, syscall.SIGTERM); err != nil {
			return fmt.Errorf("failed to close parent: %s", err)
		}
	}

	waitdone := make(chan struct{})
	go func() {
		defer close(waitdone)
		gr.wait()
	}()

	select {
	case err := <-gr.errors:
		if err == nil {
			panic("unexpected nil error")
		}
		return err
	case <-waitdone:
		if *verbose {
			log.Printf("Exiting pid %d.", os.Getpid())
		}
		return nil
	}
}

func pprintAddr(l net.Listener) []byte {
	var out bytes.Buffer
	fmt.Fprint(&out, l.Addr())
	return out.Bytes()
}
<commit_msg>增加对写pidfile的支持,配合systemd使用<commit_after>package gracegrpc

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"

	"github.com/facebookgo/grace/gracenet"
	"google.golang.org/grpc"
)

var (
	verbose    = flag.Bool("gracelog", true, "Enable logging.")
	writepid   = flag.Bool("gracepid", true, "Enable pidfile.")
	didInherit = os.Getenv("LISTEN_FDS") != ""
	ppid       = os.Getppid()
)

type graceGrpc struct {
	server   *grpc.Server
	net      *gracenet.Net
	listener net.Listener
	errors   chan error
	pidfile  string
}

func (gr *graceGrpc) InitPidFile(pid_file string) {
	gr.pidfile = pid_file
}

func NewGraceGrpc(s *grpc.Server, net, addr string) *graceGrpc {
	gr := &graceGrpc{
		server: s,
		net:    &gracenet.Net{},

		//for  StartProcess error.
		errors: make(chan error),
	}
	l, err := gr.net.Listen(net, addr)
	if err != nil {
		panic(err)
	}
	gr.listener = l
	return gr
}

func (gr *graceGrpc) serve() {
	go gr.server.Serve(gr.listener)
}

func (gr *graceGrpc) wait() {
	var wg sync.WaitGroup
	wg.Add(1)
	go gr.signalHandler(&wg)
	wg.Wait()
}

func (gr *graceGrpc) signalHandler(wg *sync.WaitGroup) {
	ch := make(chan os.Signal, 10)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR2)
	for {
		sig := <-ch
		log.Println(sig, "signal has received")
		switch sig {
		case syscall.SIGINT, syscall.SIGTERM:
			signal.Stop(ch)
			gr.server.GracefulStop()
			wg.Done()
			return
		case syscall.SIGUSR2:
			if _, err := gr.net.StartProcess(); err != nil {
				gr.errors <- err
			}
		}
	}
}

func (gr *graceGrpc) dowritepid(pid int) (err error) {
	if *writepid == false {
		return nil
	}

	pf, err := os.Create(gr.pidfile)
	defer pf.Close()
	if err != nil {
		log.Println(err)
	}

	_, err = pf.WriteString(strconv.Itoa(pid))
	if err != nil {
		log.Println(err)
	}
	return err
}

func (gr *graceGrpc) Serve() error {

	if *verbose {
		if didInherit {
			if ppid == 1 {
				log.Printf("Listening on init activated %s\n", pprintAddr(gr.listener))
			} else {
				const msg = "Graceful handoff of %s with new pid %d replace old pid %d"
				log.Printf(msg, pprintAddr(gr.listener), os.Getpid(), ppid)
			}
		} else {
			const msg = "Serving %s with pid %d\n"
			log.Printf(msg, pprintAddr(gr.listener), os.Getpid())
		}
	}

	err := gr.dowritepid(os.Getpid())
	if err != nil {
		log.Println(err)
	}

	gr.serve()

	if didInherit && ppid != 1 {
		if err := syscall.Kill(ppid, syscall.SIGTERM); err != nil {
			return fmt.Errorf("failed to close parent: %s", err)
		}
	}

	waitdone := make(chan struct{})
	go func() {
		defer close(waitdone)
		gr.wait()
	}()

	select {
	case err := <-gr.errors:
		if err == nil {
			panic("unexpected nil error")
		}
		return err
	case <-waitdone:
		if *verbose {
			log.Printf("Exiting pid %d.", os.Getpid())
		}
		return nil
	}
}

func pprintAddr(l net.Listener) []byte {
	var out bytes.Buffer
	fmt.Fprint(&out, l.Addr())
	return out.Bytes()
}
<|endoftext|>


<commit_before>package grid

import (
  "fmt"
  "math/rand"
  "time"
)

var idSeed uint= 0;

//Position structure
type pos struct {
  X int
  Y int
}

//Tiles track value and position on the grid
type tile struct {
  ID uint
  Value int
  MergeHistory []*tile
  Current pos
  Prev pos
}

func (t *tile) Move (dest *Cell) {
  fmt.Printf("Moving tile at X %d,Y %d to Cell at X %d, Y %d", t.X, t.Y, dest.X, dest.Y)
  fmt.Println()
  t.Prev = t.Current
  t.Current.X = dest.Pos
}

func (t *tile) Merge (tn *tile) {
  fmt.Printf("Merging tile %d,%d with tile %d,%d", t.X, t.Y, tn.X, tn.Y)
  fmt.Println()
  t.Value += tn.Value //Future proofs for other merge rules. Fibonacci game??? Yes please. TODO
  t.MergeHistory = append(t.MergeHistory, tn)
}

func randVal() *rand.Rand {
  r := rand.New(rand.NewSource(time.Now().UnixNano()))

  return r
}

func randTileVal() int {
  if randVal().Float64() > 0.9 {
    return 4
  } else {
    return 2
  }
}

func randPos(max int) int {
  return randVal().Int() % max }

//Grid tracks the tiles
type Grid struct {
  Size int
  StartCells int
  Tiles []*tile
  Cells [][]*Cell
}

type Cell struct {
  Tile *tile
  Pos pos
}

//Create the grid
func (g *Grid) Build () {
  g.Cells = make([][]*Cell, g.Size)
  g.Tiles = make([]*tile, 0)

  for i := range g.Cells {
    g.Cells[i] = make([]*Cell, g.Size)
    for j := range g.Cells[i] {
      cell := new(Cell)

      cell.Pos := pos{
        X: i,
        Y: j,
      }

      g.Cells[i][j] = cell
    }
  }

  start := g.StartCells

  for start != 0 {
    start = start - 1
    g.NewTile()
  }
}

func (g *Grid) Reset() {
  for i := range g.Cells {
    for j := range g.Cells[i] {
      g.Cells[i][j].Tile = nil
    }
  }
  g.Tiles = make([]*tile, 0)
}

func (g *Grid) NewTile() tile {

  avail := g.EmptyCells()
  //Two random values between 0 and Grid.Size
  i := randVal().Int() % len(avail)
  cell := avail[i]

  id := idSeed
  idSeed += 1

  newTile := tile{
    ID: id,
    Value: randTileVal(),
    MergeHistory: make([]*tile, 0),
    Current: cell.Pos,
  }

  g.Cells[cell.X][cell.Y].Tile = &newTile

  g.Tiles = append(g.Tiles, &newTile)

  return newTile
}

type vector struct {
  X int
  Y int
  Label string
}

var dMap = map[int]vector{
  1: vector{ //Up
    X: 0,
    Y: 1,
    Label: "Up",
  },
  2: vector{ //Right
    X: 1,
    Y: 0,
    Label: "Right",
  },
  3: vector{ //Down
    X: 0,
    Y: -1,
    Label: "Down",
  },
  4: vector{ //Left
    X: -1,
    Y: 0,
    Label: "Left",
  },
}

func (g *Grid) getEdge (v *vector) (max, min, delta int) {
  if(v.X == 1 || v.Y == 1) {
    max = g.Size - 1
    min = - 1
    delta = -1
  } else {
    min = g.Size
    max = 0
    delta = 1
  }

  return max, min, delta
}

//Get direction
//For table row (vector perpendicular to direction)
  //For row cell (vector opposite to direction)
    //If cell has tile
      //Add to 'new row' slice
      //If tile value matches prev tile value
        //Merge tile to prev
        //Remove tile ([2, 2, 4] .. [2 <-- 2, 4] .. [4, 4])
func (g *Grid) Shift(d int) *Grid {
  v := dMap[d]

  start, end, delta := g.getEdge(&v)

  for i:= start; i != end; i += delta {
    loc := start
    merged := false
    for j:= start; j != end; j += delta {
      var col int
      var cell, dest *Cell
      if(v.X == 0) { //Horizontal motion is zero. Direction is in Y
        col = i
        cell = g.Cells[col][j]
        dest = g.Cells[col][loc]
      } else {
        col = j
        cell = g.Cells[col][i]
        dest = g.Cells[col][loc]
      }

      if(cell.Tile != nil) { //If there is a tile for the cell we're checking
        if(loc != start) { //And if the pointer is not left at the beginning, which would mean we've moved no cells,
          potentialMerge := g.Cells[col][loc - delta].Tile; //Peek backwards to the last tile we moved

          if(!merged && potentialMerge != nil && potentialMerge.Value == cell.Tile.Value) { //If we didn't merge last time around and If the tiles match
              //Do a merge.
              potentialMerge.Merge(cell.Tile)
              cell.Tile = nil
              merged = true
              continue //Then continue
          } else {
            merged = false
          }
        }

        cell.Tile.Move(dest)
        dest.Tile = cell.Tile

        if(cell.Prev != dest.Pos) {
          cell.Tile = nil
        }

        loc += delta
      }
    }
  }

  g.NewTile()

  return g
}

func (g *Grid) EmptyCells() []*Cell {
  var ret []*Cell
  for i:=0; i < g.Size; i++ {
    for j:=0; j < g.Size; j++ {
      if(g.Cells[i][j].Tile == nil) {
        ret = append(ret, g.Cells[i][j])
      }
    }
  }

  return ret
}

func NewGrid(s, c int) Grid {
  g := Grid{
    Size: s,
    StartCells: c,
  }

  g.Build()

  return g
}
<commit_msg>grid/grid.go: Remove unused func and convert to pos type<commit_after>package grid

import (
  "fmt"
  "math/rand"
  "time"
)

var idSeed uint= 0;

//Position structure
type pos struct {
  X int
  Y int
}

//Tiles track value and position on the grid
type tile struct {
  ID uint
  Value int
  MergeHistory []*tile
  Current pos
  Prev pos
}

func (t *tile) Move (dest *Cell) {
  fmt.Printf("Moving tile at X %d,Y %d to Cell at X %d, Y %d", t.X, t.Y, dest.X, dest.Y)
  fmt.Println()
  t.Prev = t.Current
  t.Current = dest.Pos
}

func (t *tile) Merge (tn *tile) {
  fmt.Printf("Merging tile %d,%d with tile %d,%d", t.Current, tn.Current)
  fmt.Println()
  t.Value += tn.Value //Future proofs for other merge rules. Fibonacci game??? Yes please. TODO
  t.MergeHistory = append(t.MergeHistory, tn)
}

func randVal() *rand.Rand {
  r := rand.New(rand.NewSource(time.Now().UnixNano()))

  return r
}

func randTileVal() int {
  if randVal().Float64() > 0.9 {
    return 4
  } else {
    return 2
  }
}

//Grid tracks the tiles
type Grid struct {
  Size int
  StartCells int
  Tiles []*tile
  Cells [][]*Cell
}

type Cell struct {
  Tile *tile
  Pos pos
}

//Create the grid
func (g *Grid) Build () {
  g.Cells = make([][]*Cell, g.Size)
  g.Tiles = make([]*tile, 0)

  for i := range g.Cells {
    g.Cells[i] = make([]*Cell, g.Size)
    for j := range g.Cells[i] {
      cell := new(Cell)

      cell.Pos := pos{
        X: i,
        Y: j,
      }

      g.Cells[i][j] = cell
    }
  }

  start := g.StartCells

  for start != 0 {
    start = start - 1
    g.NewTile()
  }
}

func (g *Grid) Reset() {
  for i := range g.Cells {
    for j := range g.Cells[i] {
      g.Cells[i][j].Tile = nil
    }
  }
  g.Tiles = make([]*tile, 0)
}

func (g *Grid) NewTile() tile {

  avail := g.EmptyCells()
  //Two random values between 0 and Grid.Size
  i := randVal().Int() % len(avail)
  cell := avail[i]

  id := idSeed
  idSeed += 1

  newTile := tile{
    ID: id,
    Value: randTileVal(),
    MergeHistory: make([]*tile, 0),
    Current: cell.Pos,
  }

  g.Cells[cell.X][cell.Y].Tile = &newTile

  g.Tiles = append(g.Tiles, &newTile)

  return newTile
}

type vector struct {
  X int
  Y int
  Label string
}

var dMap = map[int]vector{
  1: vector{ //Up
    X: 0,
    Y: 1,
    Label: "Up",
  },
  2: vector{ //Right
    X: 1,
    Y: 0,
    Label: "Right",
  },
  3: vector{ //Down
    X: 0,
    Y: -1,
    Label: "Down",
  },
  4: vector{ //Left
    X: -1,
    Y: 0,
    Label: "Left",
  },
}

func (g *Grid) getEdge (v *vector) (max, min, delta int) {
  if(v.X == 1 || v.Y == 1) {
    max = g.Size - 1
    min = - 1
    delta = -1
  } else {
    min = g.Size
    max = 0
    delta = 1
  }

  return max, min, delta
}

//Get direction
//For table row (vector perpendicular to direction)
  //For row cell (vector opposite to direction)
    //If cell has tile
      //Add to 'new row' slice
      //If tile value matches prev tile value
        //Merge tile to prev
        //Remove tile ([2, 2, 4] .. [2 <-- 2, 4] .. [4, 4])
func (g *Grid) Shift(d int) *Grid {
  v := dMap[d]

  start, end, delta := g.getEdge(&v)

  for i:= start; i != end; i += delta {
    loc := start
    merged := false
    for j:= start; j != end; j += delta {
      var col int
      var cell, dest *Cell
      if(v.X == 0) { //Horizontal motion is zero. Direction is in Y
        col = i
        cell = g.Cells[col][j]
        dest = g.Cells[col][loc]
      } else {
        col = j
        cell = g.Cells[col][i]
        dest = g.Cells[col][loc]
      }

      if(cell.Tile != nil) { //If there is a tile for the cell we're checking
        if(loc != start) { //And if the pointer is not left at the beginning, which would mean we've moved no cells,
          potentialMerge := g.Cells[col][loc - delta].Tile; //Peek backwards to the last tile we moved

          if(!merged && potentialMerge != nil && potentialMerge.Value == cell.Tile.Value) { //If we didn't merge last time around and If the tiles match
              //Do a merge.
              potentialMerge.Merge(cell.Tile)
              cell.Tile = nil
              merged = true
              continue //Then continue
          } else {
            merged = false
          }
        }

        cell.Tile.Move(dest)
        dest.Tile = cell.Tile

        if(cell.Prev != dest.Pos) {
          cell.Tile = nil
        }

        loc += delta
      }
    }
  }

  g.NewTile()

  return g
}

func (g *Grid) EmptyCells() []*Cell {
  var ret []*Cell
  for i:=0; i < g.Size; i++ {
    for j:=0; j < g.Size; j++ {
      if(g.Cells[i][j].Tile == nil) {
        ret = append(ret, g.Cells[i][j])
      }
    }
  }

  return ret
}

func NewGrid(s, c int) Grid {
  g := Grid{
    Size: s,
    StartCells: c,
  }

  g.Build()

  return g
}
<|endoftext|>


<commit_before><commit_msg>corrected 'args' print format (added #x prefix)<commit_after><|endoftext|>


<commit_before>package autosignr

import (
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"os/exec"
	"path/filepath"
	"strings"

	log "github.com/Sirupsen/logrus"
)

// The OID for Puppet's pp_preshared_key in the Certificate Extensions
// https://docs.puppetlabs.com/puppet/latest/reference/ssl_attributes_extensions.html
const puppetPSKoid string = "1.3.6.1.4.1.34380.1.1.4"

func CertnameFromFilename(file string) string {
	base := filepath.Base(file)
	extension := filepath.Ext(base)

	return base[0 : len(base)-len(extension)]
}

func CheckCert(conf Config, file string) (bool, error) {
	name := CertnameFromFilename(file)
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return false, err
	}

	result, err := ValidateCert(conf, data, name)

	if err != nil {
		return false, err
	}

	if !result {
		log.Warningf("Unable to validate instance %s", name)
	}

	return result, nil
}

func ValidateCert(conf Config, data []byte, certname string) (bool, error) {
	if conf.CheckPSK {
		psk, err := PuppetPSKFromCSR(data)
		if err != nil {
			log.WithFields(log.Fields{
				"certname": certname,
				"err":      err,
			}).Warning("psk-extract-error")
			return false, err
		} else {
			if _, ok := conf.PresharedKeys[psk]; !ok {
				log.WithFields(log.Fields{
					"certname": certname,
					"psk":      psk,
				}).Warning("invalid-psk")
			}

			return false, errors.New("Invalid PSK")
		}
	}

	result := false
	for _, acct := range conf.Accounts {
		result = acct.Check(certname)
		if result {
			break
		}
	}

	return result, nil
}

func SignCert(conf Config, certname string) {
	cmd := fmt.Sprintf(conf.Cmdsign, certname)
	pieces := strings.Split(cmd, " ")

	cmdOut, err := exec.Command(pieces[0], pieces[1:]...).CombinedOutput()
	if err != nil {
		log.WithFields(log.Fields{
			"certname": certname,
			"err":      err,
			"output":   string(cmdOut),
		}).Error("signing-failure")
		return
	}
	log.WithFields(log.Fields{
		"certname": certname,
	}).Info("signing-success")
}

func ExistingCerts(conf Config) {

	matches, err := filepath.Glob(fmt.Sprintf("%s/*.pem", conf.Dir))
	if err != nil {
		log.Println("Glob error for: %s", err)
	}

	for _, cert := range matches {
		log.WithFields(log.Fields{
			"file": cert,
		}).Info("existing-csr")
		result, _ := CheckCert(conf, cert)

		if result {
			SignCert(conf, CertnameFromFilename(cert))
		}
	}
}

func PuppetPSKFromCSRFile(file string) (string, error) {
	var f string

	data, err := ioutil.ReadFile(file)
	if err != nil {
		return f, err
	}

	f, err = PuppetPSKFromCSR(data)

	return f, err
}

func PuppetPSKFromCSR(data []byte) (string, error) {
	var f string

	block, _ := pem.Decode(data)
	parsedcsr, err := x509.ParseCertificateRequest(block.Bytes)

	if err != nil {
		return f, err
	}

	for _, e := range parsedcsr.Extensions {
		if e.Id.String() == puppetPSKoid {
			// the first char of the trimmed string is ASCII 22,
			// synchronous idle, so remove that too
			f = strings.TrimSpace(string(e.Value))[1:]
			return f, nil
		}
	}

	return f, errors.New("Certificate did not contain a PSK")
}
<commit_msg>Logic bug in preshared key validation<commit_after>package autosignr

import (
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"os/exec"
	"path/filepath"
	"strings"

	log "github.com/Sirupsen/logrus"
)

// The OID for Puppet's pp_preshared_key in the Certificate Extensions
// https://docs.puppetlabs.com/puppet/latest/reference/ssl_attributes_extensions.html
const puppetPSKoid string = "1.3.6.1.4.1.34380.1.1.4"

func CertnameFromFilename(file string) string {
	base := filepath.Base(file)
	extension := filepath.Ext(base)

	return base[0 : len(base)-len(extension)]
}

func CheckCert(conf Config, file string) (bool, error) {
	name := CertnameFromFilename(file)
	log.Debugf("CheckCert %s", name)
	data, err := ioutil.ReadFile(file)
	if err != nil {
		return false, err
	}

	result, err := ValidateCert(conf, data, name)

	if err != nil {
		return false, err
	}

	if !result {
		log.Warningf("Unable to validate instance %s", name)
	}

	return result, nil
}

func ValidateCert(conf Config, data []byte, certname string) (bool, error) {
	log.Debugf("ValidateCert %s", certname)
	if conf.CheckPSK {
		psk, err := PuppetPSKFromCSR(data)
		if err != nil {
			log.WithFields(log.Fields{
				"certname": certname,
				"err":      err,
			}).Warning("psk-extract-error")
			return false, err
		} else {
			if _, ok := conf.PresharedKeys[psk]; !ok {
				log.WithFields(log.Fields{
					"certname": certname,
					"psk":      psk,
				}).Warning("invalid-psk")
				return false, errors.New("Invalid PSK")
			}
		}
	}

	result := false
	for _, acct := range conf.Accounts {
		result = acct.Check(certname)
		if result {
			break
		}
	}

	return result, nil
}

func SignCert(conf Config, certname string) {
	cmd := fmt.Sprintf(conf.Cmdsign, certname)
	pieces := strings.Split(cmd, " ")

	cmdOut, err := exec.Command(pieces[0], pieces[1:]...).CombinedOutput()
	if err != nil {
		log.WithFields(log.Fields{
			"certname": certname,
			"err":      err,
			"output":   string(cmdOut),
		}).Error("signing-failure")
		return
	}
	log.WithFields(log.Fields{
		"certname": certname,
	}).Info("signing-success")
}

func ExistingCerts(conf Config) {

	matches, err := filepath.Glob(fmt.Sprintf("%s/*.pem", conf.Dir))
	if err != nil {
		log.Println("Glob error for: %s", err)
	}

	for _, cert := range matches {
		log.WithFields(log.Fields{
			"file": cert,
		}).Info("existing-csr")
		result, _ := CheckCert(conf, cert)

		if result {
			SignCert(conf, CertnameFromFilename(cert))
		}
	}
}

func PuppetPSKFromCSRFile(file string) (string, error) {
	var f string

	data, err := ioutil.ReadFile(file)
	if err != nil {
		return f, err
	}

	f, err = PuppetPSKFromCSR(data)

	return f, err
}

func PuppetPSKFromCSR(data []byte) (string, error) {
	var f string

	block, _ := pem.Decode(data)
	parsedcsr, err := x509.ParseCertificateRequest(block.Bytes)

	if err != nil {
		return f, err
	}

	for _, e := range parsedcsr.Extensions {
		if e.Id.String() == puppetPSKoid {
			// the first char of the trimmed string is ASCII 22,
			// synchronous idle, so remove that too
			f = strings.TrimSpace(string(e.Value))[1:]
			return f, nil
		}
	}

	return f, errors.New("Certificate did not contain a PSK")
}
<|endoftext|>


<commit_before>package slack

import (
	"bytes"
	"context"
	"encoding/json"
	"net/url"
	"strings"

	"github.com/lestrrat/go-slack/objects"
	"github.com/pkg/errors"
)

type ChatUpdateCall struct {
	service   *ChatService
	channel   string
	timestamp string
	text      string
}

// Update returns the result of chat.update API
func (s *ChatService) Update(channel, text, ts string) *ChatUpdateCall {
	return &ChatUpdateCall{
		service:   s,
		channel:   channel,
		text:      text,
		timestamp: ts,
	}
}

func (c *ChatUpdateCall) Values() url.Values {
	v := url.Values{
		"token":   {c.service.token},
		"text":    {c.text},
		"channel": {c.channel},
		"ts":      {c.timestamp},
	}
	return v
}

func (c *ChatUpdateCall) Do(ctx context.Context) (*ChatResponse, error) {
	const endpoint = "chat.update"
	var res fullChatResponse
	if err := c.service.client.postForm(ctx, endpoint, c.Values(), &res); err != nil {
		return nil, errors.Wrap(err, `failed to post to chat.delete`)
	}

	if !res.OK {
		return nil, errors.New(res.Error.String())
	}

	return res.ChatResponse, nil
}

type ChatDeleteCall struct {
	service   *ChatService
	channel   string
	timestamp string
	asUser    bool
}

// Delete returns the result of chat.delete API
func (s *ChatService) Delete(channel, ts string) *ChatDeleteCall {
	return &ChatDeleteCall{
		service:   s,
		channel:   channel,
		timestamp: ts,
	}
}

func (c *ChatDeleteCall) Values() url.Values {
	v := url.Values{
		"token":   {c.service.token},
		"channel": {c.channel},
		"ts":      {c.timestamp},
	}

	if c.asUser {
		v.Set("as_user", "true")
	}
	return v
}

func (c *ChatDeleteCall) AsUser(b bool) *ChatDeleteCall {
	c.asUser = b
	return c
}

func (c *ChatDeleteCall) Do(ctx context.Context) (*ChatResponse, error) {
	const endpoint = "chat.delete"
	var res fullChatResponse
	if err := c.service.client.postForm(ctx, endpoint, c.Values(), &res); err != nil {
		return nil, errors.Wrap(err, `failed to post to chat.delete`)
	}

	if !res.OK {
		return nil, errors.New(res.Error.String())
	}

	return res.ChatResponse, nil
}

type ChatMeMessageCall struct {
	service *ChatService
	channel string
	text    string
}

// MeMessage returns the result of users.meMessage API
func (s *ChatService) MeMessage(channel, text string) *ChatMeMessageCall {
	return &ChatMeMessageCall{
		service: s,
		channel: channel,
		text:    text,
	}
}

func (c *ChatMeMessageCall) Values() url.Values {
	v := url.Values{
		"token":   {c.service.token},
		"channel": {c.channel},
		"text":    {c.text},
	}
	return v
}

func (c *ChatMeMessageCall) Do(ctx context.Context) (*ChatResponse, error) {
	const endpoint = "chat.meMessage"
	var res fullChatResponse
	if err := c.service.client.postForm(ctx, endpoint, c.Values(), &res); err != nil {
		return nil, errors.Wrap(err, `failed to post to chat.meMessage`)
	}

	if !res.OK {
		return nil, errors.New(res.Error.String())
	}

	return res.ChatResponse, nil
}

var replacer = strings.NewReplacer("&", "&", "<", "<", ">", ">")

func escapeMessage(s string) string {
	return replacer.Replace(s)
}

type fullChatResponse struct {
	SlackResponse
	*ChatResponse
}

type ChatPostMessageCall struct {
	service     *ChatService
	asUser      bool
	attachments objects.AttachmentList
	channel     string
	escapeText  bool
	iconEmoji   string
	iconURL     string
	linkNames   bool
	markdown    bool
	parse       string
	text        string
	unfurlLinks bool
	unfurlMedia bool
	username    string
}

// PostMessage returns the result of chat.postMessage API
func (s *ChatService) PostMessage(channel string) *ChatPostMessageCall {
	return &ChatPostMessageCall{
		service:     s,
		channel:     channel,
		escapeText:  true,
		markdown:    true,
		unfurlMedia: true,
	}
}

func (c *ChatPostMessageCall) Values() (url.Values, error) {
	v := url.Values{
		"token":   {c.service.token},
		"channel": {c.channel},
	}

	if c.asUser {
		v.Set("as_user", "true")
	}

	if len(c.attachments) > 0 {
		var buf bytes.Buffer
		if err := json.NewEncoder(&buf).Encode(c.attachments); err != nil {
			return nil, errors.Wrap(err, `failed to serialize attachments`)
		}
		v.Set("attachments", buf.String())
	}

	if len(c.iconEmoji) > 0 {
		v.Set("icon_emoji", c.iconEmoji)
	}
	if len(c.iconURL) > 0 {
		v.Set("icon_url", c.iconURL)
	}
	if c.linkNames {
		v.Set("link_names", "true")
	}

	if !c.markdown {
		v.Set("mrkdwn", "false")
	}

	if len(c.parse) > 0 {
		v.Set("parse", c.parse)
	}

	// taken from github.com/nlopes/slack:
	//    I want to send a message with explicit `as_user` `true` and
	//    `unfurl_links` `false` in request. Because setting `as_user` to
	//    `true` will change the default value for `unfurl_links` to `true`
	//    on Slack API side.
	if c.asUser && !c.unfurlLinks {
		v.Set("unfurl_link", "false")
	} else if c.unfurlLinks {
		v.Set("unfurl_link", "true")
	}

	if c.unfurlMedia {
		v.Set("unfurl_media", "true")
	}
	if len(c.username) > 0 {
		v.Set("username", c.username)
	}

	var txt = c.text
	if c.escapeText {
		txt = escapeMessage(txt)
	}
	v.Set("text", txt)

	return v, nil
}

func (c *ChatPostMessageCall) AsUser(b bool) *ChatPostMessageCall {
	c.asUser = b
	return c
}

// SetAttachments replaces the attachment list
func (c *ChatPostMessageCall) SetAttachment(l objects.AttachmentList) *ChatPostMessageCall {
	c.attachments = l
	return c
}

// Attachment appends to the attachments
func (c *ChatPostMessageCall) Attachment(a *objects.Attachment) *ChatPostMessageCall {
	c.attachments.Append(a)
	return c
}

func (c *ChatPostMessageCall) EscapeText(b bool) *ChatPostMessageCall {
	c.escapeText = b
	return c
}

func (c *ChatPostMessageCall) IconEmoji(s string) *ChatPostMessageCall {
	c.iconEmoji = s
	return c
}

func (c *ChatPostMessageCall) IconURL(s string) *ChatPostMessageCall {
	c.iconURL = s
	return c
}

func (c *ChatPostMessageCall) LinkNames(b bool) *ChatPostMessageCall {
	c.linkNames = b
	return c
}

func (c *ChatPostMessageCall) Markdown(b bool) *ChatPostMessageCall {
	c.markdown = b
	return c
}

func (c *ChatPostMessageCall) Parse(s string) *ChatPostMessageCall {
	c.parse = s
	return c
}

func (c *ChatPostMessageCall) Text(s string) *ChatPostMessageCall {
	c.text = s
	return c
}

func (c *ChatPostMessageCall) UnfurlLinks(b bool) *ChatPostMessageCall {
	c.unfurlLinks = b
	return c
}

func (c *ChatPostMessageCall) UnfurlMedia(b bool) *ChatPostMessageCall {
	c.unfurlMedia = b
	return c
}

func (c *ChatPostMessageCall) Username(s string) *ChatPostMessageCall {
	c.username = s
	return c
}

func (c *ChatPostMessageCall) Do(ctx context.Context) (*ChatResponse, error) {
	if len(c.channel) <= 0 {
		return nil, errors.New("channel not specified")
	}
	const endpoint = "chat.postMessage"

	v, err := c.Values()
	if err != nil {
		return nil, err
	}

	var res fullChatResponse
	if err := c.service.client.postForm(ctx, endpoint, v, &res); err != nil {
		return nil, errors.Wrap(err, `failed to post to chat.postMessage`)
	}

	if !res.OK {
		return nil, errors.New(res.Error.String())
	}

	return res.ChatResponse, nil
}
<commit_msg>Fix API<commit_after>package slack

import (
	"bytes"
	"context"
	"encoding/json"
	"net/url"
	"strings"

	"github.com/lestrrat/go-slack/objects"
	"github.com/pkg/errors"
)

type ChatUpdateCall struct {
	service   *ChatService
	channel   string
	timestamp string
	text      string
}

// Update returns the result of chat.update API
func (s *ChatService) Update(channel, text, ts string) *ChatUpdateCall {
	return &ChatUpdateCall{
		service:   s,
		channel:   channel,
		text:      text,
		timestamp: ts,
	}
}

func (c *ChatUpdateCall) Values() url.Values {
	v := url.Values{
		"token":   {c.service.token},
		"text":    {c.text},
		"channel": {c.channel},
		"ts":      {c.timestamp},
	}
	return v
}

func (c *ChatUpdateCall) Do(ctx context.Context) (*ChatResponse, error) {
	const endpoint = "chat.update"
	var res fullChatResponse
	if err := c.service.client.postForm(ctx, endpoint, c.Values(), &res); err != nil {
		return nil, errors.Wrap(err, `failed to post to chat.delete`)
	}

	if !res.OK {
		return nil, errors.New(res.Error.String())
	}

	return res.ChatResponse, nil
}

type ChatDeleteCall struct {
	service   *ChatService
	channel   string
	timestamp string
	asUser    bool
}

// Delete returns the result of chat.delete API
func (s *ChatService) Delete(channel, ts string) *ChatDeleteCall {
	return &ChatDeleteCall{
		service:   s,
		channel:   channel,
		timestamp: ts,
	}
}

func (c *ChatDeleteCall) Values() url.Values {
	v := url.Values{
		"token":   {c.service.token},
		"channel": {c.channel},
		"ts":      {c.timestamp},
	}

	if c.asUser {
		v.Set("as_user", "true")
	}
	return v
}

func (c *ChatDeleteCall) AsUser(b bool) *ChatDeleteCall {
	c.asUser = b
	return c
}

func (c *ChatDeleteCall) Do(ctx context.Context) (*ChatResponse, error) {
	const endpoint = "chat.delete"
	var res fullChatResponse
	if err := c.service.client.postForm(ctx, endpoint, c.Values(), &res); err != nil {
		return nil, errors.Wrap(err, `failed to post to chat.delete`)
	}

	if !res.OK {
		return nil, errors.New(res.Error.String())
	}

	return res.ChatResponse, nil
}

type ChatMeMessageCall struct {
	service *ChatService
	channel string
	text    string
}

// MeMessage returns the result of users.meMessage API
func (s *ChatService) MeMessage(channel, text string) *ChatMeMessageCall {
	return &ChatMeMessageCall{
		service: s,
		channel: channel,
		text:    text,
	}
}

func (c *ChatMeMessageCall) Values() url.Values {
	v := url.Values{
		"token":   {c.service.token},
		"channel": {c.channel},
		"text":    {c.text},
	}
	return v
}

func (c *ChatMeMessageCall) Do(ctx context.Context) (*ChatResponse, error) {
	const endpoint = "chat.meMessage"
	var res fullChatResponse
	if err := c.service.client.postForm(ctx, endpoint, c.Values(), &res); err != nil {
		return nil, errors.Wrap(err, `failed to post to chat.meMessage`)
	}

	if !res.OK {
		return nil, errors.New(res.Error.String())
	}

	return res.ChatResponse, nil
}

var replacer = strings.NewReplacer("&", "&", "<", "<", ">", ">")

func escapeMessage(s string) string {
	return replacer.Replace(s)
}

type fullChatResponse struct {
	SlackResponse
	*ChatResponse
}

type ChatPostMessageCall struct {
	service     *ChatService
	asUser      bool
	attachments objects.AttachmentList
	channel     string
	escapeText  bool
	iconEmoji   string
	iconURL     string
	linkNames   bool
	markdown    bool
	parse       string
	text        string
	unfurlLinks bool
	unfurlMedia bool
	username    string
}

// PostMessage returns the result of chat.postMessage API
func (s *ChatService) PostMessage(channel string) *ChatPostMessageCall {
	return &ChatPostMessageCall{
		service:     s,
		channel:     channel,
		escapeText:  true,
		markdown:    true,
		unfurlMedia: true,
	}
}

func (c *ChatPostMessageCall) Values() (url.Values, error) {
	v := url.Values{
		"token":   {c.service.token},
		"channel": {c.channel},
	}

	if c.asUser {
		v.Set("as_user", "true")
	}

	if len(c.attachments) > 0 {
		var buf bytes.Buffer
		if err := json.NewEncoder(&buf).Encode(c.attachments); err != nil {
			return nil, errors.Wrap(err, `failed to serialize attachments`)
		}
		v.Set("attachments", buf.String())
	}

	if len(c.iconEmoji) > 0 {
		v.Set("icon_emoji", c.iconEmoji)
	}
	if len(c.iconURL) > 0 {
		v.Set("icon_url", c.iconURL)
	}
	if c.linkNames {
		v.Set("link_names", "true")
	}

	if !c.markdown {
		v.Set("mrkdwn", "false")
	}

	if len(c.parse) > 0 {
		v.Set("parse", c.parse)
	}

	// taken from github.com/nlopes/slack:
	//    I want to send a message with explicit `as_user` `true` and
	//    `unfurl_links` `false` in request. Because setting `as_user` to
	//    `true` will change the default value for `unfurl_links` to `true`
	//    on Slack API side.
	if c.asUser && !c.unfurlLinks {
		v.Set("unfurl_link", "false")
	} else if c.unfurlLinks {
		v.Set("unfurl_link", "true")
	}

	if c.unfurlMedia {
		v.Set("unfurl_media", "true")
	}
	if len(c.username) > 0 {
		v.Set("username", c.username)
	}

	var txt = c.text
	if c.escapeText {
		txt = escapeMessage(txt)
	}
	v.Set("text", txt)

	return v, nil
}

func (c *ChatPostMessageCall) AsUser(b bool) *ChatPostMessageCall {
	c.asUser = b
	return c
}

// SetAttachments replaces the attachment list
func (c *ChatPostMessageCall) SetAttachments(l objects.AttachmentList) *ChatPostMessageCall {
	c.attachments = l
	return c
}

// Attachment appends to the attachments
func (c *ChatPostMessageCall) Attachment(a *objects.Attachment) *ChatPostMessageCall {
	c.attachments.Append(a)
	return c
}

func (c *ChatPostMessageCall) EscapeText(b bool) *ChatPostMessageCall {
	c.escapeText = b
	return c
}

func (c *ChatPostMessageCall) IconEmoji(s string) *ChatPostMessageCall {
	c.iconEmoji = s
	return c
}

func (c *ChatPostMessageCall) IconURL(s string) *ChatPostMessageCall {
	c.iconURL = s
	return c
}

func (c *ChatPostMessageCall) LinkNames(b bool) *ChatPostMessageCall {
	c.linkNames = b
	return c
}

func (c *ChatPostMessageCall) Markdown(b bool) *ChatPostMessageCall {
	c.markdown = b
	return c
}

func (c *ChatPostMessageCall) Parse(s string) *ChatPostMessageCall {
	c.parse = s
	return c
}

func (c *ChatPostMessageCall) Text(s string) *ChatPostMessageCall {
	c.text = s
	return c
}

func (c *ChatPostMessageCall) UnfurlLinks(b bool) *ChatPostMessageCall {
	c.unfurlLinks = b
	return c
}

func (c *ChatPostMessageCall) UnfurlMedia(b bool) *ChatPostMessageCall {
	c.unfurlMedia = b
	return c
}

func (c *ChatPostMessageCall) Username(s string) *ChatPostMessageCall {
	c.username = s
	return c
}

func (c *ChatPostMessageCall) Do(ctx context.Context) (*ChatResponse, error) {
	if len(c.channel) <= 0 {
		return nil, errors.New("channel not specified")
	}
	const endpoint = "chat.postMessage"

	v, err := c.Values()
	if err != nil {
		return nil, err
	}

	var res fullChatResponse
	if err := c.service.client.postForm(ctx, endpoint, v, &res); err != nil {
		return nil, errors.Wrap(err, `failed to post to chat.postMessage`)
	}

	if !res.OK {
		return nil, errors.New(res.Error.String())
	}

	return res.ChatResponse, nil
}
<|endoftext|>


<commit_before>// Package cmac implements AES-CMAC (RFC4493/NIST SP800-38b)
package cmac

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/subtle"
	"hash"
)

const (
	_Rb128 = 0x87
	_Rb64  = 0x1b
)

func shifted(x []byte) []byte {
	d := make([]byte, len(x))
	copy(d, x)
	for i := range d {
		if i > 0 {
			d[i-1] |= d[i] >> 7
		}
		d[i] <<= 1
	}
	return d
}

func gensubkey(c cipher.Block, l []byte, rb byte) []byte {
	sk := shifted(l)
	sk[len(sk)-1] ^= byte(subtle.ConstantTimeSelect(int(l[0]>>7), int(rb), 0))
	return sk
}

func gensubkeys(c cipher.Block) ([]byte, []byte) {
	var rb byte

	switch c.BlockSize() {
	case 16:
		rb = _Rb128
	case 8:
		rb = _Rb64
	default:
		panic("cmac: invalid block size")

	}

	l := make([]byte, c.BlockSize())
	c.Encrypt(l, l)

	k1 := gensubkey(c, l, rb)
	return k1, gensubkey(c, k1, rb)
}

type cmac struct {
	c  cipher.Block
	k1 []byte
	k2 []byte

	x []byte

	buf []byte
}

func newcmac(c cipher.Block) *cmac {
	k1, k2 := gensubkeys(c)
	m := &cmac{c: c, k1: k1, k2: k2}
	m.Reset()
	return m
}

func (m *cmac) block(b []byte) {
	y := make([]byte, m.c.BlockSize())
	for i := range y {
		y[i] = m.x[i] ^ b[i]
	}
	m.c.Encrypt(m.x, y)
}

func (m *cmac) Write(b []byte) (int, error) {
	d := append(m.buf, b...)

	for len(d) > m.c.BlockSize() {
		m.block(d[:m.c.BlockSize()])
		d = d[m.c.BlockSize():]
	}

	m.buf = d

	return len(b), nil
}

func (m *cmac) Sum(b []byte) []byte {
	last := make([]byte, m.c.BlockSize())

	if len(m.buf) > 0 && len(m.buf)%m.c.BlockSize() == 0 {
		for i := range last {
			last[i] = m.buf[i] ^ m.k1[i]
		}
	} else {
		copy(last, m.buf)
		last[len(m.buf)] = 0x80
		for i := range last {
			last[i] ^= m.k2[i]
		}
	}

	for i := range last {
		last[i] ^= m.x[i]
	}
	m.c.Encrypt(last, last)

	return append(b, last...)
}

func (m *cmac) Reset() {
	m.buf = nil
	m.x = make([]byte, m.c.BlockSize())
}

func (m *cmac) Size() int {
	return m.c.BlockSize()
}

func (m *cmac) BlockSize() int {
	return m.c.BlockSize()
}

// New returns a hash.Hash computing AES-CMAC.
func New(key []byte) hash.Hash {
	c, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	return NewWithCipher(c)
}

// NewWithCipher returns a hash.Hash computing CMAC using the given
// cipher.Block. The block cipher should have a block length of 8 or 16 bytes.
func NewWithCipher(c cipher.Block) hash.Hash {
	switch c.BlockSize() {
	case 8, 16:
		return newcmac(c)
	default:
		panic("cmac: invalid blocksize")
	}
}
<commit_msg>package doc<commit_after>// Package cmac implements CMAC as defined in  RFC4493 an NIST SP800-38b.
package cmac

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/subtle"
	"hash"
)

const (
	_Rb128 = 0x87
	_Rb64  = 0x1b
)

func shifted(x []byte) []byte {
	d := make([]byte, len(x))
	copy(d, x)
	for i := range d {
		if i > 0 {
			d[i-1] |= d[i] >> 7
		}
		d[i] <<= 1
	}
	return d
}

func gensubkey(c cipher.Block, l []byte, rb byte) []byte {
	sk := shifted(l)
	sk[len(sk)-1] ^= byte(subtle.ConstantTimeSelect(int(l[0]>>7), int(rb), 0))
	return sk
}

func gensubkeys(c cipher.Block) ([]byte, []byte) {
	var rb byte

	switch c.BlockSize() {
	case 16:
		rb = _Rb128
	case 8:
		rb = _Rb64
	default:
		panic("cmac: invalid block size")

	}

	l := make([]byte, c.BlockSize())
	c.Encrypt(l, l)

	k1 := gensubkey(c, l, rb)
	return k1, gensubkey(c, k1, rb)
}

type cmac struct {
	c  cipher.Block
	k1 []byte
	k2 []byte

	x []byte

	buf []byte
}

func newcmac(c cipher.Block) *cmac {
	k1, k2 := gensubkeys(c)
	m := &cmac{c: c, k1: k1, k2: k2}
	m.Reset()
	return m
}

func (m *cmac) block(b []byte) {
	y := make([]byte, m.c.BlockSize())
	for i := range y {
		y[i] = m.x[i] ^ b[i]
	}
	m.c.Encrypt(m.x, y)
}

func (m *cmac) Write(b []byte) (int, error) {
	d := append(m.buf, b...)

	for len(d) > m.c.BlockSize() {
		m.block(d[:m.c.BlockSize()])
		d = d[m.c.BlockSize():]
	}

	m.buf = d

	return len(b), nil
}

func (m *cmac) Sum(b []byte) []byte {
	last := make([]byte, m.c.BlockSize())

	if len(m.buf) > 0 && len(m.buf)%m.c.BlockSize() == 0 {
		for i := range last {
			last[i] = m.buf[i] ^ m.k1[i]
		}
	} else {
		copy(last, m.buf)
		last[len(m.buf)] = 0x80
		for i := range last {
			last[i] ^= m.k2[i]
		}
	}

	for i := range last {
		last[i] ^= m.x[i]
	}
	m.c.Encrypt(last, last)

	return append(b, last...)
}

func (m *cmac) Reset() {
	m.buf = nil
	m.x = make([]byte, m.c.BlockSize())
}

func (m *cmac) Size() int {
	return m.c.BlockSize()
}

func (m *cmac) BlockSize() int {
	return m.c.BlockSize()
}

// New returns a hash.Hash computing AES-CMAC.
func New(key []byte) hash.Hash {
	c, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	return NewWithCipher(c)
}

// NewWithCipher returns a hash.Hash computing CMAC using the given
// cipher.Block. The block cipher should have a block length of 8 or 16 bytes.
func NewWithCipher(c cipher.Block) hash.Hash {
	switch c.BlockSize() {
	case 8, 16:
		return newcmac(c)
	default:
		panic("cmac: invalid blocksize")
	}
}
<|endoftext|>


<commit_before>package epp

import (
	"bytes"
	"encoding/binary"
	"encoding/xml"
	"io"
	"net"
)

// IgnoreEOF returns err unless err == io.EOF,
// in which case it returns nil.
func IgnoreEOF(err error) error {
	if err == io.EOF {
		return nil
	}
	return err
}

// Conn represents a single connection to an EPP server.
// This implementation is not safe for concurrent use.
type Conn struct {
	net.Conn
	buf     bytes.Buffer
	decoder *xml.Decoder

	// Greeting holds the last received greeting message from the server,
	// indicating server name, status, data policy and capabilities.
	Greeting
}

// NewConn initializes an epp.Conn from a net.Conn and performs the EPP
// handshake. It reads and stores the initial EPP <greeting> message.
// https://tools.ietf.org/html/rfc5730#section-2.4
func NewConn(conn net.Conn) (*Conn, error) {
	c := newConn(conn)
	err := c.readGreeting()
	return c, err
}

// Close sends an EPP <logout> command and closes the connection c.
func (c *Conn) Close() error {
	c.Logout()
	return c.Conn.Close()
}

// newConn initializes an epp.Conn from a net.Conn.
// Used internally for testing.
func newConn(conn net.Conn) *Conn {
	c := Conn{Conn: conn}
	c.resetDecoder()
	return &c
}

// reset resets the underlying xml.Decoder and bytes.Buffer.
func (c *Conn) reset() {
	c.resetDecoder()
	c.buf.Reset()
}

// resetDecoder restores the original state of the underlying
// xml.Decoder (pos 1, line 1, stack, etc.) using a hack.
func (c *Conn) resetDecoder() {
	c.decoder = xml.NewDecoder(&c.buf)
}

// flushDataUnit writes bytes from c.buf to c using writeDataUnit.
func (c *Conn) flushDataUnit() error {
	return c.writeDataUnit(c.buf.Bytes())
}

// writeDataUnit writes a slice of bytes to c.
// Bytes written are prefixed with 32-bit header specifying the total size
// of the data unit (message + 4 byte header), in network (big-endian) order.
// http://www.ietf.org/rfc/rfc4934.txt
func (c *Conn) writeDataUnit(x []byte) error {
	logXML("<-- WRITE DATA UNIT -->", x)
	s := uint32(4 + len(x))
	err := binary.Write(c.Conn, binary.BigEndian, s)
	if err != nil {
		return err
	}
	_, err = c.Conn.Write(x)
	return err
}

// readResponse reads a single EPP response from c and parses the XML into req.
// It returns an error if the EPP response contains an error Result.
func (c *Conn) readResponse(res *response_) error {
	err := c.readDataUnit()
	if err != nil {
		return err
	}
	err = IgnoreEOF(scanResponse.Scan(c.decoder, res))
	if err != nil {
		return err
	}
	if res.Result.IsError() {
		return res.Result
	}
	return nil
}

// readDataUnit reads a single EPP message from c into
// c.buf. The bytes in c.buf are valid until the next
// call to readDataUnit.
func (c *Conn) readDataUnit() error {
	c.reset()
	var s int32
	err := binary.Read(c.Conn, binary.BigEndian, &s)
	if err != nil {
		return err
	}
	s -= 4
	if s < 0 {
		return io.ErrUnexpectedEOF
	}
	lr := io.LimitedReader{R: c.Conn, N: int64(s)}
	n, err := c.buf.ReadFrom(&lr)
	if err != nil {
		return err
	}
	if n != int64(s) || lr.N != 0 {
		return io.ErrUnexpectedEOF
	}
	logXML("<-- READ DATA UNIT -->", c.buf.Bytes())
	return nil
}

func deleteRange(s, pfx, sfx []byte) []byte {
	start := bytes.Index(s, pfx)
	if start < 0 {
		return s
	}
	end := bytes.Index(s[start+len(pfx):], sfx)
	if end < 0 {
		return s
	}
	end += start + len(pfx) + len(sfx)
	size := len(s) - (end - start)
	copy(s[start:size], s[end:])
	return s[:size]
}

func deleteBufferRange(buf *bytes.Buffer, pfx, sfx []byte) {
	v := deleteRange(buf.Bytes(), pfx, sfx)
	buf.Truncate(len(v))
}
<commit_msg>Reset buffer before creating xml.Decoder.<commit_after>package epp

import (
	"bytes"
	"encoding/binary"
	"encoding/xml"
	"io"
	"net"
)

// IgnoreEOF returns err unless err == io.EOF,
// in which case it returns nil.
func IgnoreEOF(err error) error {
	if err == io.EOF {
		return nil
	}
	return err
}

// Conn represents a single connection to an EPP server.
// This implementation is not safe for concurrent use.
type Conn struct {
	net.Conn
	buf     bytes.Buffer
	decoder *xml.Decoder

	// Greeting holds the last received greeting message from the server,
	// indicating server name, status, data policy and capabilities.
	Greeting
}

// NewConn initializes an epp.Conn from a net.Conn and performs the EPP
// handshake. It reads and stores the initial EPP <greeting> message.
// https://tools.ietf.org/html/rfc5730#section-2.4
func NewConn(conn net.Conn) (*Conn, error) {
	c := newConn(conn)
	err := c.readGreeting()
	return c, err
}

// Close sends an EPP <logout> command and closes the connection c.
func (c *Conn) Close() error {
	c.Logout()
	return c.Conn.Close()
}

// newConn initializes an epp.Conn from a net.Conn.
// Used internally for testing.
func newConn(conn net.Conn) *Conn {
	c := Conn{Conn: conn}
	c.reset()
	return &c
}

// reset resets the underlying xml.Decoder and bytes.Buffer.
func (c *Conn) reset() {
	c.buf.Reset()
	c.decoder = xml.NewDecoder(&c.buf)
}

// flushDataUnit writes bytes from c.buf to c using writeDataUnit.
func (c *Conn) flushDataUnit() error {
	return c.writeDataUnit(c.buf.Bytes())
}

// writeDataUnit writes a slice of bytes to c.
// Bytes written are prefixed with 32-bit header specifying the total size
// of the data unit (message + 4 byte header), in network (big-endian) order.
// http://www.ietf.org/rfc/rfc4934.txt
func (c *Conn) writeDataUnit(x []byte) error {
	logXML("<-- WRITE DATA UNIT -->", x)
	s := uint32(4 + len(x))
	err := binary.Write(c.Conn, binary.BigEndian, s)
	if err != nil {
		return err
	}
	_, err = c.Conn.Write(x)
	return err
}

// readResponse reads a single EPP response from c and parses the XML into req.
// It returns an error if the EPP response contains an error Result.
func (c *Conn) readResponse(res *response_) error {
	err := c.readDataUnit()
	if err != nil {
		return err
	}
	err = IgnoreEOF(scanResponse.Scan(c.decoder, res))
	if err != nil {
		return err
	}
	if res.Result.IsError() {
		return res.Result
	}
	return nil
}

// readDataUnit reads a single EPP message from c into
// c.buf. The bytes in c.buf are valid until the next
// call to readDataUnit.
func (c *Conn) readDataUnit() error {
	c.reset()
	var s int32
	err := binary.Read(c.Conn, binary.BigEndian, &s)
	if err != nil {
		return err
	}
	s -= 4
	if s < 0 {
		return io.ErrUnexpectedEOF
	}
	lr := io.LimitedReader{R: c.Conn, N: int64(s)}
	n, err := c.buf.ReadFrom(&lr)
	if err != nil {
		return err
	}
	if n != int64(s) || lr.N != 0 {
		return io.ErrUnexpectedEOF
	}
	logXML("<-- READ DATA UNIT -->", c.buf.Bytes())
	return nil
}

func deleteRange(s, pfx, sfx []byte) []byte {
	start := bytes.Index(s, pfx)
	if start < 0 {
		return s
	}
	end := bytes.Index(s[start+len(pfx):], sfx)
	if end < 0 {
		return s
	}
	end += start + len(pfx) + len(sfx)
	size := len(s) - (end - start)
	copy(s[start:size], s[end:])
	return s[:size]
}

func deleteBufferRange(buf *bytes.Buffer, pfx, sfx []byte) {
	v := deleteRange(buf.Bytes(), pfx, sfx)
	buf.Truncate(len(v))
}
<|endoftext|>


<commit_before>package main

import (
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"strconv"

	"github.com/codeskyblue/go-sh"
	"github.com/gorilla/mux"
	"github.com/gorilla/rpc"
	"github.com/gorilla/rpc/json"
)

const frontendHeaderACL = `
  acl is{{.ACL}} hdr_beg(host) {{index .HaproxyURLs #}}`
const frontendPathACL = `
  acl is{{.ACL}} path_beg -i {{index .HaproxyURLs #}}`

const frontendUse = `
  use_backend {{.Backend}} if is{{.ACL}}`

// const frontendTmpl = `
// acl is{{.ACL}} hdr_beg(host) {{.HaproxyURL}}
// use_backend {{.Backend}} if is{{.ACL}}
// `
const backendTmpl = `
backend {{.Backend}}
  server {{.Backend}} {{.Hostmachine}}:{{.Port}} check inter 10000`

const defaultBackendTmpl = `
  default_backend {{.Backend}}`

var confPath = "/usr/local/etc/haproxy/conf"
var haproxyPath = "/usr/local/etc/haproxy"

// Service to be added
type Service struct {
	// ACL to be used
	ACL string
	// URL by which service will be called
	HaproxyURLs []string
	// Backend name
	Backend string
	// storefront-services-1.myntra.com
	Hostmachine string
	// Port on Hostmachine where service runs
	Port string
}

// Services ...
type Services struct {
	Services []Service
}

// Haproxy ...
type Haproxy int

// Result ...
type Result int

// Add a frontend and backend
func (h *Haproxy) Add(r *http.Request, services *Services, result *Result) error {
	var defaultBackend interface{}
	isDefaultBackedDefined := false
	for _, service := range services.Services {
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".backend").Run()
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".frontend").Run()
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".frontendtop").Run()
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".frontendbottom").Run()

		log.Printf("Add service %s:%s", service.HaproxyURLs, service.Port)
		data := struct {
			ACL         string
			HaproxyURLs []string
			Backend     string
			Hostmachine string
			Port        string
		}{
			strings.Title(service.ACL),
			service.HaproxyURLs,
			service.Backend,
			service.Hostmachine,
			service.Port,
		}

		frontendACLs := `
  `

		frontendType := ".frontend"

		for haproxyURLId, haproxyURL := range service.HaproxyURLs {
			if haproxyURL == "default_backend" {
				defaultBackend = data
				isDefaultBackedDefined = true
			} else if strings.EqualFold(haproxyURL, "top") || strings.EqualFold(haproxyURL, "bottom") {
				frontendType += haproxyURL
			} else if string(haproxyURL[0]) == "/" {
				frontendACLs += `  `
				frontendACLs += strings.Replace(frontendPathACL, "#", strconv.Itoa(haproxyURLId), -1)
				frontendACLs += `
`
			} else {
				frontendACLs += `  `
				frontendACLs += strings.Replace(frontendHeaderACL, "#", strconv.Itoa(haproxyURLId), -1)
				frontendACLs += `
`
			}
		}

		frontendTmpl := frontendACLs + frontendUse

		log.Println(frontendTmpl)

		// Generate frontend entry
		tmpl := template.Must(template.New("frontend").Parse(frontendTmpl))
		f, err := os.OpenFile(confPath+"/"+service.ACL+frontendType, os.O_CREATE|os.O_RDWR, 0777)
		if err != nil {
			*result = 0
			return err
		}
		// fill in the template
		err = tmpl.Execute(f, data)
		if err != nil {
			*result = 0
			return err
		}

		// Generate backend entry
		tmpl = template.Must(template.New("backend").Parse(backendTmpl))
		f, err = os.OpenFile(confPath+"/"+service.ACL+".backend", os.O_CREATE|os.O_RDWR, 0777)
		if err != nil {
			*result = 0
			return err
		}
		err = tmpl.Execute(f, data)
		if err != nil {
			*result = 0
			return err
		}
	}

	// Generate default_backend if needed
	if isDefaultBackedDefined {
		tmpl := template.Must(template.New("default_backend").Parse(defaultBackendTmpl))
		f, err := os.OpenFile(confPath+"/"+".default_backend", os.O_CREATE|os.O_RDWR, 0777)
		if err != nil {
			*result = 0
			return err
		}
		err = tmpl.Execute(f, defaultBackend)
		if err != nil {
			*result = 0
			return err
		}
	}

	//join all the configs
	if err := h.generateCfg(); err != nil {
		*result = 0
		return err
	}

	if err := h.ReloadHaproxy(); err != nil {
		*result = 0
		return err
	}

	*result = 1
	return nil

}

// Remove a frontend and backend
func (h *Haproxy) Remove(r *http.Request, services *Services, result *Result) error {
	for _, service := range services.Services {
		log.Printf("Remove service %s:%s", service.HaproxyURLs, service.Port)
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".backend").Run()
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".frontend").Run()
	}

	if err := h.generateCfg(); err != nil {
		*result = 0
		return err
	}

	if err := h.ReloadHaproxy(); err != nil {
		*result = 0
		return err
	}

	*result = 1
	return nil

}

func (h *Haproxy) generateCfg() error {
	if _, err := os.Stat(haproxyPath + "/haproxy.cfg"); !os.IsNotExist(err) {
		currentTime := string(time.Now().Format("20060102150405"))
		err := os.Rename(haproxyPath+"/haproxy.cfg", haproxyPath+"/haproxy.cfg.BAK."+currentTime)
		if err != nil {
			return err
		}
	}

	var haproxyCfg []byte

	var partFunc = func(part string) {
		// walk all files in the directory
		filepath.Walk(confPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() && strings.HasSuffix(info.Name(), part) {
				b, err := ioutil.ReadFile(path)
				if err != nil {
					return err
				}
				haproxyCfg = append(haproxyCfg, b...)
			}
			return nil

		})
	}

	//append the configs in the following order
	parts := []string{".globalcfg", ".defaultcfg", ".frontendcfg", ".frontendtop", ".frontend", ".frontendbottom", ".default_backend", ".backend"}
	for i := range parts {
		partFunc(parts[i])
	}

	//write the file
	err := ioutil.WriteFile(haproxyPath+"/haproxy.cfg", haproxyCfg, 0777)
	if err != nil {
		return err
	}

	return nil

}

// StartHaproxy ...
func (h *Haproxy) StartHaproxy() error {
	// restart haproxy container
	session := sh.NewSession()
	err := session.Command("haproxy", "-p", "/var/run/haproxy.pid", "-f", "/usr/local/etc/haproxy/haproxy.cfg").Run()
	if err != nil {
		return err
	}
	return nil
}

// ReloadHaproxy ...
func (h *Haproxy) ReloadHaproxy() error {
	session := sh.NewSession()
	err := session.Command("/usr/bin/reload.sh").Run()
	if err != nil {
		return err
	}
	return nil
}

// Generate regenerates haproxy config from existing configs
func (h *Haproxy) Generate(r *http.Request, service *Service, result *Result) error {
	if err := h.generateCfg(); err != nil {
		*result = 0
		return err
	}

	*result = 1
	return nil

}

func main() {

	//if running without docker
	if os.Getenv("CONF_PATH") != "" {
		confPath = os.Getenv("CONF_PATH")
	}

	if os.Getenv("HAPROXY_PATH") != "" {
		haproxyPath = os.Getenv("HAPROXY_PATH")
	}

	s := rpc.NewServer()
	s.RegisterCodec(json.NewCodec(), "application/json")
	haproxy := new(Haproxy)
	// generate default haproxy.cfg
	log.Println("Generating default haproxy.cfg")
	haproxy.generateCfg()
	haproxy.StartHaproxy()
	s.RegisterService(haproxy, "")
	r := mux.NewRouter()
	r.Handle("/haproxy", s)
	http.ListenAndServe(":34015", r)

}
<commit_msg>proper formatting of cfg file<commit_after>package main

import (
	"html/template"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"strconv"

	"github.com/codeskyblue/go-sh"
	"github.com/gorilla/mux"
	"github.com/gorilla/rpc"
	"github.com/gorilla/rpc/json"
)

const frontendHeaderACL = `  acl is{{.ACL}} hdr_beg(host) {{index .HaproxyURLs #}}
`
const frontendPathACL = `  acl is{{.ACL}} path_beg -i {{index .HaproxyURLs #}}
`

const frontendUse = `  use_backend {{.Backend}} if is{{.ACL}}
`

// const frontendTmpl = `
// acl is{{.ACL}} hdr_beg(host) {{.HaproxyURL}}
// use_backend {{.Backend}} if is{{.ACL}}
// `
const backendTmpl = `##
backend {{.Backend}}
  server {{.Backend}} {{.Hostmachine}}:{{.Port}} check inter 10000
`

const defaultBackendTmpl = `##
  default_backend {{.Backend}}
`

var confPath = "/usr/local/etc/haproxy/conf"
var haproxyPath = "/usr/local/etc/haproxy"

// Service to be added
type Service struct {
	// ACL to be used
	ACL string
	// URL by which service will be called
	HaproxyURLs []string
	// Backend name
	Backend string
	// storefront-services-1.myntra.com
	Hostmachine string
	// Port on Hostmachine where service runs
	Port string
}

// Services ...
type Services struct {
	Services []Service
}

// Haproxy ...
type Haproxy int

// Result ...
type Result int

// Add a frontend and backend
func (h *Haproxy) Add(r *http.Request, services *Services, result *Result) error {
	var defaultBackend interface{}
	isDefaultBackedDefined := false
	for _, service := range services.Services {
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".backend").Run()
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".frontend").Run()
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".frontendtop").Run()
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".frontendbottom").Run()

		log.Printf("Add service %s:%s", service.HaproxyURLs, service.Port)
		data := struct {
			ACL         string
			HaproxyURLs []string
			Backend     string
			Hostmachine string
			Port        string
		}{
			strings.Title(service.ACL),
			service.HaproxyURLs,
			service.Backend,
			service.Hostmachine,
			service.Port,
		}

		frontendACLs := `##
`

		frontendType := ".frontend"

		for haproxyURLId, haproxyURL := range service.HaproxyURLs {
			if haproxyURL == "default_backend" {
				defaultBackend = data
				isDefaultBackedDefined = true
			} else if strings.EqualFold(haproxyURL, "top") || strings.EqualFold(haproxyURL, "bottom") {
				frontendType += haproxyURL
			} else if string(haproxyURL[0]) == "/" {
				frontendACLs += strings.Replace(frontendPathACL, "#", strconv.Itoa(haproxyURLId), -1)
			} else {
				frontendACLs += strings.Replace(frontendHeaderACL, "#", strconv.Itoa(haproxyURLId), -1)
			}
		}

		frontendTmpl := frontendACLs + frontendUse

		log.Println(frontendTmpl)

		// Generate frontend entry
		tmpl := template.Must(template.New("frontend").Parse(frontendTmpl))
		f, err := os.OpenFile(confPath+"/"+service.ACL+frontendType, os.O_CREATE|os.O_RDWR, 0777)
		if err != nil {
			*result = 0
			return err
		}
		// fill in the template
		err = tmpl.Execute(f, data)
		if err != nil {
			*result = 0
			return err
		}

		// Generate backend entry
		tmpl = template.Must(template.New("backend").Parse(backendTmpl))
		f, err = os.OpenFile(confPath+"/"+service.ACL+".backend", os.O_CREATE|os.O_RDWR, 0777)
		if err != nil {
			*result = 0
			return err
		}
		err = tmpl.Execute(f, data)
		if err != nil {
			*result = 0
			return err
		}
	}

	// Generate default_backend if needed
	if isDefaultBackedDefined {
		tmpl := template.Must(template.New("default_backend").Parse(defaultBackendTmpl))
		f, err := os.OpenFile(confPath+"/"+".default_backend", os.O_CREATE|os.O_RDWR, 0777)
		if err != nil {
			*result = 0
			return err
		}
		err = tmpl.Execute(f, defaultBackend)
		if err != nil {
			*result = 0
			return err
		}
	}

	//join all the configs
	if err := h.generateCfg(); err != nil {
		*result = 0
		return err
	}

	if err := h.ReloadHaproxy(); err != nil {
		*result = 0
		return err
	}

	*result = 1
	return nil

}

// Remove a frontend and backend
func (h *Haproxy) Remove(r *http.Request, services *Services, result *Result) error {
	for _, service := range services.Services {
		log.Printf("Remove service %s:%s", service.HaproxyURLs, service.Port)
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".backend").Run()
		sh.Command("rm", "-f", confPath+"/"+service.ACL+".frontend").Run()
	}

	if err := h.generateCfg(); err != nil {
		*result = 0
		return err
	}

	if err := h.ReloadHaproxy(); err != nil {
		*result = 0
		return err
	}

	*result = 1
	return nil

}

func (h *Haproxy) generateCfg() error {
	if _, err := os.Stat(haproxyPath + "/haproxy.cfg"); !os.IsNotExist(err) {
		currentTime := string(time.Now().Format("20060102150405"))
		err := os.Rename(haproxyPath+"/haproxy.cfg", haproxyPath+"/haproxy.cfg.BAK."+currentTime)
		if err != nil {
			return err
		}
	}

	var haproxyCfg []byte

	var partFunc = func(part string) {
		// walk all files in the directory
		filepath.Walk(confPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if !info.IsDir() && strings.HasSuffix(info.Name(), part) {
				b, err := ioutil.ReadFile(path)
				if err != nil {
					return err
				}
				haproxyCfg = append(haproxyCfg, b...)
			}
			return nil

		})
	}

	//append the configs in the following order
	parts := []string{".globalcfg", ".defaultcfg", ".frontendcfg", ".frontendtop", ".frontend", ".frontendbottom", ".default_backend", ".backend"}
	for i := range parts {
		partFunc(parts[i])
	}

	//write the file
	err := ioutil.WriteFile(haproxyPath+"/haproxy.cfg", haproxyCfg, 0777)
	if err != nil {
		return err
	}

	return nil

}

// StartHaproxy ...
func (h *Haproxy) StartHaproxy() error {
	// restart haproxy container
	session := sh.NewSession()
	err := session.Command("haproxy", "-p", "/var/run/haproxy.pid", "-f", "/usr/local/etc/haproxy/haproxy.cfg").Run()
	if err != nil {
		return err
	}
	return nil
}

// ReloadHaproxy ...
func (h *Haproxy) ReloadHaproxy() error {
	session := sh.NewSession()
	err := session.Command("/usr/bin/reload.sh").Run()
	if err != nil {
		return err
	}
	return nil
}

// Generate regenerates haproxy config from existing configs
func (h *Haproxy) Generate(r *http.Request, service *Service, result *Result) error {
	if err := h.generateCfg(); err != nil {
		*result = 0
		return err
	}

	*result = 1
	return nil

}

func main() {

	//if running without docker
	if os.Getenv("CONF_PATH") != "" {
		confPath = os.Getenv("CONF_PATH")
	}

	if os.Getenv("HAPROXY_PATH") != "" {
		haproxyPath = os.Getenv("HAPROXY_PATH")
	}

	s := rpc.NewServer()
	s.RegisterCodec(json.NewCodec(), "application/json")
	haproxy := new(Haproxy)
	// generate default haproxy.cfg
	log.Println("Generating default haproxy.cfg")
	haproxy.generateCfg()
	haproxy.StartHaproxy()
	s.RegisterService(haproxy, "")
	r := mux.NewRouter()
	r.Handle("/haproxy", s)
	http.ListenAndServe(":34015", r)

}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"log"
	"time"
)

// Question contains all the information from a question requested via HTTP,
// plus a field (filled by the reciever) for the time the request was received.
type Question struct {
	Qid          string    `schema:"qid" json:"qid"`           // The query ID
	Title        string    `schema:"title" json:"title"`       // The query text
	Body         string    `schema:"body" json:"body"`         // The body of the question (may be empty)
	Category     string    `schema:"category" json:"category"` // The category of the question (may be empty)
	ReceivedTime time.Time // The time this quesiton was recieved
}

// AsJSON returns a json representation of the Question structure
func (q *Question) AsJSON() string {
	b, err := json.Marshal(q)
	if err != nil {
		log.Println(err)
		return err.Error()
	}

	return string(b)
}

// String is a convieniencec returns the json representation of the Question structure
func (q *Question) String() string {
	return q.AsJSON()
}

// AnswerWrapper provides the xml wrapper for an answer. It's only needed at
// serilisation time
type AnswerWrapper struct {
	XMLName xml.Name `xml:"xml"`
	Answer  *Answer  `xml:"answer"`
}

// Answer provides the representation of an answer from a AnswerProcessor. It
// can be serilised to xml using xml.Marshal()
type Answer struct {
	XMLName   xml.Name `xml:"answer"`
	Answered  string   `xml:"answered,attr"`
	Pid       string   `xml:"pid,attr"`
	Qid       string   `xml:"qid,attr"`
	Time      int64    `xml:"time,attr"`
	Content   string   `xml:"content"`
	Resources string   `xml:"resources"`
}

func (a *AnswerWrapper) String() string {
	output, err := xml.MarshalIndent(a, "", "    ")
	if err != nil {
		log.Println(err)
		return err.Error()
	}

	return string(output)
}

// NewTimeOutAnswer provides the error response "answer" when a timeout has
// been hit. The timeout supplied is included in the response
func NewTimeOutAnswer(q *Question, timeout int) *Answer {
	return &Answer{
		Answered:  "no",
		Pid:       "demo-id-01",
		Qid:       q.Qid,
		Time:      int64(time.Since(q.ReceivedTime) / time.Millisecond),
		Content:   fmt.Sprintf("TIMEOUT after %d seconds: Title: %s;  Body: %s;  Category: %s", timeout, q.Title, q.Body, q.Category),
		Resources: "resource1,resource2",
	}
}
<commit_msg>Refactored the XML serialisation out of String() and into AsXML()<commit_after>package main

import (
	"encoding/json"
	"encoding/xml"
	"fmt"
	"log"
	"time"
)

// Question contains all the information from a question requested via HTTP,
// plus a field (filled by the reciever) for the time the request was received.
type Question struct {
	Qid          string    `schema:"qid" json:"qid"`           // The query ID
	Title        string    `schema:"title" json:"title"`       // The query text
	Body         string    `schema:"body" json:"body"`         // The body of the question (may be empty)
	Category     string    `schema:"category" json:"category"` // The category of the question (may be empty)
	ReceivedTime time.Time // The time this quesiton was recieved
}

// AsJSON returns a json representation of the Question structure
func (q *Question) AsJSON() string {
	b, err := json.Marshal(q)
	if err != nil {
		log.Println(err)
		return err.Error()
	}

	return string(b)
}

// String is a convieniencec returns the json representation of the Question structure
func (q *Question) String() string {
	return q.AsJSON()
}

// AnswerWrapper provides the xml wrapper for an answer. It's only needed at
// serilisation time
type AnswerWrapper struct {
	XMLName xml.Name `xml:"xml"`
	Answer  *Answer  `xml:"answer"`
}

// Answer provides the representation of an answer from a AnswerProcessor. It
// can be serilised to xml using xml.Marshal(), or the convienience method answer.AsXML()
type Answer struct {
	XMLName   xml.Name `xml:"answer"`
	Answered  string   `xml:"answered,attr"`
	Pid       string   `xml:"pid,attr"`
	Qid       string   `xml:"qid,attr"`
	Time      int64    `xml:"time,attr"`
	Content   string   `xml:"content"`
	Resources string   `xml:"resources"`
}

// AsXML returns a XML representation of the Answer structure, suitable for
// returning to the TREC question client
func (a *AnswerWrapper) AsXML() string {
	output, err := xml.MarshalIndent(a, "", "    ")
	if err != nil {
		log.Println(err)
		return err.Error()
	}

	return string(output)
}

func (a *AnswerWrapper) String() string {
	return a.AsXML()
}

// NewTimeOutAnswer provides the error response "answer" when a timeout has
// been hit. The timeout supplied is included in the response
func NewTimeOutAnswer(q *Question, timeout int) *Answer {
	return &Answer{
		Answered:  "no",
		Pid:       "demo-id-01",
		Qid:       q.Qid,
		Time:      int64(time.Since(q.ReceivedTime) / time.Millisecond),
		Content:   fmt.Sprintf("TIMEOUT after %d seconds: Title: %s;  Body: %s;  Category: %s", timeout, q.Title, q.Body, q.Category),
		Resources: "resource1,resource2",
	}
}
<|endoftext|>