<commit_before>// Copyright William Schwartz 2014. See the LICENSE file for more information.

package pigosat

import (
	"fmt"
	"reflect"
	"testing"
)

func TestlitArrayToSlice(t *testing.T) {
	// null pointers
	assertPanics(t, "null pointer", "litArrayToSlice",
		func() { litArrayToSlice(nil, 0) })
	assertPanics(t, "null pointer", "litArrayToSlice",
		func() { litArrayToSlice(nil, 1) })
	// not zero terminated
	badPtr := &cArray123[0]
	assertPanics(t, "not zero terminated - unpadded", "litArrayToSlice",
		func() { litArrayToSlice(badPtr, 3) })
	assertPanics(t, "not zero terminated - padded", "litArrayToSlice",
		func() { litArrayToSlice(badPtr, 2) })
	// zero length
	if ls := litArrayToSlice(&cZero, 0); len(ls) != 0 {
		t.Errorf("Test 0-length litArrayToSlice, maxLen==0: return value has length %d",
			len(ls))
	}
	if ls := litArrayToSlice(&cZero, 1); len(ls) != 0 {
		t.Errorf(
			"Test 0-length litArrayToSlice, maxLen>0: return value has length %d",
			len(ls))
	}
	// works correctly
	ptr := &cArray1230[0]
	expected := []Literal{1, 2, 3}
	if ls := litArrayToSlice(ptr, 3); len(ls) != 3 || !reflect.DeepEqual(ls, expected) {
		t.Errorf(
			"Test litArrayToSlice correct, maxLen==3: expected %v but got %v",
			expected, ls)
	}
	if ls := litArrayToSlice(ptr, 4); len(ls) != 3 || !reflect.DeepEqual(ls, expected) {
		t.Errorf(
			"Test litArrayToSlice correct, maxLen==4: expected %v but got %v",
			expected, ls)
	}
}

func TestAssumptionsSucceeding(t *testing.T) {
	successTests := []struct {
		assumpts  []Literal // The literals which should be assumed true/false
		solutions int       // The number of solutions that we expect to produce
	}{
		{[]Literal{1}, 10}, // We use formulaTests[0].formula below
		{[]Literal{2}, 9},
		{[]Literal{3}, 6},
		{[]Literal{4, 5}, 4},
	}

	for i, at := range successTests {
		p, _ := New(nil)
		p.AddClauses(formulaTests[0].formula)

		count := 0
		for ; ; count++ {
			for _, lit := range at.assumpts {
				p.Assume(lit)
			}
			status, sol := p.Solve()
			if status != Satisfiable {
				break
			}

			// All the UNSAT methods should give zero answers.
			if p.FailedAssumption(at.assumpts[0]) {
				t.Errorf("Test %d: FailedAssumption: expected %v not to be failed", i, at.assumpts[0])
			}
			if r := p.FailedAssumptions(); !reflect.DeepEqual(r, []Literal{}) {
				t.Errorf("Test %d: FailedAssumptions: expected [], got %v", i, r)
			}

			p.BlockSolution(sol)
		}
		if count != at.solutions {
			t.Errorf("Test %d: Expected %d solution(s) for assumptions %v; got %d",
				i, at.solutions, at.assumpts, count)
		}
	}
}

func TestAssumptionsFailing(t *testing.T) {
	p, _ := New(nil)
	p.AddClauses(formulaTests[0].formula)
	p.Assume(3)
	p.Assume(4)
	p.Assume(5)
	p.Solve()

	failed := []Literal{3, 4}
	if actual := p.FailedAssumptions(); !reflect.DeepEqual(failed, actual) {
		t.Errorf("Expected failed assumptions %v != %v actual", failed, actual)
	}
	for _, f := range failed {
		if !p.FailedAssumption(f) {
			t.Errorf("Expected literal %v to be a failed assumption", f)
		}
	}

	i, wanted := 0, [][]Literal{{3, 5}, {3, 5}, {5, 4}, {}}
	for a := p.MaxSatisfiableAssumptions(); len(a) > 0; i++ {
		if !reflect.DeepEqual(a, wanted[i]) {
			t.Errorf("(Next)MaxSat'Assumpt's: Got %v wanted %v", a, wanted[i])
		}
		a = p.NextMaxSatisfiableAssumptions()
	}
	if a := p.MaxSatisfiableAssumptions(); !reflect.DeepEqual(a, []Literal{}) {
		t.Errorf("MaxSatisfiableAssumptions: CNF inconsistent, so wanted []Literal{}, but got %v", a)
	}
}

// ExampleAssume demonstrates how to use Assume and related methods.
func ExamplePigosat_Assume() {
	var formula = Formula{{1, 2, 3}, {1, 2}, {2, 3}}
	p, _ := New(nil)
	p.AddClauses(formula)
	fmt.Println("Formula:", formula)
	status, solution := p.Solve()
	fmt.Println("No assumptions:", status, "solution ==", solution)

	// Satisfiable assumptions
	fmt.Println()
	fmt.Println("**** SATISFIABLE ASSUMPTIONS ****")
	p.Assume(1)
	p.Assume(-2)
	// Assumptions do not change the number of clauses.
	fmt.Println("Assume  1, -2 : Number of clauses:", p.AddedOriginalClauses())
	status, solution = p.Solve()
	fmt.Println("               ", status, "solution ==", solution)

	// Calls to p.AddClauses or p.Assume cancel assumptions 1 and -2
	// immediately, or a second call to p.Solve also cancels the assumptions.
	p.Assume(-3)
	status, solution = p.Solve()
	fmt.Println("Assume      -3:", status, "solution ==", solution)

	// Unsatisfiable assumptions
	fmt.Println()
	fmt.Println("**** UNSATISFIABLE ASSUMPTIONS ****")
	p.Assume(-1)                 // assume unit clause Clause{-1, 0}
	p.Assume(-2)                 // assume unit clause Clause{-2, 0}
	status, solution = p.Solve() // assumes -1 and -2 hold
	fmt.Println("Assume -1, -2 :", status, "solution ==", solution)
	// Assumptions -1 and -2 caused unsatisfiability.
	fmt.Println("                Failed assumptions:", p.FailedAssumptions())
	fmt.Println("                Assumption -1 failed:", p.FailedAssumption(-1))
	fmt.Println("                Assumption -2 failed:", p.FailedAssumption(-2))
	// Not every subset of the assumptions causes unsatisfiability.
	// p.MaxSatisfiableAssumptions would return the same as
	// p.NextMaxSatisfiableAssumptions, but the return value wouldn't change
	// with each call like it does with p.NextMaxSatisfiableAssumptions.
	fmt.Println("                Maximal satisfiable subset of assumptions 1:", p.NextMaxSatisfiableAssumptions())
	fmt.Println("                Maximal satisfiable subset of assumptions 2:", p.NextMaxSatisfiableAssumptions())
	fmt.Println("                Maximal satisfiable subset of assumptions 3:", p.NextMaxSatisfiableAssumptions())
	// p.NextMaxSatisfiableAssumptions does add clauses.
	fmt.Println("                Number of clauses:", p.AddedOriginalClauses())

	// Unknown status
	// Assumptions are valid but p.Solve returns no Solution assignment. The
	// solver knowns the status is Unknown until a call to p.Assume,
	// p.AddClauses, or p.Solve resets the assumptions.

	// Output:
	// Formula: [[1 2 3] [1 2] [2 3]]
	// No assumptions: Satisfiable solution == {1:true , 2:true , 3:true}

	// **** SATISFIABLE ASSUMPTIONS ****
	// Assume  1, -2 : Number of clauses: 3
	//                 Satisfiable solution == {1:true , 2:false, 3:true}
	// Assume      -3: Satisfiable solution == {1:true , 2:true , 3:false}

	// **** UNSATISFIABLE ASSUMPTIONS ****
	// Assume -1, -2 : Unsatisfiable solution == {}
	//                 Failed assumptions: [-1 -2]
	//                 Assumption -1 failed: true
	//                 Assumption -2 failed: true
	//                 Maximal satisfiable subset of assumptions 1: [-1]
	//                 Maximal satisfiable subset of assumptions 2: [-2]
	//                 Maximal satisfiable subset of assumptions 3: []
	//                 Number of clauses: 5
}
<commit_msg>Add test that demonstrates crash in use of Assume and FailedAssumption after UNSAT<commit_after>// Copyright William Schwartz 2014. See the LICENSE file for more information.

package pigosat

import (
	"fmt"
	"reflect"
	"testing"
)

func TestlitArrayToSlice(t *testing.T) {
	// null pointers
	assertPanics(t, "null pointer", "litArrayToSlice",
		func() { litArrayToSlice(nil, 0) })
	assertPanics(t, "null pointer", "litArrayToSlice",
		func() { litArrayToSlice(nil, 1) })
	// not zero terminated
	badPtr := &cArray123[0]
	assertPanics(t, "not zero terminated - unpadded", "litArrayToSlice",
		func() { litArrayToSlice(badPtr, 3) })
	assertPanics(t, "not zero terminated - padded", "litArrayToSlice",
		func() { litArrayToSlice(badPtr, 2) })
	// zero length
	if ls := litArrayToSlice(&cZero, 0); len(ls) != 0 {
		t.Errorf("Test 0-length litArrayToSlice, maxLen==0: return value has length %d",
			len(ls))
	}
	if ls := litArrayToSlice(&cZero, 1); len(ls) != 0 {
		t.Errorf(
			"Test 0-length litArrayToSlice, maxLen>0: return value has length %d",
			len(ls))
	}
	// works correctly
	ptr := &cArray1230[0]
	expected := []Literal{1, 2, 3}
	if ls := litArrayToSlice(ptr, 3); len(ls) != 3 || !reflect.DeepEqual(ls, expected) {
		t.Errorf(
			"Test litArrayToSlice correct, maxLen==3: expected %v but got %v",
			expected, ls)
	}
	if ls := litArrayToSlice(ptr, 4); len(ls) != 3 || !reflect.DeepEqual(ls, expected) {
		t.Errorf(
			"Test litArrayToSlice correct, maxLen==4: expected %v but got %v",
			expected, ls)
	}
}

func TestAssumptionsSucceeding(t *testing.T) {
	successTests := []struct {
		assumpts  []Literal // The literals which should be assumed true/false
		solutions int       // The number of solutions that we expect to produce
	}{
		{[]Literal{1}, 10}, // We use formulaTests[0].formula below
		{[]Literal{2}, 9},
		{[]Literal{3}, 6},
		{[]Literal{4, 5}, 4},
	}

	for i, at := range successTests {
		p, _ := New(nil)
		p.AddClauses(formulaTests[0].formula)

		count := 0
		for ; ; count++ {
			for _, lit := range at.assumpts {
				p.Assume(lit)
			}
			status, sol := p.Solve()
			if status != Satisfiable {
				break
			}

			// All the UNSAT methods should give zero answers.
			if p.FailedAssumption(at.assumpts[0]) {
				t.Errorf("Test %d: FailedAssumption: expected %v not to be failed", i, at.assumpts[0])
			}
			if r := p.FailedAssumptions(); !reflect.DeepEqual(r, []Literal{}) {
				t.Errorf("Test %d: FailedAssumptions: expected [], got %v", i, r)
			}

			p.BlockSolution(sol)
		}
		if count != at.solutions {
			t.Errorf("Test %d: Expected %d solution(s) for assumptions %v; got %d",
				i, at.solutions, at.assumpts, count)
		}
	}
}

func TestAssumptionsFailing(t *testing.T) {
	p, _ := New(nil)
	p.AddClauses(formulaTests[0].formula)
	p.Assume(3)
	p.Assume(4)
	p.Assume(5)
	p.Solve()

	failed := []Literal{3, 4}
	if actual := p.FailedAssumptions(); !reflect.DeepEqual(failed, actual) {
		t.Errorf("Expected failed assumptions %v != %v actual", failed, actual)
	}
	for _, f := range failed {
		if !p.FailedAssumption(f) {
			t.Errorf("Expected literal %v to be a failed assumption", f)
		}
	}

	i, wanted := 0, [][]Literal{{3, 5}, {3, 5}, {5, 4}, {}}
	for a := p.MaxSatisfiableAssumptions(); len(a) > 0; i++ {
		if !reflect.DeepEqual(a, wanted[i]) {
			t.Errorf("(Next)MaxSat'Assumpt's: Got %v wanted %v", a, wanted[i])
		}
		a = p.NextMaxSatisfiableAssumptions()
	}
	if a := p.MaxSatisfiableAssumptions(); !reflect.DeepEqual(a, []Literal{}) {
		t.Errorf("MaxSatisfiableAssumptions: CNF inconsistent, so wanted []Literal{}, but got %v", a)
	}
}

