<commit_before>package server

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

	"github.com/bouk/httprouter"
	"github.com/influxdata/chronograf"
)

type dbLinks struct {
	Self string `json:"self"`              // Self link mapping to this resource
	RPs  string `json:"retentionPolicies"` // URL for retention policies for this database
}

type dbResponse struct {
	Name          string       `json:"name"`                        // a unique string identifier for the database
	Duration      string       `json:"duration,omitempty"`          // the duration (when creating a default retention policy)
	Replication   int32        `json:"replication,omitempty"`       // the replication factor (when creating a default retention policy)
	ShardDuration string       `json:"shardDuration,omitempty"`     // the shard duration (when creating a default retention policy)
	RPs           []rpResponse `json:"retentionPolicies,omitempty"` // RPs are the retention policies for a database
	Links         dbLinks      `json:"links"`                       // Links are URI locations related to the database
}

// newDBResponse creates the response for the /databases endpoint
func newDBResponse(srcID int, name string, rps []rpResponse) dbResponse {
	base := "/chronograf/v1/sources"
	return dbResponse{
		Name: name,
		RPs:  rps,
		Links: dbLinks{
			Self: fmt.Sprintf("%s/%d/dbs/%s", base, srcID, name),
			RPs:  fmt.Sprintf("%s/%d/dbs/%s/rps", base, srcID, name),
		},
	}
}

type dbsResponse struct {
	Databases []dbResponse `json:"databases"`
}

type rpLinks struct {
	Self string `json:"self"` // Self link mapping to this resource
}

type rpResponse struct {
	Name          string  `json:"name"`          // a unique string identifier for the retention policy
	Duration      string  `json:"duration"`      // the duration
	Replication   int32   `json:"replication"`   // the replication factor
	ShardDuration string  `json:"shardDuration"` // the shard duration
	Default       bool    `json:"default"`       // whether the RP should be the default
	Links         rpLinks `json:"links"`         // Links are URI locations related to the database
}

// WithLinks adds links to an rpResponse in place
func (r *rpResponse) WithLinks(srcID int, dbName string) {
	base := "/chronograf/v1/sources"
	r.Links = rpLinks{
		Self: fmt.Sprintf("%s/%d/dbs/%s/rps/%s", base, srcID, dbName, r.Name),
	}
}

type rpsResponse struct {
	RetentionPolicies []rpResponse `json:"retentionPolicies"`
}

// GetDatabases queries the list of all databases for a source
func (h *Service) GetDatabases(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	databases, err := db.AllDB(ctx)
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}

	dbs := make([]dbResponse, len(databases))
	for i, d := range databases {
		rps, err := h.allRPs(ctx, db, srcID, d.Name)
		if err != nil {
			Error(w, http.StatusBadRequest, err.Error(), h.Logger)
			return
		}
		dbs[i] = newDBResponse(srcID, d.Name, rps)
	}

	res := dbsResponse{
		Databases: dbs,
	}

	encodeJSON(w, http.StatusOK, res, h.Logger)
}

// NewDatabase creates a new database within the datastore
func (h *Service) NewDatabase(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases

	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	postedDB := &chronograf.Database{}
	if err := json.NewDecoder(r.Body).Decode(postedDB); err != nil {
		invalidJSON(w, h.Logger)
		return
	}

	if err := ValidDatabaseRequest(postedDB); err != nil {
		invalidData(w, err, h.Logger)
		return
	}

	database, err := db.CreateDB(ctx, postedDB)
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}

	rps, err := h.allRPs(ctx, db, srcID, database.Name)
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}
	res := newDBResponse(srcID, database.Name, rps)
	encodeJSON(w, http.StatusCreated, res, h.Logger)
}

