<commit_before>package http

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

	platform "github.com/influxdata/influxdb"
)

const (
	// PlatformErrorCodeHeader shows the error code of platform error.
	PlatformErrorCodeHeader = "X-Platform-Error-Code"
)

// AuthzError is returned for authorization errors. When this error type is returned,
// the user can be presented with a generic "authorization failed" error, but
// the system can log the underlying AuthzError() so that operators have insight
// into what actually failed with authorization.
type AuthzError interface {
	error
	AuthzError() error
}

// CheckErrorStatus for status and any error in the response.
func CheckErrorStatus(code int, res *http.Response) error {
	err := CheckError(res)
	if err != nil {
		return err
	}

	if res.StatusCode != code {
		return fmt.Errorf("unexpected status code: %s", res.Status)
	}

	return nil
}

// CheckError reads the http.Response and returns an error if one exists.
// It will automatically recognize the errors returned by Influx services
// and decode the error into an internal error type. If the error cannot
// be determined in that way, it will create a generic error message.
//
// If there is no error, then this returns nil.
func CheckError(resp *http.Response) (err error) {
	switch resp.StatusCode / 100 {
	case 4, 5:
		// We will attempt to parse this error outside of this block.
	case 2:
		return nil
	default:
		// TODO(jsternberg): Figure out what to do here?
		return &platform.Error{
			Code: platform.EInternal,
			Msg:  fmt.Sprintf("unexpected status code: %d %s", resp.StatusCode, resp.Status),
		}
	}
	pe := new(platform.Error)
	parseErr := json.NewDecoder(resp.Body).Decode(pe)
	if parseErr != nil {
		return parseErr
	}
	return pe
}

// EncodeError encodes err with the appropriate status code and format,
// sets the X-Platform-Error-Code headers on the response.
// We're no longer using X-Influx-Error and X-Influx-Reference.
// and sets the response status to the corresponding status code.
func EncodeError(ctx context.Context, err error, w http.ResponseWriter) {
	if err == nil {
		return
	}

	code := platform.ErrorCode(err)
	httpCode, ok := statusCodePlatformError[code]
	if !ok {
		httpCode = http.StatusBadRequest
	}
	w.Header().Set(PlatformErrorCodeHeader, code)
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(httpCode)
	var e error
	if pe, ok := err.(*platform.Error); ok {
		e = &platform.Error{
			Code: code,
			Op:   platform.ErrorOp(err),
			Msg:  platform.ErrorMessage(err),
			Err:  pe.Err,
		}
	} else {
		e = &platform.Error{
			Code: platform.EInternal,
			Err:  err,
		}
	}
	b, _ := json.Marshal(e)
	_, _ = w.Write(b)
}

// UnauthorizedError encodes a error message and status code for unauthorized access.
func UnauthorizedError(ctx context.Context, w http.ResponseWriter) {
	EncodeError(ctx, &platform.Error{
		Code: platform.EUnauthorized,
		Msg:  "unauthorized access",
	}, w)
}