// TestAssumptionsNoCrash tests that where Picosat crashes when using the
// *assumption* methods when the solver's state is not UNSAT, Pigosat should
// just give a simple zero answer.
func TestAssumptionsNoCrash(t *testing.T) {
	p, _ := New(nil)
	p.AddClauses(formulaTests[0].formula)
	p.Assume(3)
	p.Assume(4)
	p.Assume(5)
	p.Solve()

	if r := p.Res(); r != Unsatisfiable {
		t.Fatalf("Expected %v, got %v", Unsatisfiable, r)
	}
	if !p.FailedAssumption(3) {
		t.Fatalf("Expected literal 3 to be a failed assumption")
	}
	if r := p.Res(); r != Unsatisfiable {
		t.Fatalf("Expected %v, got %v", Unsatisfiable, r)
	}

	p.Assume(3)
	if r := p.Res(); r != Unsatisfiable {
		t.Fatalf("Expected %v, got %v", Unsatisfiable, r)
	}

	if p.FailedAssumption(3) {
		t.Errorf("FailedAssumption: Expected false, got true")
	}
	if r := p.FailedAssumptions(); len(r) != 0 {
		t.Errorf("FailedAssumptions: Expected []Literal{}, got %v", r)
	}
}

// ExampleAssume demonstrates how to use Assume and related methods.
func ExamplePigosat_Assume() {
	var formula = Formula{{1, 2, 3}, {1, 2}, {2, 3}}
	p, _ := New(nil)
	p.AddClauses(formula)
	fmt.Println("Formula:", formula)
	status, solution := p.Solve()
	fmt.Println("No assumptions:", status, "solution ==", solution)

	// Satisfiable assumptions
	fmt.Println()
	fmt.Println("**** SATISFIABLE ASSUMPTIONS ****")
	p.Assume(1)
	p.Assume(-2)
	// Assumptions do not change the number of clauses.
	fmt.Println("Assume  1, -2 : Number of clauses:", p.AddedOriginalClauses())
	status, solution = p.Solve()
	fmt.Println("               ", status, "solution ==", solution)

	// Calls to p.AddClauses or p.Assume cancel assumptions 1 and -2
	// immediately, or a second call to p.Solve also cancels the assumptions.
	p.Assume(-3)
	status, solution = p.Solve()
	fmt.Println("Assume      -3:", status, "solution ==", solution)

	// Unsatisfiable assumptions
	fmt.Println()
	fmt.Println("**** UNSATISFIABLE ASSUMPTIONS ****")
	p.Assume(-1)                 // assume unit clause Clause{-1, 0}
	p.Assume(-2)                 // assume unit clause Clause{-2, 0}
	status, solution = p.Solve() // assumes -1 and -2 hold
	fmt.Println("Assume -1, -2 :", status, "solution ==", solution)
	// Assumptions -1 and -2 caused unsatisfiability.
	fmt.Println("                Failed assumptions:", p.FailedAssumptions())
	fmt.Println("                Assumption -1 failed:", p.FailedAssumption(-1))
	fmt.Println("                Assumption -2 failed:", p.FailedAssumption(-2))
	// Not every subset of the assumptions causes unsatisfiability.
	// p.MaxSatisfiableAssumptions would return the same as
	// p.NextMaxSatisfiableAssumptions, but the return value wouldn't change
	// with each call like it does with p.NextMaxSatisfiableAssumptions.
	fmt.Println("                Maximal satisfiable subset of assumptions 1:", p.NextMaxSatisfiableAssumptions())
	fmt.Println("                Maximal satisfiable subset of assumptions 2:", p.NextMaxSatisfiableAssumptions())
	fmt.Println("                Maximal satisfiable subset of assumptions 3:", p.NextMaxSatisfiableAssumptions())
	// p.NextMaxSatisfiableAssumptions does add clauses.
	fmt.Println("                Number of clauses:", p.AddedOriginalClauses())

	// Unknown status
	// Assumptions are valid but p.Solve returns no Solution assignment. The
	// solver knowns the status is Unknown until a call to p.Assume,
	// p.AddClauses, or p.Solve resets the assumptions.

	// Output:
	// Formula: [[1 2 3] [1 2] [2 3]]
	// No assumptions: Satisfiable solution == {1:true , 2:true , 3:true}

	// **** SATISFIABLE ASSUMPTIONS ****
	// Assume  1, -2 : Number of clauses: 3
	//                 Satisfiable solution == {1:true , 2:false, 3:true}
	// Assume      -3: Satisfiable solution == {1:true , 2:true , 3:false}

	// **** UNSATISFIABLE ASSUMPTIONS ****
	// Assume -1, -2 : Unsatisfiable solution == {}
	//                 Failed assumptions: [-1 -2]
	//                 Assumption -1 failed: true
	//                 Assumption -2 failed: true
	//                 Maximal satisfiable subset of assumptions 1: [-1]
	//                 Maximal satisfiable subset of assumptions 2: [-2]
	//                 Maximal satisfiable subset of assumptions 3: []
	//                 Number of clauses: 5
}
<|endoftext|>


<commit_before>package astilectron

import (
	"net"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/asticode/go-astilog"
	"github.com/asticode/go-astitools/context"
	"github.com/asticode/go-astitools/exec"
	"github.com/pkg/errors"
)

// Versions
const (
	VersionAstilectron = "0.15.0"
	VersionElectron    = "1.8.1"
)

// Misc vars
var (
	validOSes = map[string]bool{
		"darwin":  true,
		"linux":   true,
		"windows": true,
	}
)

// App event names
const (
	EventNameAppClose         = "app.close"
	EventNameAppCmdQuit       = "app.cmd.quit" // Sends an event to Electron to properly quit the app
	EventNameAppCmdStop       = "app.cmd.stop" // Cancel the context which results in exiting abruptly Electron's app
	EventNameAppCrash         = "app.crash"
	EventNameAppErrorAccept   = "app.error.accept"
	EventNameAppEventReady    = "app.event.ready"
	EventNameAppNoAccept      = "app.no.accept"
	EventNameAppTooManyAccept = "app.too.many.accept"
)

// Astilectron represents an object capable of interacting with Astilectron
type Astilectron struct {
	canceller    *asticontext.Canceller
	channelQuit  chan bool
	closeOnce    sync.Once
	dispatcher   *dispatcher
	displayPool  *displayPool
	identifier   *identifier
	listener     net.Listener
	options      Options
	paths        *Paths
	provisioner  Provisioner
	reader       *reader
	stderrWriter *astiexec.StdWriter
	stdoutWriter *astiexec.StdWriter
	writer       *writer
}

// Options represents Astilectron options
type Options struct {
	AppName            string
	AppIconDarwinPath  string // Darwin systems requires a specific .icns file
	AppIconDefaultPath string
	BaseDirectoryPath  string
	ElectronSwitches   []string
}

// New creates a new Astilectron instance
func New(o Options) (a *Astilectron, err error) {
	// Validate the OS
	if !IsValidOS(runtime.GOOS) {
		err = errors.Wrapf(err, "OS %s is invalid")
		return
	}

	// Init
	a = &Astilectron{
		canceller:   asticontext.NewCanceller(),
		channelQuit: make(chan bool),
		dispatcher:  newDispatcher(),
		displayPool: newDisplayPool(),
		identifier:  newIdentifier(),
		options:     o,
		provisioner: DefaultProvisioner,
	}

	// Set paths
	if a.paths, err = newPaths(runtime.GOOS, runtime.GOARCH, o); err != nil {
		err = errors.Wrap(err, "creating new paths failed")
		return
	}

	// Add default listeners
	a.On(EventNameAppCmdStop, func(e Event) (deleteListener bool) {
		a.Stop()
		return
	})
	a.On(EventNameDisplayEventAdded, func(e Event) (deleteListener bool) {
		a.displayPool.update(e.Displays)
		return
	})
	a.On(EventNameDisplayEventMetricsChanged, func(e Event) (deleteListener bool) {
		a.displayPool.update(e.Displays)
		return
	})
	a.On(EventNameDisplayEventRemoved, func(e Event) (deleteListener bool) {
		a.displayPool.update(e.Displays)
		return
	})
	return
}

// IsValidOS validates the OS
func IsValidOS(os string) (ok bool) {
	_, ok = validOSes[os]
	return
}

// SetProvisioner sets the provisioner
func (a *Astilectron) SetProvisioner(p Provisioner) *Astilectron {
	a.provisioner = p
	return a
}

// On implements the Listenable interface
func (a *Astilectron) On(eventName string, l Listener) {
	a.dispatcher.addListener(mainTargetID, eventName, l)
}

// Start starts Astilectron
func (a *Astilectron) Start() (err error) {
	// Log
	astilog.Debug("Starting...")

	// Provision
	if err = a.provision(); err != nil {
		return errors.Wrap(err, "provisioning failed")
	}

	// Unfortunately communicating with Electron through stdin/stdout doesn't work on Windows so all communications
	// will be done through TCP
	if err = a.listenTCP(); err != nil {
		return errors.Wrap(err, "listening failed")
	}

	// Execute
	if err = a.execute(); err != nil {
		err = errors.Wrap(err, "executing failed")
		return
	}
	return
}

// provision provisions Astilectron
func (a *Astilectron) provision() error {
	astilog.Debug("Provisioning...")
	var ctx, _ = a.canceller.NewContext()
	return a.provisioner.Provision(ctx, a.options.AppName, runtime.GOOS, runtime.GOARCH, *a.paths)
}

// listenTCP listens to the first TCP connection coming its way (this should be Astilectron)
func (a *Astilectron) listenTCP() (err error) {
	// Log
	astilog.Debug("Listening...")

	// Listen
	if a.listener, err = net.Listen("tcp", "127.0.0.1:"); err != nil {
		return errors.Wrap(err, "tcp net.Listen failed")
	}

	// Check a connection has been accepted quickly enough
	var chanAccepted = make(chan bool)
	go a.watchNoAccept(30*time.Second, chanAccepted)

	// Accept connections
	go a.acceptTCP(chanAccepted)
	return
}

// watchNoAccept checks whether a TCP connection is accepted quickly enough
func (a *Astilectron) watchNoAccept(timeout time.Duration, chanAccepted chan bool) {
	var t = time.NewTimer(timeout)
	defer t.Stop()
	for {
		select {
		case <-chanAccepted:
			return
		case <-t.C:
			astilog.Errorf("No TCP connection has been accepted in the past %s", timeout)
			a.dispatcher.dispatch(Event{Name: EventNameAppNoAccept, TargetID: mainTargetID})
			a.dispatcher.dispatch(Event{Name: EventNameAppCmdStop, TargetID: mainTargetID})
			return
		}
	}
}

// watchAcceptTCP accepts TCP connections
func (a *Astilectron) acceptTCP(chanAccepted chan bool) {
	for i := 0; i <= 1; i++ {
		// Accept
		var conn net.Conn
		var err error
		if conn, err = a.listener.Accept(); err != nil {
			astilog.Errorf("%s while TCP accepting", err)
			a.dispatcher.dispatch(Event{Name: EventNameAppErrorAccept, TargetID: mainTargetID})
			a.dispatcher.dispatch(Event{Name: EventNameAppCmdStop, TargetID: mainTargetID})
			return
		}

		// We only accept the first connection which should be Astilectron, close the next one and stop
		// the app
		if i > 0 {
			astilog.Errorf("Too many TCP connections")
			a.dispatcher.dispatch(Event{Name: EventNameAppTooManyAccept, TargetID: mainTargetID})
			a.dispatcher.dispatch(Event{Name: EventNameAppCmdStop, TargetID: mainTargetID})
			conn.Close()
			return
		}

		// Let the timer know a connection has been accepted
		chanAccepted <- true

		// Create reader and writer
		a.writer = newWriter(conn)
		a.reader = newReader(a.dispatcher, conn)
		go a.reader.read()
	}
}

// execute executes Astilectron in Electron
func (a *Astilectron) execute() (err error) {
	// Log
	astilog.Debug("Executing...")

	// Create command
	var ctx, _ = a.canceller.NewContext()
	var cmd = exec.CommandContext(ctx, a.paths.AppExecutable(), append([]string{a.paths.AstilectronApplication(), a.listener.Addr().String()}, a.options.ElectronSwitches...)...)
	a.stderrWriter = astiexec.NewStdWriter(func(i []byte) { astilog.Debugf("Stderr says: %s", i) })
	a.stdoutWriter = astiexec.NewStdWriter(func(i []byte) { astilog.Debugf("Stdout says: %s", i) })
	cmd.Stderr = a.stderrWriter
	cmd.Stdout = a.stdoutWriter

	// Execute command
	if err = a.executeCmd(cmd); err != nil {
		return errors.Wrap(err, "executing cmd failed")
	}
	return
}