// DropDatabase removes a database from a data source
func (h *Service) DropDatabase(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases

	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")

	dropErr := db.DropDB(ctx, dbID)
	if dropErr != nil {
		Error(w, http.StatusBadRequest, dropErr.Error(), h.Logger)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

// RetentionPolicies lists retention policies within a database
func (h *Service) RetentionPolicies(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")
	res, err := h.allRPs(ctx, db, srcID, dbID)
	if err != nil {
		msg := fmt.Sprintf("Unable to connect get RPs %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}
	encodeJSON(w, http.StatusOK, res, h.Logger)
}

func (h *Service) allRPs(ctx context.Context, db chronograf.Databases, srcID int, dbID string) ([]rpResponse, error) {
	allRP, err := db.AllRP(ctx, dbID)
	if err != nil {
		return nil, err
	}

	rps := make([]rpResponse, len(allRP))
	for i, rp := range allRP {
		rp := rpResponse{
			Name:          rp.Name,
			Duration:      rp.Duration,
			Replication:   rp.Replication,
			ShardDuration: rp.ShardDuration,
			Default:       rp.Default,
		}
		rp.WithLinks(srcID, dbID)
		rps[i] = rp
	}
	return rps, nil
}

// NewRetentionPolicy creates a new retention policy for a database
func (h *Service) NewRetentionPolicy(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	postedRP := &chronograf.RetentionPolicy{}
	if err := json.NewDecoder(r.Body).Decode(postedRP); err != nil {
		invalidJSON(w, h.Logger)
		return
	}
	if err := ValidRetentionPolicyRequest(postedRP); err != nil {
		invalidData(w, err, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")
	rp, err := db.CreateRP(ctx, dbID, postedRP)
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}
	res := rpResponse{
		Name:          rp.Name,
		Duration:      rp.Duration,
		Replication:   rp.Replication,
		ShardDuration: rp.ShardDuration,
		Default:       rp.Default,
	}
	res.WithLinks(srcID, dbID)
	encodeJSON(w, http.StatusCreated, res, h.Logger)
}

// UpdateRetentionPolicy modifies an existing retention policy for a database
func (h *Service) UpdateRetentionPolicy(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	postedRP := &chronograf.RetentionPolicy{}
	if err := json.NewDecoder(r.Body).Decode(postedRP); err != nil {
		invalidJSON(w, h.Logger)
		return
	}
	if err := ValidRetentionPolicyRequest(postedRP); err != nil {
		invalidData(w, err, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")
	rpID := httprouter.GetParamFromContext(ctx, "rpid")
	rp, err := db.UpdateRP(ctx, dbID, rpID, postedRP)

	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}

	res := rpResponse{
		Name:          rp.Name,
		Duration:      rp.Duration,
		Replication:   rp.Replication,
		ShardDuration: rp.ShardDuration,
		Default:       rp.Default,
	}
	res.WithLinks(srcID, dbID)
	encodeJSON(w, http.StatusCreated, res, h.Logger)
}

// DropRetentionPolicy removes a retention policy from a database
func (h *Service) DropRetentionPolicy(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")
	rpID := httprouter.GetParamFromContext(ctx, "rpid")
	dropErr := db.DropRP(ctx, dbID, rpID)
	if dropErr != nil {
		Error(w, http.StatusBadRequest, dropErr.Error(), h.Logger)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

// ValidDatabaseRequest checks if the database posted is valid
func ValidDatabaseRequest(d *chronograf.Database) error {
	if len(d.Name) == 0 {
		return fmt.Errorf("name is required")
	}
	return nil
}

// ValidRetentionPolicyRequest checks if a retention policy is valid on POST
func ValidRetentionPolicyRequest(rp *chronograf.RetentionPolicy) error {
	if len(rp.Name) == 0 {
		return fmt.Errorf("name is required")
	}
	if len(rp.Duration) == 0 {
		return fmt.Errorf("duration is required")
	}
	if rp.Replication == 0 {
		return fmt.Errorf("replication factor is invalid")
	}
	return nil
}
<commit_msg>Update rpResponse default to isDefault<commit_after>package server

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

	"github.com/bouk/httprouter"
	"github.com/influxdata/chronograf"
)

type dbLinks struct {
	Self string `json:"self"`              // Self link mapping to this resource
	RPs  string `json:"retentionPolicies"` // URL for retention policies for this database
}

type dbResponse struct {
	Name          string       `json:"name"`                        // a unique string identifier for the database
	Duration      string       `json:"duration,omitempty"`          // the duration (when creating a default retention policy)
	Replication   int32        `json:"replication,omitempty"`       // the replication factor (when creating a default retention policy)
	ShardDuration string       `json:"shardDuration,omitempty"`     // the shard duration (when creating a default retention policy)
	RPs           []rpResponse `json:"retentionPolicies,omitempty"` // RPs are the retention policies for a database
	Links         dbLinks      `json:"links"`                       // Links are URI locations related to the database
}

// newDBResponse creates the response for the /databases endpoint
func newDBResponse(srcID int, name string, rps []rpResponse) dbResponse {
	base := "/chronograf/v1/sources"
	return dbResponse{
		Name: name,
		RPs:  rps,
		Links: dbLinks{
			Self: fmt.Sprintf("%s/%d/dbs/%s", base, srcID, name),
			RPs:  fmt.Sprintf("%s/%d/dbs/%s/rps", base, srcID, name),
		},
	}
}

type dbsResponse struct {
	Databases []dbResponse `json:"databases"`
}

type rpLinks struct {
	Self string `json:"self"` // Self link mapping to this resource
}

type rpResponse struct {
	Name          string  `json:"name"`          // a unique string identifier for the retention policy
	Duration      string  `json:"duration"`      // the duration
	Replication   int32   `json:"replication"`   // the replication factor
	ShardDuration string  `json:"shardDuration"` // the shard duration
	Default       bool    `json:"isDefault"`     // whether the RP should be the default
	Links         rpLinks `json:"links"`         // Links are URI locations related to the database
}

// WithLinks adds links to an rpResponse in place
func (r *rpResponse) WithLinks(srcID int, dbName string) {
	base := "/chronograf/v1/sources"
	r.Links = rpLinks{
		Self: fmt.Sprintf("%s/%d/dbs/%s/rps/%s", base, srcID, dbName, r.Name),
	}
}

type rpsResponse struct {
	RetentionPolicies []rpResponse `json:"retentionPolicies"`
}

// GetDatabases queries the list of all databases for a source
func (h *Service) GetDatabases(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	databases, err := db.AllDB(ctx)
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}

	dbs := make([]dbResponse, len(databases))
	for i, d := range databases {
		rps, err := h.allRPs(ctx, db, srcID, d.Name)
		if err != nil {
			Error(w, http.StatusBadRequest, err.Error(), h.Logger)
			return
		}
		dbs[i] = newDBResponse(srcID, d.Name, rps)
	}

	res := dbsResponse{
		Databases: dbs,
	}

	encodeJSON(w, http.StatusOK, res, h.Logger)
}

// NewDatabase creates a new database within the datastore
func (h *Service) NewDatabase(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases

	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	postedDB := &chronograf.Database{}
	if err := json.NewDecoder(r.Body).Decode(postedDB); err != nil {
		invalidJSON(w, h.Logger)
		return
	}

	if err := ValidDatabaseRequest(postedDB); err != nil {
		invalidData(w, err, h.Logger)
		return
	}

	database, err := db.CreateDB(ctx, postedDB)
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}

	rps, err := h.allRPs(ctx, db, srcID, database.Name)
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}
	res := newDBResponse(srcID, database.Name, rps)
	encodeJSON(w, http.StatusCreated, res, h.Logger)
}

// DropDatabase removes a database from a data source
func (h *Service) DropDatabase(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases

	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")

	dropErr := db.DropDB(ctx, dbID)
	if dropErr != nil {
		Error(w, http.StatusBadRequest, dropErr.Error(), h.Logger)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

// RetentionPolicies lists retention policies within a database
func (h *Service) RetentionPolicies(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")
	res, err := h.allRPs(ctx, db, srcID, dbID)
	if err != nil {
		msg := fmt.Sprintf("Unable to connect get RPs %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}
	encodeJSON(w, http.StatusOK, res, h.Logger)
}

func (h *Service) allRPs(ctx context.Context, db chronograf.Databases, srcID int, dbID string) ([]rpResponse, error) {
	allRP, err := db.AllRP(ctx, dbID)
	if err != nil {
		return nil, err
	}

	rps := make([]rpResponse, len(allRP))
	for i, rp := range allRP {
		rp := rpResponse{
			Name:          rp.Name,
			Duration:      rp.Duration,
			Replication:   rp.Replication,
			ShardDuration: rp.ShardDuration,
			Default:       rp.Default,
		}
		rp.WithLinks(srcID, dbID)
		rps[i] = rp
	}
	return rps, nil
}

// NewRetentionPolicy creates a new retention policy for a database
func (h *Service) NewRetentionPolicy(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	postedRP := &chronograf.RetentionPolicy{}
	if err := json.NewDecoder(r.Body).Decode(postedRP); err != nil {
		invalidJSON(w, h.Logger)
		return
	}
	if err := ValidRetentionPolicyRequest(postedRP); err != nil {
		invalidData(w, err, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")
	rp, err := db.CreateRP(ctx, dbID, postedRP)
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}
	res := rpResponse{
		Name:          rp.Name,
		Duration:      rp.Duration,
		Replication:   rp.Replication,
		ShardDuration: rp.ShardDuration,
		Default:       rp.Default,
	}
	res.WithLinks(srcID, dbID)
	encodeJSON(w, http.StatusCreated, res, h.Logger)
}

// UpdateRetentionPolicy modifies an existing retention policy for a database
func (h *Service) UpdateRetentionPolicy(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	postedRP := &chronograf.RetentionPolicy{}
	if err := json.NewDecoder(r.Body).Decode(postedRP); err != nil {
		invalidJSON(w, h.Logger)
		return
	}
	if err := ValidRetentionPolicyRequest(postedRP); err != nil {
		invalidData(w, err, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")
	rpID := httprouter.GetParamFromContext(ctx, "rpid")
	rp, err := db.UpdateRP(ctx, dbID, rpID, postedRP)

	if err != nil {
		Error(w, http.StatusBadRequest, err.Error(), h.Logger)
		return
	}

	res := rpResponse{
		Name:          rp.Name,
		Duration:      rp.Duration,
		Replication:   rp.Replication,
		ShardDuration: rp.ShardDuration,
		Default:       rp.Default,
	}
	res.WithLinks(srcID, dbID)
	encodeJSON(w, http.StatusCreated, res, h.Logger)
}

// DropRetentionPolicy removes a retention policy from a database
func (h *Service) DropRetentionPolicy(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	srcID, err := paramID("id", r)
	if err != nil {
		Error(w, http.StatusUnprocessableEntity, err.Error(), h.Logger)
		return
	}

	src, err := h.SourcesStore.Get(ctx, srcID)
	if err != nil {
		notFound(w, srcID, h.Logger)
		return
	}

	db := h.Databases
	if err = db.Connect(ctx, &src); err != nil {
		msg := fmt.Sprintf("Unable to connect to source %d: %v", srcID, err)
		Error(w, http.StatusBadRequest, msg, h.Logger)
		return
	}

	dbID := httprouter.GetParamFromContext(ctx, "dbid")
	rpID := httprouter.GetParamFromContext(ctx, "rpid")
	dropErr := db.DropRP(ctx, dbID, rpID)
	if dropErr != nil {
		Error(w, http.StatusBadRequest, dropErr.Error(), h.Logger)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

// ValidDatabaseRequest checks if the database posted is valid
func ValidDatabaseRequest(d *chronograf.Database) error {
	if len(d.Name) == 0 {
		return fmt.Errorf("name is required")
	}
	return nil
}

// ValidRetentionPolicyRequest checks if a retention policy is valid on POST
func ValidRetentionPolicyRequest(rp *chronograf.RetentionPolicy) error {
	if len(rp.Name) == 0 {
		return fmt.Errorf("name is required")
	}
	if len(rp.Duration) == 0 {
		return fmt.Errorf("duration is required")
	}
	if rp.Replication == 0 {
		return fmt.Errorf("replication factor is invalid")
	}
	return nil
}
<|endoftext|>


<commit_before>package cli

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/user"
	"reflect"
	"strconv"
	"strings"
)

var defaultConfig = `{
	"ping" : {
		"timeout" : "2s",
		"interval": "1s",
		"count":	4
	},
	"hping" : {
		"timeout"  : "2s",
		"method"   : "HEAD",
		"data"	   : "mylg",
		"count"	   : 5
	},
	"web" : {
		"port"	   : 8080,
		"address"  : "127.0.0.1"
	},
	"scan" : {
		"port"     : "1-500"
	},
	"trace" : {
		"wait"  : "2s"
	},
	"snmp" : {
		"community" : "public",
		"timeout" : "1s",
		"version" : "2c",
		"retries" : 1,
		"port"    : 161
	}
}`

// Config represents configuration
type Config struct {
	Ping  Ping  `json:"ping"`
	Hping HPing `json:"hping"`
	Web   Web   `json:"web"`
	Scan  Scan  `json:"scan"`
	Trace Trace `json:"trace"`
	Snmp  SNMP  `json:"snmp"`
}

// Ping represents ping command options
type Ping struct {
	Timeout  string `json:"timeout"`
	Interval string `json:"interval"`
	Count    int    `json:"count"`
}

// HPing represents ping command options
type HPing struct {
	Timeout string `json:"timeout"`
	Method  string `json:"method"`
	Data    string `json:"data"`
	Count   int    `json:"count"`
}

// Web represents web command options
type Web struct {
	Port    int    `json:"port"`
	Address string `json:"address"`
}

// Scan represents scan command options
type Scan struct {
	Port string `json:"port"`
}

// Trace represents trace command options
type Trace struct {
	Wait string `json:"wait"`
}

// SNMP represents nms command options
type SNMP struct {
	Community string `json:"community"`
	Timeout   string `json:"timeout"`
	Version   string `json:"version"`
	Retries   int    `json:"retries"`
	Port      int    `json:"port"`
}

// WriteConfig write config to disk
func WriteConfig(cfg Config) error {
	f, err := cfgFile()
	if err != nil {
		return err
	}

	h, err := os.Create(f)
	if err != nil {
		return err
	}

	b, err := json.Marshal(cfg)
	if err != nil {
		return err
	}
	_, err = h.Write(bytes.ToLower(b))
	if err != nil {
		return err
	}
	h.Close()

	return nil
}

// GetOptions returns option(s)/value(s) for specific command
func GetOptions(s interface{}, key string) ([]string, []interface{}) {
	var (
		opts []string
		vals []interface{}
	)
	v := reflect.ValueOf(s)
	t := v.Type()

	for i := 0; i < v.NumField(); i++ {
		if t.Field(i).Name == key {
			f := v.Field(i)
			ft := f.Type()
			for j := 0; j < f.NumField(); j++ {
				vals = append(vals, f.Field(j))
				opts = append(opts, ft.Field(j).Name)
			}
			break
		}
	}
	return opts, vals
}

// GetCMDNames returns command line names
func GetCMDNames(s interface{}) []string {
	var fields []string

	v := reflect.ValueOf(s)
	t := v.Type()

	for i := 0; i < t.NumField(); i++ {
		fields = append(fields, t.Field(i).Name)
	}
	return fields
}

// UpgradeConfig adds / removes new command(s)/option(s)
func UpgradeConfig(cfg *Config) error {
	var (
		conf  map[string]interface{}
		cConf Config
	)

	b := make([]byte, 2048)
	f, err := cfgFile()
	if err != nil {
		return err
	}

	h, err := os.Open(f)
	if err != nil {
		return err
	}
	n, _ := h.Read(b)
	b = b[:n]
	// load saved/old config to conf
	json.Unmarshal(b, &conf)
	// load default config to cConf
	json.Unmarshal([]byte(defaultConfig), &cConf)

	for _, cmd := range GetCMDNames(cConf) {
		opts, vals := GetOptions(cConf, cmd)
		for i, opt := range opts {
			if v, ok := conf[strings.ToLower(cmd)].(interface{}); ok {
				if _, ok = v.(map[string]interface{})[strings.ToLower(opt)]; !ok {
					args := fmt.Sprintf("%s %s %v", cmd, opt, vals[i])
					SetConfig(args, cfg)
				}
			} else {
				// there is new command
				args := fmt.Sprintf("%s %s %v", cmd, opt, vals[i])
				SetConfig(args, cfg)
			}
		}
	}
	return nil
}

// LoadConfig loads configuration
func LoadConfig() Config {
	var cfg Config

	cfg = ReadConfig()
	UpgradeConfig(&cfg)
	return cfg
}

// InitConfig creates new config file
func InitConfig(f string) ([]byte, error) {
	h, err := os.Create(f)
	if err != nil {
		return []byte(""), err
	}

	h.Chmod(os.FileMode(int(0600)))
	h.WriteString(defaultConfig)
	h.Close()

	return []byte(defaultConfig), nil
}

// ReadConfig reads configuration from existing
// or default configuration
func ReadConfig() Config {
	var (
		b    = make([]byte, 2048)
		conf Config
		err  error
	)
	f, err := cfgFile()
	if err != nil {

	}

	h, err := os.Open(f)

	if err != nil {
		switch {
		case os.IsNotExist(err):
			if b, err = InitConfig(f); err != nil {
				println(err.Error())
			}
		case os.IsPermission(err):
			println("cannot read configuration file due to insufficient permissions")
			b = []byte(defaultConfig)
		default:
			println(err.Error())
			b = []byte(defaultConfig)
		}
	} else {
		n, _ := h.Read(b)
		b = b[:n]
	}

	err = json.Unmarshal(b, &conf)
	if err != nil {
		println(err.Error())
		b = []byte(defaultConfig)
		json.Unmarshal(b, &conf)
	}

	return conf
}

// ReadDefaultConfig returns default configuration
func ReadDefaultConfig() (Config, error) {
	var (
		b    = make([]byte, 2048)
		conf Config
	)
	b = []byte(defaultConfig)
	err := json.Unmarshal(b, &conf)
	return conf, err
}

// cfgFile returns config file
func cfgFile() (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", err
	}
	return user.HomeDir + "/.mylg.config", nil
}

//
func optionType(v reflect.Value, opt string) string {
	opt = strings.Title(opt)
	for i := 0; i < v.NumField(); i++ {
		if v.Type().Field(i).Name == opt {
			return v.Type().Field(i).Type.Name()
		}
	}
	return "nan"
}

// SetConfig handles update option's value
func SetConfig(args string, s *Config) error {
	var (
		v       reflect.Value
		integer int64
		float   float64
		err     error
	)

	args = strings.ToLower(args)
	f := strings.Fields(args)
	if len(f) < 3 {
		helpSet()
		return fmt.Errorf("syntax error")
	}

	cmd, opt, val := f[0], f[1], f[2]
	opt = strings.Title(opt)

	v = reflect.ValueOf(s)
	v = reflect.Indirect(v)
	v = v.FieldByName(strings.Title(cmd))

	if !v.IsValid() {
		return fmt.Errorf("invalid command")
	}

	switch optionType(v, opt) {
	case "string":
		// string
		err = SetValue(v.Addr(), opt, fmt.Sprintf("%v", val))
	case "int":
		// integer
		if integer, err = strconv.ParseInt(val, 10, 64); err == nil {
			err = SetValue(v.Addr(), strings.Title(opt), integer)
		} else {
			err = fmt.Errorf("the value should be integer")
		}
	case "float":
		// float
		if float, err = strconv.ParseFloat(val, 64); err == nil {
			err = SetValue(v.Addr(), opt, float)
		} else {
			err = fmt.Errorf("the value should be float")
		}
	default:
		err = fmt.Errorf("invalid option")
	}

	if err != nil {
		return err
	}

	// save config
	err = WriteConfig(*s)
	return err
}

// SetValue set optioni's value
func SetValue(v reflect.Value, rec string, val interface{}) error {

	if v.Kind() != reflect.Ptr {
		return fmt.Errorf("not a pointer value")
	}

	v = reflect.Indirect(v)
	switch v.Kind() {
	case reflect.Int:
		if value, ok := val.(int64); ok {
			v.SetInt(value)
		} else {
			return fmt.Errorf("the value should be integer")
		}
	case reflect.Float64:
		if value, ok := val.(float64); ok {
			v.SetFloat(value)
		} else {
			return fmt.Errorf("the value should be float")
		}
	case reflect.String:
		if value, ok := val.(string); ok {
			v.SetString(value)
		} else {
			return fmt.Errorf("the value shouldn't be number")
		}
	case reflect.Struct:
		for i := 0; i < v.NumField(); i++ {
			if v.Type().Field(i).Name == rec {
				err := SetValue(v.Field(i).Addr(), rec, val)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// ShowConfig prints the configuration
func ShowConfig(s *Config) {
	var v reflect.Value

	v = reflect.ValueOf(s)
	v = reflect.Indirect(v)

	for i := 0; i < v.NumField(); i++ {
		cmd := v.Type().Field(i).Name
		cmd = strings.ToLower(cmd)

		vv := v.Field(i).Addr()
		vv = reflect.Indirect(vv)

		for j := 0; j < vv.NumField(); j++ {
			subCmd := vv.Type().Field(j).Name
			subCmd = strings.ToLower(subCmd)
			value := vv.Field(j)
			fmt.Printf("set %-8s %-10s %v\n", cmd, subCmd, value)
		}
	}
}

// helpSet shows set command
func helpSet() {
	println(`
          usage:
               set command option value
          example:
               set ping timeout 2s
	`)

}
<commit_msg>add snmp v3 params<commit_after>package cli

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"os/user"
	"reflect"
	"strconv"
	"strings"
)

var defaultConfig = `{
	"ping" : {
		"timeout" : "2s",
		"interval": "1s",
		"count":	4
	},
	"hping" : {
		"timeout"  : "2s",
		"method"   : "HEAD",
		"data"	   : "mylg",
		"count"	   : 5
	},
	"web" : {
		"port"	   : 8080,
		"address"  : "127.0.0.1"
	},
	"scan" : {
		"port"     : "1-500"
	},
	"trace" : {
		"wait"  : "2s"
	},
	"snmp" : {
		"community"     : "public",
		"timeout"       : "1s",
		"version"       : "2c",
		"retries"       : 1,
		"port"          : 161,
		"securitylevel" : "noauthnopriv",
		"authpass"      : "nopass",
		"authproto"     : "sha",
		"Privacypass"   : "nopass",
		"Privacyproto"  : "aes"
	}
}`

// Config represents configuration
type Config struct {
	Ping  Ping  `json:"ping"`
	Hping HPing `json:"hping"`
	Web   Web   `json:"web"`
	Scan  Scan  `json:"scan"`
	Trace Trace `json:"trace"`
	Snmp  SNMP  `json:"snmp"`
}

// Ping represents ping command options
type Ping struct {
	Timeout  string `json:"timeout"`
	Interval string `json:"interval"`
	Count    int    `json:"count"`
}

// HPing represents ping command options
type HPing struct {
	Timeout string `json:"timeout"`
	Method  string `json:"method"`
	Data    string `json:"data"`
	Count   int    `json:"count"`
}

// Web represents web command options
type Web struct {
	Port    int    `json:"port"`
	Address string `json:"address"`
}

// Scan represents scan command options
type Scan struct {
	Port string `json:"port"`
}

// Trace represents trace command options
type Trace struct {
	Wait string `json:"wait"`
}

// SNMP represents nms command options
type SNMP struct {
	Community     string `json:"community"`
	Timeout       string `json:"timeout"`
	Version       string `json:"version"`
	Retries       int    `json:"retries"`
	Port          int    `json:"port"`
	Securitylevel string `json:"securitylevel"`
	Authpass      string `json:"authpass"`
	Authproto     string `json:"authproto"`
	Privacypass   string `json:"privacypass"`
	Privacyproto  string `json:"privacyproto"`
}

// WriteConfig write config to disk
func WriteConfig(cfg Config) error {
	f, err := cfgFile()
	if err != nil {
		return err
	}

	h, err := os.Create(f)
	if err != nil {
		return err
	}

	b, err := json.Marshal(cfg)
	if err != nil {
		return err
	}
	_, err = h.Write(bytes.ToLower(b))
	if err != nil {
		return err
	}
	h.Close()

	return nil
}

// GetOptions returns option(s)/value(s) for specific command
func GetOptions(s interface{}, key string) ([]string, []interface{}) {
	var (
		opts []string
		vals []interface{}
	)
	v := reflect.ValueOf(s)
	t := v.Type()

	for i := 0; i < v.NumField(); i++ {
		if t.Field(i).Name == key {
			f := v.Field(i)
			ft := f.Type()
			for j := 0; j < f.NumField(); j++ {
				vals = append(vals, f.Field(j))
				opts = append(opts, ft.Field(j).Name)
			}
			break
		}
	}
	return opts, vals
}

// GetCMDNames returns command line names
func GetCMDNames(s interface{}) []string {
	var fields []string

	v := reflect.ValueOf(s)
	t := v.Type()

	for i := 0; i < t.NumField(); i++ {
		fields = append(fields, t.Field(i).Name)
	}
	return fields
}

// UpgradeConfig adds / removes new command(s)/option(s)
func UpgradeConfig(cfg *Config) error {
	var (
		conf  map[string]interface{}
		cConf Config
	)

	b := make([]byte, 2048)
	f, err := cfgFile()
	if err != nil {
		return err
	}

	h, err := os.Open(f)
	if err != nil {
		return err
	}
	n, _ := h.Read(b)
	b = b[:n]
	// load saved/old config to conf
	json.Unmarshal(b, &conf)
	// load default config to cConf
	json.Unmarshal([]byte(defaultConfig), &cConf)

	for _, cmd := range GetCMDNames(cConf) {
		opts, vals := GetOptions(cConf, cmd)
		for i, opt := range opts {
			if v, ok := conf[strings.ToLower(cmd)].(interface{}); ok {
				if _, ok = v.(map[string]interface{})[strings.ToLower(opt)]; !ok {
					args := fmt.Sprintf("%s %s %v", cmd, opt, vals[i])
					SetConfig(args, cfg)
				}
			} else {
				// there is new command
				args := fmt.Sprintf("%s %s %v", cmd, opt, vals[i])
				SetConfig(args, cfg)
			}
		}
	}
	return nil
}

// LoadConfig loads configuration
func LoadConfig() Config {
	var cfg Config

	cfg = ReadConfig()
	UpgradeConfig(&cfg)
	return cfg
}

// InitConfig creates new config file
func InitConfig(f string) ([]byte, error) {
	h, err := os.Create(f)
	if err != nil {
		return []byte(""), err
	}

	h.Chmod(os.FileMode(int(0600)))
	h.WriteString(defaultConfig)
	h.Close()

	return []byte(defaultConfig), nil
}

// ReadConfig reads configuration from existing
// or default configuration
func ReadConfig() Config {
	var (
		b    = make([]byte, 2048)
		conf Config
		err  error
	)
	f, err := cfgFile()
	if err != nil {

	}

	h, err := os.Open(f)

	if err != nil {
		switch {
		case os.IsNotExist(err):
			if b, err = InitConfig(f); err != nil {
				println(err.Error())
			}
		case os.IsPermission(err):
			println("cannot read configuration file due to insufficient permissions")
			b = []byte(defaultConfig)
		default:
			println(err.Error())
			b = []byte(defaultConfig)
		}
	} else {
		n, _ := h.Read(b)
		b = b[:n]
	}

	err = json.Unmarshal(b, &conf)
	if err != nil {
		println(err.Error())
		b = []byte(defaultConfig)
		json.Unmarshal(b, &conf)
	}

	return conf
}

// ReadDefaultConfig returns default configuration
func ReadDefaultConfig() (Config, error) {
	var (
		b    = make([]byte, 2048)
		conf Config
	)
	b = []byte(defaultConfig)
	err := json.Unmarshal(b, &conf)
	return conf, err
}

// cfgFile returns config file
func cfgFile() (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", err
	}
	return user.HomeDir + "/.mylg.config", nil
}

//
func optionType(v reflect.Value, opt string) string {
	opt = strings.Title(opt)
	for i := 0; i < v.NumField(); i++ {
		if v.Type().Field(i).Name == opt {
			return v.Type().Field(i).Type.Name()
		}
	}
	return "nan"
}

// SetConfig handles update option's value
func SetConfig(args string, s *Config) error {
	var (
		v       reflect.Value
		integer int64
		float   float64
		err     error
	)

	args = strings.ToLower(args)
	f := strings.Fields(args)
	if len(f) < 3 {
		helpSet()
		return fmt.Errorf("syntax error")
	}

	cmd, opt, val := f[0], f[1], f[2]
	opt = strings.Title(opt)

	v = reflect.ValueOf(s)
	v = reflect.Indirect(v)
	v = v.FieldByName(strings.Title(cmd))

	if !v.IsValid() {
		return fmt.Errorf("invalid command")
	}

	switch optionType(v, opt) {
	case "string":
		// string
		err = SetValue(v.Addr(), opt, fmt.Sprintf("%v", val))
	case "int":
		// integer
		if integer, err = strconv.ParseInt(val, 10, 64); err == nil {
			err = SetValue(v.Addr(), strings.Title(opt), integer)
		} else {
			err = fmt.Errorf("the value should be integer")
		}
	case "float":
		// float
		if float, err = strconv.ParseFloat(val, 64); err == nil {
			err = SetValue(v.Addr(), opt, float)
		} else {
			err = fmt.Errorf("the value should be float")
		}
	default:
		err = fmt.Errorf("invalid option")
	}

	if err != nil {
		return err
	}

	// save config
	err = WriteConfig(*s)
	return err
}

// SetValue set optioni's value
func SetValue(v reflect.Value, rec string, val interface{}) error {

	if v.Kind() != reflect.Ptr {
		return fmt.Errorf("not a pointer value")
	}

	v = reflect.Indirect(v)
	switch v.Kind() {
	case reflect.Int:
		if value, ok := val.(int64); ok {
			v.SetInt(value)
		} else {
			return fmt.Errorf("the value should be integer")
		}
	case reflect.Float64:
		if value, ok := val.(float64); ok {
			v.SetFloat(value)
		} else {
			return fmt.Errorf("the value should be float")
		}
	case reflect.String:
		if value, ok := val.(string); ok {
			v.SetString(value)
		} else {
			return fmt.Errorf("the value shouldn't be number")
		}
	case reflect.Struct:
		for i := 0; i < v.NumField(); i++ {
			if v.Type().Field(i).Name == rec {
				err := SetValue(v.Field(i).Addr(), rec, val)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// ShowConfig prints the configuration
func ShowConfig(s *Config) {
	var v reflect.Value

	v = reflect.ValueOf(s)
	v = reflect.Indirect(v)

	for i := 0; i < v.NumField(); i++ {
		cmd := v.Type().Field(i).Name
		cmd = strings.ToLower(cmd)

		vv := v.Field(i).Addr()
		vv = reflect.Indirect(vv)

		for j := 0; j < vv.NumField(); j++ {
			subCmd := vv.Type().Field(j).Name
			subCmd = strings.ToLower(subCmd)
			value := vv.Field(j)
			fmt.Printf("set %-8s %-15s %v\n", cmd, subCmd, value)
		}
	}
}

// helpSet shows set command
func helpSet() {
	println(`
          usage:
               set command option value
          example:
               set ping timeout 2s
	`)

}
<|endoftext|>


<commit_before>// This example declares a durable Exchange, an ephemeral (auto-delete) Queue,
// binds the Queue to the Exchange with a routing key, and consumes every
// message published to that Exchange with that routing key.
//
package main

import (
	"flag"
	"fmt"
	"github.com/streadway/amqp"
	"log"
	"time"
)

var (
	uri          *string = flag.String("uri", "amqp://guest:guest@localhost:5672/", "AMQP URI")
	exchangeName *string = flag.String("exchange", "test-exchange", "AMQP exchange name")
	queueName    *string = flag.String("queue", "test-queue", "AMQP queue name")
	routingKey   *string = flag.String("routing-key", "test-key", "AMQP routing key")
	lifetimeStr  *string = flag.String("lifetime", "5s", "lifetime of process before shutdown (0s=infinite)")
	lifetime     time.Duration
)

func init() {
	flag.Parse()
	var err error
	if lifetime, err = time.ParseDuration(*lifetimeStr); err != nil {
		log.Fatalf("%s: invalid -lifetime", *lifetimeStr)
	}
}

func main() {
	c, err := NewConsumer(*uri, *exchangeName, *queueName, *routingKey)
	if err != nil {
		log.Fatalf("%s", err)
	}

	if lifetime == 0 {
		log.Printf("running forever")
		select {}
	} else {
		log.Printf("running for %s", lifetime)
		time.Sleep(lifetime)
	}

	log.Printf("shutting down")
	c.Shutdown()
}

type Consumer struct {
	Channel *amqp.Channel
	quit    chan chan bool
}

func NewConsumer(amqpURI, exchange, queue, routing string) (*Consumer, error) {
	c := &Consumer{
		Channel: nil,
		quit:    make(chan chan bool),
	}

	log.Printf("dialing %s", amqpURI)
	connection, err := amqp.Dial(amqpURI)
	if err != nil {
		return nil, fmt.Errorf("Dial: %s", err)
	}

	log.Printf("got Connection, getting Channel")
	channel, err := connection.Channel()
	if err != nil {
		return nil, fmt.Errorf("Channel: %s", err)
	}
	c.Channel = channel

	log.Printf("got Channel, declaring Exchange (%s)", exchange)
	noArgs := amqp.Table{}
	if err := c.Channel.ExchangeDeclare(
		exchange,          // name of the exchange
		amqp.UntilDeleted, // lifetime = durable
		amqp.Direct,       // type
		false,             // internal
		false,             // noWait
		noArgs,            // arguments
	); err != nil {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Exchange Declare: %s", err)
	}

	log.Printf("declared Exchange, declaring Queue (%s)", queue)
	queueState, err := channel.QueueDeclare(
		queue,            // name of the queue
		amqp.UntilUnused, // lifetime = auto-delete
		false,            // exclusive
		false,            // noWait
		noArgs,           // arguments
	)
	if err != nil {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Queue Declare: %s", err)
	}
	if !queueState.Declared {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Queue Declare: somehow Undeclared")
	}

	log.Printf("declared Queue, binding to Exchange (routing '%s')", routing)
	if err := channel.QueueBind(
		queue,    // name of the queue
		routing,  // routingKey
		exchange, // sourceExchange
		false,    // noWait
		noArgs,   // arguments
	); err != nil {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Queue Bind: %s", err)
	}

	log.Printf("Queue bound to Exchange, starting Consume")
	deliveries, err := channel.Consume(
		queue,  // name
		false,  // noAck
		false,  // exclusive
		false,  // noLocal
		false,  // noWait
		"",     // consumerTag,
		noArgs, // arguments
		nil,    // deliveries (ie. create a deliveries channel for me)
	)
	if err != nil {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Queue Consume: %s", err)
	}

	go c.handle(deliveries)
	return c, nil
}

func (c *Consumer) Shutdown() {
	q := make(chan bool)
	c.quit <- q
	<-q
}

func (c *Consumer) handle(deliveries chan amqp.Delivery) {
	for {
		select {
		case q := <-c.quit:
			closeChannel(c.Channel)
			q <- true
			return
		case d := <-deliveries:
			log.Printf(
				"got %dB delivery: [%v] %s",
				len(d.Body),
				d.DeliveryTag,
				d.Body,
			)
		}
	}
}

func closeChannel(c *amqp.Channel) {
	if err := c.Close(); err != nil {
		log.Printf("AMQP Channel Close error: %s", err)
	}
}
<commit_msg>Shutdown when the delivery channel closes, likely on channel closure or Cancel<commit_after>// This example declares a durable Exchange, an ephemeral (auto-delete) Queue,
// binds the Queue to the Exchange with a routing key, and consumes every
// message published to that Exchange with that routing key.
//
package main

import (
	"flag"
	"fmt"
	"github.com/streadway/amqp"
	"log"
	"time"
)

var (
	uri          *string = flag.String("uri", "amqp://guest:guest@localhost:5672/", "AMQP URI")
	exchangeName *string = flag.String("exchange", "test-exchange", "AMQP exchange name")
	queueName    *string = flag.String("queue", "test-queue", "AMQP queue name")
	routingKey   *string = flag.String("routing-key", "test-key", "AMQP routing key")
	lifetimeStr  *string = flag.String("lifetime", "5s", "lifetime of process before shutdown (0s=infinite)")
	lifetime     time.Duration
)

func init() {
	flag.Parse()
	var err error
	if lifetime, err = time.ParseDuration(*lifetimeStr); err != nil {
		log.Fatalf("%s: invalid -lifetime", *lifetimeStr)
	}
}

func main() {
	c, err := NewConsumer(*uri, *exchangeName, *queueName, *routingKey)
	if err != nil {
		log.Fatalf("%s", err)
	}

	if lifetime == 0 {
		log.Printf("running forever")
		select {}
	} else {
		log.Printf("running for %s", lifetime)
		time.Sleep(lifetime)
	}

	log.Printf("shutting down")
	c.Shutdown()
}

type Consumer struct {
	Channel *amqp.Channel
	quit    chan chan bool
}

func NewConsumer(amqpURI, exchange, queue, routing string) (*Consumer, error) {
	c := &Consumer{
		Channel: nil,
		quit:    make(chan chan bool),
	}

	log.Printf("dialing %s", amqpURI)
	connection, err := amqp.Dial(amqpURI)
	if err != nil {
		return nil, fmt.Errorf("Dial: %s", err)
	}

	log.Printf("got Connection, getting Channel")
	channel, err := connection.Channel()
	if err != nil {
		return nil, fmt.Errorf("Channel: %s", err)
	}
	c.Channel = channel

	log.Printf("got Channel, declaring Exchange (%s)", exchange)
	noArgs := amqp.Table{}
	if err := c.Channel.ExchangeDeclare(
		exchange,          // name of the exchange
		amqp.UntilDeleted, // lifetime = durable
		amqp.Direct,       // type
		false,             // internal
		false,             // noWait
		noArgs,            // arguments
	); err != nil {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Exchange Declare: %s", err)
	}

	log.Printf("declared Exchange, declaring Queue (%s)", queue)
	queueState, err := channel.QueueDeclare(
		queue,            // name of the queue
		amqp.UntilUnused, // lifetime = auto-delete
		false,            // exclusive
		false,            // noWait
		noArgs,           // arguments
	)
	if err != nil {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Queue Declare: %s", err)
	}
	if !queueState.Declared {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Queue Declare: somehow Undeclared")
	}

	log.Printf("declared Queue, binding to Exchange (routing '%s')", routing)
	if err := channel.QueueBind(
		queue,    // name of the queue
		routing,  // routingKey
		exchange, // sourceExchange
		false,    // noWait
		noArgs,   // arguments
	); err != nil {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Queue Bind: %s", err)
	}

	log.Printf("Queue bound to Exchange, starting Consume")
	deliveries, err := channel.Consume(
		queue,  // name
		false,  // noAck
		false,  // exclusive
		false,  // noLocal
		false,  // noWait
		"",     // consumerTag,
		noArgs, // arguments
		nil,    // deliveries (ie. create a deliveries channel for me)
	)
	if err != nil {
		closeChannel(c.Channel)
		return nil, fmt.Errorf("Queue Consume: %s", err)
	}

	go c.handle(deliveries)
	return c, nil
}

func (c *Consumer) Shutdown() {
	q := make(chan bool)
	c.quit <- q
	<-q
}

func (c *Consumer) handle(deliveries chan amqp.Delivery) {
	for {
		select {
		case q := <-c.quit:
			closeChannel(c.Channel)
			q <- true
			return
		case d, ok := <-deliveries:
			if ok {
				log.Printf(
					"got %dB delivery: [%v] %s",
					len(d.Body),
					d.DeliveryTag,
					d.Body,
				)
			} else {
				log.Printf("delivery channel closed")
				c.Shutdown()
				return
			}
		}
	}
}

func closeChannel(c *amqp.Channel) {
	if err := c.Close(); err != nil {
		log.Printf("AMQP Channel Close error: %s", err)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"sync"
	"testing"

	"github.com/buchgr/bazel-remote/cache"
	"github.com/buchgr/bazel-remote/cache/disk"
	"github.com/buchgr/bazel-remote/utils"

	pb "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
	"github.com/golang/protobuf/proto"
)

func TestDownloadFile(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	blobSize := int64(1024)

	hash, err := testutils.CreateCacheFile(filepath.Join(cacheDir, "cas"), blobSize)
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, blobSize, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")

	req, err := http.NewRequest("GET", "/cas/"+hash, bytes.NewReader([]byte{}))
	if err != nil {
		t.Fatal(err)
	}
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.CacheHandler)
	handler.ServeHTTP(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Fatal("Handler returned wrong status code",
			"expected", http.StatusOK,
			"got", status,
		)
	}

	rsp := rr.Result()
	if contentLen := rsp.ContentLength; contentLen != blobSize {
		t.Error("GET request returned wrong content length",
			"expected", blobSize,
			"got", contentLen)
	}

	hasher := sha256.New()
	io.Copy(hasher, rsp.Body)
	if actualHash := hex.EncodeToString(hasher.Sum(nil)); actualHash != hash {
		t.Error("Received the wrong content.",
			"expected hash", hash,
			"actualHash", actualHash,
		)
	}

	req, err = http.NewRequest("HEAD", "/cas/"+hash, bytes.NewReader([]byte{}))
	if err != nil {
		t.Fatal(err)
	}
	handler.ServeHTTP(rr, req)
	rsp = rr.Result()
	if contentLen := rsp.ContentLength; contentLen != blobSize {
		t.Error("HEAD request returned wrong content length",
			"expected", blobSize,
			"got", contentLen)
	}
}

func TestUploadFilesConcurrently(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	const NumUploads = 1000

	var requests [NumUploads]*http.Request
	for i := 0; i < NumUploads; i++ {
		data, hash := testutils.RandomDataAndHash(1024)
		r, err := http.NewRequest("PUT", "/cas/"+hash, bytes.NewReader(data))
		if err != nil {
			t.Fatal(err)
		}
		requests[i] = r
	}

	c := disk.New(cacheDir, 1000*1024, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	handler := http.HandlerFunc(h.CacheHandler)

	var wg sync.WaitGroup
	wg.Add(len(requests))
	for _, request := range requests {
		go func(request *http.Request) {
			defer wg.Done()
			rr := httptest.NewRecorder()
			handler.ServeHTTP(rr, request)

			if status := rr.Code; status != http.StatusOK {
				t.Error("Handler returned wrong status code",
					"expected", http.StatusOK,
					"got", status,
				)
			}
		}(request)
	}

	wg.Wait()

	f, err := os.Open(cacheDir)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	files, err := f.Readdir(-1)
	if err != nil {
		t.Fatal(err)
	}
	var totalSize int64
	for _, fileinfo := range files {
		if !fileinfo.IsDir() {
			size := fileinfo.Size()
			if size != 1024 {
				t.Error("Expected all files to be 1024 bytes.", "Got", size)
			}
			totalSize += fileinfo.Size()
		}
	}

	// Test that purging worked and kept cache size in check.
	if upperBound := int64(NumUploads) * 1024 * 1024; totalSize > upperBound {
		t.Error("Cache size too big. Expected at most", upperBound, "got",
			totalSize)
	}
}

func TestUploadSameFileConcurrently(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	data, hash := testutils.RandomDataAndHash(1024)

	c := disk.New(cacheDir, 1024, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	handler := http.HandlerFunc(h.CacheHandler)

	var wg sync.WaitGroup
	wg.Add(100)
	for i := 0; i < 100; i++ {
		go func() {
			defer wg.Done()

			rr := httptest.NewRecorder()

			request, err := http.NewRequest("PUT", "/cas/"+hash, bytes.NewReader(data))
			if err != nil {
				t.Error(err)
				return
			}

			handler.ServeHTTP(rr, request)

			if status := rr.Code; status != http.StatusOK {
				resp, _ := ioutil.ReadAll(rr.Body)
				t.Error("Handler returned wrong status code",
					"expected", http.StatusOK,
					"got", status,
					string(resp),
				)
			}
		}()
	}

	wg.Wait()
}

func TestUploadCorruptedFile(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	data, hash := testutils.RandomDataAndHash(1024)
	corruptedData := data[:999]

	r, err := http.NewRequest("PUT", "/cas/"+hash, bytes.NewReader(corruptedData))
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, 2048, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.CacheHandler)
	handler.ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusInternalServerError {
		t.Error("Handler returned wrong status code",
			"expected", http.StatusInternalServerError,
			"got", status)
	}

	// Check that no file was saved in the cache
	f, err := os.Open(cacheDir)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	entries, err := f.Readdir(-1)
	if err != nil {
		t.Fatal(err)
	}
	for _, fileEntry := range entries {
		if !fileEntry.IsDir() {
			t.Error("Unexpected file in the cache", fileEntry.Name())
		}
	}
}

func TestUploadEmptyActionResult(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	data, hash := testutils.RandomDataAndHash(0)

	r, err := http.NewRequest("PUT", "/ac/"+hash, bytes.NewReader(data))
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, 2048, nil)
	validate := true
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), validate, "")
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.CacheHandler)
	handler.ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatal("Handler returned wrong status code",
			"expected", http.StatusOK,
			"got", status)
	}

	cacheFile := filepath.Join(cacheDir, "ac", hash[:2], hash)
	cachedData, err := ioutil.ReadFile(cacheFile)
	if err != nil {
		t.Fatal(err)
	}
	if len(cachedData) == 0 {
		t.Fatal("expected non-zero length ActionResult to be cached")
	}

	ar := pb.ActionResult{}
	err = proto.Unmarshal(cachedData, &ar)
	if err != nil {
		t.Fatal(err)
	}
	if ar.ExecutionMetadata == nil {
		t.Fatal("expected non-nil ExecutionMetadata")
	}
	ar.ExecutionMetadata = nil

	remarshaled, err := proto.Marshal(&ar)
	if err != nil {
		t.Fatal(err)
	}
	if len(remarshaled) != 0 {
		t.Fatal("expected zero-length blob once the metadata is stripped")
	}
}

func TestEmptyBlobAvailable(t *testing.T) {
	testEmptyBlobAvailable(t, "HEAD")
	testEmptyBlobAvailable(t, "GET")
}

func testEmptyBlobAvailable(t *testing.T, method string) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	data, hash := testutils.RandomDataAndHash(0)
	r, err := http.NewRequest(method, "/cas/"+hash, bytes.NewReader(data))
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, 2048, nil)
	validate := true
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), validate, "")
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.CacheHandler)
	handler.ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatal("Handler returned wrong status code for",
			"expected", http.StatusOK,
			"got", status)
	}
}

func TestStatusPage(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	r, err := http.NewRequest("GET", "/status", bytes.NewReader([]byte{}))
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, 2048, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.StatusPageHandler)
	handler.ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Error("StatusPageHandler returned wrong status code",
			"expected", http.StatusOK,
			"got", status)
	}

	var data statusPageData
	err = json.Unmarshal(rr.Body.Bytes(), &data)
	if err != nil {
		t.Fatal(err)
	}

	if numFiles := data.NumFiles; numFiles != 0 {
		t.Error("StatusPageHandler returned wrong number of files",
			"expected", 0,
			"got", numFiles)
	}
}

func TestParseRequestURL(t *testing.T) {
	{
		_, _, err := parseRequestURL("invalid/url", true)
		if err == nil {
			t.Error("Failed to reject an invalid URL")
		}
	}

	const aSha256sum = "fec3be77b8aa0d307ed840581ded3d114c86f36d4914c81e33a72877020c0603"

	{
		kind, hash, err := parseRequestURL("cas/"+aSha256sum, true)
		if err != nil {
			t.Error("Failed to parse a valid CAS URL")
		}
		if hash != aSha256sum {
			t.Error("Cache key parsed incorrectly")
		}
		if kind != cache.CAS {
			t.Errorf("Expected kind CAS but got %s", kind)
		}
	}

	{
		kind, hash, err := parseRequestURL("ac/"+aSha256sum, true)
		if err != nil {
			t.Error("Failed to parse a valid AC URL")
		}
		if hash != aSha256sum {
			t.Error("Cache key parsed incorrectly")
		}
		if kind != cache.AC {
			t.Errorf("Expected kind AC but got %s", kind)
		}
	}

	{
		kind, hash, err := parseRequestURL("prefix/ac/"+aSha256sum, true)
		if err != nil {
			t.Error("Failed to parse a valid AC URL with prefix")
		}
		if hash != aSha256sum {
			t.Error("Cache key parsed incorrectly")
		}
		if kind != cache.AC {
			t.Errorf("Expected kind AC but got %s", kind)
		}
	}

	{
		kind, hash, err := parseRequestURL("prefix/ac/"+aSha256sum, false)
		if err != nil {
			t.Error("Failed to parse a valid AC URL with prefix")
		}
		if hash != aSha256sum {
			t.Error("Cache key parsed incorrectly")
		}
		if kind != cache.RAW {
			t.Errorf("Expected kind RAW but got %s", kind)
		}
	}
}

type fakeResponseWriter struct {
	statusCode *int
	response   string
}

func (r fakeResponseWriter) Header() http.Header {
	return http.Header{}
}

func (r fakeResponseWriter) Write(data []byte) (int, error) {
	r.response = string(data)
	return 0, nil
}

func (r fakeResponseWriter) WriteHeader(statusCode int) {
	*r.statusCode = statusCode
}

func TestRemoteReturnsNotFound(t *testing.T) {
	cacheDir, err := ioutil.TempDir("", "bazel-remote")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(cacheDir)
	emptyCache := disk.New(cacheDir, 1024, nil)

	h := NewHTTPCache(emptyCache, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	// create a fake http.Request
	_, hash := testutils.RandomDataAndHash(1024)
	url, _ := url.Parse(fmt.Sprintf("http://localhost:8080/ac/%s", hash))
	reader := bytes.NewReader([]byte{})
	body := ioutil.NopCloser(reader)
	req := &http.Request{
		Method:     "GET",
		URL:        url,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Body:       body,
	}
	statusCode := 0
	respWriter := fakeResponseWriter{
		statusCode: &statusCode,
	}
	h.CacheHandler(respWriter, req)
	if statusCode != http.StatusNotFound {
		t.Errorf("Wrong status code, expected %d, got %d", http.StatusNotFound, statusCode)
	}
}
<commit_msg>also log the method in testEmptyBlobAvailable<commit_after>package server

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"sync"
	"testing"

	"github.com/buchgr/bazel-remote/cache"
	"github.com/buchgr/bazel-remote/cache/disk"
	"github.com/buchgr/bazel-remote/utils"

	pb "github.com/bazelbuild/remote-apis/build/bazel/remote/execution/v2"
	"github.com/golang/protobuf/proto"
)

func TestDownloadFile(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	blobSize := int64(1024)

	hash, err := testutils.CreateCacheFile(filepath.Join(cacheDir, "cas"), blobSize)
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, blobSize, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")

	req, err := http.NewRequest("GET", "/cas/"+hash, bytes.NewReader([]byte{}))
	if err != nil {
		t.Fatal(err)
	}
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.CacheHandler)
	handler.ServeHTTP(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Fatal("Handler returned wrong status code",
			"expected", http.StatusOK,
			"got", status,
		)
	}

	rsp := rr.Result()
	if contentLen := rsp.ContentLength; contentLen != blobSize {
		t.Error("GET request returned wrong content length",
			"expected", blobSize,
			"got", contentLen)
	}

	hasher := sha256.New()
	io.Copy(hasher, rsp.Body)
	if actualHash := hex.EncodeToString(hasher.Sum(nil)); actualHash != hash {
		t.Error("Received the wrong content.",
			"expected hash", hash,
			"actualHash", actualHash,
		)
	}

	req, err = http.NewRequest("HEAD", "/cas/"+hash, bytes.NewReader([]byte{}))
	if err != nil {
		t.Fatal(err)
	}
	handler.ServeHTTP(rr, req)
	rsp = rr.Result()
	if contentLen := rsp.ContentLength; contentLen != blobSize {
		t.Error("HEAD request returned wrong content length",
			"expected", blobSize,
			"got", contentLen)
	}
}

func TestUploadFilesConcurrently(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	const NumUploads = 1000

	var requests [NumUploads]*http.Request
	for i := 0; i < NumUploads; i++ {
		data, hash := testutils.RandomDataAndHash(1024)
		r, err := http.NewRequest("PUT", "/cas/"+hash, bytes.NewReader(data))
		if err != nil {
			t.Fatal(err)
		}
		requests[i] = r
	}

	c := disk.New(cacheDir, 1000*1024, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	handler := http.HandlerFunc(h.CacheHandler)

	var wg sync.WaitGroup
	wg.Add(len(requests))
	for _, request := range requests {
		go func(request *http.Request) {
			defer wg.Done()
			rr := httptest.NewRecorder()
			handler.ServeHTTP(rr, request)

			if status := rr.Code; status != http.StatusOK {
				t.Error("Handler returned wrong status code",
					"expected", http.StatusOK,
					"got", status,
				)
			}
		}(request)
	}

	wg.Wait()

	f, err := os.Open(cacheDir)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	files, err := f.Readdir(-1)
	if err != nil {
		t.Fatal(err)
	}
	var totalSize int64
	for _, fileinfo := range files {
		if !fileinfo.IsDir() {
			size := fileinfo.Size()
			if size != 1024 {
				t.Error("Expected all files to be 1024 bytes.", "Got", size)
			}
			totalSize += fileinfo.Size()
		}
	}

	// Test that purging worked and kept cache size in check.
	if upperBound := int64(NumUploads) * 1024 * 1024; totalSize > upperBound {
		t.Error("Cache size too big. Expected at most", upperBound, "got",
			totalSize)
	}
}

func TestUploadSameFileConcurrently(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	data, hash := testutils.RandomDataAndHash(1024)

	c := disk.New(cacheDir, 1024, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	handler := http.HandlerFunc(h.CacheHandler)

	var wg sync.WaitGroup
	wg.Add(100)
	for i := 0; i < 100; i++ {
		go func() {
			defer wg.Done()

			rr := httptest.NewRecorder()

			request, err := http.NewRequest("PUT", "/cas/"+hash, bytes.NewReader(data))
			if err != nil {
				t.Error(err)
				return
			}

			handler.ServeHTTP(rr, request)

			if status := rr.Code; status != http.StatusOK {
				resp, _ := ioutil.ReadAll(rr.Body)
				t.Error("Handler returned wrong status code",
					"expected", http.StatusOK,
					"got", status,
					string(resp),
				)
			}
		}()
	}

	wg.Wait()
}

func TestUploadCorruptedFile(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	data, hash := testutils.RandomDataAndHash(1024)
	corruptedData := data[:999]

	r, err := http.NewRequest("PUT", "/cas/"+hash, bytes.NewReader(corruptedData))
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, 2048, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.CacheHandler)
	handler.ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusInternalServerError {
		t.Error("Handler returned wrong status code",
			"expected", http.StatusInternalServerError,
			"got", status)
	}

	// Check that no file was saved in the cache
	f, err := os.Open(cacheDir)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	entries, err := f.Readdir(-1)
	if err != nil {
		t.Fatal(err)
	}
	for _, fileEntry := range entries {
		if !fileEntry.IsDir() {
			t.Error("Unexpected file in the cache", fileEntry.Name())
		}
	}
}

func TestUploadEmptyActionResult(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	data, hash := testutils.RandomDataAndHash(0)

	r, err := http.NewRequest("PUT", "/ac/"+hash, bytes.NewReader(data))
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, 2048, nil)
	validate := true
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), validate, "")
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.CacheHandler)
	handler.ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatal("Handler returned wrong status code",
			"expected", http.StatusOK,
			"got", status)
	}

	cacheFile := filepath.Join(cacheDir, "ac", hash[:2], hash)
	cachedData, err := ioutil.ReadFile(cacheFile)
	if err != nil {
		t.Fatal(err)
	}
	if len(cachedData) == 0 {
		t.Fatal("expected non-zero length ActionResult to be cached")
	}

	ar := pb.ActionResult{}
	err = proto.Unmarshal(cachedData, &ar)
	if err != nil {
		t.Fatal(err)
	}
	if ar.ExecutionMetadata == nil {
		t.Fatal("expected non-nil ExecutionMetadata")
	}
	ar.ExecutionMetadata = nil

	remarshaled, err := proto.Marshal(&ar)
	if err != nil {
		t.Fatal(err)
	}
	if len(remarshaled) != 0 {
		t.Fatal("expected zero-length blob once the metadata is stripped")
	}
}