// statusCodePlatformError is the map convert platform.Error to error
var statusCodePlatformError = map[string]int{
	platform.EInternal:            http.StatusInternalServerError,
	platform.EInvalid:             http.StatusBadRequest,
	platform.EUnprocessableEntity: http.StatusUnprocessableEntity,
	platform.EEmptyValue:          http.StatusBadRequest,
	platform.EConflict:            http.StatusUnprocessableEntity,
	platform.ENotFound:            http.StatusNotFound,
	platform.EUnavailable:         http.StatusServiceUnavailable,
	platform.EForbidden:           http.StatusForbidden,
	platform.ETooManyRequests:     http.StatusTooManyRequests,
	platform.EUnauthorized:        http.StatusUnauthorized,
	platform.EMethodNotAllowed:    http.StatusMethodNotAllowed,
}
<commit_msg>fix(http): do not discard non-json encoded errors when using `CheckError` (#13844)<commit_after>package http

import (
	"bytes"
	"context"
	"encoding/json"
	stderrors "errors"
	"fmt"
	"io"
	"net/http"

	platform "github.com/influxdata/influxdb"
	"github.com/pkg/errors"
)

const (
	// PlatformErrorCodeHeader shows the error code of platform error.
	PlatformErrorCodeHeader = "X-Platform-Error-Code"
)

// AuthzError is returned for authorization errors. When this error type is returned,
// the user can be presented with a generic "authorization failed" error, but
// the system can log the underlying AuthzError() so that operators have insight
// into what actually failed with authorization.
type AuthzError interface {
	error
	AuthzError() error
}

// CheckErrorStatus for status and any error in the response.
func CheckErrorStatus(code int, res *http.Response) error {
	err := CheckError(res)
	if err != nil {
		return err
	}

	if res.StatusCode != code {
		return fmt.Errorf("unexpected status code: %s", res.Status)
	}

	return nil
}

// CheckError reads the http.Response and returns an error if one exists.
// It will automatically recognize the errors returned by Influx services
// and decode the error into an internal error type. If the error cannot
// be determined in that way, it will create a generic error message.
//
// If there is no error, then this returns nil.
func CheckError(resp *http.Response) (err error) {
	switch resp.StatusCode / 100 {
	case 4, 5:
		// We will attempt to parse this error outside of this block.
	case 2:
		return nil
	default:
		// TODO(jsternberg): Figure out what to do here?
		return &platform.Error{
			Code: platform.EInternal,
			Msg:  fmt.Sprintf("unexpected status code: %d %s", resp.StatusCode, resp.Status),
		}
	}
	pe := new(platform.Error)

	var buf bytes.Buffer
	if _, err := io.Copy(&buf, resp.Body); err != nil {
		return &platform.Error{
			Code: platform.EInternal,
			Msg:  err.Error(),
		}
	}
	parseErr := json.Unmarshal(buf.Bytes(), pe)
	if parseErr != nil {
		return errors.Wrap(stderrors.New(buf.String()), parseErr.Error())
	}
	return pe
}

// EncodeError encodes err with the appropriate status code and format,
// sets the X-Platform-Error-Code headers on the response.
// We're no longer using X-Influx-Error and X-Influx-Reference.
// and sets the response status to the corresponding status code.
func EncodeError(ctx context.Context, err error, w http.ResponseWriter) {
	if err == nil {
		return
	}

	code := platform.ErrorCode(err)
	httpCode, ok := statusCodePlatformError[code]
	if !ok {
		httpCode = http.StatusBadRequest
	}
	w.Header().Set(PlatformErrorCodeHeader, code)
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(httpCode)
	var e error
	if pe, ok := err.(*platform.Error); ok {
		e = &platform.Error{
			Code: code,
			Op:   platform.ErrorOp(err),
			Msg:  platform.ErrorMessage(err),
			Err:  pe.Err,
		}
	} else {
		e = &platform.Error{
			Code: platform.EInternal,
			Err:  err,
		}
	}
	b, _ := json.Marshal(e)
	_, _ = w.Write(b)
}

// UnauthorizedError encodes a error message and status code for unauthorized access.
func UnauthorizedError(ctx context.Context, w http.ResponseWriter) {
	EncodeError(ctx, &platform.Error{
		Code: platform.EUnauthorized,
		Msg:  "unauthorized access",
	}, w)
}

// statusCodePlatformError is the map convert platform.Error to error
var statusCodePlatformError = map[string]int{
	platform.EInternal:            http.StatusInternalServerError,
	platform.EInvalid:             http.StatusBadRequest,
	platform.EUnprocessableEntity: http.StatusUnprocessableEntity,
	platform.EEmptyValue:          http.StatusBadRequest,
	platform.EConflict:            http.StatusUnprocessableEntity,
	platform.ENotFound:            http.StatusNotFound,
	platform.EUnavailable:         http.StatusServiceUnavailable,
	platform.EForbidden:           http.StatusForbidden,
	platform.ETooManyRequests:     http.StatusTooManyRequests,
	platform.EUnauthorized:        http.StatusUnauthorized,
	platform.EMethodNotAllowed:    http.StatusMethodNotAllowed,
}
<|endoftext|>


<commit_before>package eventsocket

import (
	"net/http"

	"github.com/gorilla/mux"
)

type httpServer struct {
	router *mux.Router
}

// install the http server's router
func (h *httpServer) route() error {
	// log.Info("Initializing EventSocket Router")

	// instantiate a new controller
	C, err := newHttpController()
	if err != nil {
		// log.Error(fmt.Sprintf("Encountered error while instantiating new HttpController: %s", err.Error()))
		return err
	}

	// get a new router
	h.router = mux.NewRouter()
	s := h.router.PathPrefix("/v1").Subrouter()

	// handle routes
	// s.HandleFunc("/clients/{cid}/ws", C.Client.GetWs).Methods("GET")
	// s.HandleFunc("/clients/{cid}", C.Client.Get).Methods("GET")
	s.HandleFunc("/clients", C.Client.Create).Methods("POST")

	return nil
}

// handle and serve the api
func (h *httpServer) listen(listenAddr string) error {
	http.Handle("/", h.router)

	err := http.ListenAndServe(listenAddr, nil)
	if err != nil {
		// log.Error(fmt.Sprintf("%s: %s", "ListenAndServe Error", err.Error()))
		return err
	}

	return nil
}
<commit_msg>httpserver.listen: convert ListenAndServe to net.Listen and http.Serve<commit_after>package eventsocket

import (
	"fmt"
	"net"
	"net/http"

	"github.com/LiftMe/glip/log"
	"github.com/gorilla/mux"
)

type httpServer struct {
	router *mux.Router
}

// install the http server's router
func (h *httpServer) route() error {
	// log.Info("Initializing EventSocket Router")

	// instantiate a new controller
	C, err := newHttpController()
	if err != nil {
		// log.Error(fmt.Sprintf("Encountered error while instantiating new HttpController: %s", err.Error()))
		return err
	}

	// get a new router
	h.router = mux.NewRouter()
	s := h.router.PathPrefix("/v1").Subrouter()

	// handle routes
	// s.HandleFunc("/clients/{cid}/ws", C.Client.GetWs).Methods("GET")
	// s.HandleFunc("/clients/{cid}", C.Client.Get).Methods("GET")
	s.HandleFunc("/clients", C.Client.Create).Methods("POST")

	return nil
}

// handle and serve the api
func (h *httpServer) listen(listenAddr string) error {
	http.Handle("/", h.router)

	l, err := net.Listen("tcp", listenAddr)
	if err != nil {
		log.Error(fmt.Sprintf("%s: %s", "Listen Error", err.Error()))
		return err
	}

	err = http.Serve(l, nil)
	if err != nil {
		log.Error(fmt.Sprintf("%s: %s", "Serve Error", err.Error()))
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package sudoku

import (
	"fmt"
	"log"
	"math"
	"strings"
)

//The actual techniques are intialized in hs_techniques.go, and actually defined in hst_*.go files.
//Techniques is ALL technies. CheapTechniques is techniques that are reasonably cheap to compute.
//ExpensiveTechniques is techniques that should only be used if all else has failed.
var Techniques []SolveTechnique
var CheapTechniques []SolveTechnique
var ExpensiveTechniques []SolveTechnique

var GuessTechnique SolveTechnique

//EVERY technique, even the weird one like Guess
var AllTechniques []SolveTechnique

//Worst case scenario, how many times we'd call HumanSolve to get a difficulty.
const MAX_DIFFICULTY_ITERATIONS = 50

//This number is the 'Constant' term from the multiple linear regression to learn the weights.
var difficultyConstant float64

//How close we have to get to the average to feel comfortable our difficulty is converging.
const DIFFICULTY_CONVERGENCE = 0.0005

type SolveDirections []*SolveStep

type SolveStep struct {
	//The cells that will be affected by the techinque
	TargetCells CellList
	//The cells that together lead the techinque to being valid
	PointerCells CellList
	//The numbers we will remove (or, in the case of Fill, add)
	//TODO: shouldn't this be renamed TargetNums?
	TargetNums IntSlice
	//The numbers in pointerCells that lead us to remove TargetNums from TargetCells.
	//This is only very rarely needed (at this time only for hiddenSubset techniques)
	PointerNums IntSlice
	//The general technique that underlies this step.
	Technique SolveTechnique
}

func (self *SolveStep) IsUseful(grid *Grid) bool {
	//Returns true IFF calling Apply with this step and the given grid would result in some useful work. Does not modify the gri.d

	//All of this logic is substantially recreated in Apply.

	if self.Technique == nil {
		return false
	}

	//TODO: test this.
	if self.Technique.IsFill() {
		if len(self.TargetCells) == 0 || len(self.TargetNums) == 0 {
			return false
		}
		cell := self.TargetCells[0].InGrid(grid)
		return self.TargetNums[0] != cell.Number()
	} else {
		useful := false
		for _, cell := range self.TargetCells {
			gridCell := cell.InGrid(grid)
			for _, exclude := range self.TargetNums {
				//It's right to use Possible because it includes the logic of "it's not possible if there's a number in there already"
				//TODO: ensure the comment above is correct logically.
				if gridCell.Possible(exclude) {
					useful = true
				}
			}
		}
		return useful
	}
}

func (self *SolveStep) Apply(grid *Grid) {
	//All of this logic is substantially recreated in IsUseful.
	if self.Technique.IsFill() {
		if len(self.TargetCells) == 0 || len(self.TargetNums) == 0 {
			return
		}
		cell := self.TargetCells[0].InGrid(grid)
		cell.SetNumber(self.TargetNums[0])
	} else {
		for _, cell := range self.TargetCells {
			gridCell := cell.InGrid(grid)
			for _, exclude := range self.TargetNums {
				gridCell.setExcluded(exclude, true)
			}
		}
	}
}

func (self *SolveStep) Description() string {
	result := ""
	if self.Technique.IsFill() {
		result += fmt.Sprintf("We put %s in cell %s ", self.TargetNums.Description(), self.TargetCells.Description())
	} else {
		//TODO: pluralize based on length of lists.
		result += fmt.Sprintf("We remove the possibilities %s from cells %s ", self.TargetNums.Description(), self.TargetCells.Description())
	}
	result += "because " + self.Technique.Description(self) + "."
	return result
}

func (self *SolveStep) normalize() {
	//Puts the solve step in its normal status. In practice this means that the various slices are sorted, so that the Description of them is stable.
	self.PointerCells.Sort()
	self.TargetCells.Sort()
	self.TargetNums.Sort()
	self.PointerNums.Sort()
}

func (self SolveDirections) Stats() []string {
	//TODO: test this.
	techniqueCount := make(map[string]int)
	for _, step := range self {
		techniqueCount[step.Technique.Name()] += 1
	}
	var result []string

	//TODO: use a standard divider across the codebase
	divider := "-------------------------"

	result = append(result, divider)
	result = append(result, fmt.Sprintf("Difficulty : %f", self.Difficulty()))
	result = append(result, divider)
	result = append(result, fmt.Sprintf("Step count: %d", len(self)))
	result = append(result, divider)

	//We want a stable ordering for technique counts.
	for _, technique := range AllTechniques {
		result = append(result, fmt.Sprintf("%s : %d", technique.Name(), techniqueCount[technique.Name()]))
	}

	result = append(result, divider)

	return result
}

func (self SolveDirections) Description() []string {

	if len(self) == 0 {
		return []string{""}
	}

	descriptions := make([]string, len(self))

	for i, step := range self {
		intro := ""
		switch i {
		case 0:
			intro = "First, "
		case len(self) - 1:
			intro = "Finally, "
		default:
			//TODO: switch between "then" and "next" randomly.
			intro = "Next, "
		}
		descriptions[i] = intro + strings.ToLower(step.Description())

	}
	return descriptions
}

func (self SolveDirections) Difficulty() float64 {
	//How difficult the solve directions described are. The measure of difficulty we use is
	//just summing up weights we see; this captures:
	//* Number of steps
	//* Average difficulty of steps
	//* Number of hard steps
	//* (kind of) the hardest step: because the difficulties go up expontentionally.

	//This method assumes the weights have been calibrated empirically to give scores between 0.0 and 1.0
	//without normalization here.

	if len(self) == 0 {
		//The puzzle was not able to be solved, apparently.
		return 1.0
	}

	accum := difficultyConstant
	for _, step := range self {
		accum += step.Technique.Difficulty()
	}

	if accum < 0.0 {
		log.Println("Accumuldated difficulty snapped to 0.0:", accum)
		accum = 0.0
	}

	if accum > 1.0 {
		log.Println("Accumulated difficulty snapped to 1.0:", accum)
		accum = 1.0
	}

	return accum
}

func (self SolveDirections) Walkthrough(grid *Grid) string {

	//TODO: test this.

	clone := grid.Copy()
	defer clone.Done()

	DIVIDER := "\n\n--------------------------------------------\n\n"

	intro := fmt.Sprintf("This will take %d steps to solve.", len(self))

	intro += "\nWhen you start, your grid looks like this:\n"

	intro += clone.Diagram()

	intro += "\n"

	intro += DIVIDER

	descriptions := self.Description()

	results := make([]string, len(self))

	for i, description := range descriptions {

		result := description + "\n"
		result += "After doing that, your grid will look like: \n\n"

		self[i].Apply(clone)

		result += clone.Diagram()

		results[i] = result
	}

	return intro + strings.Join(results, DIVIDER) + DIVIDER + "Now the puzzle is solved."
}

func (self *Grid) HumanWalkthrough() string {
	steps := self.HumanSolution()
	return steps.Walkthrough(self)
}

func (self *Grid) HumanSolution() SolveDirections {
	clone := self.Copy()
	defer clone.Done()
	return clone.HumanSolve()
}

type branchPoint struct {
	//The point at which we branched for a guess. non-nil if we are in a branch.
	grid *Grid
	//The step we'll apply to get us into the branch point.
	step *SolveStep
	//Ther other numbers to try from the branch point.
	otherNums IntSlice
	//The steps we've taken since the branch point.
	branchSteps []*SolveStep
	//The earlier branch point
	previousBranchPoint *branchPoint
	nextBranchPoint     *branchPoint
}

func (self *Grid) HumanSolve() SolveDirections {

	var results []*SolveStep

	var branch *branchPoint

	//Note: trying these all in parallel is much slower (~15x) than doing them in sequence.
	//The reason is that in sequence we bailed early as soon as we found one step; now we try them all.

	for !self.Solved() {

		var possibilities []*SolveStep

		if branch != nil && self.Invalid() {
			//We're in a branch, and got to a point where we found an invalidity.
			//We chose the wrong branch. We should unwind to the branch point and go down the other branch.

			//Unwind ourselves
			self.Load(branch.grid.DataString())

			//Throwout the steps down the wrong branch we took.
			branch.branchSteps = nil

			if len(branch.otherNums) > 0 {

				//Pop off the nextNum to do
				nextNum := branch.otherNums[0]
				branch.otherNums = branch.otherNums[1:]

				//Stuff it into the TargetNums for the branch step.
				branch.step.TargetNums = IntSlice{nextNum}

				//Stuff the possibility list with the mangled branch.step.
				possibilities = []*SolveStep{branch.step}
			} else {
				//Well, crap. We're out of luck, nothing more for us to do.
				//TODO: pick a DIFFERENT guess operation at this grid state.
				//TODO: we could also unravel this sub-branch and go up to a higher branching level.
			}

		} else {

			//Normal operation; get potential solve steps by running them all.
			possibilities = runTechniques(CheapTechniques, self)
		}

		//TODO: provide hints to the techniques of where to look based on the last filled cell

		//Now pick one to apply.
		if len(possibilities) == 0 {
			//Okay, let's try the ExpensiveTechniques, as a hail mary.
			possibilities = runTechniques(ExpensiveTechniques, self)
			if len(possibilities) == 0 {
				//Hmm, didn't find any possivbilities. We're getting to be out of options...

				//Try to guess as a hail mary
				possibilities = runTechniques([]SolveTechnique{GuessTechnique}, self)

				if len(possibilities) == 0 {
					//Okay, we're well and truly done--not even any guesses came up with something. Nothing we can do.
					break
				} else {
					//Yay, found something! remember the branch point, so we can jump back to it.

					//Push new branch point onto the doubly-linked list of branch points
					newBranch := &branchPoint{
						previousBranchPoint: branch,
					}
					if branch != nil {
						branch.nextBranchPoint = newBranch
					}
					branch = newBranch

					//We're just going to choose the first one.
					possibilities = possibilities[0:1]

					//TODO: this doesn't hold our special excludes, which we might have worked quite a bit to set up.
					//Ideally we'd have a way to keep those overrides.
					//The worst case is that we have a few unnecessary Cull steps just before the branch point.
					branch.grid = self.Copy()
					branch.step = possibilities[0]
					branch.otherNums = branch.step.PointerNums

					//Null out the branchPointStep's pointerNums; their only point was to communicate out the other possibilities.
					//And from now on they'll just be confusing.
					branch.step.PointerNums = nil

				}
			}
		}

		//TODO: consider if we should stop picking techniques based on their weight here.
		//Now that Find returns a slice instead of a single, we're already much more likely to select an "easy" technique. ... Right?

		possibilitiesWeights := make([]float64, len(possibilities))
		for i, possibility := range possibilities {
			possibilitiesWeights[i] = possibility.Technique.Difficulty()
		}
		step := possibilities[randomIndexWithInvertedWeights(possibilitiesWeights)]

		if branch == nil {
			results = append(results, step)
		} else {
			//We're in a branch point; we don't know if it's the RIGHT branch
			//So keep the steps somewhere else so we can throw them out if we unwind.
			branch.branchSteps = append(branch.branchSteps, step)
		}
		step.Apply(self)

	}

	if !self.Solved() {
		//We couldn't solve the puzzle.
		return nil
	}

	if branch != nil {
		//Apparently we're in the branch where the solution acutally lay. commit those steps and return them.

		//Walk up to the first branch.
		currentBranch := branch
		for {
			if currentBranch.previousBranchPoint == nil {
				break
			}
			currentBranch = currentBranch.previousBranchPoint
		}

		//now currentBranch is the earliest branch point
		//Walk down the list and copy in all of those steps
		for currentBranch != nil {
			results = append(results, currentBranch.branchSteps...)
			currentBranch = currentBranch.nextBranchPoint
		}
	}

	return results
}

func runTechniques(techniques []SolveTechnique, grid *Grid) []*SolveStep {
	numTechniques := len(techniques)
	possibilitiesChan := make(chan []*SolveStep)

	var possibilities []*SolveStep

	for _, technique := range techniques {
		go func(theTechnique SolveTechnique) {
			possibilitiesChan <- theTechnique.Find(grid)
		}(technique)
	}

	//Collect all of the results
	for i := 0; i < numTechniques; i++ {
		for _, possibility := range <-possibilitiesChan {
			possibilities = append(possibilities, possibility)
		}
	}

	return possibilities
}

func (self *Grid) Difficulty() float64 {
	//This can be an extremely expensive method. Do not call repeatedly!
	//returns the difficulty of the grid, which is a number between 0.0 and 1.0.
	//This is a probabilistic measure; repeated calls may return different numbers, although generally we wait for the results to converge.

	//We solve the same puzzle N times, then ask each set of steps for their difficulty, and combine those to come up with the overall difficulty.

	accum := 0.0
	average := 0.0
	lastAverage := 0.0

	for i := 0; i < MAX_DIFFICULTY_ITERATIONS; i++ {
		grid := self.Copy()
		steps := grid.HumanSolve()
		difficulty := steps.Difficulty()

		accum += difficulty
		average = accum / (float64(i) + 1.0)

		if math.Abs(average-lastAverage) < DIFFICULTY_CONVERGENCE {
			//Okay, we've already converged. Just return early!
			return average
		}

		lastAverage = average
	}

	//We weren't converging... oh well!
	return average

}
<commit_msg>TESTS FAIL. Panic if we have too many nested branch points.<commit_after>package sudoku

import (
	"fmt"
	"log"
	"math"
	"strings"
)

//The actual techniques are intialized in hs_techniques.go, and actually defined in hst_*.go files.
//Techniques is ALL technies. CheapTechniques is techniques that are reasonably cheap to compute.
//ExpensiveTechniques is techniques that should only be used if all else has failed.
var Techniques []SolveTechnique
var CheapTechniques []SolveTechnique
var ExpensiveTechniques []SolveTechnique

var GuessTechnique SolveTechnique

//EVERY technique, even the weird one like Guess
var AllTechniques []SolveTechnique

//Worst case scenario, how many times we'd call HumanSolve to get a difficulty.
const MAX_DIFFICULTY_ITERATIONS = 50

//This number is the 'Constant' term from the multiple linear regression to learn the weights.
var difficultyConstant float64

//How close we have to get to the average to feel comfortable our difficulty is converging.
const DIFFICULTY_CONVERGENCE = 0.0005

type SolveDirections []*SolveStep

type SolveStep struct {
	//The cells that will be affected by the techinque
	TargetCells CellList
	//The cells that together lead the techinque to being valid
	PointerCells CellList
	//The numbers we will remove (or, in the case of Fill, add)
	//TODO: shouldn't this be renamed TargetNums?
	TargetNums IntSlice
	//The numbers in pointerCells that lead us to remove TargetNums from TargetCells.
	//This is only very rarely needed (at this time only for hiddenSubset techniques)
	PointerNums IntSlice
	//The general technique that underlies this step.
	Technique SolveTechnique
}

func (self *SolveStep) IsUseful(grid *Grid) bool {
	//Returns true IFF calling Apply with this step and the given grid would result in some useful work. Does not modify the gri.d

	//All of this logic is substantially recreated in Apply.

	if self.Technique == nil {
		return false
	}

	//TODO: test this.
	if self.Technique.IsFill() {
		if len(self.TargetCells) == 0 || len(self.TargetNums) == 0 {
			return false
		}
		cell := self.TargetCells[0].InGrid(grid)
		return self.TargetNums[0] != cell.Number()
	} else {
		useful := false
		for _, cell := range self.TargetCells {
			gridCell := cell.InGrid(grid)
			for _, exclude := range self.TargetNums {
				//It's right to use Possible because it includes the logic of "it's not possible if there's a number in there already"
				//TODO: ensure the comment above is correct logically.
				if gridCell.Possible(exclude) {
					useful = true
				}
			}
		}
		return useful
	}
}

func (self *SolveStep) Apply(grid *Grid) {
	//All of this logic is substantially recreated in IsUseful.
	if self.Technique.IsFill() {
		if len(self.TargetCells) == 0 || len(self.TargetNums) == 0 {
			return
		}
		cell := self.TargetCells[0].InGrid(grid)
		cell.SetNumber(self.TargetNums[0])
	} else {
		for _, cell := range self.TargetCells {
			gridCell := cell.InGrid(grid)
			for _, exclude := range self.TargetNums {
				gridCell.setExcluded(exclude, true)
			}
		}
	}
}

func (self *SolveStep) Description() string {
	result := ""
	if self.Technique.IsFill() {
		result += fmt.Sprintf("We put %s in cell %s ", self.TargetNums.Description(), self.TargetCells.Description())
	} else {
		//TODO: pluralize based on length of lists.
		result += fmt.Sprintf("We remove the possibilities %s from cells %s ", self.TargetNums.Description(), self.TargetCells.Description())
	}
	result += "because " + self.Technique.Description(self) + "."
	return result
}

func (self *SolveStep) normalize() {
	//Puts the solve step in its normal status. In practice this means that the various slices are sorted, so that the Description of them is stable.
	self.PointerCells.Sort()
	self.TargetCells.Sort()
	self.TargetNums.Sort()
	self.PointerNums.Sort()
}

func (self SolveDirections) Stats() []string {
	//TODO: test this.
	techniqueCount := make(map[string]int)
	for _, step := range self {
		techniqueCount[step.Technique.Name()] += 1
	}
	var result []string

	//TODO: use a standard divider across the codebase
	divider := "-------------------------"

	result = append(result, divider)
	result = append(result, fmt.Sprintf("Difficulty : %f", self.Difficulty()))
	result = append(result, divider)
	result = append(result, fmt.Sprintf("Step count: %d", len(self)))
	result = append(result, divider)

	//We want a stable ordering for technique counts.
	for _, technique := range AllTechniques {
		result = append(result, fmt.Sprintf("%s : %d", technique.Name(), techniqueCount[technique.Name()]))
	}

	result = append(result, divider)

	return result
}

func (self SolveDirections) Description() []string {

	if len(self) == 0 {
		return []string{""}
	}

	descriptions := make([]string, len(self))

	for i, step := range self {
		intro := ""
		switch i {
		case 0:
			intro = "First, "
		case len(self) - 1:
			intro = "Finally, "
		default:
			//TODO: switch between "then" and "next" randomly.
			intro = "Next, "
		}
		descriptions[i] = intro + strings.ToLower(step.Description())

	}
	return descriptions
}

func (self SolveDirections) Difficulty() float64 {
	//How difficult the solve directions described are. The measure of difficulty we use is
	//just summing up weights we see; this captures:
	//* Number of steps
	//* Average difficulty of steps
	//* Number of hard steps
	//* (kind of) the hardest step: because the difficulties go up expontentionally.

	//This method assumes the weights have been calibrated empirically to give scores between 0.0 and 1.0
	//without normalization here.

	if len(self) == 0 {
		//The puzzle was not able to be solved, apparently.
		return 1.0
	}

	accum := difficultyConstant
	for _, step := range self {
		accum += step.Technique.Difficulty()
	}

	if accum < 0.0 {
		log.Println("Accumuldated difficulty snapped to 0.0:", accum)
		accum = 0.0
	}

	if accum > 1.0 {
		log.Println("Accumulated difficulty snapped to 1.0:", accum)
		accum = 1.0
	}

	return accum
}

func (self SolveDirections) Walkthrough(grid *Grid) string {

	//TODO: test this.

	clone := grid.Copy()
	defer clone.Done()

	DIVIDER := "\n\n--------------------------------------------\n\n"

	intro := fmt.Sprintf("This will take %d steps to solve.", len(self))

	intro += "\nWhen you start, your grid looks like this:\n"

	intro += clone.Diagram()

	intro += "\n"

	intro += DIVIDER

	descriptions := self.Description()

	results := make([]string, len(self))

	for i, description := range descriptions {

		result := description + "\n"
		result += "After doing that, your grid will look like: \n\n"

		self[i].Apply(clone)

		result += clone.Diagram()

		results[i] = result
	}

	return intro + strings.Join(results, DIVIDER) + DIVIDER + "Now the puzzle is solved."
}

func (self *Grid) HumanWalkthrough() string {
	steps := self.HumanSolution()
	return steps.Walkthrough(self)
}

func (self *Grid) HumanSolution() SolveDirections {
	clone := self.Copy()
	defer clone.Done()
	return clone.HumanSolve()
}

type branchPoint struct {
	//The point at which we branched for a guess. non-nil if we are in a branch.
	grid *Grid
	//The step we'll apply to get us into the branch point.
	step *SolveStep
	//Ther other numbers to try from the branch point.
	otherNums IntSlice
	//The steps we've taken since the branch point.
	branchSteps []*SolveStep
	//The earlier branch point
	previousBranchPoint *branchPoint
	nextBranchPoint     *branchPoint
}

func (self *Grid) HumanSolve() SolveDirections {

	var results []*SolveStep

	var branch *branchPoint

	numBranches := 0

	//Note: trying these all in parallel is much slower (~15x) than doing them in sequence.
	//The reason is that in sequence we bailed early as soon as we found one step; now we try them all.

	for !self.Solved() {

		var possibilities []*SolveStep

		if branch != nil && self.Invalid() {
			//We're in a branch, and got to a point where we found an invalidity.
			//We chose the wrong branch. We should unwind to the branch point and go down the other branch.

			//Unwind ourselves
			self.Load(branch.grid.DataString())

			//Throwout the steps down the wrong branch we took.
			branch.branchSteps = nil

			if len(branch.otherNums) > 0 {

				//Pop off the nextNum to do
				nextNum := branch.otherNums[0]
				branch.otherNums = branch.otherNums[1:]

				//Stuff it into the TargetNums for the branch step.
				branch.step.TargetNums = IntSlice{nextNum}

				//Stuff the possibility list with the mangled branch.step.
				possibilities = []*SolveStep{branch.step}
			} else {
				//Well, crap. We're out of luck, nothing more for us to do.
				//TODO: pick a DIFFERENT guess operation at this grid state.
				//TODO: we could also unravel this sub-branch and go up to a higher branching level.
			}

		} else {

			//Normal operation; get potential solve steps by running them all.
			possibilities = runTechniques(CheapTechniques, self)
		}

		//TODO: provide hints to the techniques of where to look based on the last filled cell

		//Now pick one to apply.
		if len(possibilities) == 0 {
			//Okay, let's try the ExpensiveTechniques, as a hail mary.
			possibilities = runTechniques(ExpensiveTechniques, self)
			if len(possibilities) == 0 {
				//Hmm, didn't find any possivbilities. We're getting to be out of options...

				//Try to guess as a hail mary
				possibilities = runTechniques([]SolveTechnique{GuessTechnique}, self)

				if len(possibilities) == 0 {
					//Okay, we're well and truly done--not even any guesses came up with something. Nothing we can do.
					break
				} else {
					//Yay, found something! remember the branch point, so we can jump back to it.

					//If we branch more than a few times, things are probably REALLY wrong.
					numBranches++
					if numBranches >= 10 {
						panic("Too many branches")
					}

					//Push new branch point onto the doubly-linked list of branch points
					newBranch := &branchPoint{
						previousBranchPoint: branch,
					}
					if branch != nil {
						branch.nextBranchPoint = newBranch
					}
					branch = newBranch

					//We're just going to choose the first one.
					possibilities = possibilities[0:1]

					//TODO: this doesn't hold our special excludes, which we might have worked quite a bit to set up.
					//Ideally we'd have a way to keep those overrides.
					//The worst case is that we have a few unnecessary Cull steps just before the branch point.
					branch.grid = self.Copy()
					branch.step = possibilities[0]
					branch.otherNums = branch.step.PointerNums

					//Null out the branchPointStep's pointerNums; their only point was to communicate out the other possibilities.
					//And from now on they'll just be confusing.
					branch.step.PointerNums = nil

				}
			}
		}

		//TODO: consider if we should stop picking techniques based on their weight here.
		//Now that Find returns a slice instead of a single, we're already much more likely to select an "easy" technique. ... Right?

		possibilitiesWeights := make([]float64, len(possibilities))
		for i, possibility := range possibilities {
			possibilitiesWeights[i] = possibility.Technique.Difficulty()
		}
		step := possibilities[randomIndexWithInvertedWeights(possibilitiesWeights)]

		if branch == nil {
			results = append(results, step)
		} else {
			//We're in a branch point; we don't know if it's the RIGHT branch
			//So keep the steps somewhere else so we can throw them out if we unwind.
			branch.branchSteps = append(branch.branchSteps, step)
		}
		step.Apply(self)

	}

	if !self.Solved() {
		//We couldn't solve the puzzle.
		return nil
	}

	if branch != nil {
		//Apparently we're in the branch where the solution acutally lay. commit those steps and return them.

		//Walk up to the first branch.
		currentBranch := branch
		for {
			if currentBranch.previousBranchPoint == nil {
				break
			}
			currentBranch = currentBranch.previousBranchPoint
		}

		//now currentBranch is the earliest branch point
		//Walk down the list and copy in all of those steps
		for currentBranch != nil {
			results = append(results, currentBranch.branchSteps...)
			currentBranch = currentBranch.nextBranchPoint
		}
	}

	return results
}

func runTechniques(techniques []SolveTechnique, grid *Grid) []*SolveStep {
	numTechniques := len(techniques)
	possibilitiesChan := make(chan []*SolveStep)

	var possibilities []*SolveStep

	for _, technique := range techniques {
		go func(theTechnique SolveTechnique) {
			possibilitiesChan <- theTechnique.Find(grid)
		}(technique)
	}

	//Collect all of the results
	for i := 0; i < numTechniques; i++ {
		for _, possibility := range <-possibilitiesChan {
			possibilities = append(possibilities, possibility)
		}
	}

	return possibilities
}

func (self *Grid) Difficulty() float64 {
	//This can be an extremely expensive method. Do not call repeatedly!
	//returns the difficulty of the grid, which is a number between 0.0 and 1.0.
	//This is a probabilistic measure; repeated calls may return different numbers, although generally we wait for the results to converge.

	//We solve the same puzzle N times, then ask each set of steps for their difficulty, and combine those to come up with the overall difficulty.

	accum := 0.0
	average := 0.0
	lastAverage := 0.0

	for i := 0; i < MAX_DIFFICULTY_ITERATIONS; i++ {
		grid := self.Copy()
		steps := grid.HumanSolve()
		difficulty := steps.Difficulty()

		accum += difficulty
		average = accum / (float64(i) + 1.0)

		if math.Abs(average-lastAverage) < DIFFICULTY_CONVERGENCE {
			//Okay, we've already converged. Just return early!
			return average
		}

		lastAverage = average
	}

	//We weren't converging... oh well!
	return average

}
<|endoftext|>


<commit_before>package jsonrpc

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/cenkalti/rpc2"
)

type jsonCodec struct {
	dec *json.Decoder // for reading JSON values
	enc *json.Encoder // for writing JSON values
	c   io.Closer

	// temporary work space
	msg            message
	serverRequest  serverRequest
	clientRequest  clientRequest
	clientResponse clientResponse

	// JSON-RPC clients can use arbitrary json values as request IDs.
	// Package rpc expects uint64 request IDs.
	// We assign uint64 sequence numbers to incoming requests
	// but save the original request ID in the pending map.
	// When rpc responds, we use the sequence number in
	// the response to find the original request ID.
	mutext  sync.Mutex // protects seq, pending
	pending map[uint64]*json.RawMessage
	seq     uint64
}

func NewJSONCodec(conn io.ReadWriteCloser) rpc2.Codec {
	return &jsonCodec{
		dec:     json.NewDecoder(conn),
		enc:     json.NewEncoder(conn),
		c:       conn,
		pending: make(map[uint64]*json.RawMessage),
	}
}

type clientRequest struct {
	Method string         `json:"method"`
	Params [1]interface{} `json:"params"`
	Id     *uint64        `json:"id"`
}
type serverRequest struct {
	Method string           `json:"method"`
	Params *json.RawMessage `json:"params"`
	Id     *json.RawMessage `json:"id"`
}

type clientResponse struct {
	Id     uint64           `json:"id"`
	Result *json.RawMessage `json:"result"`
	Error  interface{}      `json:"error"`
}
type serverResponse struct {
	Id     *json.RawMessage `json:"id"`
	Result interface{}      `json:"result"`
	Error  interface{}      `json:"error"`
}

type message struct {
	Method string           `json:"method"`
	Params *json.RawMessage `json:"params"`
	Id     *json.RawMessage `json:"id"`
	Result *json.RawMessage `json:"result"`
	Error  interface{}      `json:"error"`
}

func (c *jsonCodec) ReadHeader(req *rpc2.Request, resp *rpc2.Response) error {
	c.msg = message{}
	if err := c.dec.Decode(&c.msg); err != nil {
		return err
	}

	if c.msg.Method != "" {
		// We are server and read a request from client.
		c.serverRequest.Id = c.msg.Id
		c.serverRequest.Method = c.msg.Method
		c.serverRequest.Params = c.msg.Params

		req.Method = c.serverRequest.Method

		// JSON request id can be any JSON value;
		// RPC package expects uint64.  Translate to
		// internal uint64 and save JSON on the side.
		if c.serverRequest.Id == nil {
			// Notification
		} else {
			c.mutext.Lock()
			c.seq++
			c.pending[c.seq] = c.serverRequest.Id
			c.serverRequest.Id = nil
			req.Seq = c.seq
			c.mutext.Unlock()
		}

		return nil

	} else if c.msg.Result != nil {
		// We are client and read a response from server.
		err := json.Unmarshal(*c.msg.Id, &c.clientResponse.Id)
		if err != nil {
			return err
		}
		c.clientResponse.Result = c.msg.Result
		c.clientResponse.Error = c.msg.Error

		resp.Error = ""
		resp.Seq = c.clientResponse.Id
		if c.clientResponse.Error != nil || c.clientResponse.Result == nil {
			x, ok := c.clientResponse.Error.(string)
			if !ok {
				return fmt.Errorf("invalid error %v", c.clientResponse.Error)
			}
			if x == "" {
				x = "unspecified error"
			}
			resp.Error = x
		}
		return nil
	}
	return errors.New("cannot determine message type")
}

var errMissingParams = errors.New("jsonrpc: request body missing params")

func (c *jsonCodec) ReadRequestBody(x interface{}) error {
	if x == nil {
		return nil
	}
	if c.serverRequest.Params == nil {
		return errMissingParams
	}
	// JSON params is array value.
	// RPC params is struct.
	// Unmarshal into array containing struct for now.
	// Should think about making RPC more general.
	var params [1]interface{}
	params[0] = x
	return json.Unmarshal(*c.serverRequest.Params, &params)

}

func (c *jsonCodec) ReadResponseBody(x interface{}) error {
	if x == nil {
		return nil
	}
	return json.Unmarshal(*c.clientResponse.Result, x)
}

func (c *jsonCodec) WriteRequest(r *rpc2.Request, param interface{}) error {
	c.clientRequest.Method = r.Method
	c.clientRequest.Params[0] = param
	if r.Seq == 0 {
		// Notification
		c.clientRequest.Id = nil
	} else {
		seq := r.Seq
		c.clientRequest.Id = &seq
	}
	return c.enc.Encode(&c.clientRequest)
}

var null = json.RawMessage([]byte("null"))

func (c *jsonCodec) WriteResponse(r *rpc2.Response, x interface{}) error {
	var resp serverResponse
	c.mutext.Lock()
	b, ok := c.pending[r.Seq]
	if !ok {
		c.mutext.Unlock()
		return errors.New("invalid sequence number in response")
	}
	delete(c.pending, r.Seq)
	c.mutext.Unlock()

	if b == nil {
		// Invalid request so no id.  Use JSON null.
		b = &null
	}
	resp.Id = b
	resp.Result = x
	if r.Error == "" {
		resp.Error = nil
	} else {
		resp.Error = r.Error
	}
	return c.enc.Encode(resp)
}

func (c *jsonCodec) Close() error {
	return c.c.Close()
}
<commit_msg>json: fix error responses<commit_after>package jsonrpc

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/cenkalti/rpc2"
)

type jsonCodec struct {
	dec *json.Decoder // for reading JSON values
	enc *json.Encoder // for writing JSON values
	c   io.Closer

	// temporary work space
	msg            message
	serverRequest  serverRequest
	clientRequest  clientRequest
	clientResponse clientResponse

	// JSON-RPC clients can use arbitrary json values as request IDs.
	// Package rpc expects uint64 request IDs.
	// We assign uint64 sequence numbers to incoming requests
	// but save the original request ID in the pending map.
	// When rpc responds, we use the sequence number in
	// the response to find the original request ID.
	mutext  sync.Mutex // protects seq, pending
	pending map[uint64]*json.RawMessage
	seq     uint64
}

func NewJSONCodec(conn io.ReadWriteCloser) rpc2.Codec {
	return &jsonCodec{
		dec:     json.NewDecoder(conn),
		enc:     json.NewEncoder(conn),
		c:       conn,
		pending: make(map[uint64]*json.RawMessage),
	}
}

type clientRequest struct {
	Method string         `json:"method"`
	Params [1]interface{} `json:"params"`
	Id     *uint64        `json:"id"`
}
type serverRequest struct {
	Method string           `json:"method"`
	Params *json.RawMessage `json:"params"`
	Id     *json.RawMessage `json:"id"`
}

type clientResponse struct {
	Id     uint64           `json:"id"`
	Result *json.RawMessage `json:"result"`
	Error  interface{}      `json:"error"`
}
type serverResponse struct {
	Id     *json.RawMessage `json:"id"`
	Result interface{}      `json:"result"`
	Error  interface{}      `json:"error"`
}

type message struct {
	Method string           `json:"method"`
	Params *json.RawMessage `json:"params"`
	Id     *json.RawMessage `json:"id"`
	Result *json.RawMessage `json:"result"`
	Error  interface{}      `json:"error"`
}

func (c *jsonCodec) ReadHeader(req *rpc2.Request, resp *rpc2.Response) error {
	c.msg = message{}
	if err := c.dec.Decode(&c.msg); err != nil {
		return err
	}

	if c.msg.Method != "" {
		// We are server and read a request from client.
		c.serverRequest.Id = c.msg.Id
		c.serverRequest.Method = c.msg.Method
		c.serverRequest.Params = c.msg.Params

		req.Method = c.serverRequest.Method

		// JSON request id can be any JSON value;
		// RPC package expects uint64.  Translate to
		// internal uint64 and save JSON on the side.
		if c.serverRequest.Id == nil {
			// Notification
		} else {
			c.mutext.Lock()
			c.seq++
			c.pending[c.seq] = c.serverRequest.Id
			c.serverRequest.Id = nil
			req.Seq = c.seq
			c.mutext.Unlock()
		}

		return nil

	} else if c.msg.Result != nil {
		// We are client and read a response from server.
		err := json.Unmarshal(*c.msg.Id, &c.clientResponse.Id)
		if err != nil {
			return err
		}
		c.clientResponse.Result = c.msg.Result
		c.clientResponse.Error = c.msg.Error

		resp.Error = ""
		resp.Seq = c.clientResponse.Id
		if c.clientResponse.Error != nil || c.clientResponse.Result == nil {
			x, ok := c.clientResponse.Error.(string)
			if !ok {
				return fmt.Errorf("invalid error %v", c.clientResponse.Error)
			}
			if x == "" {
				x = "unspecified error"
			}
			resp.Error = x
		}
		return nil
	}
	return errors.New("cannot determine message type")
}

var errMissingParams = errors.New("jsonrpc: request body missing params")

func (c *jsonCodec) ReadRequestBody(x interface{}) error {
	if x == nil {
		return nil
	}
	if c.serverRequest.Params == nil {
		return errMissingParams
	}
	// JSON params is array value.
	// RPC params is struct.
	// Unmarshal into array containing struct for now.
	// Should think about making RPC more general.
	var params [1]interface{}
	params[0] = x
	return json.Unmarshal(*c.serverRequest.Params, &params)

}

func (c *jsonCodec) ReadResponseBody(x interface{}) error {
	if x == nil {
		return nil
	}
	return json.Unmarshal(*c.clientResponse.Result, x)
}

func (c *jsonCodec) WriteRequest(r *rpc2.Request, param interface{}) error {
	c.clientRequest.Method = r.Method
	c.clientRequest.Params[0] = param
	if r.Seq == 0 {
		// Notification
		c.clientRequest.Id = nil
	} else {
		seq := r.Seq
		c.clientRequest.Id = &seq
	}
	return c.enc.Encode(&c.clientRequest)
}

var null = json.RawMessage([]byte("null"))

func (c *jsonCodec) WriteResponse(r *rpc2.Response, x interface{}) error {
	c.mutext.Lock()
	b, ok := c.pending[r.Seq]
	if !ok {
		c.mutext.Unlock()
		return errors.New("invalid sequence number in response")
	}
	delete(c.pending, r.Seq)
	c.mutext.Unlock()

	if b == nil {
		// Invalid request so no id.  Use JSON null.
		b = &null
	}
	resp := serverResponse{Id: b}
	if r.Error == "" {
		resp.Result = x
	} else {
		resp.Error = r.Error
	}
	return c.enc.Encode(resp)
}

func (c *jsonCodec) Close() error {
	return c.c.Close()
}
<|endoftext|>


<commit_before>// Based on net/rpc/jsonrpc by:
// Copyright 2010 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package jsonrpc2

import (
	"encoding/json"
	"errors"
	"io"
	"math"
	"net"
	"net/rpc"
	"reflect"
	"sync"
)

const seqNotify = math.MaxUint64

type clientCodec struct {
	dec *json.Decoder // for reading JSON values
	enc *json.Encoder // for writing JSON values
	c   io.Closer

	// temporary work space
	resp clientResponse

	// JSON-RPC responses include the request id but not the request method.
	// Package rpc expects both.
	// We save the request method in pending when sending a request
	// and then look it up by request ID when filling out the rpc Response.
	mutex   sync.Mutex        // protects pending
	pending map[uint64]string // map request id to method name
}

// NewClientCodec returns a new rpc.ClientCodec using JSON-RPC 2.0 on conn.
func NewClientCodec(conn io.ReadWriteCloser) rpc.ClientCodec {
	return &clientCodec{
		dec:     json.NewDecoder(conn),
		enc:     json.NewEncoder(conn),
		c:       conn,
		pending: make(map[uint64]string),
	}
}

type clientRequest struct {
	Version string      `json:"jsonrpc"`
	Method  string      `json:"method"`
	Params  interface{} `json:"params,omitempty"`
	ID      *uint64     `json:"id,omitempty"`
}

func (c *clientCodec) WriteRequest(r *rpc.Request, param interface{}) error {
	// If return error: it will be returned as is for this call.
	// Allow param to be only Array, Slice, Map or Struct.
	// When param is nil or uninitialized Map or Slice - omit "params".
	if param != nil {
		switch k := reflect.TypeOf(param).Kind(); k {
		case reflect.Map:
			if reflect.TypeOf(param).Key().Kind() == reflect.String {
				if reflect.ValueOf(param).IsNil() {
					param = nil
				}
			}
		case reflect.Slice:
			if reflect.ValueOf(param).IsNil() {
				param = nil
			}
		case reflect.Array, reflect.Struct:
		case reflect.Ptr:
			switch k := reflect.TypeOf(param).Elem().Kind(); k {
			case reflect.Map:
				if reflect.TypeOf(param).Elem().Key().Kind() == reflect.String {
					if reflect.ValueOf(param).Elem().IsNil() {
						param = nil
					}
				}
			case reflect.Slice:
				if reflect.ValueOf(param).Elem().IsNil() {
					param = nil
				}
			case reflect.Array, reflect.Struct:
			default:
				return NewError(errInternal.Code, "unsupported param type: Ptr to "+k.String())
			}
		default:
			return NewError(errInternal.Code, "unsupported param type: "+k.String())
		}
	}

	var req clientRequest
	if r.Seq != seqNotify {
		c.mutex.Lock()
		c.pending[r.Seq] = r.ServiceMethod
		c.mutex.Unlock()
		req.ID = &r.Seq
	}
	req.Version = "2.0"
	req.Method = r.ServiceMethod
	req.Params = param
	if err := c.enc.Encode(&req); err != nil {
		return NewError(errInternal.Code, err.Error())
	}
	return nil
}

type clientResponse struct {
	Version string           `json:"jsonrpc"`
	ID      *uint64          `json:"id"`
	Result  *json.RawMessage `json:"result,omitempty"`
	Error   *Error           `json:"error,omitempty"`
}

func (r *clientResponse) reset() {
	r.Version = ""
	r.ID = nil
	r.Result = nil
	r.Error = nil
}

func (r *clientResponse) UnmarshalJSON(raw []byte) error {
	r.reset()
	type resp *clientResponse
	if err := json.Unmarshal(raw, resp(r)); err != nil {
		return errors.New("bad response: " + string(raw))
	}

	var o = make(map[string]*json.RawMessage)
	if err := json.Unmarshal(raw, &o); err != nil {
		return errors.New("bad response: " + string(raw))
	}
	_, okVer := o["jsonrpc"]
	_, okID := o["id"]
	_, okRes := o["result"]
	_, okErr := o["error"]
	if !okVer || !okID || !(okRes || okErr) || (okRes && okErr) || len(o) > 3 {
		return errors.New("bad response: " + string(raw))
	}
	if r.Version != "2.0" {
		return errors.New("bad response: " + string(raw))
	}
	if okRes && r.Result == nil {
		r.Result = &null
	}
	if okErr {
		if o["error"] == nil {
			return errors.New("bad response: " + string(raw))
		}
		oe := make(map[string]*json.RawMessage)
		if err := json.Unmarshal(*o["error"], &oe); err != nil {
			return errors.New("bad response: " + string(raw))
		}
		if oe["code"] == nil || oe["message"] == nil {
			return errors.New("bad response: " + string(raw))
		}
		if _, ok := oe["data"]; (!ok && len(oe) > 2) || len(oe) > 3 {
			return errors.New("bad response: " + string(raw))
		}
	}
	if o["id"] == nil && !okErr {
		return errors.New("bad response: " + string(raw))
	}

	return nil
}

func (c *clientCodec) ReadResponseHeader(r *rpc.Response) error {
	// If return err:
	// - io.EOF will became ErrShutdown or io.ErrUnexpectedEOF
	// - it will be returned as is for all pending calls
	// - client will be shutdown
	// So, return io.EOF as is, return *Error for all other errors.
	if err := c.dec.Decode(&c.resp); err != nil {
		if err == io.EOF {
			return err
		}
		return NewError(errInternal.Code, err.Error())
	}
	if c.resp.ID == nil {
		return c.resp.Error
	}

	c.mutex.Lock()
	r.ServiceMethod = c.pending[*c.resp.ID]
	delete(c.pending, *c.resp.ID)
	c.mutex.Unlock()

	r.Error = ""
	r.Seq = *c.resp.ID
	if c.resp.Error != nil {
		r.Error = c.resp.Error.Error()
	}
	return nil
}

func (c *clientCodec) ReadResponseBody(x interface{}) error {
	// If x!=nil and return error e:
	// - this call get e.Error() appended to "reading body "
	// - other pending calls get error as is XXX actually other calls
	//   shouldn't be affected by this error at all, so let's at least
	//   provide different error message for other calls
	if x == nil {
		return nil
	}
	if err := json.Unmarshal(*c.resp.Result, x); err != nil {
		e := NewError(errInternal.Code, err.Error())
		e.Data = NewError(errInternal.Code, "some other Call failed to unmarshal Reply")
		return e
	}
	return nil
}

func (c *clientCodec) Close() error {
	return c.c.Close()
}

// Client represents a JSON RPC 2.0 Client.
// There may be multiple outstanding Calls associated
// with a single Client, and a Client may be used by
// multiple goroutines simultaneously.
//
// It also provides all methods of net/rpc Client.
type Client struct {
	*rpc.Client
	codec *clientCodec
}

// Notify try to invoke the named function. It return error only in case
// it wasn't able to send request.
func (c Client) Notify(serviceMethod string, args interface{}) error {
	req := &rpc.Request{
		ServiceMethod: serviceMethod,
		Seq:           seqNotify,
	}
	return c.codec.WriteRequest(req, args)
}

// NewClient returns a new Client to handle requests to the
// set of services at the other end of the connection.
func NewClient(conn io.ReadWriteCloser) *Client {
	codec := NewClientCodec(conn)
	client := rpc.NewClientWithCodec(codec)
	return &Client{client, codec.(*clientCodec)}
}

// Dial connects to a JSON-RPC 2.0 server at the specified network address.
func Dial(network, address string) (*Client, error) {
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return NewClient(conn), err
}
<commit_msg>allow use of a client codec other than the one supplied this package for client.<commit_after>// Based on net/rpc/jsonrpc by:
// Copyright 2010 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package jsonrpc2

import (
	"encoding/json"
	"errors"
	"io"
	"math"
	"net"
	"net/rpc"
	"reflect"
	"sync"
)

const seqNotify = math.MaxUint64

type clientCodec struct {
	dec *json.Decoder // for reading JSON values
	enc *json.Encoder // for writing JSON values
	c   io.Closer

	// temporary work space
	resp clientResponse

	// JSON-RPC responses include the request id but not the request method.
	// Package rpc expects both.
	// We save the request method in pending when sending a request
	// and then look it up by request ID when filling out the rpc Response.
	mutex   sync.Mutex        // protects pending
	pending map[uint64]string // map request id to method name
}

// NewClientCodec returns a new rpc.ClientCodec using JSON-RPC 2.0 on conn.
func NewClientCodec(conn io.ReadWriteCloser) rpc.ClientCodec {
	return &clientCodec{
		dec:     json.NewDecoder(conn),
		enc:     json.NewEncoder(conn),
		c:       conn,
		pending: make(map[uint64]string),
	}
}

type clientRequest struct {
	Version string      `json:"jsonrpc"`
	Method  string      `json:"method"`
	Params  interface{} `json:"params,omitempty"`
	ID      *uint64     `json:"id,omitempty"`
}

func (c *clientCodec) WriteRequest(r *rpc.Request, param interface{}) error {
	// If return error: it will be returned as is for this call.
	// Allow param to be only Array, Slice, Map or Struct.
	// When param is nil or uninitialized Map or Slice - omit "params".
	if param != nil {
		switch k := reflect.TypeOf(param).Kind(); k {
		case reflect.Map:
			if reflect.TypeOf(param).Key().Kind() == reflect.String {
				if reflect.ValueOf(param).IsNil() {
					param = nil
				}
			}
		case reflect.Slice:
			if reflect.ValueOf(param).IsNil() {
				param = nil
			}
		case reflect.Array, reflect.Struct:
		case reflect.Ptr:
			switch k := reflect.TypeOf(param).Elem().Kind(); k {
			case reflect.Map:
				if reflect.TypeOf(param).Elem().Key().Kind() == reflect.String {
					if reflect.ValueOf(param).Elem().IsNil() {
						param = nil
					}
				}
			case reflect.Slice:
				if reflect.ValueOf(param).Elem().IsNil() {
					param = nil
				}
			case reflect.Array, reflect.Struct:
			default:
				return NewError(errInternal.Code, "unsupported param type: Ptr to "+k.String())
			}
		default:
			return NewError(errInternal.Code, "unsupported param type: "+k.String())
		}
	}

	var req clientRequest
	if r.Seq != seqNotify {
		c.mutex.Lock()
		c.pending[r.Seq] = r.ServiceMethod
		c.mutex.Unlock()
		req.ID = &r.Seq
	}
	req.Version = "2.0"
	req.Method = r.ServiceMethod
	req.Params = param
	if err := c.enc.Encode(&req); err != nil {
		return NewError(errInternal.Code, err.Error())
	}
	return nil
}

type clientResponse struct {
	Version string           `json:"jsonrpc"`
	ID      *uint64          `json:"id"`
	Result  *json.RawMessage `json:"result,omitempty"`
	Error   *Error           `json:"error,omitempty"`
}

func (r *clientResponse) reset() {
	r.Version = ""
	r.ID = nil
	r.Result = nil
	r.Error = nil
}

func (r *clientResponse) UnmarshalJSON(raw []byte) error {
	r.reset()
	type resp *clientResponse
	if err := json.Unmarshal(raw, resp(r)); err != nil {
		return errors.New("bad response: " + string(raw))
	}

	var o = make(map[string]*json.RawMessage)
	if err := json.Unmarshal(raw, &o); err != nil {
		return errors.New("bad response: " + string(raw))
	}
	_, okVer := o["jsonrpc"]
	_, okID := o["id"]
	_, okRes := o["result"]
	_, okErr := o["error"]
	if !okVer || !okID || !(okRes || okErr) || (okRes && okErr) || len(o) > 3 {
		return errors.New("bad response: " + string(raw))
	}
	if r.Version != "2.0" {
		return errors.New("bad response: " + string(raw))
	}
	if okRes && r.Result == nil {
		r.Result = &null
	}
	if okErr {
		if o["error"] == nil {
			return errors.New("bad response: " + string(raw))
		}
		oe := make(map[string]*json.RawMessage)
		if err := json.Unmarshal(*o["error"], &oe); err != nil {
			return errors.New("bad response: " + string(raw))
		}
		if oe["code"] == nil || oe["message"] == nil {
			return errors.New("bad response: " + string(raw))
		}
		if _, ok := oe["data"]; (!ok && len(oe) > 2) || len(oe) > 3 {
			return errors.New("bad response: " + string(raw))
		}
	}
	if o["id"] == nil && !okErr {
		return errors.New("bad response: " + string(raw))
	}

	return nil
}

func (c *clientCodec) ReadResponseHeader(r *rpc.Response) error {
	// If return err:
	// - io.EOF will became ErrShutdown or io.ErrUnexpectedEOF
	// - it will be returned as is for all pending calls
	// - client will be shutdown
	// So, return io.EOF as is, return *Error for all other errors.
	if err := c.dec.Decode(&c.resp); err != nil {
		if err == io.EOF {
			return err
		}
		return NewError(errInternal.Code, err.Error())
	}
	if c.resp.ID == nil {
		return c.resp.Error
	}

	c.mutex.Lock()
	r.ServiceMethod = c.pending[*c.resp.ID]
	delete(c.pending, *c.resp.ID)
	c.mutex.Unlock()

	r.Error = ""
	r.Seq = *c.resp.ID
	if c.resp.Error != nil {
		r.Error = c.resp.Error.Error()
	}
	return nil
}

func (c *clientCodec) ReadResponseBody(x interface{}) error {
	// If x!=nil and return error e:
	// - this call get e.Error() appended to "reading body "
	// - other pending calls get error as is XXX actually other calls
	//   shouldn't be affected by this error at all, so let's at least
	//   provide different error message for other calls
	if x == nil {
		return nil
	}
	if err := json.Unmarshal(*c.resp.Result, x); err != nil {
		e := NewError(errInternal.Code, err.Error())
		e.Data = NewError(errInternal.Code, "some other Call failed to unmarshal Reply")
		return e
	}
	return nil
}

func (c *clientCodec) Close() error {
	return c.c.Close()
}

// Client represents a JSON RPC 2.0 Client.
// There may be multiple outstanding Calls associated
// with a single Client, and a Client may be used by
// multiple goroutines simultaneously.
//
// It also provides all methods of net/rpc Client.
type Client struct {
	*rpc.Client
	codec rpc.ClientCodec
}

// Notify try to invoke the named function. It return error only in case
// it wasn't able to send request.
func (c Client) Notify(serviceMethod string, args interface{}) error {
	req := &rpc.Request{
		ServiceMethod: serviceMethod,
		Seq:           seqNotify,
	}
	return c.codec.WriteRequest(req, args)
}

// NewClient returns a new Client to handle requests to the
// set of services at the other end of the connection.
func NewClient(conn io.ReadWriteCloser) *Client {
	codec := NewClientCodec(conn)
	client := rpc.NewClientWithCodec(codec)
	return &Client{client, codec}
}

// NewClientWithCodec returns a new Client using the given rpc.ClientCodec.
func NewClientWithCodec(codec rpc.ClientCodec) *Client {
	client := rpc.NewClientWithCodec(codec)
	return &Client{client, codec}
}

// Dial connects to a JSON-RPC 2.0 server at the specified network address.
func Dial(network, address string) (*Client, error) {
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return NewClient(conn), err
}
<|endoftext|>


<commit_before>/* Copyright 2014 Ooyala, Inc. All rights reserved.
 *
 * This file is 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 monitor

import (
	"atlantis/supervisor/containers/serialize"
	"atlantis/supervisor/rpc/types"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/jigish/go-flags"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

const (
	OK = iota
	Warning
	Critical
	Uknown
)

type Config struct {
	ContainerFile   string `toml:"container_file"`
	ContainersDir   string `toml:"container_dir"`
	InventoryDir    string `toml:"inventory_dir"`
	SSHIdentity     string `toml:"ssh_identity"`
	SSHUser         string `toml:"ssh_user"`
	CheckName       string `toml:"check_name"`
	CheckDir        string `toml:"check_dir"`
	DefaultGroup    string `toml:"default_group"`
	TimeoutDuration uint   `toml:"timeout_duration"`
	Verbose         bool   `toml:"verbose"`
}

type Opts struct {
	ContainerFile   string `short:"f" long:"container-file" description:"file to get container information"`
	ContainersDir   string `short:"s" long:"containers-dir" description:"directory containing configs for each container"`
	SSHIdentity     string `short:"i" long:"ssh-identity" description:"file containing the SSH key for all containers"`
	SSHUser         string `short:"u" long:"ssh-user" description:"user account to ssh into containers"`
	CheckName       string `short:"n" long:"check-name" description:"service name that will appear in Nagios for the monitor"`
	CheckDir        string `short:"d" long:"check-dir" description:"directory containing all the scripts for the monitoring checks"`
	DefaultGroup    string `short:"g" long:"default-group" description:"default contact group to use if there is no valid group provided"`
	Config          string `short:"c" long:"config-file" default:"/etc/atlantis/supervisor/monitor.toml" description:"the config file to use"`
	TimeoutDuration uint   `short:"t" long:"timeout-duration" description:"max number of seconds to wait for a monitoring check to finish"`
	Verbose         bool   `short:"v" long:"verbose" default:false description:"print verbose debug information"`
}

type ServiceCheck struct {
	Service  string
	User     string
	Identity string
	Host     string
	Port     uint16
	Script   string
}

//TODO(mchandra):Need defaults defined by constants
var config = &Config{
	ContainerFile:   "/etc/atlantis/supervisor/save/containers",
	ContainersDir:   "/etc/atlantis/containers",
	InventoryDir:    "/etc/atlantis/supervisor/inventory",
	SSHIdentity:     "/opt/atlantis/supervisor/master_id_rsa",
	SSHUser:         "root",
	CheckName:       "ContainerMonitor",
	CheckDir:        "/check_mk_checks",
	DefaultGroup:    "atlantis_orphan_apps",
	TimeoutDuration: 11,
	Verbose:         false,
}

func (s *ServiceCheck) cmd() *exec.Cmd {
	return silentSshCmd(s.User, s.Identity, s.Host, s.Script, s.Port)
}

func (s *ServiceCheck) timeOutMsg() string {
	return fmt.Sprintf("%d %s - Timeout occured during check\n", Critical, s.Service)
}

func (s *ServiceCheck) errMsg(err error) string {
	if err != nil {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, err.Error())
	} else {
		return fmt.Sprintf("%d %s - Error encountered while monitoring the service\n", Critical, s.Service)
	}
}

func (s *ServiceCheck) validate(msg string) string {
	m := strings.SplitN(msg, " ", 4)
	if len(m) < 2 {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, "Check validation failed; no service name found in check")
	}
	if m[1] != s.Service {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, "Check validation failed; found service "+m[1]+", expected "+s.Service)
	}

	return msg
}

func (s *ServiceCheck) runCheck(done chan bool) {
	out, err := s.cmd().Output()
	if err != nil {
		fmt.Print(s.errMsg(err))
	} else {
		fmt.Print(s.validate(string(out)))
	}
	done <- true
}

func (s *ServiceCheck) checkWithTimeout(results chan bool, d time.Duration) {
	done := make(chan bool, 1)
	go s.runCheck(done)
	select {
	case <-done:
		results <- true
	case <-time.After(d):
		fmt.Print(s.timeOutMsg())
		results <- true
	}
}

type ContainerCheck struct {
	Name         string
	User         string
	Identity     string
	Directory    string
	Inventory    string
	ContactGroup string
	container    *types.Container
}

type ContainerConfig struct {
	Dependencies map[string]interface{}
}

func (c *ContainerCheck) verifyContactGroup(group string) bool {
	output, err := exec.Command("/usr/bin/cmk_admin", "-l").Output()
	if err != nil {
		fmt.Printf("%d %s - Error listing existing contact_groups for validation, please try again later! Error: %s\n", Critical, c.Name, err.Error())
		return false
	}
	for _, l := range strings.Split(string(output), "\n") {
		cg := strings.TrimSpace(strings.TrimPrefix(l, "*"))
		if cg == c.ContactGroup {
			return true
		}
	}
	return false
}

func (c *ContainerCheck) parseContactGroup() {
	c.ContactGroup = config.DefaultGroup
	config_file := filepath.Join(config.ContainersDir, c.container.ID, "config.json")
	var cont_config ContainerConfig
	if err := serialize.RetrieveObject(config_file, &cont_config); err != nil {
		fmt.Printf("%d %s - Could not retrieve container config %s: %s\n", Critical, c.Name, config_file, err)
	} else {
		dep, ok := cont_config.Dependencies["cmk"]
		if !ok {
			fmt.Printf("%d %s - cmk dep not present, defaulting to %s contact group!\n", OK, c.Name, config.DefaultGroup)
			return
		}
		cmk_dep, ok := dep.(map[string]interface{})
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but value is not map[string]string!\n", Critical, c.Name)
			return
		}
		val, ok := cmk_dep["contact_group"]
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but no contact_group key!\n", Critical, c.Name)
			return
		}
		group, ok := val.(string)
		if ok {
			group = strings.ToLower(group)
			if c.verifyContactGroup(group) {
				c.ContactGroup = group
			} else {
				fmt.Printf("%d %s - Specified contact_group does not exist in cmk! Falling back to default group %s.\n", Critical, c.Name, config.DefaultGroup)
			}
		} else {
			fmt.Printf("%d %s - Value for contact_group key of cmk dep is not a string!\n", Critical, c.Name)
		}
	}
}

func (c *ContainerCheck) updateContactGroup(name string) {
	if len(c.ContactGroup) == 0 {
		c.parseContactGroup()
	}
	inventoryPath := path.Join(c.Inventory, name)
	if _, err := os.Stat(inventoryPath); os.IsNotExist(err) {
		output, err := exec.Command("/usr/bin/cmk_admin", "-s", name, "-a", c.ContactGroup).CombinedOutput()
		if err != nil {
			fmt.Printf("%d %s - Failure to update contact group for service %s. Error: %s\n", OK, c.Name, name, err.Error())
		} else {
			os.Create(inventoryPath)
		}
		if config.Verbose {
			fmt.Printf("\n/usr/bin/cmk_admin -s %s -a %s\n%s\n\n", name, c.ContactGroup, output)
		}
	}
}

func (c *ContainerCheck) Run(t time.Duration, done chan bool) {
	c.updateContactGroup(c.Name)
	defer func() { done <- true }()
	o, err := silentSshCmd(c.User, c.Identity, c.container.Host, "ls "+c.Directory, c.container.SSHPort).Output()
	if err != nil {
		fmt.Printf("%d %s - Error getting checks for container: %s\n", Critical, c.Name, err.Error())
		return
	}
	fmt.Printf("%d %s - Got checks for container\n", OK, c.Name)
	scripts := strings.Split(strings.TrimSpace(string(o)), "\n")
	if len(scripts) == 0 || len(scripts[0]) == 0 {
		// nothing to check on this container, exit
		return
	}
	c.checkAll(scripts, t)
}

func (c *ContainerCheck) checkAll(scripts []string, t time.Duration) {
	results := make(chan bool, len(scripts))
	for _, s := range scripts {
		serviceName := fmt.Sprintf("%s_%s", strings.Split(s, ".")[0], c.container.ID)
		c.updateContactGroup(serviceName)
		go c.serviceCheck(s).checkWithTimeout(results, t)
	}
	for _ = range scripts {
		<-results
	}
}

func (c *ContainerCheck) serviceCheck(script string) *ServiceCheck {
	// The full path to the script is required
	command := fmt.Sprintf("%s/%s %d %s", c.Directory, script, c.container.PrimaryPort, c.container.ID)
	// The service name is obtained be removing the file extension from the script and appending the container
	// id
	serviceName := fmt.Sprintf("%s_%s", strings.Split(script, ".")[0], c.container.ID)
	return &ServiceCheck{serviceName, c.User, c.Identity, c.container.Host, c.container.SSHPort, command}
}

func silentSshCmd(user, identity, host, cmd string, port uint16) *exec.Cmd {
	args := []string{"-q", user + "@" + host, "-i", identity, "-p", fmt.Sprintf("%d", port), "-o", "StrictHostKeyChecking=no", cmd}
	return exec.Command("ssh", args...)
}

func overlayConfig() {
	opts := &Opts{}
	flags.Parse(opts)
	if opts.Config != "" {
		_, err := toml.DecodeFile(opts.Config, config)
		if err != nil {
			// no need to panic here. we have reasonable defaults.
		}
	}
	if opts.ContainerFile != "" {
		config.ContainerFile = opts.ContainerFile
	}
	if opts.ContainersDir != "" {
		config.ContainersDir = opts.ContainersDir
	}
	if opts.SSHIdentity != "" {
		config.SSHIdentity = opts.SSHIdentity
	}
	if opts.SSHUser != "" {
		config.SSHUser = opts.SSHUser
	}
	if opts.CheckDir != "" {
		config.CheckDir = opts.CheckDir
	}
	if opts.CheckName != "" {
		config.CheckName = opts.CheckName
	}
	if opts.DefaultGroup != "" {
		config.DefaultGroup = opts.DefaultGroup
	}
	if opts.TimeoutDuration != 0 {
		config.TimeoutDuration = opts.TimeoutDuration
	}
	if opts.Verbose {
		config.Verbose = true
	}
}

//file containing containers and service name to show in Nagios for the monitor itself
func Run() {
	overlayConfig()
	var contMap map[string]*types.Container
	//Check if folder exists
	_, err := os.Stat(config.ContainerFile)
	if os.IsNotExist(err) {
		fmt.Printf("%d %s - Container file does not exists %s. Likely no live containers present.\n", OK, config.CheckName, config.ContainerFile)
		return
	}
	if err := serialize.RetrieveObject(config.ContainerFile, &contMap); err != nil {
		fmt.Printf("%d %s - Error retrieving %s: %s\n", Critical, config.CheckName, config.ContainerFile, err)
		return
	}
	done := make(chan bool, len(contMap))
	config.SSHIdentity = strings.Replace(config.SSHIdentity, "~", os.Getenv("HOME"), 1)
	for _, c := range contMap {
		if c.Host == "" {
			c.Host = "localhost"
		}
		check := &ContainerCheck{config.CheckName + "_" + c.ID, config.SSHUser, config.SSHIdentity, config.CheckDir, config.InventoryDir, "", c}
		go check.Run(time.Duration(config.TimeoutDuration)*time.Second, done)
	}
	for _ = range contMap {
		<-done
	}
	// Clean up inventories from containers that no longer exist
	err = filepath.Walk(config.InventoryDir, func(path string, _ os.FileInfo, _ error) error {
		if path == config.InventoryDir {
			return nil
		}
		var err error
		split := strings.Split(path, "_")
		cont := split[len(split)-1]
		if _, ok := contMap[cont]; !ok {
			err = os.Remove(path)
		}
		return err
	})
	if err != nil {
		fmt.Printf("%d %s - Error iterating over inventory to delete obsolete markers. Error: %s\n", OK, config.CheckName, err.Error())
	}
}
<commit_msg>don't run checks until after contact group has been set<commit_after>/* Copyright 2014 Ooyala, Inc. All rights reserved.
 *
 * This file is 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 monitor

import (
	"atlantis/supervisor/containers/serialize"
	"atlantis/supervisor/rpc/types"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/jigish/go-flags"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

const (
	OK = iota
	Warning
	Critical
	Uknown
)

type Config struct {
	ContainerFile   string `toml:"container_file"`
	ContainersDir   string `toml:"container_dir"`
	InventoryDir    string `toml:"inventory_dir"`
	SSHIdentity     string `toml:"ssh_identity"`
	SSHUser         string `toml:"ssh_user"`
	CheckName       string `toml:"check_name"`
	CheckDir        string `toml:"check_dir"`
	DefaultGroup    string `toml:"default_group"`
	TimeoutDuration uint   `toml:"timeout_duration"`
	Verbose         bool   `toml:"verbose"`
}

type Opts struct {
	ContainerFile   string `short:"f" long:"container-file" description:"file to get container information"`
	ContainersDir   string `short:"s" long:"containers-dir" description:"directory containing configs for each container"`
	SSHIdentity     string `short:"i" long:"ssh-identity" description:"file containing the SSH key for all containers"`
	SSHUser         string `short:"u" long:"ssh-user" description:"user account to ssh into containers"`
	CheckName       string `short:"n" long:"check-name" description:"service name that will appear in Nagios for the monitor"`
	CheckDir        string `short:"d" long:"check-dir" description:"directory containing all the scripts for the monitoring checks"`
	DefaultGroup    string `short:"g" long:"default-group" description:"default contact group to use if there is no valid group provided"`
	Config          string `short:"c" long:"config-file" default:"/etc/atlantis/supervisor/monitor.toml" description:"the config file to use"`
	TimeoutDuration uint   `short:"t" long:"timeout-duration" description:"max number of seconds to wait for a monitoring check to finish"`
	Verbose         bool   `short:"v" long:"verbose" default:false description:"print verbose debug information"`
}

type ServiceCheck struct {
	Service  string
	User     string
	Identity string
	Host     string
	Port     uint16
	Script   string
}

//TODO(mchandra):Need defaults defined by constants
var config = &Config{
	ContainerFile:   "/etc/atlantis/supervisor/save/containers",
	ContainersDir:   "/etc/atlantis/containers",
	InventoryDir:    "/etc/atlantis/supervisor/inventory",
	SSHIdentity:     "/opt/atlantis/supervisor/master_id_rsa",
	SSHUser:         "root",
	CheckName:       "ContainerMonitor",
	CheckDir:        "/check_mk_checks",
	DefaultGroup:    "atlantis_orphan_apps",
	TimeoutDuration: 11,
	Verbose:         false,
}

func (s *ServiceCheck) cmd() *exec.Cmd {
	return silentSshCmd(s.User, s.Identity, s.Host, s.Script, s.Port)
}

func (s *ServiceCheck) timeOutMsg() string {
	return fmt.Sprintf("%d %s - Timeout occured during check\n", Critical, s.Service)
}

func (s *ServiceCheck) errMsg(err error) string {
	if err != nil {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, err.Error())
	} else {
		return fmt.Sprintf("%d %s - Error encountered while monitoring the service\n", Critical, s.Service)
	}
}

func (s *ServiceCheck) validate(msg string) string {
	m := strings.SplitN(msg, " ", 4)
	if len(m) < 2 {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, "Check validation failed; no service name found in check")
	}
	if m[1] != s.Service {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, "Check validation failed; found service "+m[1]+", expected "+s.Service)
	}

	return msg
}

func (s *ServiceCheck) runCheck(done chan bool) {
	out, err := s.cmd().Output()
	if err != nil {
		fmt.Print(s.errMsg(err))
	} else {
		fmt.Print(s.validate(string(out)))
	}
	done <- true
}

func (s *ServiceCheck) checkWithTimeout(results chan bool, d time.Duration) {
	done := make(chan bool, 1)
	go s.runCheck(done)
	select {
	case <-done:
		results <- true
	case <-time.After(d):
		fmt.Print(s.timeOutMsg())
		results <- true
	}
}

type ContainerCheck struct {
	Name         string
	User         string
	Identity     string
	Directory    string
	Inventory    string
	ContactGroup string
	container    *types.Container
}

type ContainerConfig struct {
	Dependencies map[string]interface{}
}

func (c *ContainerCheck) verifyContactGroup(group string) bool {
	output, err := exec.Command("/usr/bin/cmk_admin", "-l").Output()
	if err != nil {
		fmt.Printf("%d %s - Error listing existing contact_groups for validation, please try again later! Error: %s\n", Critical, c.Name, err.Error())
		return false
	}
	for _, l := range strings.Split(string(output), "\n") {
		cg := strings.TrimSpace(strings.TrimPrefix(l, "*"))
		if cg == c.ContactGroup {
			return true
		}
	}
	return false
}

func (c *ContainerCheck) parseContactGroup() {
	c.ContactGroup = config.DefaultGroup
	config_file := filepath.Join(config.ContainersDir, c.container.ID, "config.json")
	var cont_config ContainerConfig
	if err := serialize.RetrieveObject(config_file, &cont_config); err != nil {
		fmt.Printf("%d %s - Could not retrieve container config %s: %s\n", Critical, c.Name, config_file, err)
	} else {
		dep, ok := cont_config.Dependencies["cmk"]
		if !ok {
			fmt.Printf("%d %s - cmk dep not present, defaulting to %s contact group!\n", OK, c.Name, config.DefaultGroup)
			return
		}
		cmk_dep, ok := dep.(map[string]interface{})
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but value is not map[string]string!\n", Critical, c.Name)
			return
		}
		val, ok := cmk_dep["contact_group"]
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but no contact_group key!\n", Critical, c.Name)
			return
		}
		group, ok := val.(string)
		if ok {
			group = strings.ToLower(group)
			if c.verifyContactGroup(group) {
				c.ContactGroup = group
			} else {
				fmt.Printf("%d %s - Specified contact_group does not exist in cmk! Falling back to default group %s.\n", Critical, c.Name, config.DefaultGroup)
			}
		} else {
			fmt.Printf("%d %s - Value for contact_group key of cmk dep is not a string!\n", Critical, c.Name)
		}
	}
}

func (c *ContainerCheck) updateContactGroup(name string) (updated bool) {
	if len(c.ContactGroup) == 0 {
		c.parseContactGroup()
	}
	inventoryPath := path.Join(c.Inventory, name)
	if _, err := os.Stat(inventoryPath); os.IsNotExist(err) {
		output, err := exec.Command("/usr/bin/cmk_admin", "-s", name, "-a", c.ContactGroup).CombinedOutput()
		if err != nil {
			fmt.Printf("%d %s - Failure to update contact group for service %s. Error: %s\n", OK, c.Name, name, err.Error())
		} else {
			os.Create(inventoryPath)
			updated = true
		}
		if config.Verbose {
			fmt.Printf("\n/usr/bin/cmk_admin -s %s -a %s\n%s\n\n", name, c.ContactGroup, output)
		}
	}
	return
}

func (c *ContainerCheck) Run(t time.Duration, done chan bool) {
	defer func() { done <- true }()
	if c.updateContactGroup(c.Name) {
		return
	}
	o, err := silentSshCmd(c.User, c.Identity, c.container.Host, "ls "+c.Directory, c.container.SSHPort).Output()
	if err != nil {
		fmt.Printf("%d %s - Error getting checks for container: %s\n", Critical, c.Name, err.Error())
		return
	}
	fmt.Printf("%d %s - Got checks for container\n", OK, c.Name)
	scripts := strings.Split(strings.TrimSpace(string(o)), "\n")
	if len(scripts) == 0 || len(scripts[0]) == 0 {
		// nothing to check on this container, exit
		return
	}
	c.checkAll(scripts, t)
}

func (c *ContainerCheck) checkAll(scripts []string, t time.Duration) {
	results := make(chan bool, len(scripts))
	for _, s := range scripts {
		serviceName := fmt.Sprintf("%s_%s", strings.Split(s, ".")[0], c.container.ID)
		if c.updateContactGroup(serviceName) {
			results <- true
		} else {
			go c.serviceCheck(s).checkWithTimeout(results, t)
		}
	}
	for _ = range scripts {
		<-results
	}
}

func (c *ContainerCheck) serviceCheck(script string) *ServiceCheck {
	// The full path to the script is required
	command := fmt.Sprintf("%s/%s %d %s", c.Directory, script, c.container.PrimaryPort, c.container.ID)
	// The service name is obtained be removing the file extension from the script and appending the container
	// id
	serviceName := fmt.Sprintf("%s_%s", strings.Split(script, ".")[0], c.container.ID)
	return &ServiceCheck{serviceName, c.User, c.Identity, c.container.Host, c.container.SSHPort, command}
}

func silentSshCmd(user, identity, host, cmd string, port uint16) *exec.Cmd {
	args := []string{"-q", user + "@" + host, "-i", identity, "-p", fmt.Sprintf("%d", port), "-o", "StrictHostKeyChecking=no", cmd}
	return exec.Command("ssh", args...)
}

func overlayConfig() {
	opts := &Opts{}
	flags.Parse(opts)
	if opts.Config != "" {
		_, err := toml.DecodeFile(opts.Config, config)
		if err != nil {
			// no need to panic here. we have reasonable defaults.
		}
	}
	if opts.ContainerFile != "" {
		config.ContainerFile = opts.ContainerFile
	}
	if opts.ContainersDir != "" {
		config.ContainersDir = opts.ContainersDir
	}
	if opts.SSHIdentity != "" {
		config.SSHIdentity = opts.SSHIdentity
	}
	if opts.SSHUser != "" {
		config.SSHUser = opts.SSHUser
	}
	if opts.CheckDir != "" {
		config.CheckDir = opts.CheckDir
	}
	if opts.CheckName != "" {
		config.CheckName = opts.CheckName
	}
	if opts.DefaultGroup != "" {
		config.DefaultGroup = opts.DefaultGroup
	}
	if opts.TimeoutDuration != 0 {
		config.TimeoutDuration = opts.TimeoutDuration
	}
	if opts.Verbose {
		config.Verbose = true
	}
}

//file containing containers and service name to show in Nagios for the monitor itself
func Run() {
	overlayConfig()
	var contMap map[string]*types.Container
	//Check if folder exists
	_, err := os.Stat(config.ContainerFile)
	if os.IsNotExist(err) {
		fmt.Printf("%d %s - Container file does not exists %s. Likely no live containers present.\n", OK, config.CheckName, config.ContainerFile)
		return
	}
	if err := serialize.RetrieveObject(config.ContainerFile, &contMap); err != nil {
		fmt.Printf("%d %s - Error retrieving %s: %s\n", Critical, config.CheckName, config.ContainerFile, err)
		return
	}
	done := make(chan bool, len(contMap))
	config.SSHIdentity = strings.Replace(config.SSHIdentity, "~", os.Getenv("HOME"), 1)
	for _, c := range contMap {
		if c.Host == "" {
			c.Host = "localhost"
		}
		check := &ContainerCheck{config.CheckName + "_" + c.ID, config.SSHUser, config.SSHIdentity, config.CheckDir, config.InventoryDir, "", c}
		go check.Run(time.Duration(config.TimeoutDuration)*time.Second, done)
	}
	for _ = range contMap {
		<-done
	}
	// Clean up inventories from containers that no longer exist
	err = filepath.Walk(config.InventoryDir, func(path string, _ os.FileInfo, _ error) error {
		if path == config.InventoryDir {
			return nil
		}
		var err error
		split := strings.Split(path, "_")
		cont := split[len(split)-1]
		if _, ok := contMap[cont]; !ok {
			err = os.Remove(path)
		}
		return err
	})
	if err != nil {
		fmt.Printf("%d %s - Error iterating over inventory to delete obsolete markers. Error: %s\n", OK, config.CheckName, err.Error())
	}
}
<|endoftext|>


<commit_before>package rabbithole

import (
	"encoding/json"
	"net/http"
	"net/url"
)

//
// GET /api/permissions
//

// Example response:
//
// [{"user":"guest","vhost":"/","configure":".*","write":".*","read":".*"}]

type PermissionInfo struct {
	User  string `json:"user"`
	Vhost string `json:"vhost"`

	// Configuration permissions
	Configure string `json:"configure"`
	// Write permissions
	Write     string `json:"write"`
	// Read permissions
	Read      string `json:"read"`
}

// Returns permissions for all users and virtual hosts.
func (c *Client) ListPermissions() (rec []PermissionInfo, err error) {
	req, err := newGETRequest(c, "permissions/")
	if err != nil {
		return []PermissionInfo{}, err
	}

	if err = executeAndParseRequest(req, &rec); err != nil {
		return []PermissionInfo{}, err
	}

	return rec, nil
}

//
// GET /api/users/{user}/permissions
//

// Returns permissions of a specific user.
func (c *Client) ListPermissionsOf(username string) (rec []PermissionInfo, err error) {
	req, err := newGETRequest(c, "users/"+url.QueryEscape(username)+"/permissions")
	if err != nil {
		return []PermissionInfo{}, err
	}

	if err = executeAndParseRequest(req, &rec); err != nil {
		return []PermissionInfo{}, err
	}

	return rec, nil
}

//
// GET /api/permissions/{vhost}/{user}
//

// Returns permissions of user in virtual host.
func (c *Client) GetPermissionsIn(vhost, username string) (rec PermissionInfo, err error) {
	req, err := newGETRequest(c, "permissions/"+url.QueryEscape(vhost)+"/"+url.QueryEscape(username))
	if err != nil {
		return PermissionInfo{}, err
	}

	if err = executeAndParseRequest(req, &rec); err != nil {
		return PermissionInfo{}, err
	}

	return rec, nil
}

//
// PUT /api/permissions/{vhost}/{user}
//

type Permissions struct {
	Configure string `json:"configure"`
	Write     string `json:"write"`
	Read      string `json:"read"`
}

// Updates permissions of user in virtual host.
func (c *Client) UpdatePermissionsIn(vhost, username string, permissions Permissions) (res *http.Response, err error) {
	body, err := json.Marshal(permissions)
	if err != nil {
		return nil, err
	}

	req, err := newRequestWithBody(c, "PUT", "permissions/"+url.QueryEscape(vhost)+"/"+url.QueryEscape(username), body)
	if err != nil {
		return nil, err
	}

	res, err = executeRequest(c, req)
	if err != nil {
		return nil, err
	}

	return res, nil
}

//
// DELETE /api/permissions/{vhost}/{user}
//

func (c *Client) ClearPermissionsIn(vhost, username string) (res *http.Response, err error) {
	req, err := newRequestWithBody(c, "DELETE", "permissions/"+url.QueryEscape(vhost)+"/"+url.QueryEscape(username), nil)
	if err != nil {
		return nil, err
	}

	res, err = executeRequest(c, req)
	if err != nil {
		return nil, err
	}

	return res, nil
}
<commit_msg>Ditto<commit_after>package rabbithole

import (
	"encoding/json"
	"net/http"
	"net/url"
)

//
// GET /api/permissions
//

// Example response:
//
// [{"user":"guest","vhost":"/","configure":".*","write":".*","read":".*"}]

type PermissionInfo struct {
	User  string `json:"user"`
	Vhost string `json:"vhost"`

	// Configuration permissions
	Configure string `json:"configure"`
	// Write permissions
	Write     string `json:"write"`
	// Read permissions
	Read      string `json:"read"`
}

// Returns permissions for all users and virtual hosts.
func (c *Client) ListPermissions() (rec []PermissionInfo, err error) {
	req, err := newGETRequest(c, "permissions/")
	if err != nil {
		return []PermissionInfo{}, err
	}

	if err = executeAndParseRequest(req, &rec); err != nil {
		return []PermissionInfo{}, err
	}

	return rec, nil
}

//
// GET /api/users/{user}/permissions
//

// Returns permissions of a specific user.
func (c *Client) ListPermissionsOf(username string) (rec []PermissionInfo, err error) {
	req, err := newGETRequest(c, "users/"+url.QueryEscape(username)+"/permissions")
	if err != nil {
		return []PermissionInfo{}, err
	}

	if err = executeAndParseRequest(req, &rec); err != nil {
		return []PermissionInfo{}, err
	}

	return rec, nil
}

//
// GET /api/permissions/{vhost}/{user}
//

// Returns permissions of user in virtual host.
func (c *Client) GetPermissionsIn(vhost, username string) (rec PermissionInfo, err error) {
	req, err := newGETRequest(c, "permissions/"+url.QueryEscape(vhost)+"/"+url.QueryEscape(username))
	if err != nil {
		return PermissionInfo{}, err
	}

	if err = executeAndParseRequest(req, &rec); err != nil {
		return PermissionInfo{}, err
	}

	return rec, nil
}

//
// PUT /api/permissions/{vhost}/{user}
//

type Permissions struct {
	Configure string `json:"configure"`
	Write     string `json:"write"`
	Read      string `json:"read"`
}

// Updates permissions of user in virtual host.
func (c *Client) UpdatePermissionsIn(vhost, username string, permissions Permissions) (res *http.Response, err error) {
	body, err := json.Marshal(permissions)
	if err != nil {
		return nil, err
	}

	req, err := newRequestWithBody(c, "PUT", "permissions/"+url.QueryEscape(vhost)+"/"+url.QueryEscape(username), body)
	if err != nil {
		return nil, err
	}

	res, err = executeRequest(c, req)
	if err != nil {
		return nil, err
	}

	return res, nil
}

//
// DELETE /api/permissions/{vhost}/{user}
//

// Clears (deletes) permissions of user in virtual host.
func (c *Client) ClearPermissionsIn(vhost, username string) (res *http.Response, err error) {
	req, err := newRequestWithBody(c, "DELETE", "permissions/"+url.QueryEscape(vhost)+"/"+url.QueryEscape(username), nil)
	if err != nil {
		return nil, err
	}

	res, err = executeRequest(c, req)
	if err != nil {
		return nil, err
	}

	return res, nil
}
<|endoftext|>


<commit_before>package container

import (
	"fmt"

	containertypes "github.com/docker/docker/api/types/container"
	mounttypes "github.com/docker/docker/api/types/mount"
	networktypes "github.com/docker/docker/api/types/network"
	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/go-connections/nat"
)

// WithName sets the name of the container
func WithName(name string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Name = name
	}
}

// WithLinks sets the links of the container
func WithLinks(links ...string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.HostConfig.Links = links
	}
}

// WithImage sets the image of the container
func WithImage(image string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.Image = image
	}
}

// WithCmd sets the comannds of the container
func WithCmd(cmds ...string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.Cmd = strslice.StrSlice(cmds)
	}
}

// WithNetworkMode sets the network mode of the container
func WithNetworkMode(mode string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.HostConfig.NetworkMode = containertypes.NetworkMode(mode)
	}
}

// WithExposedPorts sets the exposed ports of the container
func WithExposedPorts(ports ...string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.ExposedPorts = map[nat.Port]struct{}{}
		for _, port := range ports {
			c.Config.ExposedPorts[nat.Port(port)] = struct{}{}
		}
	}
}

// WithTty sets the TTY mode of the container
func WithTty(tty bool) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.Tty = tty
	}
}

// WithWorkingDir sets the working dir of the container
func WithWorkingDir(dir string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.WorkingDir = dir
	}
}

// WithMount adds an mount
func WithMount(m mounttypes.Mount) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.HostConfig.Mounts = append(c.HostConfig.Mounts, m)
	}
}

// WithVolume sets the volume of the container
func WithVolume(name string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		if c.Config.Volumes == nil {
			c.Config.Volumes = map[string]struct{}{}
		}
		c.Config.Volumes[name] = struct{}{}
	}
}

// WithBind sets the bind mount of the container
func WithBind(src, target string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.HostConfig.Binds = append(c.HostConfig.Binds, fmt.Sprintf("%s:%s", src, target))
	}
}

// WithIPv4 sets the specified ip for the specified network of the container
func WithIPv4(network, ip string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		if c.NetworkingConfig.EndpointsConfig == nil {
			c.NetworkingConfig.EndpointsConfig = map[string]*networktypes.EndpointSettings{}
		}
		if v, ok := c.NetworkingConfig.EndpointsConfig[network]; !ok || v == nil {
			c.NetworkingConfig.EndpointsConfig[network] = &networktypes.EndpointSettings{}
		}
		if c.NetworkingConfig.EndpointsConfig[network].IPAMConfig == nil {
			c.NetworkingConfig.EndpointsConfig[network].IPAMConfig = &networktypes.EndpointIPAMConfig{}
		}
		c.NetworkingConfig.EndpointsConfig[network].IPAMConfig.IPv4Address = ip
	}
}

// WithIPv6 sets the specified ip6 for the specified network of the container
func WithIPv6(network, ip string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		if c.NetworkingConfig.EndpointsConfig == nil {
			c.NetworkingConfig.EndpointsConfig = map[string]*networktypes.EndpointSettings{}
		}
		if v, ok := c.NetworkingConfig.EndpointsConfig[network]; !ok || v == nil {
			c.NetworkingConfig.EndpointsConfig[network] = &networktypes.EndpointSettings{}
		}
		if c.NetworkingConfig.EndpointsConfig[network].IPAMConfig == nil {
			c.NetworkingConfig.EndpointsConfig[network].IPAMConfig = &networktypes.EndpointIPAMConfig{}
		}
		c.NetworkingConfig.EndpointsConfig[network].IPAMConfig.IPv6Address = ip
	}
}

// WithLogDriver sets the log driver to use for the container
func WithLogDriver(driver string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		if c.HostConfig == nil {
			c.HostConfig = &containertypes.HostConfig{}
		}
		c.HostConfig.LogConfig.Type = driver
	}
}

// WithAutoRemove sets the container to be removed on exit
func WithAutoRemove(c *TestContainerConfig) {
	if c.HostConfig == nil {
		c.HostConfig = &containertypes.HostConfig{}
	}
	c.HostConfig.AutoRemove = true
}
<commit_msg>integration/internal/container/ops: rm unused code<commit_after>package container

import (
	"fmt"

	containertypes "github.com/docker/docker/api/types/container"
	mounttypes "github.com/docker/docker/api/types/mount"
	networktypes "github.com/docker/docker/api/types/network"
	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/go-connections/nat"
)

// WithName sets the name of the container
func WithName(name string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Name = name
	}
}

// WithLinks sets the links of the container
func WithLinks(links ...string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.HostConfig.Links = links
	}
}

// WithImage sets the image of the container
func WithImage(image string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.Image = image
	}
}

// WithCmd sets the comannds of the container
func WithCmd(cmds ...string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.Cmd = strslice.StrSlice(cmds)
	}
}

// WithNetworkMode sets the network mode of the container
func WithNetworkMode(mode string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.HostConfig.NetworkMode = containertypes.NetworkMode(mode)
	}
}

// WithExposedPorts sets the exposed ports of the container
func WithExposedPorts(ports ...string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.ExposedPorts = map[nat.Port]struct{}{}
		for _, port := range ports {
			c.Config.ExposedPorts[nat.Port(port)] = struct{}{}
		}
	}
}

// WithTty sets the TTY mode of the container
func WithTty(tty bool) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.Tty = tty
	}
}

// WithWorkingDir sets the working dir of the container
func WithWorkingDir(dir string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.Config.WorkingDir = dir
	}
}

// WithMount adds an mount
func WithMount(m mounttypes.Mount) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.HostConfig.Mounts = append(c.HostConfig.Mounts, m)
	}
}

// WithVolume sets the volume of the container
func WithVolume(name string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		if c.Config.Volumes == nil {
			c.Config.Volumes = map[string]struct{}{}
		}
		c.Config.Volumes[name] = struct{}{}
	}
}

// WithBind sets the bind mount of the container
func WithBind(src, target string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.HostConfig.Binds = append(c.HostConfig.Binds, fmt.Sprintf("%s:%s", src, target))
	}
}

// WithIPv4 sets the specified ip for the specified network of the container
func WithIPv4(network, ip string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		if c.NetworkingConfig.EndpointsConfig == nil {
			c.NetworkingConfig.EndpointsConfig = map[string]*networktypes.EndpointSettings{}
		}
		if v, ok := c.NetworkingConfig.EndpointsConfig[network]; !ok || v == nil {
			c.NetworkingConfig.EndpointsConfig[network] = &networktypes.EndpointSettings{}
		}
		if c.NetworkingConfig.EndpointsConfig[network].IPAMConfig == nil {
			c.NetworkingConfig.EndpointsConfig[network].IPAMConfig = &networktypes.EndpointIPAMConfig{}
		}
		c.NetworkingConfig.EndpointsConfig[network].IPAMConfig.IPv4Address = ip
	}
}

// WithIPv6 sets the specified ip6 for the specified network of the container
func WithIPv6(network, ip string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		if c.NetworkingConfig.EndpointsConfig == nil {
			c.NetworkingConfig.EndpointsConfig = map[string]*networktypes.EndpointSettings{}
		}
		if v, ok := c.NetworkingConfig.EndpointsConfig[network]; !ok || v == nil {
			c.NetworkingConfig.EndpointsConfig[network] = &networktypes.EndpointSettings{}
		}
		if c.NetworkingConfig.EndpointsConfig[network].IPAMConfig == nil {
			c.NetworkingConfig.EndpointsConfig[network].IPAMConfig = &networktypes.EndpointIPAMConfig{}
		}
		c.NetworkingConfig.EndpointsConfig[network].IPAMConfig.IPv6Address = ip
	}
}

// WithLogDriver sets the log driver to use for the container
func WithLogDriver(driver string) func(*TestContainerConfig) {
	return func(c *TestContainerConfig) {
		c.HostConfig.LogConfig.Type = driver
	}
}

// WithAutoRemove sets the container to be removed on exit
func WithAutoRemove(c *TestContainerConfig) {
	c.HostConfig.AutoRemove = true
}
<|endoftext|>