// executeCmd executes the command
func (a *Astilectron) executeCmd(cmd *exec.Cmd) (err error) {
	var e = synchronousFunc(a.canceller, a, func() {
		// Start command
		astilog.Debugf("Starting cmd %s", strings.Join(cmd.Args, " "))
		if err = cmd.Start(); err != nil {
			err = errors.Wrapf(err, "starting cmd %s failed", strings.Join(cmd.Args, " "))
			return
		}

		// Watch command
		go a.watchCmd(cmd)
	}, EventNameAppEventReady)

	// Update display pool
	if e.Displays != nil {
		a.displayPool.update(e.Displays)
	}
	return
}

// watchCmd watches the cmd execution
func (a *Astilectron) watchCmd(cmd *exec.Cmd) {
	// Wait
	cmd.Wait()

	// Check the canceller to check whether it was a crash
	if !a.canceller.Cancelled() {
		astilog.Debug("App has crashed")
		a.dispatcher.dispatch(Event{Name: EventNameAppCrash, TargetID: mainTargetID})
	} else {
		astilog.Debug("App has closed")
		a.dispatcher.dispatch(Event{Name: EventNameAppClose, TargetID: mainTargetID})
	}
	a.dispatcher.dispatch(Event{Name: EventNameAppCmdStop, TargetID: mainTargetID})
}

// Close closes Astilectron properly
func (a *Astilectron) Close() {
	astilog.Debug("Closing...")
	a.canceller.Cancel()
	if a.listener != nil {
		a.listener.Close()
	}
	if a.reader != nil {
		a.reader.close()
	}
	if a.stderrWriter != nil {
		a.stderrWriter.Close()
	}
	if a.stdoutWriter != nil {
		a.stdoutWriter.Close()
	}
	if a.writer != nil {
		a.writer.close()
	}
}

// HandleSignals handles signals
func (a *Astilectron) HandleSignals() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGABRT, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM)
	go func() {
		for sig := range ch {
			astilog.Debugf("Received signal %s", sig)
			a.Stop()
		}
	}()
}

// Stop orders Astilectron to stop
func (a *Astilectron) Stop() {
	astilog.Debug("Stopping...")
	a.canceller.Cancel()
	a.closeOnce.Do(func() {
		close(a.channelQuit)
	})
}

// Wait is a blocking pattern
func (a *Astilectron) Wait() {
	<-a.channelQuit
}

// Quit quits the app
func (a *Astilectron) Quit() error {
	return a.writer.write(Event{Name: EventNameAppCmdQuit})
}

// Paths returns the paths
func (a *Astilectron) Paths() Paths {
	return *a.paths
}

// Displays returns the displays
func (a *Astilectron) Displays() []*Display {
	return a.displayPool.all()
}

// PrimaryDisplay returns the primary display
func (a *Astilectron) PrimaryDisplay() *Display {
	return a.displayPool.primary()
}

// NewMenu creates a new app menu
func (a *Astilectron) NewMenu(i []*MenuItemOptions) *Menu {
	return newMenu(nil, mainTargetID, i, a.canceller, a.dispatcher, a.identifier, a.writer)
}

// NewWindow creates a new window
func (a *Astilectron) NewWindow(url string, o *WindowOptions) (*Window, error) {
	return newWindow(a.options, url, o, a.canceller, a.dispatcher, a.identifier, a.writer)
}

// NewWindowInDisplay creates a new window in a specific display
// This overrides the center attribute
func (a *Astilectron) NewWindowInDisplay(d *Display, url string, o *WindowOptions) (*Window, error) {
	if o.X != nil {
		*o.X += d.Bounds().X
	} else {
		o.X = PtrInt(d.Bounds().X)
	}
	if o.Y != nil {
		*o.Y += d.Bounds().Y
	} else {
		o.Y = PtrInt(d.Bounds().Y)
	}
	return newWindow(a.options, url, o, a.canceller, a.dispatcher, a.identifier, a.writer)
}

// NewTray creates a new tray
func (a *Astilectron) NewTray(o *TrayOptions) *Tray {
	return newTray(o, a.canceller, a.dispatcher, a.identifier, a.writer)
}
<commit_msg>update astilectron version<commit_after>package astilectron

import (
	"net"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/asticode/go-astilog"
	"github.com/asticode/go-astitools/context"
	"github.com/asticode/go-astitools/exec"
	"github.com/pkg/errors"
)

// Versions
const (
	VersionAstilectron = "0.16.0"
	VersionElectron    = "1.8.1"
)

// Misc vars
var (
	validOSes = map[string]bool{
		"darwin":  true,
		"linux":   true,
		"windows": true,
	}
)

// App event names
const (
	EventNameAppClose         = "app.close"
	EventNameAppCmdQuit       = "app.cmd.quit" // Sends an event to Electron to properly quit the app
	EventNameAppCmdStop       = "app.cmd.stop" // Cancel the context which results in exiting abruptly Electron's app
	EventNameAppCrash         = "app.crash"
	EventNameAppErrorAccept   = "app.error.accept"
	EventNameAppEventReady    = "app.event.ready"
	EventNameAppNoAccept      = "app.no.accept"
	EventNameAppTooManyAccept = "app.too.many.accept"
)

// Astilectron represents an object capable of interacting with Astilectron
type Astilectron struct {
	canceller    *asticontext.Canceller
	channelQuit  chan bool
	closeOnce    sync.Once
	dispatcher   *dispatcher
	displayPool  *displayPool
	identifier   *identifier
	listener     net.Listener
	options      Options
	paths        *Paths
	provisioner  Provisioner
	reader       *reader
	stderrWriter *astiexec.StdWriter
	stdoutWriter *astiexec.StdWriter
	writer       *writer
}

// Options represents Astilectron options
type Options struct {
	AppName            string
	AppIconDarwinPath  string // Darwin systems requires a specific .icns file
	AppIconDefaultPath string
	BaseDirectoryPath  string
	ElectronSwitches   []string
}

// New creates a new Astilectron instance
func New(o Options) (a *Astilectron, err error) {
	// Validate the OS
	if !IsValidOS(runtime.GOOS) {
		err = errors.Wrapf(err, "OS %s is invalid")
		return
	}

	// Init
	a = &Astilectron{
		canceller:   asticontext.NewCanceller(),
		channelQuit: make(chan bool),
		dispatcher:  newDispatcher(),
		displayPool: newDisplayPool(),
		identifier:  newIdentifier(),
		options:     o,
		provisioner: DefaultProvisioner,
	}

	// Set paths
	if a.paths, err = newPaths(runtime.GOOS, runtime.GOARCH, o); err != nil {
		err = errors.Wrap(err, "creating new paths failed")
		return
	}

	// Add default listeners
	a.On(EventNameAppCmdStop, func(e Event) (deleteListener bool) {
		a.Stop()
		return
	})
	a.On(EventNameDisplayEventAdded, func(e Event) (deleteListener bool) {
		a.displayPool.update(e.Displays)
		return
	})
	a.On(EventNameDisplayEventMetricsChanged, func(e Event) (deleteListener bool) {
		a.displayPool.update(e.Displays)
		return
	})
	a.On(EventNameDisplayEventRemoved, func(e Event) (deleteListener bool) {
		a.displayPool.update(e.Displays)
		return
	})
	return
}

// IsValidOS validates the OS
func IsValidOS(os string) (ok bool) {
	_, ok = validOSes[os]
	return
}

// SetProvisioner sets the provisioner
func (a *Astilectron) SetProvisioner(p Provisioner) *Astilectron {
	a.provisioner = p
	return a
}

// On implements the Listenable interface
func (a *Astilectron) On(eventName string, l Listener) {
	a.dispatcher.addListener(mainTargetID, eventName, l)
}

// Start starts Astilectron
func (a *Astilectron) Start() (err error) {
	// Log
	astilog.Debug("Starting...")

	// Provision
	if err = a.provision(); err != nil {
		return errors.Wrap(err, "provisioning failed")
	}

	// Unfortunately communicating with Electron through stdin/stdout doesn't work on Windows so all communications
	// will be done through TCP
	if err = a.listenTCP(); err != nil {
		return errors.Wrap(err, "listening failed")
	}

	// Execute
	if err = a.execute(); err != nil {
		err = errors.Wrap(err, "executing failed")
		return
	}
	return
}

// provision provisions Astilectron
func (a *Astilectron) provision() error {
	astilog.Debug("Provisioning...")
	var ctx, _ = a.canceller.NewContext()
	return a.provisioner.Provision(ctx, a.options.AppName, runtime.GOOS, runtime.GOARCH, *a.paths)
}

// listenTCP listens to the first TCP connection coming its way (this should be Astilectron)
func (a *Astilectron) listenTCP() (err error) {
	// Log
	astilog.Debug("Listening...")

	// Listen
	if a.listener, err = net.Listen("tcp", "127.0.0.1:"); err != nil {
		return errors.Wrap(err, "tcp net.Listen failed")
	}

	// Check a connection has been accepted quickly enough
	var chanAccepted = make(chan bool)
	go a.watchNoAccept(30*time.Second, chanAccepted)

	// Accept connections
	go a.acceptTCP(chanAccepted)
	return
}

// watchNoAccept checks whether a TCP connection is accepted quickly enough
func (a *Astilectron) watchNoAccept(timeout time.Duration, chanAccepted chan bool) {
	var t = time.NewTimer(timeout)
	defer t.Stop()
	for {
		select {
		case <-chanAccepted:
			return
		case <-t.C:
			astilog.Errorf("No TCP connection has been accepted in the past %s", timeout)
			a.dispatcher.dispatch(Event{Name: EventNameAppNoAccept, TargetID: mainTargetID})
			a.dispatcher.dispatch(Event{Name: EventNameAppCmdStop, TargetID: mainTargetID})
			return
		}
	}
}

// watchAcceptTCP accepts TCP connections
func (a *Astilectron) acceptTCP(chanAccepted chan bool) {
	for i := 0; i <= 1; i++ {
		// Accept
		var conn net.Conn
		var err error
		if conn, err = a.listener.Accept(); err != nil {
			astilog.Errorf("%s while TCP accepting", err)
			a.dispatcher.dispatch(Event{Name: EventNameAppErrorAccept, TargetID: mainTargetID})
			a.dispatcher.dispatch(Event{Name: EventNameAppCmdStop, TargetID: mainTargetID})
			return
		}

		// We only accept the first connection which should be Astilectron, close the next one and stop
		// the app
		if i > 0 {
			astilog.Errorf("Too many TCP connections")
			a.dispatcher.dispatch(Event{Name: EventNameAppTooManyAccept, TargetID: mainTargetID})
			a.dispatcher.dispatch(Event{Name: EventNameAppCmdStop, TargetID: mainTargetID})
			conn.Close()
			return
		}

		// Let the timer know a connection has been accepted
		chanAccepted <- true

		// Create reader and writer
		a.writer = newWriter(conn)
		a.reader = newReader(a.dispatcher, conn)
		go a.reader.read()
	}
}

// execute executes Astilectron in Electron
func (a *Astilectron) execute() (err error) {
	// Log
	astilog.Debug("Executing...")

	// Create command
	var ctx, _ = a.canceller.NewContext()
	var cmd = exec.CommandContext(ctx, a.paths.AppExecutable(), append([]string{a.paths.AstilectronApplication(), a.listener.Addr().String()}, a.options.ElectronSwitches...)...)
	a.stderrWriter = astiexec.NewStdWriter(func(i []byte) { astilog.Debugf("Stderr says: %s", i) })
	a.stdoutWriter = astiexec.NewStdWriter(func(i []byte) { astilog.Debugf("Stdout says: %s", i) })
	cmd.Stderr = a.stderrWriter
	cmd.Stdout = a.stdoutWriter

	// Execute command
	if err = a.executeCmd(cmd); err != nil {
		return errors.Wrap(err, "executing cmd failed")
	}
	return
}

// executeCmd executes the command
func (a *Astilectron) executeCmd(cmd *exec.Cmd) (err error) {
	var e = synchronousFunc(a.canceller, a, func() {
		// Start command
		astilog.Debugf("Starting cmd %s", strings.Join(cmd.Args, " "))
		if err = cmd.Start(); err != nil {
			err = errors.Wrapf(err, "starting cmd %s failed", strings.Join(cmd.Args, " "))
			return
		}

		// Watch command
		go a.watchCmd(cmd)
	}, EventNameAppEventReady)

	// Update display pool
	if e.Displays != nil {
		a.displayPool.update(e.Displays)
	}
	return
}

// watchCmd watches the cmd execution
func (a *Astilectron) watchCmd(cmd *exec.Cmd) {
	// Wait
	cmd.Wait()

	// Check the canceller to check whether it was a crash
	if !a.canceller.Cancelled() {
		astilog.Debug("App has crashed")
		a.dispatcher.dispatch(Event{Name: EventNameAppCrash, TargetID: mainTargetID})
	} else {
		astilog.Debug("App has closed")
		a.dispatcher.dispatch(Event{Name: EventNameAppClose, TargetID: mainTargetID})
	}
	a.dispatcher.dispatch(Event{Name: EventNameAppCmdStop, TargetID: mainTargetID})
}