func TestEmptyBlobAvailable(t *testing.T) {
	testEmptyBlobAvailable(t, "HEAD")
	testEmptyBlobAvailable(t, "GET")
}

func testEmptyBlobAvailable(t *testing.T, method string) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	data, hash := testutils.RandomDataAndHash(0)
	r, err := http.NewRequest(method, "/cas/"+hash, bytes.NewReader(data))
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, 2048, nil)
	validate := true
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), validate, "")
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.CacheHandler)
	handler.ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Fatal("Handler returned wrong status code for method",
			method,
			"expected", http.StatusOK,
			"got", status)
	}
}

func TestStatusPage(t *testing.T) {
	cacheDir := testutils.TempDir(t)
	defer os.RemoveAll(cacheDir)

	r, err := http.NewRequest("GET", "/status", bytes.NewReader([]byte{}))
	if err != nil {
		t.Fatal(err)
	}

	c := disk.New(cacheDir, 2048, nil)
	h := NewHTTPCache(c, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	rr := httptest.NewRecorder()
	handler := http.HandlerFunc(h.StatusPageHandler)
	handler.ServeHTTP(rr, r)

	if status := rr.Code; status != http.StatusOK {
		t.Error("StatusPageHandler returned wrong status code",
			"expected", http.StatusOK,
			"got", status)
	}

	var data statusPageData
	err = json.Unmarshal(rr.Body.Bytes(), &data)
	if err != nil {
		t.Fatal(err)
	}

	if numFiles := data.NumFiles; numFiles != 0 {
		t.Error("StatusPageHandler returned wrong number of files",
			"expected", 0,
			"got", numFiles)
	}
}

func TestParseRequestURL(t *testing.T) {
	{
		_, _, err := parseRequestURL("invalid/url", true)
		if err == nil {
			t.Error("Failed to reject an invalid URL")
		}
	}

	const aSha256sum = "fec3be77b8aa0d307ed840581ded3d114c86f36d4914c81e33a72877020c0603"

	{
		kind, hash, err := parseRequestURL("cas/"+aSha256sum, true)
		if err != nil {
			t.Error("Failed to parse a valid CAS URL")
		}
		if hash != aSha256sum {
			t.Error("Cache key parsed incorrectly")
		}
		if kind != cache.CAS {
			t.Errorf("Expected kind CAS but got %s", kind)
		}
	}

	{
		kind, hash, err := parseRequestURL("ac/"+aSha256sum, true)
		if err != nil {
			t.Error("Failed to parse a valid AC URL")
		}
		if hash != aSha256sum {
			t.Error("Cache key parsed incorrectly")
		}
		if kind != cache.AC {
			t.Errorf("Expected kind AC but got %s", kind)
		}
	}

	{
		kind, hash, err := parseRequestURL("prefix/ac/"+aSha256sum, true)
		if err != nil {
			t.Error("Failed to parse a valid AC URL with prefix")
		}
		if hash != aSha256sum {
			t.Error("Cache key parsed incorrectly")
		}
		if kind != cache.AC {
			t.Errorf("Expected kind AC but got %s", kind)
		}
	}

	{
		kind, hash, err := parseRequestURL("prefix/ac/"+aSha256sum, false)
		if err != nil {
			t.Error("Failed to parse a valid AC URL with prefix")
		}
		if hash != aSha256sum {
			t.Error("Cache key parsed incorrectly")
		}
		if kind != cache.RAW {
			t.Errorf("Expected kind RAW but got %s", kind)
		}
	}
}

type fakeResponseWriter struct {
	statusCode *int
	response   string
}

func (r fakeResponseWriter) Header() http.Header {
	return http.Header{}
}

func (r fakeResponseWriter) Write(data []byte) (int, error) {
	r.response = string(data)
	return 0, nil
}

func (r fakeResponseWriter) WriteHeader(statusCode int) {
	*r.statusCode = statusCode
}

func TestRemoteReturnsNotFound(t *testing.T) {
	cacheDir, err := ioutil.TempDir("", "bazel-remote")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(cacheDir)
	emptyCache := disk.New(cacheDir, 1024, nil)

	h := NewHTTPCache(emptyCache, testutils.NewSilentLogger(), testutils.NewSilentLogger(), true, "")
	// create a fake http.Request
	_, hash := testutils.RandomDataAndHash(1024)
	url, _ := url.Parse(fmt.Sprintf("http://localhost:8080/ac/%s", hash))
	reader := bytes.NewReader([]byte{})
	body := ioutil.NopCloser(reader)
	req := &http.Request{
		Method:     "GET",
		URL:        url,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Body:       body,
	}
	statusCode := 0
	respWriter := fakeResponseWriter{
		statusCode: &statusCode,
	}
	h.CacheHandler(respWriter, req)
	if statusCode != http.StatusNotFound {
		t.Errorf("Wrong status code, expected %d, got %d", http.StatusNotFound, statusCode)
	}
}
<|endoftext|>


<commit_before>package bridge

import (
	"errors"
	"log"
	"net"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"

	dockerapi "github.com/fsouza/go-dockerclient"
)

var serviceIDPattern = regexp.MustCompile(`^(.+?):([a-zA-Z0-9][a-zA-Z0-9_.-]+):[0-9]+(?::udp)?$`)

var serviceNamePattern = regexp.MustCompile(`_[0-9]+$`)

type Bridge struct {
	sync.Mutex
	registry       RegistryAdapter
	docker         *dockerapi.Client
	services       map[string][]*Service
	deadContainers map[string]*DeadContainer
	config         Config
}

func New(docker *dockerapi.Client, adapterUri string, config Config) (*Bridge, error) {
	uri, err := url.Parse(adapterUri)
	if err != nil {
		return nil, errors.New("bad adapter uri: " + adapterUri)
	}
	factory, found := AdapterFactories.Lookup(uri.Scheme)
	if !found {
		return nil, errors.New("unrecognized adapter: " + adapterUri)
	}

	log.Println("Using", uri.Scheme, "adapter:", uri)
	return &Bridge{
		docker:         docker,
		config:         config,
		registry:       factory.New(uri),
		services:       make(map[string][]*Service),
		deadContainers: make(map[string]*DeadContainer),
	}, nil
}

func (b *Bridge) Ping() error {
	return b.registry.Ping()
}

func (b *Bridge) Add(containerId string) {
	b.Lock()
	defer b.Unlock()
	b.add(containerId, false)
}

func (b *Bridge) Remove(containerId string) {
	b.remove(containerId, true)
}

func (b *Bridge) RemoveOnExit(containerId string) {
	b.remove(containerId, b.shouldRemove(containerId))
}

func (b *Bridge) Refresh() {
	b.Lock()
	defer b.Unlock()

	for containerId, deadContainer := range b.deadContainers {
		deadContainer.TTL -= b.config.RefreshInterval
		if deadContainer.TTL <= 0 {
			delete(b.deadContainers, containerId)
		}
	}

	for containerId, services := range b.services {
		for _, service := range services {
			err := b.registry.Refresh(service)
			if err != nil {
				log.Println("refresh failed:", service.ID, err)
				continue
			}
			log.Println("refreshed:", containerId[:12], service.ID)
		}
	}
}

func (b *Bridge) Sync(quiet bool) {
	b.Lock()
	defer b.Unlock()

	containers, err := b.docker.ListContainers(dockerapi.ListContainersOptions{})
	if err != nil && quiet {
		log.Println("error listing containers, skipping sync")
		return
	} else if err != nil && !quiet {
		log.Fatal(err)
	}

	log.Printf("Syncing services on %d containers", len(containers))

	// NOTE: This assumes reregistering will do the right thing, i.e. nothing..
	for _, listing := range containers {
		services := b.services[listing.ID]
		if services == nil {
			b.add(listing.ID, quiet)
		} else {
			for _, service := range services {
				err := b.registry.Register(service)
				if err != nil {
					log.Println("sync register failed:", service, err)
				}
			}
		}
	}

	// Clean up services that were registered previously, but aren't
	// acknowledged within registrator
	if b.config.Cleanup {
		// Remove services if its corresponding container is not running
		log.Println("Listing non-exited containers")
		filters := map[string][]string{"status": {"created", "restarting", "running", "paused"}}
		nonExitedContainers, err := b.docker.ListContainers(dockerapi.ListContainersOptions{Filters: filters})
		if err != nil {
			log.Println("error listing nonExitedContainers, skipping sync", err)
			return
		}
		for listingId, _ := range b.services {
			found := false
			for _, container := range nonExitedContainers {
				if listingId == container.ID {
					found = true
					break
				}
			}
			// This is a container that does not exist
			if !found {
				log.Printf("stale: Removing service %s because it does not exist", listingId)
				go b.RemoveOnExit(listingId)
			}
		}

		log.Println("Cleaning up dangling services")
		extServices, err := b.registry.Services()
		if err != nil {
			log.Println("cleanup failed:", err)
			return
		}

	Outer:
		for _, extService := range extServices {
			matches := serviceIDPattern.FindStringSubmatch(extService.ID)
			if len(matches) != 3 {
				// There's no way this was registered by us, so leave it
				continue
			}
			serviceHostname := matches[1]
			if serviceHostname != Hostname {
				// ignore because registered on a different host
				continue
			}
			serviceContainerName := matches[2]
			for _, listing := range b.services {
				for _, service := range listing {
					if service.Name == extService.Name && serviceContainerName == service.Origin.container.Name[1:] {
						continue Outer
					}
				}
			}
			log.Println("dangling:", extService.ID)
			err := b.registry.Deregister(extService)
			if err != nil {
				log.Println("deregister failed:", extService.ID, err)
				continue
			}
			log.Println(extService.ID, "removed")
		}
	}
}

func (b *Bridge) add(containerId string, quiet bool) {
	if d := b.deadContainers[containerId]; d != nil {
		b.services[containerId] = d.Services
		delete(b.deadContainers, containerId)
	}

	if b.services[containerId] != nil {
		log.Println("container, ", containerId[:12], ", already exists, ignoring")
		// Alternatively, remove and readd or resubmit.
		return
	}

	container, err := b.docker.InspectContainer(containerId)
	if err != nil {
		log.Println("unable to inspect container:", containerId[:12], err)
		return
	}

	ports := make(map[string]ServicePort)

	// Extract configured host port mappings, relevant when using --net=host
	for port, _ := range container.Config.ExposedPorts {
		published := []dockerapi.PortBinding{ {"0.0.0.0", port.Port()}, }
		ports[string(port)] = servicePort(container, port, published, b.config.Network, quiet)
	}

	// Extract runtime port mappings, relevant when using --net=bridge
	for port, published := range container.NetworkSettings.Ports {
		ports[string(port)] = servicePort(container, port, published, b.config.Network, quiet)
	}

	if len(ports) == 0 && !quiet {
		log.Println("ignored:", container.ID[:12], "no published ports")
		return
	}

	servicePorts := make(map[string]ServicePort)
	for key, port := range ports {
		// If internal and can't evaluate the exposed IP (will happen on wrong network), ignore
		if b.config.Internal == true && port.ExposedIP == "" {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "internal IP not found on network", b.config.Network)
			}
			continue
		}

		if b.config.Internal != true && port.HostPort == "" {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "port", port.ExposedPort, "not published on host")
			}
			continue
		}
		servicePorts[key] = port
	}

	for _, port := range servicePorts {
		service := b.newService(port)
		if service == nil {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "service on port", port.ExposedPort)
			}
			continue
		}
		err := b.registry.Register(service)
		if err != nil {
			log.Println("register failed:", service, err)
			continue
		}
		b.services[container.ID] = append(b.services[container.ID], service)
		log.Println("added:", container.ID[:12], service.ID)
	}
}

func (b *Bridge) newService(port ServicePort) *Service {
	container := port.container
	defaultName := strings.Replace(serviceNamePattern.ReplaceAllString(strings.TrimPrefix(container.Name, "/"), ""), "_", "-", -1)

	// not sure about this logic. kind of want to remove it.
	hostname := Hostname
	if hostname == "" {
		hostname = port.HostIP
	}
	if port.HostIP == "0.0.0.0" {
		ip, err := net.ResolveIPAddr("ip", hostname)
		if err == nil {
			port.HostIP = ip.String()
		}
	}

	if b.config.HostIp != "" {
		port.HostIP = b.config.HostIp
	}

	metadata := serviceMetaData(container.Config, port.ExposedPort)

	ignore := mapDefault(metadata, "ignore", "")
	if ignore != "" {
		return nil
	}

	service := new(Service)
	service.Origin = port
	service.Name = mapDefault(metadata, "name", defaultName)
	service.ID = container.ID + ":" + service.Name + ":" + port.ExposedPort

	var p int

	if b.config.Internal == true {
		service.IP = port.ExposedIP
		p, _ = strconv.Atoi(port.ExposedPort)
	} else {
		service.IP = port.HostIP
		p, _ = strconv.Atoi(port.HostPort)
	}
	service.Port = p

	if b.config.UseIpFromLabel != "" {
		containerIp := container.Config.Labels[b.config.UseIpFromLabel]
		if containerIp != "" {
			slashIndex := strings.LastIndex(containerIp, "/")
			if slashIndex > -1 {
				service.IP = containerIp[:slashIndex]
			} else {
				service.IP = containerIp
			}
			log.Println("using container IP " + service.IP + " from label '" +
				b.config.UseIpFromLabel  + "'")
		} else {
			log.Println("Label '" + b.config.UseIpFromLabel +
				"' not found in container configuration")
		}
	}

	// NetworkMode can point to another container (kuberenetes pods)
	networkMode := container.HostConfig.NetworkMode
	if networkMode != "" {
		if strings.HasPrefix(networkMode, "container:") {
			networkContainerId := strings.Split(networkMode, ":")[1]
			log.Println(service.Name + ": detected container NetworkMode, linked to: " + networkContainerId[:12])
			networkContainer, err := b.docker.InspectContainer(networkContainerId)
			if err != nil {
				log.Println("unable to inspect network container:", networkContainerId[:12], err)
			} else {
				service.IP = networkContainer.NetworkSettings.IPAddress
				log.Println(service.Name + ": using network container IP " + service.IP)
			}
		}
	}

	if port.PortType == "udp" {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags, "udp")
		service.ID = service.ID + ":udp"
	} else {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags)
		service.ID = service.ID + ":tcp"
	}

	id := mapDefault(metadata, "id", "")
	if id != "" {
		service.ID = id
	}

	delete(metadata, "id")
	delete(metadata, "tags")
	delete(metadata, "name")
	service.Attrs = metadata
	service.TTL = b.config.RefreshTtl

	return service
}