// Close closes Astilectron properly
func (a *Astilectron) Close() {
	astilog.Debug("Closing...")
	a.canceller.Cancel()
	if a.listener != nil {
		a.listener.Close()
	}
	if a.reader != nil {
		a.reader.close()
	}
	if a.stderrWriter != nil {
		a.stderrWriter.Close()
	}
	if a.stdoutWriter != nil {
		a.stdoutWriter.Close()
	}
	if a.writer != nil {
		a.writer.close()
	}
}

// HandleSignals handles signals
func (a *Astilectron) HandleSignals() {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, syscall.SIGABRT, syscall.SIGINT, syscall.SIGQUIT, syscall.SIGTERM)
	go func() {
		for sig := range ch {
			astilog.Debugf("Received signal %s", sig)
			a.Stop()
		}
	}()
}

// Stop orders Astilectron to stop
func (a *Astilectron) Stop() {
	astilog.Debug("Stopping...")
	a.canceller.Cancel()
	a.closeOnce.Do(func() {
		close(a.channelQuit)
	})
}

// Wait is a blocking pattern
func (a *Astilectron) Wait() {
	<-a.channelQuit
}

// Quit quits the app
func (a *Astilectron) Quit() error {
	return a.writer.write(Event{Name: EventNameAppCmdQuit})
}

// Paths returns the paths
func (a *Astilectron) Paths() Paths {
	return *a.paths
}

// Displays returns the displays
func (a *Astilectron) Displays() []*Display {
	return a.displayPool.all()
}

// PrimaryDisplay returns the primary display
func (a *Astilectron) PrimaryDisplay() *Display {
	return a.displayPool.primary()
}

// NewMenu creates a new app menu
func (a *Astilectron) NewMenu(i []*MenuItemOptions) *Menu {
	return newMenu(nil, mainTargetID, i, a.canceller, a.dispatcher, a.identifier, a.writer)
}

// NewWindow creates a new window
func (a *Astilectron) NewWindow(url string, o *WindowOptions) (*Window, error) {
	return newWindow(a.options, url, o, a.canceller, a.dispatcher, a.identifier, a.writer)
}

// NewWindowInDisplay creates a new window in a specific display
// This overrides the center attribute
func (a *Astilectron) NewWindowInDisplay(d *Display, url string, o *WindowOptions) (*Window, error) {
	if o.X != nil {
		*o.X += d.Bounds().X
	} else {
		o.X = PtrInt(d.Bounds().X)
	}
	if o.Y != nil {
		*o.Y += d.Bounds().Y
	} else {
		o.Y = PtrInt(d.Bounds().Y)
	}
	return newWindow(a.options, url, o, a.canceller, a.dispatcher, a.identifier, a.writer)
}

// NewTray creates a new tray
func (a *Astilectron) NewTray(o *TrayOptions) *Tray {
	return newTray(o, a.canceller, a.dispatcher, a.identifier, a.writer)
}
<|endoftext|>


<commit_before>package etcd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"path"
	"strings"
	"sync"
	"time"
)

// Errors introduced by handling requests
var (
	ErrRequestCancelled = errors.New("sending request is cancelled")
)

type RawRequest struct {
	Method       string
	RelativePath string
	Values       url.Values
	Cancel       <-chan bool
}

// NewRawRequest returns a new RawRequest
func NewRawRequest(method, relativePath string, values url.Values, cancel <-chan bool) *RawRequest {
	return &RawRequest{
		Method:       method,
		RelativePath: relativePath,
		Values:       values,
		Cancel:       cancel,
	}
}