func (b *Bridge) remove(containerId string, deregister bool) {
	b.Lock()
	defer b.Unlock()

	if deregister {
		deregisterAll := func(services []*Service) {
			for _, service := range services {
				err := b.registry.Deregister(service)
				if err != nil {
					log.Println("deregister failed:", service.ID, err)
					continue
				}
				log.Println("removed:", containerId[:12], service.ID)
			}
		}
		deregisterAll(b.services[containerId])
		if d := b.deadContainers[containerId]; d != nil {
			deregisterAll(d.Services)
			delete(b.deadContainers, containerId)
		}
	} else if b.config.RefreshTtl != 0 && b.services[containerId] != nil {
		// need to stop the refreshing, but can't delete it yet
		b.deadContainers[containerId] = &DeadContainer{b.config.RefreshTtl, b.services[containerId]}
	}
	delete(b.services, containerId)
}

// bit set on ExitCode if it represents an exit via a signal
var dockerSignaledBit = 128

func (b *Bridge) shouldRemove(containerId string) bool {
	if b.config.DeregisterCheck == "always" {
		return true
	}
	container, err := b.docker.InspectContainer(containerId)
	if _, ok := err.(*dockerapi.NoSuchContainer); ok {
		// the container has already been removed from Docker
		// e.g. probabably run with "--rm" to remove immediately
		// so its exit code is not accessible
		log.Printf("registrator: container %v was removed, could not fetch exit code", containerId[:12])
		return true
	}

	switch {
	case err != nil:
		log.Printf("registrator: error fetching status for container %v on \"die\" event: %v\n", containerId[:12], err)
		return false
	case container.State.Running:
		log.Printf("registrator: not removing container %v, still running", containerId[:12])
		return false
	case container.State.ExitCode == 0:
		return true
	case container.State.ExitCode&dockerSignaledBit == dockerSignaledBit:
		return true
	}
	return false
}

var Hostname string

func init() {
	// It's ok for Hostname to ultimately be an empty string
	// An empty string will fall back to trying to make a best guess
	Hostname, _ = os.Hostname()
}
<commit_msg>do not add UDP tags by default<commit_after>package bridge

import (
	"errors"
	"log"
	"net"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"

	dockerapi "github.com/fsouza/go-dockerclient"
)

var serviceIDPattern = regexp.MustCompile(`^(.+?):([a-zA-Z0-9][a-zA-Z0-9_.-]+):[0-9]+(?::udp)?$`)

var serviceNamePattern = regexp.MustCompile(`_[0-9]+$`)

type Bridge struct {
	sync.Mutex
	registry       RegistryAdapter
	docker         *dockerapi.Client
	services       map[string][]*Service
	deadContainers map[string]*DeadContainer
	config         Config
}

func New(docker *dockerapi.Client, adapterUri string, config Config) (*Bridge, error) {
	uri, err := url.Parse(adapterUri)
	if err != nil {
		return nil, errors.New("bad adapter uri: " + adapterUri)
	}
	factory, found := AdapterFactories.Lookup(uri.Scheme)
	if !found {
		return nil, errors.New("unrecognized adapter: " + adapterUri)
	}

	log.Println("Using", uri.Scheme, "adapter:", uri)
	return &Bridge{
		docker:         docker,
		config:         config,
		registry:       factory.New(uri),
		services:       make(map[string][]*Service),
		deadContainers: make(map[string]*DeadContainer),
	}, nil
}