// getCancelable issues a cancelable GET request
func (c *Client) getCancelable(key string, options Options,
	cancel <-chan bool) (*RawResponse, error) {
	logger.Debugf("get %s [%s]", key, c.cluster.Leader)
	p := keyToPath(key)

	// If consistency level is set to STRONG, append
	// the `consistent` query string.
	if c.config.Consistency == STRONG_CONSISTENCY {
		options["consistent"] = true
	}

	str, err := options.toParameters(VALID_GET_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("GET", p, nil, cancel)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// get issues a GET request
func (c *Client) get(key string, options Options) (*RawResponse, error) {
	return c.getCancelable(key, options, nil)
}

// put issues a PUT request
func (c *Client) put(key string, value string, ttl uint64,
	options Options) (*RawResponse, error) {

	logger.Debugf("put %s, %s, ttl: %d, [%s]", key, value, ttl, c.cluster.Leader)
	p := keyToPath(key)

	str, err := options.toParameters(VALID_PUT_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("PUT", p, buildValues(value, ttl), nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// post issues a POST request
func (c *Client) post(key string, value string, ttl uint64) (*RawResponse, error) {
	logger.Debugf("post %s, %s, ttl: %d, [%s]", key, value, ttl, c.cluster.Leader)
	p := keyToPath(key)

	req := NewRawRequest("POST", p, buildValues(value, ttl), nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// delete issues a DELETE request
func (c *Client) delete(key string, options Options) (*RawResponse, error) {
	logger.Debugf("delete %s [%s]", key, c.cluster.Leader)
	p := keyToPath(key)

	str, err := options.toParameters(VALID_DELETE_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("DELETE", p, nil, nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// SendRequest sends a HTTP request and returns a Response as defined by etcd
func (c *Client) SendRequest(rr *RawRequest) (*RawResponse, error) {

	var req *http.Request
	var resp *http.Response
	var httpPath string
	var err error
	var respBody []byte

	var numReqs = 1

	checkRetry := c.CheckRetry
	if checkRetry == nil {
		checkRetry = DefaultCheckRetry
	}

	cancelled := make(chan bool, 1)
	reqLock := new(sync.Mutex)

	if rr.Cancel != nil {
		cancelRoutine := make(chan bool)
		defer close(cancelRoutine)

		go func() {
			select {
			case <-rr.Cancel:
				cancelled <- true
				logger.Debug("send.request is cancelled")
			case <-cancelRoutine:
				return
			}

			// Repeat canceling request until this thread is stopped
			// because we have no idea about whether it succeeds.
			for {
				reqLock.Lock()
				c.httpClient.Transport.(*http.Transport).CancelRequest(req)
				reqLock.Unlock()

				select {
				case <-time.After(100 * time.Millisecond):
				case <-cancelRoutine:
					return
				}
			}
		}()
	}

	// If we connect to a follower and consistency is required, retry until
	// we connect to a leader
	sleep := 25 * time.Millisecond
	maxSleep := time.Second
	for attempt := 0; ; attempt++ {
		if attempt > 0 {
			select {
			case <-cancelled:
				return nil, ErrRequestCancelled
			case <-time.After(sleep):
				sleep = sleep * 2
				if sleep > maxSleep {
					sleep = maxSleep
				}
			}
		}

		logger.Debug("Connecting to etcd: attempt", attempt+1, "for", rr.RelativePath)

		if rr.Method == "GET" && c.config.Consistency == WEAK_CONSISTENCY {
			// If it's a GET and consistency level is set to WEAK,
			// then use a random machine.
			httpPath = c.getHttpPath(true, rr.RelativePath)
		} else {
			// Else use the leader.
			httpPath = c.getHttpPath(false, rr.RelativePath)
		}

		// Return a cURL command if curlChan is set
		if c.cURLch != nil {
			command := fmt.Sprintf("curl -X %s %s", rr.Method, httpPath)
			for key, value := range rr.Values {
				command += fmt.Sprintf(" -d %s=%s", key, value[0])
			}
			c.sendCURL(command)
		}

		logger.Debug("send.request.to ", httpPath, " | method ", rr.Method)

		reqLock.Lock()
		if rr.Values == nil {
			if req, err = http.NewRequest(rr.Method, httpPath, nil); err != nil {
				return nil, err
			}
		} else {
			body := strings.NewReader(rr.Values.Encode())
			if req, err = http.NewRequest(rr.Method, httpPath, body); err != nil {
				return nil, err
			}

			req.Header.Set("Content-Type",
				"application/x-www-form-urlencoded; param=value")
		}
		reqLock.Unlock()

		resp, err = c.httpClient.Do(req)
		defer func() {
			if resp != nil {
				resp.Body.Close()
			}
		}()

		// If the request was cancelled, return ErrRequestCancelled directly
		select {
		case <-cancelled:
			return nil, ErrRequestCancelled
		default:
		}

		numReqs++

		// network error, change a machine!
		if err != nil {
			logger.Debug("network error:", err.Error())
			lastResp := http.Response{}
			if checkErr := checkRetry(c.cluster, numReqs, lastResp, err); checkErr != nil {
				return nil, checkErr
			}

			c.cluster.switchLeader(attempt % len(c.cluster.Machines))
			continue
		}

		// if there is no error, it should receive response
		logger.Debug("recv.response.from", httpPath)

		if validHttpStatusCode[resp.StatusCode] {
			// try to read byte code and break the loop
			respBody, err = ioutil.ReadAll(resp.Body)
			if err == nil {
				logger.Debug("recv.success.", httpPath)
				break
			}
			// ReadAll error may be caused due to cancel request
			select {
			case <-cancelled:
				return nil, ErrRequestCancelled
			default:
			}
		}

		// if resp is TemporaryRedirect, set the new leader and retry
		if resp.StatusCode == http.StatusTemporaryRedirect {
			u, err := resp.Location()

			if err != nil {
				logger.Warning(err)
			} else {
				// Update cluster leader based on redirect location
				// because it should point to the leader address
				c.cluster.updateLeaderFromURL(u)
				logger.Debug("recv.response.relocate", u.String())
			}
			resp.Body.Close()
			continue
		}

		if checkErr := checkRetry(c.cluster, numReqs, *resp,
			errors.New("Unexpected HTTP status code")); checkErr != nil {
			return nil, checkErr
		}
		resp.Body.Close()
	}

	r := &RawResponse{
		StatusCode: resp.StatusCode,
		Body:       respBody,
		Header:     resp.Header,
	}

	return r, nil
}

// DefaultCheckRetry defines the retrying behaviour for bad HTTP requests
// If we have retried 2 * machine number, stop retrying.
// If status code is InternalServerError, sleep for 200ms.
func DefaultCheckRetry(cluster *Cluster, numReqs int, lastResp http.Response,
	err error) error {

	if numReqs >= 2*len(cluster.Machines) {
		return newError(ErrCodeEtcdNotReachable,
			"Tried to connect to each peer twice and failed", 0)
	}

	code := lastResp.StatusCode
	if code == http.StatusInternalServerError {
		time.Sleep(time.Millisecond * 200)

	}

	logger.Warning("bad response status code", code)
	return nil
}

func (c *Client) getHttpPath(random bool, s ...string) string {
	var machine string
	if random {
		machine = c.cluster.Machines[rand.Intn(len(c.cluster.Machines))]
	} else {
		machine = c.cluster.Leader
	}

	fullPath := machine + "/" + version
	for _, seg := range s {
		fullPath = fullPath + "/" + seg
	}

	return fullPath
}

// buildValues builds a url.Values map according to the given value and ttl
func buildValues(value string, ttl uint64) url.Values {
	v := url.Values{}

	if value != "" {
		v.Set("value", value)
	}

	if ttl > 0 {
		v.Set("ttl", fmt.Sprintf("%v", ttl))
	}

	return v
}

// convert key string to http path exclude version
// for example: key[foo] -> path[keys/foo]
// key[/] -> path[keys/]
func keyToPath(key string) string {
	p := path.Join("keys", key)

	// corner case: if key is "/" or "//" ect
	// path join will clear the tailing "/"
	// we need to add it back
	if p == "keys" {
		p = "keys/"
	}

	return p
}
<commit_msg>Use defer on reqLock.Unlock() to guarantee it gets unlocked on error<commit_after>package etcd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"path"
	"strings"
	"sync"
	"time"
)

// Errors introduced by handling requests
var (
	ErrRequestCancelled = errors.New("sending request is cancelled")
)

type RawRequest struct {
	Method       string
	RelativePath string
	Values       url.Values
	Cancel       <-chan bool
}

// NewRawRequest returns a new RawRequest
func NewRawRequest(method, relativePath string, values url.Values, cancel <-chan bool) *RawRequest {
	return &RawRequest{
		Method:       method,
		RelativePath: relativePath,
		Values:       values,
		Cancel:       cancel,
	}
}

// getCancelable issues a cancelable GET request
func (c *Client) getCancelable(key string, options Options,
	cancel <-chan bool) (*RawResponse, error) {
	logger.Debugf("get %s [%s]", key, c.cluster.Leader)
	p := keyToPath(key)

	// If consistency level is set to STRONG, append
	// the `consistent` query string.
	if c.config.Consistency == STRONG_CONSISTENCY {
		options["consistent"] = true
	}

	str, err := options.toParameters(VALID_GET_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("GET", p, nil, cancel)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// get issues a GET request
func (c *Client) get(key string, options Options) (*RawResponse, error) {
	return c.getCancelable(key, options, nil)
}

// put issues a PUT request
func (c *Client) put(key string, value string, ttl uint64,
	options Options) (*RawResponse, error) {

	logger.Debugf("put %s, %s, ttl: %d, [%s]", key, value, ttl, c.cluster.Leader)
	p := keyToPath(key)

	str, err := options.toParameters(VALID_PUT_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("PUT", p, buildValues(value, ttl), nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// post issues a POST request
func (c *Client) post(key string, value string, ttl uint64) (*RawResponse, error) {
	logger.Debugf("post %s, %s, ttl: %d, [%s]", key, value, ttl, c.cluster.Leader)
	p := keyToPath(key)

	req := NewRawRequest("POST", p, buildValues(value, ttl), nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// delete issues a DELETE request
func (c *Client) delete(key string, options Options) (*RawResponse, error) {
	logger.Debugf("delete %s [%s]", key, c.cluster.Leader)
	p := keyToPath(key)

	str, err := options.toParameters(VALID_DELETE_OPTIONS)
	if err != nil {
		return nil, err
	}
	p += str

	req := NewRawRequest("DELETE", p, nil, nil)
	resp, err := c.SendRequest(req)

	if err != nil {
		return nil, err
	}

	return resp, nil
}

// SendRequest sends a HTTP request and returns a Response as defined by etcd
func (c *Client) SendRequest(rr *RawRequest) (*RawResponse, error) {

	var req *http.Request
	var resp *http.Response
	var httpPath string
	var err error
	var respBody []byte

	var numReqs = 1

	checkRetry := c.CheckRetry
	if checkRetry == nil {
		checkRetry = DefaultCheckRetry
	}

	cancelled := make(chan bool, 1)
	reqLock := new(sync.Mutex)

	if rr.Cancel != nil {
		cancelRoutine := make(chan bool)
		defer close(cancelRoutine)

		go func() {
			select {
			case <-rr.Cancel:
				cancelled <- true
				logger.Debug("send.request is cancelled")
			case <-cancelRoutine:
				return
			}

			// Repeat canceling request until this thread is stopped
			// because we have no idea about whether it succeeds.
			for {
				reqLock.Lock()
				c.httpClient.Transport.(*http.Transport).CancelRequest(req)
				reqLock.Unlock()

				select {
				case <-time.After(100 * time.Millisecond):
				case <-cancelRoutine:
					return
				}
			}
		}()
	}

	// If we connect to a follower and consistency is required, retry until
	// we connect to a leader
	sleep := 25 * time.Millisecond
	maxSleep := time.Second
	for attempt := 0; ; attempt++ {
		if attempt > 0 {
			select {
			case <-cancelled:
				return nil, ErrRequestCancelled
			case <-time.After(sleep):
				sleep = sleep * 2
				if sleep > maxSleep {
					sleep = maxSleep
				}
			}
		}

		logger.Debug("Connecting to etcd: attempt", attempt+1, "for", rr.RelativePath)

		if rr.Method == "GET" && c.config.Consistency == WEAK_CONSISTENCY {
			// If it's a GET and consistency level is set to WEAK,
			// then use a random machine.
			httpPath = c.getHttpPath(true, rr.RelativePath)
		} else {
			// Else use the leader.
			httpPath = c.getHttpPath(false, rr.RelativePath)
		}

		// Return a cURL command if curlChan is set
		if c.cURLch != nil {
			command := fmt.Sprintf("curl -X %s %s", rr.Method, httpPath)
			for key, value := range rr.Values {
				command += fmt.Sprintf(" -d %s=%s", key, value[0])
			}
			c.sendCURL(command)
		}

		logger.Debug("send.request.to ", httpPath, " | method ", rr.Method)

		req, err := func() (*http.Request, error) {
			reqLock.Lock()
			defer reqLock.Unlock()

			if rr.Values == nil {
				if req, err = http.NewRequest(rr.Method, httpPath, nil); err != nil {
					return nil, err
				}
			} else {
				body := strings.NewReader(rr.Values.Encode())
				if req, err = http.NewRequest(rr.Method, httpPath, body); err != nil {
					return nil, err
				}

				req.Header.Set("Content-Type",
					"application/x-www-form-urlencoded; param=value")
			}
			return req, nil
		}()

		if err != nil {
			return nil, err
		}

		resp, err = c.httpClient.Do(req)
		defer func() {
			if resp != nil {
				resp.Body.Close()
			}
		}()

		// If the request was cancelled, return ErrRequestCancelled directly
		select {
		case <-cancelled:
			return nil, ErrRequestCancelled
		default:
		}

		numReqs++

		// network error, change a machine!
		if err != nil {
			logger.Debug("network error:", err.Error())
			lastResp := http.Response{}
			if checkErr := checkRetry(c.cluster, numReqs, lastResp, err); checkErr != nil {
				return nil, checkErr
			}

			c.cluster.switchLeader(attempt % len(c.cluster.Machines))
			continue
		}

		// if there is no error, it should receive response
		logger.Debug("recv.response.from", httpPath)

		if validHttpStatusCode[resp.StatusCode] {
			// try to read byte code and break the loop
			respBody, err = ioutil.ReadAll(resp.Body)
			if err == nil {
				logger.Debug("recv.success.", httpPath)
				break
			}
			// ReadAll error may be caused due to cancel request
			select {
			case <-cancelled:
				return nil, ErrRequestCancelled
			default:
			}
		}

		// if resp is TemporaryRedirect, set the new leader and retry
		if resp.StatusCode == http.StatusTemporaryRedirect {
			u, err := resp.Location()

			if err != nil {
				logger.Warning(err)
			} else {
				// Update cluster leader based on redirect location
				// because it should point to the leader address
				c.cluster.updateLeaderFromURL(u)
				logger.Debug("recv.response.relocate", u.String())
			}
			resp.Body.Close()
			continue
		}

		if checkErr := checkRetry(c.cluster, numReqs, *resp,
			errors.New("Unexpected HTTP status code")); checkErr != nil {
			return nil, checkErr
		}
		resp.Body.Close()
	}

	r := &RawResponse{
		StatusCode: resp.StatusCode,
		Body:       respBody,
		Header:     resp.Header,
	}

	return r, nil
}

// DefaultCheckRetry defines the retrying behaviour for bad HTTP requests
// If we have retried 2 * machine number, stop retrying.
// If status code is InternalServerError, sleep for 200ms.
func DefaultCheckRetry(cluster *Cluster, numReqs int, lastResp http.Response,
	err error) error {

	if numReqs >= 2*len(cluster.Machines) {
		return newError(ErrCodeEtcdNotReachable,
			"Tried to connect to each peer twice and failed", 0)
	}

	code := lastResp.StatusCode
	if code == http.StatusInternalServerError {
		time.Sleep(time.Millisecond * 200)

	}

	logger.Warning("bad response status code", code)
	return nil
}

func (c *Client) getHttpPath(random bool, s ...string) string {
	var machine string
	if random {
		machine = c.cluster.Machines[rand.Intn(len(c.cluster.Machines))]
	} else {
		machine = c.cluster.Leader
	}

	fullPath := machine + "/" + version
	for _, seg := range s {
		fullPath = fullPath + "/" + seg
	}

	return fullPath
}

// buildValues builds a url.Values map according to the given value and ttl
func buildValues(value string, ttl uint64) url.Values {
	v := url.Values{}

	if value != "" {
		v.Set("value", value)
	}

	if ttl > 0 {
		v.Set("ttl", fmt.Sprintf("%v", ttl))
	}

	return v
}

// convert key string to http path exclude version
// for example: key[foo] -> path[keys/foo]
// key[/] -> path[keys/]
func keyToPath(key string) string {
	p := path.Join("keys", key)

	// corner case: if key is "/" or "//" ect
	// path join will clear the tailing "/"
	// we need to add it back
	if p == "keys" {
		p = "keys/"
	}

	return p
}
<|endoftext|>


<commit_before>package slack

import "encoding/json"

// AttachmentField contains information for an attachment field
// An Attachment can contain multiple of these
type AttachmentField struct {
	Title string `json:"title"`
	Value string `json:"value"`
	Short bool   `json:"short"`
}

// AttachmentAction is a button or menu to be included in the attachment. Required when
// using message buttons or menus and otherwise not useful. A maximum of 5 actions may be
// provided per attachment.
type AttachmentAction struct {
	Name            string                        `json:"name"`                       // Required.
	Text            string                        `json:"text"`                       // Required.
	Style           string                        `json:"style,omitempty"`            // Optional. Allowed values: "default", "primary", "danger".
	Type            string                        `json:"type"`                       // Required. Must be set to "button" or "select".
	Value           string                        `json:"value,omitempty"`            // Optional.
	DataSource      string                        `json:"data_source,omitempty"`      // Optional.
	MinQueryLength  int                           `json:"min_query_length,omitempty"` // Optional. Default value is 1.
	Options         []AttachmentActionOption      `json:"options,omitempty"`          // Optional. Maximum of 100 options can be provided in each menu.
	SelectedOptions []AttachmentActionOption      `json:"selected_options,omitempty"` // Optional. The first element of this array will be set as the pre-selected option for this menu.
	OptionGroups    []AttachmentActionOptionGroup `json:"option_groups,omitempty"`    // Optional.
	Confirm         *ConfirmationField            `json:"confirm,omitempty"`          // Optional.
	URL             string                        `json:"url,omitempty"`              // Optional.
}

// AttachmentActionOption the individual option to appear in action menu.
type AttachmentActionOption struct {
	Text        string `json:"text"`                  // Required.
	Value       string `json:"value"`                 // Required.
	Description string `json:"description,omitempty"` // Optional. Up to 30 characters.
}

// AttachmentActionOptionGroup is a semi-hierarchal way to list available options to appear in action menu.
type AttachmentActionOptionGroup struct {
	Text    string                   `json:"text"`    // Required.
	Options []AttachmentActionOption `json:"options"` // Required.
}

// AttachmentActionCallback is sent from Slack when a user clicks a button in an interactive message (aka AttachmentAction)
type AttachmentActionCallback struct {
	Actions    []AttachmentAction `json:"actions"`
	CallbackID string             `json:"callback_id"`
	Team       Team               `json:"team"`
	Channel    Channel            `json:"channel"`
	User       User               `json:"user"`

	Name  string `json:"name"`
	Value string `json:"value"`

	OriginalMessage Message `json:"original_message"`

	ActionTs     string `json:"action_ts"`
	MessageTs    string `json:"message_ts"`
	AttachmentID string `json:"attachment_id"`
	Token        string `json:"token"`
	ResponseURL  string `json:"response_url"`
}

// ConfirmationField are used to ask users to confirm actions
type ConfirmationField struct {
	Title       string `json:"title,omitempty"`        // Optional.
	Text        string `json:"text"`                   // Required.
	OkText      string `json:"ok_text,omitempty"`      // Optional. Defaults to "Okay"
	DismissText string `json:"dismiss_text,omitempty"` // Optional. Defaults to "Cancel"
}

// Attachment contains all the information for an attachment
type Attachment struct {
	Color    string `json:"color,omitempty"`
	Fallback string `json:"fallback"`

	CallbackID string `json:"callback_id,omitempty"`

	AuthorName    string `json:"author_name,omitempty"`
	AuthorSubname string `json:"author_subname,omitempty"`
	AuthorLink    string `json:"author_link,omitempty"`
	AuthorIcon    string `json:"author_icon,omitempty"`

	Title     string `json:"title,omitempty"`
	TitleLink string `json:"title_link,omitempty"`
	Pretext   string `json:"pretext,omitempty"`
	Text      string `json:"text"`

	ImageURL string `json:"image_url,omitempty"`
	ThumbURL string `json:"thumb_url,omitempty"`

	Fields     []AttachmentField  `json:"fields,omitempty"`
	Actions    []AttachmentAction `json:"actions,omitempty"`
	MarkdownIn []string           `json:"mrkdwn_in,omitempty"`

	Footer     string `json:"footer,omitempty"`
	FooterIcon string `json:"footer_icon,omitempty"`

	Ts json.Number `json:"ts,omitempty"`
}
<commit_msg>Add id field to attachments (#219)<commit_after>package slack

import "encoding/json"

// AttachmentField contains information for an attachment field
// An Attachment can contain multiple of these
type AttachmentField struct {
	Title string `json:"title"`
	Value string `json:"value"`
	Short bool   `json:"short"`
}

// AttachmentAction is a button or menu to be included in the attachment. Required when
// using message buttons or menus and otherwise not useful. A maximum of 5 actions may be
// provided per attachment.
type AttachmentAction struct {
	Name            string                        `json:"name"`                       // Required.
	Text            string                        `json:"text"`                       // Required.
	Style           string                        `json:"style,omitempty"`            // Optional. Allowed values: "default", "primary", "danger".
	Type            string                        `json:"type"`                       // Required. Must be set to "button" or "select".
	Value           string                        `json:"value,omitempty"`            // Optional.
	DataSource      string                        `json:"data_source,omitempty"`      // Optional.
	MinQueryLength  int                           `json:"min_query_length,omitempty"` // Optional. Default value is 1.
	Options         []AttachmentActionOption      `json:"options,omitempty"`          // Optional. Maximum of 100 options can be provided in each menu.
	SelectedOptions []AttachmentActionOption      `json:"selected_options,omitempty"` // Optional. The first element of this array will be set as the pre-selected option for this menu.
	OptionGroups    []AttachmentActionOptionGroup `json:"option_groups,omitempty"`    // Optional.
	Confirm         *ConfirmationField            `json:"confirm,omitempty"`          // Optional.
	URL             string                        `json:"url,omitempty"`              // Optional.
}

// AttachmentActionOption the individual option to appear in action menu.
type AttachmentActionOption struct {
	Text        string `json:"text"`                  // Required.
	Value       string `json:"value"`                 // Required.
	Description string `json:"description,omitempty"` // Optional. Up to 30 characters.
}

// AttachmentActionOptionGroup is a semi-hierarchal way to list available options to appear in action menu.
type AttachmentActionOptionGroup struct {
	Text    string                   `json:"text"`    // Required.
	Options []AttachmentActionOption `json:"options"` // Required.
}

// AttachmentActionCallback is sent from Slack when a user clicks a button in an interactive message (aka AttachmentAction)
type AttachmentActionCallback struct {
	Actions    []AttachmentAction `json:"actions"`
	CallbackID string             `json:"callback_id"`
	Team       Team               `json:"team"`
	Channel    Channel            `json:"channel"`
	User       User               `json:"user"`

	Name  string `json:"name"`
	Value string `json:"value"`

	OriginalMessage Message `json:"original_message"`

	ActionTs     string `json:"action_ts"`
	MessageTs    string `json:"message_ts"`
	AttachmentID string `json:"attachment_id"`
	Token        string `json:"token"`
	ResponseURL  string `json:"response_url"`
}

// ConfirmationField are used to ask users to confirm actions
type ConfirmationField struct {
	Title       string `json:"title,omitempty"`        // Optional.
	Text        string `json:"text"`                   // Required.
	OkText      string `json:"ok_text,omitempty"`      // Optional. Defaults to "Okay"
	DismissText string `json:"dismiss_text,omitempty"` // Optional. Defaults to "Cancel"
}

// Attachment contains all the information for an attachment
type Attachment struct {
	Color    string `json:"color,omitempty"`
	Fallback string `json:"fallback"`

	CallbackID string `json:"callback_id,omitempty"`
	ID         int    `json:"id,omitempty"`

	AuthorName    string `json:"author_name,omitempty"`
	AuthorSubname string `json:"author_subname,omitempty"`
	AuthorLink    string `json:"author_link,omitempty"`
	AuthorIcon    string `json:"author_icon,omitempty"`

	Title     string `json:"title,omitempty"`
	TitleLink string `json:"title_link,omitempty"`
	Pretext   string `json:"pretext,omitempty"`
	Text      string `json:"text"`

	ImageURL string `json:"image_url,omitempty"`
	ThumbURL string `json:"thumb_url,omitempty"`

	Fields     []AttachmentField  `json:"fields,omitempty"`
	Actions    []AttachmentAction `json:"actions,omitempty"`
	MarkdownIn []string           `json:"mrkdwn_in,omitempty"`

	Footer     string `json:"footer,omitempty"`
	FooterIcon string `json:"footer_icon,omitempty"`

	Ts json.Number `json:"ts,omitempty"`
}
<|endoftext|>


<commit_before>package ethutil

import (
	"bytes"
	"encoding/binary"
	"fmt"
)

// Number to bytes
//
// Returns the number in bytes with the specified base
func NumberToBytes(num interface{}, bits int) []byte {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, num)
	if err != nil {
		fmt.Println("NumberToBytes failed:", err)
	}

	return buf.Bytes()[buf.Len()-(bits/8):]
}

// Bytes to number
//
// Attempts to cast a byte slice to a unsigned integer
func BytesToNumber(b []byte) uint64 {
	var number uint64

	// Make sure the buffer is 64bits
	data := make([]byte, 8)
	data = append(data[:len(b)], b...)

	buf := bytes.NewReader(data)
	err := binary.Read(buf, binary.BigEndian, &number)
	if err != nil {
		fmt.Println("BytesToNumber failed:", err)
	}

	return number
}

// Read variable int
//
// Read a variable length number in big endian byte order
func ReadVarint(reader *bytes.Reader) (ret uint64) {
	if reader.Len() == 8 {
		var num uint64
		binary.Read(reader, binary.BigEndian, &num)
		ret = uint64(num)
	} else if reader.Len() == 4 {
		var num uint32
		binary.Read(reader, binary.BigEndian, &num)
		ret = uint64(num)
	} else if reader.Len() == 2 {
		var num uint16
		binary.Read(reader, binary.BigEndian, &num)
		ret = uint64(num)
	} else {
		var num uint8
		binary.Read(reader, binary.BigEndian, &num)
		ret = uint64(num)
	}

	return ret
}

// Binary length
//
// Returns the true binary length of the given number
func BinaryLength(num int) int {
	if num == 0 {
		return 0
	}

	return 1 + BinaryLength(num>>8)
}

// Copy bytes
//
// Returns an exact copy of the provided bytes
func CopyBytes(b []byte) (copiedBytes []byte) {
	copiedBytes = make([]byte, len(b))
	copy(copiedBytes, b)

	return
}

func IsHex(str string) bool {
	l := len(str)
	return l >= 4 && l%2 == 0 && str[0:2] == "0x"
}

func StringToByteFunc(str string, cb func(str string) []byte) (ret []byte) {
	if str[0:2] == "0x" {
		ret = FromHex(str[2:])
	} else {
		ret = cb(str)
	}

	return
}
<commit_msg>Length checking<commit_after>package ethutil

import (
	"bytes"
	"encoding/binary"
	"fmt"
)

// Number to bytes
//
// Returns the number in bytes with the specified base
func NumberToBytes(num interface{}, bits int) []byte {
	buf := new(bytes.Buffer)
	err := binary.Write(buf, binary.BigEndian, num)
	if err != nil {
		fmt.Println("NumberToBytes failed:", err)
	}

	return buf.Bytes()[buf.Len()-(bits/8):]
}

// Bytes to number
//
// Attempts to cast a byte slice to a unsigned integer
func BytesToNumber(b []byte) uint64 {
	var number uint64

	// Make sure the buffer is 64bits
	data := make([]byte, 8)
	data = append(data[:len(b)], b...)

	buf := bytes.NewReader(data)
	err := binary.Read(buf, binary.BigEndian, &number)
	if err != nil {
		fmt.Println("BytesToNumber failed:", err)
	}

	return number
}

// Read variable int
//
// Read a variable length number in big endian byte order
func ReadVarint(reader *bytes.Reader) (ret uint64) {
	if reader.Len() == 8 {
		var num uint64
		binary.Read(reader, binary.BigEndian, &num)
		ret = uint64(num)
	} else if reader.Len() == 4 {
		var num uint32
		binary.Read(reader, binary.BigEndian, &num)
		ret = uint64(num)
	} else if reader.Len() == 2 {
		var num uint16
		binary.Read(reader, binary.BigEndian, &num)
		ret = uint64(num)
	} else {
		var num uint8
		binary.Read(reader, binary.BigEndian, &num)
		ret = uint64(num)
	}

	return ret
}

// Binary length
//
// Returns the true binary length of the given number
func BinaryLength(num int) int {
	if num == 0 {
		return 0
	}

	return 1 + BinaryLength(num>>8)
}

// Copy bytes
//
// Returns an exact copy of the provided bytes
func CopyBytes(b []byte) (copiedBytes []byte) {
	copiedBytes = make([]byte, len(b))
	copy(copiedBytes, b)

	return
}

func IsHex(str string) bool {
	l := len(str)
	return l >= 4 && l%2 == 0 && str[0:2] == "0x"
}

func StringToByteFunc(str string, cb func(str string) []byte) (ret []byte) {
	if len(str) > 1 && str[0:2] == "0x" {
		ret = FromHex(str[2:])
	} else {
		ret = cb(str)
	}

	return
}
<|endoftext|>


<commit_before>package cli

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/99designs/aws-vault/prompt"
	"github.com/99designs/aws-vault/server"
	"github.com/99designs/aws-vault/vault"
	"github.com/99designs/keyring"
	"gopkg.in/alecthomas/kingpin.v2"
)

type ExecCommandInput struct {
	Profile          string
	Command          string
	Args             []string
	Keyring          keyring.Keyring
	Duration         time.Duration
	RoleDuration     time.Duration
	MfaToken         string
	MfaPrompt        prompt.PromptFunc
	StartServer      bool
	CredentialHelper bool
	Signals          chan os.Signal
	NoSession        bool
}

// json metadata for AWS credential process. Ref: https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes
type AwsCredentialHelperData struct {
	Version         int    `json:"Version"`
	AccessKeyID     string `json:"AccessKeyId"`
	SecretAccessKey string `json:"SecretAccessKey"`
	SessionToken    string `json:"SessionToken"`
	Expiration      string `json:"Expiration,omitempty"`
}

func ConfigureExecCommand(app *kingpin.Application) {
	input := ExecCommandInput{}

	cmd := app.Command("exec", "Executes a command with AWS credentials in the environment")
	cmd.Flag("no-session", "Use root credentials, no session created").
		Short('n').
		BoolVar(&input.NoSession)

	cmd.Flag("session-ttl", "Expiration time for aws session").
		Default("4h").
		OverrideDefaultFromEnvar("AWS_SESSION_TTL").
		Short('t').
		DurationVar(&input.Duration)

	cmd.Flag("assume-role-ttl", "Expiration time for aws assumed role").
		Default("15m").
		OverrideDefaultFromEnvar("AWS_ASSUME_ROLE_TTL").
		DurationVar(&input.RoleDuration)

	cmd.Flag("mfa-token", "The mfa token to use").
		Short('m').
		StringVar(&input.MfaToken)

	cmd.Flag("json", "AWS credential helper. Ref: https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes").
		Short('j').
		BoolVar(&input.CredentialHelper)

	cmd.Flag("server", "Run the server in the background for credentials").
		Short('s').
		BoolVar(&input.StartServer)

	cmd.Arg("profile", "Name of the profile").
		Required().
		HintAction(ProfileNames).
		StringVar(&input.Profile)

	cmd.Arg("cmd", "Command to execute").
		Default(os.Getenv("SHELL")).
		StringVar(&input.Command)

	cmd.Arg("args", "Command arguments").
		StringsVar(&input.Args)

	cmd.Action(func(c *kingpin.ParseContext) error {
		input.Keyring = keyringImpl
		input.MfaPrompt = prompt.Method(GlobalFlags.PromptDriver)
		input.Signals = make(chan os.Signal)
		ExecCommand(app, input)
		return nil
	})
}

func ExecCommand(app *kingpin.Application, input ExecCommandInput) {
	if os.Getenv("AWS_VAULT") != "" {
		app.Fatalf("aws-vault sessions should be nested with care, unset $AWS_VAULT to force")
		return
	}

	var setEnv = true

	if input.NoSession && input.StartServer {
		app.Fatalf("Can't start a credential server without a session")
		return
	}

	creds, err := vault.NewVaultCredentials(input.Keyring, input.Profile, vault.VaultOptions{
		SessionDuration:    input.Duration,
		AssumeRoleDuration: input.RoleDuration,
		MfaToken:           input.MfaToken,
		MfaPrompt:          input.MfaPrompt,
		NoSession:          input.NoSession,
		Config:             awsConfig,
	})
	if err != nil {
		app.Fatalf("%v", err)
	}

	val, err := creds.Get()
	if err != nil {
		app.Fatalf(awsConfig.FormatCredentialError(err, input.Profile))
	}

	if input.StartServer {
		if err := server.StartCredentialsServer(creds); err != nil {
			app.Fatalf("Failed to start credential server: %v", err)
		} else {
			setEnv = false
		}
	}

	if input.CredentialHelper {
		credentialData := AwsCredentialHelperData{
			Version:         1,
			AccessKeyID:     val.AccessKeyID,
			SecretAccessKey: val.SecretAccessKey,
			SessionToken:    val.SessionToken,
		}
		if !input.NoSession {
			credentialData.Expiration = time.Now().Add(input.Duration).Format("2006-01-02T15:04:05Z")
		}
		json, err := json.Marshal(&credentialData)
		if err != nil {
			app.Fatalf("Error creating credential json")
		}
		fmt.Printf(string(json))
	} else {

		env := environ(os.Environ())
		env.Set("AWS_VAULT", input.Profile)

		env.Unset("AWS_ACCESS_KEY_ID")
		env.Unset("AWS_SECRET_ACCESS_KEY")
		env.Unset("AWS_CREDENTIAL_FILE")
		env.Unset("AWS_DEFAULT_PROFILE")
		env.Unset("AWS_PROFILE")

		if profile, _ := awsConfig.Profile(input.Profile); profile.Region != "" {
			log.Printf("Setting subprocess env: AWS_DEFAULT_REGION=%s, AWS_REGION=%s", profile.Region, profile.Region)
			env.Set("AWS_DEFAULT_REGION", profile.Region)
			env.Set("AWS_REGION", profile.Region)
		}

		if setEnv {
			log.Println("Setting subprocess env: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY")
			env.Set("AWS_ACCESS_KEY_ID", val.AccessKeyID)
			env.Set("AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

			if val.SessionToken != "" {
				log.Println("Setting subprocess env: AWS_SESSION_TOKEN, AWS_SECURITY_TOKEN")
				env.Set("AWS_SESSION_TOKEN", val.SessionToken)
				env.Set("AWS_SECURITY_TOKEN", val.SessionToken)
			}
		}

		cmd := exec.Command(input.Command, input.Args...)
		cmd.Env = env
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		signal.Notify(input.Signals, os.Interrupt, os.Kill)

		if err := cmd.Start(); err != nil {
			app.Fatalf("%v", err)
		}
		// wait for the command to finish
		waitCh := make(chan error, 1)
		go func() {
			waitCh <- cmd.Wait()
			close(waitCh)
		}()

		for {
			select {
			case sig := <-input.Signals:
				if err = cmd.Process.Signal(sig); err != nil {
					app.Errorf("%v", err)
					break
				}
			case err := <-waitCh:
				var waitStatus syscall.WaitStatus
				if exitError, ok := err.(*exec.ExitError); ok {
					waitStatus = exitError.Sys().(syscall.WaitStatus)
					os.Exit(waitStatus.ExitStatus())
				}
				if err != nil {
					app.Fatalf("%v", err)
				}
				return
			}
		}
	}
}

// environ is a slice of strings representing the environment, in the form "key=value".
type environ []string

// Unset an environment variable by key
func (e *environ) Unset(key string) {
	for i := range *e {
		if strings.HasPrefix((*e)[i], key+"=") {
			(*e)[i] = (*e)[len(*e)-1]
			*e = (*e)[:len(*e)-1]
			break
		}
	}
}

// Set adds an environment variable, replacing any existing ones of the same key
func (e *environ) Set(key, val string) {
	e.Unset(key)
	*e = append(*e, key+"="+val)
}

// ProfileNames returns a slice of profile names from the AWS config
func ProfileNames() []string {
	var profileNames []string
	for _, profile := range awsConfig.Profiles() {
		profileNames = append(profileNames, profile.Name)
	}
	return profileNames
}
<commit_msg>Update to use UTC<commit_after>package cli

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/99designs/aws-vault/prompt"
	"github.com/99designs/aws-vault/server"
	"github.com/99designs/aws-vault/vault"
	"github.com/99designs/keyring"
	"gopkg.in/alecthomas/kingpin.v2"
)

type ExecCommandInput struct {
	Profile          string
	Command          string
	Args             []string
	Keyring          keyring.Keyring
	Duration         time.Duration
	RoleDuration     time.Duration
	MfaToken         string
	MfaPrompt        prompt.PromptFunc
	StartServer      bool
	CredentialHelper bool
	Signals          chan os.Signal
	NoSession        bool
}

// json metadata for AWS credential process. Ref: https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes
type AwsCredentialHelperData struct {
	Version         int    `json:"Version"`
	AccessKeyID     string `json:"AccessKeyId"`
	SecretAccessKey string `json:"SecretAccessKey"`
	SessionToken    string `json:"SessionToken"`
	Expiration      string `json:"Expiration,omitempty"`
}

func ConfigureExecCommand(app *kingpin.Application) {
	input := ExecCommandInput{}

	cmd := app.Command("exec", "Executes a command with AWS credentials in the environment")
	cmd.Flag("no-session", "Use root credentials, no session created").
		Short('n').
		BoolVar(&input.NoSession)

	cmd.Flag("session-ttl", "Expiration time for aws session").
		Default("4h").
		OverrideDefaultFromEnvar("AWS_SESSION_TTL").
		Short('t').
		DurationVar(&input.Duration)

	cmd.Flag("assume-role-ttl", "Expiration time for aws assumed role").
		Default("15m").
		OverrideDefaultFromEnvar("AWS_ASSUME_ROLE_TTL").
		DurationVar(&input.RoleDuration)

	cmd.Flag("mfa-token", "The mfa token to use").
		Short('m').
		StringVar(&input.MfaToken)

	cmd.Flag("json", "AWS credential helper. Ref: https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes").
		Short('j').
		BoolVar(&input.CredentialHelper)

	cmd.Flag("server", "Run the server in the background for credentials").
		Short('s').
		BoolVar(&input.StartServer)

	cmd.Arg("profile", "Name of the profile").
		Required().
		HintAction(ProfileNames).
		StringVar(&input.Profile)

	cmd.Arg("cmd", "Command to execute").
		Default(os.Getenv("SHELL")).
		StringVar(&input.Command)

	cmd.Arg("args", "Command arguments").
		StringsVar(&input.Args)

	cmd.Action(func(c *kingpin.ParseContext) error {
		input.Keyring = keyringImpl
		input.MfaPrompt = prompt.Method(GlobalFlags.PromptDriver)
		input.Signals = make(chan os.Signal)
		ExecCommand(app, input)
		return nil
	})
}

func ExecCommand(app *kingpin.Application, input ExecCommandInput) {
	if os.Getenv("AWS_VAULT") != "" {
		app.Fatalf("aws-vault sessions should be nested with care, unset $AWS_VAULT to force")
		return
	}

	var setEnv = true

	if input.NoSession && input.StartServer {
		app.Fatalf("Can't start a credential server without a session")
		return
	}

	creds, err := vault.NewVaultCredentials(input.Keyring, input.Profile, vault.VaultOptions{
		SessionDuration:    input.Duration,
		AssumeRoleDuration: input.RoleDuration,
		MfaToken:           input.MfaToken,
		MfaPrompt:          input.MfaPrompt,
		NoSession:          input.NoSession,
		Config:             awsConfig,
	})
	if err != nil {
		app.Fatalf("%v", err)
	}

	val, err := creds.Get()
	if err != nil {
		app.Fatalf(awsConfig.FormatCredentialError(err, input.Profile))
	}

	if input.StartServer {
		if err := server.StartCredentialsServer(creds); err != nil {
			app.Fatalf("Failed to start credential server: %v", err)
		} else {
			setEnv = false
		}
	}

	if input.CredentialHelper {
		credentialData := AwsCredentialHelperData{
			Version:         1,
			AccessKeyID:     val.AccessKeyID,
			SecretAccessKey: val.SecretAccessKey,
			SessionToken:    val.SessionToken,
		}
		if !input.NoSession {
			credentialData.Expiration = time.Now().UTC().Add(input.Duration).Format("2006-01-02T15:04:05Z")
		}
		json, err := json.Marshal(&credentialData)
		if err != nil {
			app.Fatalf("Error creating credential json")
		}
		fmt.Printf(string(json))
	} else {

		env := environ(os.Environ())
		env.Set("AWS_VAULT", input.Profile)

		env.Unset("AWS_ACCESS_KEY_ID")
		env.Unset("AWS_SECRET_ACCESS_KEY")
		env.Unset("AWS_CREDENTIAL_FILE")
		env.Unset("AWS_DEFAULT_PROFILE")
		env.Unset("AWS_PROFILE")

		if profile, _ := awsConfig.Profile(input.Profile); profile.Region != "" {
			log.Printf("Setting subprocess env: AWS_DEFAULT_REGION=%s, AWS_REGION=%s", profile.Region, profile.Region)
			env.Set("AWS_DEFAULT_REGION", profile.Region)
			env.Set("AWS_REGION", profile.Region)
		}

		if setEnv {
			log.Println("Setting subprocess env: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY")
			env.Set("AWS_ACCESS_KEY_ID", val.AccessKeyID)
			env.Set("AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

			if val.SessionToken != "" {
				log.Println("Setting subprocess env: AWS_SESSION_TOKEN, AWS_SECURITY_TOKEN")
				env.Set("AWS_SESSION_TOKEN", val.SessionToken)
				env.Set("AWS_SECURITY_TOKEN", val.SessionToken)
			}
		}

		cmd := exec.Command(input.Command, input.Args...)
		cmd.Env = env
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		signal.Notify(input.Signals, os.Interrupt, os.Kill)

		if err := cmd.Start(); err != nil {
			app.Fatalf("%v", err)
		}
		// wait for the command to finish
		waitCh := make(chan error, 1)
		go func() {
			waitCh <- cmd.Wait()
			close(waitCh)
		}()

		for {
			select {
			case sig := <-input.Signals:
				if err = cmd.Process.Signal(sig); err != nil {
					app.Errorf("%v", err)
					break
				}
			case err := <-waitCh:
				var waitStatus syscall.WaitStatus
				if exitError, ok := err.(*exec.ExitError); ok {
					waitStatus = exitError.Sys().(syscall.WaitStatus)
					os.Exit(waitStatus.ExitStatus())
				}
				if err != nil {
					app.Fatalf("%v", err)
				}
				return
			}
		}
	}
}

// environ is a slice of strings representing the environment, in the form "key=value".
type environ []string

// Unset an environment variable by key
func (e *environ) Unset(key string) {
	for i := range *e {
		if strings.HasPrefix((*e)[i], key+"=") {
			(*e)[i] = (*e)[len(*e)-1]
			*e = (*e)[:len(*e)-1]
			break
		}
	}
}

// Set adds an environment variable, replacing any existing ones of the same key
func (e *environ) Set(key, val string) {
	e.Unset(key)
	*e = append(*e, key+"="+val)
}

// ProfileNames returns a slice of profile names from the AWS config
func ProfileNames() []string {
	var profileNames []string
	for _, profile := range awsConfig.Profiles() {
		profileNames = append(profileNames, profile.Name)
	}
	return profileNames
}
<|endoftext|>


<commit_before>/*Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package event

import (
	"fmt"
	"github.com/skyrings/skyring-common/conf"
	"github.com/skyrings/skyring-common/db"
	"github.com/skyrings/skyring-common/models"
	"github.com/skyrings/skyring-common/tools/logger"
	"github.com/skyrings/skyring-common/tools/task"
	"github.com/skyrings/skyring/apps/skyring"
	"github.com/skyrings/skyring/nodemanager/saltnodemanager"
	"gopkg.in/mgo.v2/bson"
	"os"
	"time"
)

var (
	curr_hostname, err = os.Hostname()
)

var handlermap = map[string]interface{}{
	"skyring/dbus/node/*/generic/storage/block/added":           block_add_handler,
	"skyring/dbus/node/*/generic/storage/block/removed":         block_remove_handler,
	"skyring/dbus/node/*/generic/storage/block/changed":         block_change_handler,
	"skyring/dbus/node/*/generic/storage/mount/changed":         mount_change_handler,
	"skyring/dbus/node/*/generic/storage/drive/added":           drive_add_handler,
	"skyring/dbus/node/*/generic/storage/drive/removed":         drive_remove_handler,
	"skyring/dbus/node/*/generic/storage/drive/possibleFailure": drive_remove_handler,
	"skyring/dbus/node/*/generic/service/collectd":              collectd_status_handler,
	"salt/node/appeared":                                        node_appeared_handler,
	"salt/node/lost":                                            node_lost_handler,
	"skyring/collectd/node/*/threshold/*/*":                     collectd_threshold_handler,
}

// ALL HANDLERS ARE JUST WRITING THE EVENTS TO DB. OTHER HANDLING AND CORRELATION
// IS TO BE DONE

func block_add_handler(event models.Event) error {
	return nil
}

func block_remove_handler(event models.Event) error {
	return nil
}

func block_change_handler(event models.Event) error {
	return nil
}

func mount_change_handler(event models.Event) error {
	return nil
}

func drive_add_handler(event models.Event) error {
	return nil
}

func drive_remove_handler(event models.Event) error {
	return nil
}

func collectd_status_handler(event models.Event) error {
	return nil
}

func update_node_status(nodeStatus string, event models.Event) error {
	sessionCopy := db.GetDatastore().Copy()
	defer sessionCopy.Close()
	coll := sessionCopy.DB(conf.SystemConfig.DBConfig.Database).C(models.COLL_NAME_STORAGE_NODES)
	if err := coll.Update(bson.M{"nodeid": event.NodeId}, bson.M{"$set": bson.M{"status": nodeStatus}}); err != nil {
		logger.Get().Error("Error updating the node status: %s", err)
		return err
	}
	return nil
}

func node_appeared_handler(event models.Event) error {
	if err := update_node_status(models.STATUS_UP, event); err != nil {
		return err
	}
	return nil
}

func node_lost_handler(event models.Event) error {
	if err := update_node_status(models.STATUS_DOWN, event); err != nil {
		return err
	}
	return nil
}

func collectd_threshold_handler(event models.Event) error {
	return nil
}

func handle_node_start_event(node string) error {
	sessionCopy := db.GetDatastore().Copy()
	defer sessionCopy.Close()
	coll := sessionCopy.DB(conf.SystemConfig.DBConfig.Database).C(models.COLL_NAME_STORAGE_NODES)

	var storage_node models.Node

	_ = coll.Find(bson.M{"hostname": node}).One(&storage_node)
	if storage_node.State != models.NODE_STATE_INITIALIZING {
		logger.Get().Warning(fmt.Sprintf("Node with name: %s not in activating state to update other details", node))
		return nil
	}
	asyncTask := func(t *task.Task) {
		for {
			select {
			case <-t.StopCh:
				return
			default:
				t.UpdateStatus("started the task for InitializeNode: %s", t.ID)
				// Process the request
				if err := initializeStorageNode(node, t); err != nil {
					t.UpdateStatus("Failed")
					t.Done(models.TASK_STATUS_FAILURE)
				} else {
					t.UpdateStatus("Success")
					t.Done(models.TASK_STATUS_SUCCESS)
				}
				return
			}
		}
	}
	if taskId, err := skyring.GetApp().GetTaskManager().Run(fmt.Sprintf("Initialize Node: %s", node), asyncTask, 120*time.Second, nil, nil, nil); err != nil {
		logger.Get().Error("Unable to create the task for Initialize Node: %s. error: %v", node, err)
	} else {
		logger.Get().Debug("Task created for initialize node. Task id: %s", taskId.String())
	}
	return nil
}

func initializeStorageNode(node string, t *task.Task) error {
	sProfiles, err := skyring.GetDbProvider().StorageProfileInterface().StorageProfiles(nil, models.QueryOps{})
	if err != nil {
		logger.Get().Error("Unable to get the storage profiles. May not be able to apply storage profiles for node: %v err:%v", node, err)
	}
	if storage_node, ok := saltnodemanager.GetStorageNodeInstance(node, sProfiles); ok {
		if err := updateStorageNodeToDB(*storage_node); err != nil {
			logger.Get().Error("Unable to add details of node: %s to DB. error: %v", node, err)
			t.UpdateStatus("Unable to add details of node: %s to DB. error: %v", node, err)
			updateNodeState(node, models.NODE_STATE_FAILED)
			return err
		}
		if nodeErrorMap, configureError := skyring.GetCoreNodeManager().SetUpMonitoring(node, curr_hostname); configureError != nil && len(nodeErrorMap) != 0 {
			if len(nodeErrorMap) != 0 {
				logger.Get().Error("Unable to setup collectd on %s because of %v", node, nodeErrorMap)
				t.UpdateStatus("Unable to setup collectd on %s because of %v", node, nodeErrorMap)
				updateNodeState(node, models.NODE_STATE_FAILED)
				return err
			} else {
				logger.Get().Error("Config Error during monitoring setup for node:%s Error:%v", node, configureError)
				t.UpdateStatus("Config Error during monitoring setup for node:%s Error:%v", node, configureError)
				updateNodeState(node, models.NODE_STATE_FAILED)
				return err
			}
		}
		return nil
	} else {
		logger.Get().Critical("Error getting the details for node: %s", node)
		t.UpdateStatus("Error getting the details for node: %s", node)
		updateNodeState(node, models.NODE_STATE_FAILED)
		return fmt.Errorf("Error getting the details for node: %s", node)
	}
}

func updateNodeState(node string, state int) error {
	sessionCopy := db.GetDatastore().Copy()
	defer sessionCopy.Close()
	coll := sessionCopy.DB(conf.SystemConfig.DBConfig.Database).C(models.COLL_NAME_STORAGE_NODES)
	if err := coll.Update(bson.M{"hostname": node}, bson.M{"$set": bson.M{"state": state}}); err != nil {
		logger.Get().Critical("Error updating the node state for node: %s. error: %v", node, err)
		return err
	}
	return nil
}

func updateStorageNodeToDB(storage_node models.Node) error {
	// Add the node details to the DB
	sessionCopy := db.GetDatastore().Copy()
	defer sessionCopy.Close()
	coll := sessionCopy.DB(conf.SystemConfig.DBConfig.Database).C(models.COLL_NAME_STORAGE_NODES)

	storage_node.State = models.NODE_STATE_ACTIVE
	if err := coll.Update(bson.M{"hostname": storage_node.Hostname}, storage_node); err != nil {
		logger.Get().Critical("Error Updating the node: %s. error: %v", storage_node.Hostname, err)
		return err
	}
	return nil
}
<commit_msg>sample-notification: patch for sample notification.<commit_after>/*Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package event

import (
	"fmt"
	"github.com/skyrings/skyring-common/conf"
	"github.com/skyrings/skyring-common/db"
	"github.com/skyrings/skyring-common/models"
	"github.com/skyrings/skyring-common/notifier"
	"github.com/skyrings/skyring-common/tools/logger"
	"github.com/skyrings/skyring-common/tools/task"
	"github.com/skyrings/skyring/apps/skyring"
	"github.com/skyrings/skyring/nodemanager/saltnodemanager"
	"gopkg.in/mgo.v2/bson"
	"os"
	"time"
)

var (
	curr_hostname, err = os.Hostname()
)

var handlermap = map[string]interface{}{
	"skyring/dbus/node/*/generic/storage/block/added":           block_add_handler,
	"skyring/dbus/node/*/generic/storage/block/removed":         block_remove_handler,
	"skyring/dbus/node/*/generic/storage/block/changed":         block_change_handler,
	"skyring/dbus/node/*/generic/storage/mount/changed":         mount_change_handler,
	"skyring/dbus/node/*/generic/storage/drive/added":           drive_add_handler,
	"skyring/dbus/node/*/generic/storage/drive/removed":         drive_remove_handler,
	"skyring/dbus/node/*/generic/storage/drive/possibleFailure": drive_remove_handler,
	"skyring/dbus/node/*/generic/service/collectd":              collectd_status_handler,
	"salt/node/appeared":                                        node_appeared_handler,
	"salt/node/lost":                                            node_lost_handler,
	"skyring/collectd/node/*/threshold/*/*":                     collectd_threshold_handler,
}

// ALL HANDLERS ARE JUST WRITING THE EVENTS TO DB. OTHER HANDLING AND CORRELATION
// IS TO BE DONE

func block_add_handler(event models.Event) error {
	return nil
}

func block_remove_handler(event models.Event) error {
	return nil
}

func block_change_handler(event models.Event) error {
	return nil
}

func mount_change_handler(event models.Event) error {
	return nil
}

func drive_add_handler(event models.Event) error {
	return nil
}

func drive_remove_handler(event models.Event) error {
	sessionCopy := db.GetDatastore().Copy()
	defer sessionCopy.Close()
	var node models.Node
	coll := sessionCopy.DB(conf.SystemConfig.DBConfig.Database).C(models.COLL_NAME_STORAGE_NODES)
	if err := coll.Find(bson.M{"nodeid": event.NodeId}).One(&node); err != nil {
		logger.Get().Error("Node information read from DB failed for node: %s. error: %v", event.NodeId, err)
		return nil
	}
	subject := fmt.Sprintf("Warning: Storage Drive Removed from %s", node.Hostname)
	body := event.Message
	notifier.MailNotify(subject, body)
	return nil
}

func collectd_status_handler(event models.Event) error {
	return nil
}

func update_node_status(nodeStatus string, event models.Event) error {
	sessionCopy := db.GetDatastore().Copy()
	defer sessionCopy.Close()
	coll := sessionCopy.DB(conf.SystemConfig.DBConfig.Database).C(models.COLL_NAME_STORAGE_NODES)
	if err := coll.Update(bson.M{"nodeid": event.NodeId}, bson.M{"$set": bson.M{"status": nodeStatus}}); err != nil {
		logger.Get().Error("Error updating the node status: %s", err)
		return err
	}
	return nil
}

func node_appeared_handler(event models.Event) error {
	if err := update_node_status(models.STATUS_UP, event); err != nil {
		return err
	}
	return nil
}

func node_lost_handler(event models.Event) error {
	if err := update_node_status(models.STATUS_DOWN, event); err != nil {
		return err
	}
	return nil
}

func collectd_threshold_handler(event models.Event) error {
	return nil
}

func handle_node_start_event(node string) error {
	sessionCopy := db.GetDatastore().Copy()
	defer sessionCopy.Close()
	coll := sessionCopy.DB(conf.SystemConfig.DBConfig.Database).C(models.COLL_NAME_STORAGE_NODES)

	var storage_node models.Node

	_ = coll.Find(bson.M{"hostname": node}).One(&storage_node)
	if storage_node.State != models.NODE_STATE_INITIALIZING {
		logger.Get().Warning(fmt.Sprintf("Node with name: %s not in activating state to update other details", node))
		return nil
	}
	asyncTask := func(t *task.Task) {
		for {
			select {
			case <-t.StopCh:
				return
			default:
				t.UpdateStatus("started the task for InitializeNode: %s", t.ID)
				// Process the request
				if err := initializeStorageNode(node, t); err != nil {
					t.UpdateStatus("Failed")
					t.Done(models.TASK_STATUS_FAILURE)
				} else {
					t.UpdateStatus("Success")
					t.Done(models.TASK_STATUS_SUCCESS)
				}
				return
			}
		}
	}
	if taskId, err := skyring.GetApp().GetTaskManager().Run(fmt.Sprintf("Initialize Node: %s", node), asyncTask, 120*time.Second, nil, nil, nil); err != nil {
		logger.Get().Error("Unable to create the task for Initialize Node: %s. error: %v", node, err)
	} else {
		logger.Get().Debug("Task created for initialize node. Task id: %s", taskId.String())
	}
	return nil
}

func initializeStorageNode(node string, t *task.Task) error {
	sProfiles, err := skyring.GetDbProvider().StorageProfileInterface().StorageProfiles(nil, models.QueryOps{})
	if err != nil {
		logger.Get().Error("Unable to get the storage profiles. May not be able to apply storage profiles for node: %v err:%v", node, err)
	}
	if storage_node, ok := saltnodemanager.GetStorageNodeInstance(node, sProfiles); ok {
		if err := updateStorageNodeToDB(*storage_node); err != nil {
			logger.Get().Error("Unable to add details of node: %s to DB. error: %v", node, err)
			t.UpdateStatus("Unable to add details of node: %s to DB. error: %v", node, err)
			updateNodeState(node, models.NODE_STATE_FAILED)
			return err
		}
		if nodeErrorMap, configureError := skyring.GetCoreNodeManager().SetUpMonitoring(node, curr_hostname); configureError != nil && len(nodeErrorMap) != 0 {
			if len(nodeErrorMap) != 0 {
				logger.Get().Error("Unable to setup collectd on %s because of %v", node, nodeErrorMap)
				t.UpdateStatus("Unable to setup collectd on %s because of %v", node, nodeErrorMap)
				updateNodeState(node, models.NODE_STATE_FAILED)
				return err
			} else {
				logger.Get().Error("Config Error during monitoring setup for node:%s Error:%v", node, configureError)
				t.UpdateStatus("Config Error during monitoring setup for node:%s Error:%v", node, configureError)
				updateNodeState(node, models.NODE_STATE_FAILED)
				return err
			}
		}
		return nil
	} else {
		logger.Get().Critical("Error getting the details for node: %s", node)
		t.UpdateStatus("Error getting the details for node: %s", node)
		updateNodeState(node, models.NODE_STATE_FAILED)
		return fmt.Errorf("Error getting the details for node: %s", node)
	}
}

func updateNodeState(node string, state int) error {
	sessionCopy := db.GetDatastore().Copy()
	defer sessionCopy.Close()
	coll := sessionCopy.DB(conf.SystemConfig.DBConfig.Database).C(models.COLL_NAME_STORAGE_NODES)
	if err := coll.Update(bson.M{"hostname": node}, bson.M{"$set": bson.M{"state": state}}); err != nil {
		logger.Get().Critical("Error updating the node state for node: %s. error: %v", node, err)
		return err
	}
	return nil
}

func updateStorageNodeToDB(storage_node models.Node) error {
	// Add the node details to the DB
	sessionCopy := db.GetDatastore().Copy()
	defer sessionCopy.Close()
	coll := sessionCopy.DB(conf.SystemConfig.DBConfig.Database).C(models.COLL_NAME_STORAGE_NODES)

	storage_node.State = models.NODE_STATE_ACTIVE
	if err := coll.Update(bson.M{"hostname": storage_node.Hostname}, storage_node); err != nil {
		logger.Get().Critical("Error Updating the node: %s. error: %v", storage_node.Hostname, err)
		return err
	}
	return nil
}
<|endoftext|>