func (b *Bridge) Ping() error {
	return b.registry.Ping()
}

func (b *Bridge) Add(containerId string) {
	b.Lock()
	defer b.Unlock()
	b.add(containerId, false)
}

func (b *Bridge) Remove(containerId string) {
	b.remove(containerId, true)
}

func (b *Bridge) RemoveOnExit(containerId string) {
	b.remove(containerId, b.shouldRemove(containerId))
}

func (b *Bridge) Refresh() {
	b.Lock()
	defer b.Unlock()

	for containerId, deadContainer := range b.deadContainers {
		deadContainer.TTL -= b.config.RefreshInterval
		if deadContainer.TTL <= 0 {
			delete(b.deadContainers, containerId)
		}
	}

	for containerId, services := range b.services {
		for _, service := range services {
			err := b.registry.Refresh(service)
			if err != nil {
				log.Println("refresh failed:", service.ID, err)
				continue
			}
			log.Println("refreshed:", containerId[:12], service.ID)
		}
	}
}

func (b *Bridge) Sync(quiet bool) {
	b.Lock()
	defer b.Unlock()

	containers, err := b.docker.ListContainers(dockerapi.ListContainersOptions{})
	if err != nil && quiet {
		log.Println("error listing containers, skipping sync")
		return
	} else if err != nil && !quiet {
		log.Fatal(err)
	}

	log.Printf("Syncing services on %d containers", len(containers))

	// NOTE: This assumes reregistering will do the right thing, i.e. nothing..
	for _, listing := range containers {
		services := b.services[listing.ID]
		if services == nil {
			b.add(listing.ID, quiet)
		} else {
			for _, service := range services {
				err := b.registry.Register(service)
				if err != nil {
					log.Println("sync register failed:", service, err)
				}
			}
		}
	}

	// Clean up services that were registered previously, but aren't
	// acknowledged within registrator
	if b.config.Cleanup {
		// Remove services if its corresponding container is not running
		log.Println("Listing non-exited containers")
		filters := map[string][]string{"status": {"created", "restarting", "running", "paused"}}
		nonExitedContainers, err := b.docker.ListContainers(dockerapi.ListContainersOptions{Filters: filters})
		if err != nil {
			log.Println("error listing nonExitedContainers, skipping sync", err)
			return
		}
		for listingId, _ := range b.services {
			found := false
			for _, container := range nonExitedContainers {
				if listingId == container.ID {
					found = true
					break
				}
			}
			// This is a container that does not exist
			if !found {
				log.Printf("stale: Removing service %s because it does not exist", listingId)
				go b.RemoveOnExit(listingId)
			}
		}

		log.Println("Cleaning up dangling services")
		extServices, err := b.registry.Services()
		if err != nil {
			log.Println("cleanup failed:", err)
			return
		}

	Outer:
		for _, extService := range extServices {
			matches := serviceIDPattern.FindStringSubmatch(extService.ID)
			if len(matches) != 3 {
				// There's no way this was registered by us, so leave it
				continue
			}
			serviceHostname := matches[1]
			if serviceHostname != Hostname {
				// ignore because registered on a different host
				continue
			}
			serviceContainerName := matches[2]
			for _, listing := range b.services {
				for _, service := range listing {
					if service.Name == extService.Name && serviceContainerName == service.Origin.container.Name[1:] {
						continue Outer
					}
				}
			}
			log.Println("dangling:", extService.ID)
			err := b.registry.Deregister(extService)
			if err != nil {
				log.Println("deregister failed:", extService.ID, err)
				continue
			}
			log.Println(extService.ID, "removed")
		}
	}
}

func (b *Bridge) add(containerId string, quiet bool) {
	if d := b.deadContainers[containerId]; d != nil {
		b.services[containerId] = d.Services
		delete(b.deadContainers, containerId)
	}

	if b.services[containerId] != nil {
		log.Println("container, ", containerId[:12], ", already exists, ignoring")
		// Alternatively, remove and readd or resubmit.
		return
	}

	container, err := b.docker.InspectContainer(containerId)
	if err != nil {
		log.Println("unable to inspect container:", containerId[:12], err)
		return
	}

	ports := make(map[string]ServicePort)

	// Extract configured host port mappings, relevant when using --net=host
	for port, _ := range container.Config.ExposedPorts {
		published := []dockerapi.PortBinding{ {"0.0.0.0", port.Port()}, }
		ports[string(port)] = servicePort(container, port, published, b.config.Network, quiet)
	}

	// Extract runtime port mappings, relevant when using --net=bridge
	for port, published := range container.NetworkSettings.Ports {
		ports[string(port)] = servicePort(container, port, published, b.config.Network, quiet)
	}

	if len(ports) == 0 && !quiet {
		log.Println("ignored:", container.ID[:12], "no published ports")
		return
	}

	servicePorts := make(map[string]ServicePort)
	for key, port := range ports {
		// If internal and can't evaluate the exposed IP (will happen on wrong network), ignore
		if b.config.Internal == true && port.ExposedIP == "" {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "internal IP not found on network", b.config.Network)
			}
			continue
		}

		if b.config.Internal != true && port.HostPort == "" {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "port", port.ExposedPort, "not published on host")
			}
			continue
		}
		servicePorts[key] = port
	}

	for _, port := range servicePorts {
		service := b.newService(port)
		if service == nil {
			if !quiet {
				log.Println("ignored:", container.ID[:12], "service on port", port.ExposedPort)
			}
			continue
		}
		err := b.registry.Register(service)
		if err != nil {
			log.Println("register failed:", service, err)
			continue
		}
		b.services[container.ID] = append(b.services[container.ID], service)
		log.Println("added:", container.ID[:12], service.ID)
	}
}

func (b *Bridge) newService(port ServicePort) *Service {
	container := port.container
	defaultName := strings.Replace(serviceNamePattern.ReplaceAllString(strings.TrimPrefix(container.Name, "/"), ""), "_", "-", -1)

	// not sure about this logic. kind of want to remove it.
	hostname := Hostname
	if hostname == "" {
		hostname = port.HostIP
	}
	if port.HostIP == "0.0.0.0" {
		ip, err := net.ResolveIPAddr("ip", hostname)
		if err == nil {
			port.HostIP = ip.String()
		}
	}

	if b.config.HostIp != "" {
		port.HostIP = b.config.HostIp
	}

	metadata := serviceMetaData(container.Config, port.ExposedPort)

	ignore := mapDefault(metadata, "ignore", "")
	if ignore != "" {
		return nil
	}

	service := new(Service)
	service.Origin = port
	service.Name = mapDefault(metadata, "name", defaultName)
	service.ID = container.ID + ":" + service.Name + ":" + port.ExposedPort

	var p int

	if b.config.Internal == true {
		service.IP = port.ExposedIP
		p, _ = strconv.Atoi(port.ExposedPort)
	} else {
		service.IP = port.HostIP
		p, _ = strconv.Atoi(port.HostPort)
	}
	service.Port = p

	if b.config.UseIpFromLabel != "" {
		containerIp := container.Config.Labels[b.config.UseIpFromLabel]
		if containerIp != "" {
			slashIndex := strings.LastIndex(containerIp, "/")
			if slashIndex > -1 {
				service.IP = containerIp[:slashIndex]
			} else {
				service.IP = containerIp
			}
			log.Println("using container IP " + service.IP + " from label '" +
				b.config.UseIpFromLabel  + "'")
		} else {
			log.Println("Label '" + b.config.UseIpFromLabel +
				"' not found in container configuration")
		}
	}

	// NetworkMode can point to another container (kuberenetes pods)
	networkMode := container.HostConfig.NetworkMode
	if networkMode != "" {
		if strings.HasPrefix(networkMode, "container:") {
			networkContainerId := strings.Split(networkMode, ":")[1]
			log.Println(service.Name + ": detected container NetworkMode, linked to: " + networkContainerId[:12])
			networkContainer, err := b.docker.InspectContainer(networkContainerId)
			if err != nil {
				log.Println("unable to inspect network container:", networkContainerId[:12], err)
			} else {
				service.IP = networkContainer.NetworkSettings.IPAddress
				log.Println(service.Name + ": using network container IP " + service.IP)
			}
		}
	}

	if port.PortType == "udp" {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags)
		service.ID = service.ID + ":udp"
	} else {
		service.Tags = combineTags(
			mapDefault(metadata, "tags", ""), b.config.ForceTags)
		service.ID = service.ID + ":tcp"
	}

	id := mapDefault(metadata, "id", "")
	if id != "" {
		service.ID = id
	}

	delete(metadata, "id")
	delete(metadata, "tags")
	delete(metadata, "name")
	service.Attrs = metadata
	service.TTL = b.config.RefreshTtl

	return service
}

func (b *Bridge) remove(containerId string, deregister bool) {
	b.Lock()
	defer b.Unlock()

	if deregister {
		deregisterAll := func(services []*Service) {
			for _, service := range services {
				err := b.registry.Deregister(service)
				if err != nil {
					log.Println("deregister failed:", service.ID, err)
					continue
				}
				log.Println("removed:", containerId[:12], service.ID)
			}
		}
		deregisterAll(b.services[containerId])
		if d := b.deadContainers[containerId]; d != nil {
			deregisterAll(d.Services)
			delete(b.deadContainers, containerId)
		}
	} else if b.config.RefreshTtl != 0 && b.services[containerId] != nil {
		// need to stop the refreshing, but can't delete it yet
		b.deadContainers[containerId] = &DeadContainer{b.config.RefreshTtl, b.services[containerId]}
	}
	delete(b.services, containerId)
}

// bit set on ExitCode if it represents an exit via a signal
var dockerSignaledBit = 128

func (b *Bridge) shouldRemove(containerId string) bool {
	if b.config.DeregisterCheck == "always" {
		return true
	}
	container, err := b.docker.InspectContainer(containerId)
	if _, ok := err.(*dockerapi.NoSuchContainer); ok {
		// the container has already been removed from Docker
		// e.g. probabably run with "--rm" to remove immediately
		// so its exit code is not accessible
		log.Printf("registrator: container %v was removed, could not fetch exit code", containerId[:12])
		return true
	}

	switch {
	case err != nil:
		log.Printf("registrator: error fetching status for container %v on \"die\" event: %v\n", containerId[:12], err)
		return false
	case container.State.Running:
		log.Printf("registrator: not removing container %v, still running", containerId[:12])
		return false
	case container.State.ExitCode == 0:
		return true
	case container.State.ExitCode&dockerSignaledBit == dockerSignaledBit:
		return true
	}
	return false
}

var Hostname string

func init() {
	// It's ok for Hostname to ultimately be an empty string
	// An empty string will fall back to trying to make a best guess
	Hostname, _ = os.Hostname()
}
<|endoftext|>