<commit_before>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package local

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"

	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/etcd/etcdserver"
	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/etcd/etcdserver/etcdhttp"
	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/etcd/pkg/types"
)

const (
	memberName  = "simple"
	clusterName = "simple-cluster"
	tempPrefix  = "simple-etcd-"

	// No peer URL exists but etcd doesn't allow the value to be empty.
	peerURL    = "http://localhost:0"
	clusterCfg = memberName + "=" + peerURL
)

// SimpleEtcd provides a single node etcd server.
type SimpleEtcd struct {
	Port     int
	listener net.Listener
	server   *etcdserver.EtcdServer
	dataDir  string
}

func NewSimpleEtcd() (*SimpleEtcd, error) {
	var err error
	se := &SimpleEtcd{}
	se.listener, err = net.Listen("tcp", ":0")
	if err != nil {
		return nil, err
	}

	se.Port = se.listener.Addr().(*net.TCPAddr).Port
	clientURLs, err := interfaceURLs(se.Port)
	if err != nil {
		se.Destroy()
		return nil, err
	}

	se.dataDir, err = ioutil.TempDir("", tempPrefix)
	if err != nil {
		se.Destroy()
		return nil, err
	}

	peerURLs, err := types.NewURLs([]string{peerURL})
	if err != nil {
		se.Destroy()
		return nil, err
	}

	cfg := &etcdserver.ServerConfig{
		Name:       memberName,
		ClientURLs: clientURLs,
		PeerURLs:   peerURLs,
		DataDir:    se.dataDir,
		InitialPeerURLsMap: types.URLsMap{
			memberName: peerURLs,
		},
		NewCluster:    true,
		TickMs:        100,
		ElectionTicks: 10,
	}

	se.server, err = etcdserver.NewServer(cfg)
	if err != nil {
		return nil, err
	}

	se.server.Start()
	go http.Serve(se.listener,
		etcdhttp.NewClientHandler(se.server, cfg.ReqTimeout()))

	return se, nil
}

func (se *SimpleEtcd) Destroy() error {
	var err error
	firstErr := func(e error) {
		if e != nil && err == nil {
			err = e
		}
	}

	if se.listener != nil {
		firstErr(se.listener.Close())
	}

	if se.server != nil {
		se.server.Stop()
	}

	if se.dataDir != "" {
		firstErr(os.RemoveAll(se.dataDir))
	}

	return err
}

// Generate all publishable URLs for a given HTTP port.
func interfaceURLs(port int) (types.URLs, error) {
	allAddrs, err := net.InterfaceAddrs()
	if err != nil {
		return []url.URL{}, err
	}

	var allURLs types.URLs
	for _, a := range allAddrs {
		ip, ok := a.(*net.IPNet)
		if !ok || !ip.IP.IsGlobalUnicast() {
			continue
		}

		tcp := net.TCPAddr{
			IP:   ip.IP,
			Port: port,
		}

		u := url.URL{
			Scheme: "http",
			Host:   tcp.String(),
		}
		allURLs = append(allURLs, u)
	}

	if len(allAddrs) == 0 {
		return []url.URL{}, fmt.Errorf("no publishable addresses")
	}

	return allURLs, nil
}
<commit_msg>platform/local: update name of etcd v2 API package<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package local

import (
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"net/url"
	"os"

	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/etcd/etcdserver"
	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/etcd/etcdserver/api/v2http"
	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/etcd/pkg/types"
)

const (
	memberName  = "simple"
	clusterName = "simple-cluster"
	tempPrefix  = "simple-etcd-"

	// No peer URL exists but etcd doesn't allow the value to be empty.
	peerURL    = "http://localhost:0"
	clusterCfg = memberName + "=" + peerURL
)

// SimpleEtcd provides a single node etcd server.
type SimpleEtcd struct {
	Port     int
	listener net.Listener
	server   *etcdserver.EtcdServer
	dataDir  string
}

func NewSimpleEtcd() (*SimpleEtcd, error) {
	var err error
	se := &SimpleEtcd{}
	se.listener, err = net.Listen("tcp", ":0")
	if err != nil {
		return nil, err
	}

	se.Port = se.listener.Addr().(*net.TCPAddr).Port
	clientURLs, err := interfaceURLs(se.Port)
	if err != nil {
		se.Destroy()
		return nil, err
	}

	se.dataDir, err = ioutil.TempDir("", tempPrefix)
	if err != nil {
		se.Destroy()
		return nil, err
	}

	peerURLs, err := types.NewURLs([]string{peerURL})
	if err != nil {
		se.Destroy()
		return nil, err
	}

	cfg := &etcdserver.ServerConfig{
		Name:       memberName,
		ClientURLs: clientURLs,
		PeerURLs:   peerURLs,
		DataDir:    se.dataDir,
		InitialPeerURLsMap: types.URLsMap{
			memberName: peerURLs,
		},
		NewCluster:    true,
		TickMs:        100,
		ElectionTicks: 10,
	}

	se.server, err = etcdserver.NewServer(cfg)
	if err != nil {
		return nil, err
	}

	se.server.Start()
	go http.Serve(se.listener,
		v2http.NewClientHandler(se.server, cfg.ReqTimeout()))

	return se, nil
}

func (se *SimpleEtcd) Destroy() error {
	var err error
	firstErr := func(e error) {
		if e != nil && err == nil {
			err = e
		}
	}

	if se.listener != nil {
		firstErr(se.listener.Close())
	}

	if se.server != nil {
		se.server.Stop()
	}

	if se.dataDir != "" {
		firstErr(os.RemoveAll(se.dataDir))
	}

	return err
}

// Generate all publishable URLs for a given HTTP port.
func interfaceURLs(port int) (types.URLs, error) {
	allAddrs, err := net.InterfaceAddrs()
	if err != nil {
		return []url.URL{}, err
	}

	var allURLs types.URLs
	for _, a := range allAddrs {
		ip, ok := a.(*net.IPNet)
		if !ok || !ip.IP.IsGlobalUnicast() {
			continue
		}

		tcp := net.TCPAddr{
			IP:   ip.IP,
			Port: port,
		}

		u := url.URL{
			Scheme: "http",
			Host:   tcp.String(),
		}
		allURLs = append(allURLs, u)
	}

	if len(allAddrs) == 0 {
		return []url.URL{}, fmt.Errorf("no publishable addresses")
	}

	return allURLs, nil
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package aggregatorservice contains the functions needed for handling the aggregation requests.
package aggregatorservice

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os/exec"
	"strings"

	log "github.com/golang/glog"
	"cloud.google.com/go/pubsub"
	"cloud.google.com/go/storage"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/ioutils"
	"github.com/google/privacy-sandbox-aggregation-service/service/query"
	"github.com/google/privacy-sandbox-aggregation-service/service/utils"
)

// DataflowCfg contains parameters necessary for running pipelines on Dataflow.
type DataflowCfg struct {
	Project         string
	Region          string
	TempLocation    string
	StagingLocation string
}

// ServerCfg contains file URIs necessary for the service.
type ServerCfg struct {
	PrivateKeyParamsURI             string
	DpfAggregatePartialReportBinary string
	WorkspaceURI                    string
}

// SharedInfoHandler handles HTTP requests for the information shared with other helpers.
type SharedInfoHandler struct {
	SharedInfo *query.HelperSharedInfo
}

func (h *SharedInfoHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	err := json.NewEncoder(w).Encode(h.SharedInfo)
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// QueryHandler handles the request in the pubsub messages.
type QueryHandler struct {
	ServerCfg                 ServerCfg
	PipelineRunner            string
	DataflowCfg               DataflowCfg
	Origin                    string
	SharedDir                 string
	RequestPubSubTopic        string
	RequestPubsubSubscription string

	PubSubTopicClient, PubSubSubscriptionClient *pubsub.Client
	GCSClient                                   *storage.Client
}

// Setup creates the cloud API clients.
func (h *QueryHandler) Setup(ctx context.Context) error {
	topicProject, _, err := utils.ParsePubSubResourceName(h.RequestPubSubTopic)
	if err != nil {
		return err
	}
	h.PubSubTopicClient, err = pubsub.NewClient(ctx, topicProject)
	if err != nil {
		return err
	}

	subscriptionProject, _, err := utils.ParsePubSubResourceName(h.RequestPubsubSubscription)
	if err != nil {
		return err
	}

	if subscriptionProject == topicProject {
		h.PubSubSubscriptionClient = h.PubSubTopicClient
	} else {
		h.PubSubSubscriptionClient, err = pubsub.NewClient(ctx, subscriptionProject)
		if err != nil {
			return err
		}

	}

	h.GCSClient, err = storage.NewClient(ctx)
	return err
}

// Close closes the cloud API clients.
func (h *QueryHandler) Close() {
	h.PubSubTopicClient.Close()
	h.PubSubSubscriptionClient.Close()
	h.GCSClient.Close()
}

// SetupPullRequests gets ready to pull requests contained in a PubSub message subscription, and handles the request.
func (h *QueryHandler) SetupPullRequests(ctx context.Context) error {
	_, subID, err := utils.ParsePubSubResourceName(h.RequestPubsubSubscription)
	if err != nil {
		return err
	}
	sub := h.PubSubSubscriptionClient.Subscription(subID)

	// Only allow pulling one message at a time to avoid overloading the memory.
	sub.ReceiveSettings.Synchronous = true
	sub.ReceiveSettings.MaxOutstandingMessages = 1
	return sub.Receive(ctx, func(ctx context.Context, msg *pubsub.Message) {
		request := &query.AggregateRequest{}
		err := json.Unmarshal(msg.Data, request)
		if err != nil {
			log.Error(err)
			msg.Nack()
			return
		}
		if err := h.aggregatePartialReportHierarchy(ctx, request); err != nil {
			log.Error(err)
			msg.Nack()
			return
		}
		msg.Ack()
	})
}

func getFinalPartialResultURI(resultDir, queryID, origin string) string {
	return ioutils.JoinPath(resultDir, fmt.Sprintf("%s_%s", queryID, strings.ReplaceAll(origin, ".", "_")))
}

func (h *QueryHandler) aggregatePartialReportHierarchy(ctx context.Context, request *query.AggregateRequest) error {
	config, err := query.ReadExpansionConfigFile(ctx, request.ExpandConfigURI)
	if err != nil {
		return nil
	}

	finalLevel := int32(len(config.PrefixLengths)) - 1
	if request.QueryLevel > finalLevel {
		return fmt.Errorf("expect request level <= finalLevel %d, got %d", finalLevel, request.QueryLevel)
	}

	partialReportURI := request.PartialReportURI
	outputDecryptedReportURI := ""
	if request.QueryLevel > 0 {
		// If it is not the first-level aggregation, check if the result from the partner helper is ready for the previous level.
		exist, err := utils.IsGCSObjectExist(ctx, h.GCSClient,
			query.GetRequestPartialResultURI(request.PartnerSharedInfo.SharedDir, request.QueryID, request.QueryLevel-1),
		)
		if err != nil {
			return err
		}
		if !exist {
			// When the partial result from the partner helper is not ready, nack the message with an error.
			return fmt.Errorf("result from %s for level %d of query %s is not ready", request.PartnerSharedInfo.Origin, request.QueryLevel-1, request.QueryID)
		}

		// If it is not the first-level aggregation, the pipeline should read the decrypted reports instead of the original encrypted ones.
		partialReportURI = query.GetRequestDecryptedReportURI(h.ServerCfg.WorkspaceURI, request.QueryID)
	} else {
		outputDecryptedReportURI = query.GetRequestDecryptedReportURI(h.ServerCfg.WorkspaceURI, request.QueryID)
	}

	expandParamsURI, err := query.GetRequestExpandParamsURI(ctx, config, request,
		h.ServerCfg.WorkspaceURI,
		h.SharedDir,
		request.PartnerSharedInfo.SharedDir,
	)
	if err != nil {
		return err
	}

	var outputResultURI string
	// The final-level results are not supposed to be shared with the partner helpers.
	if request.QueryLevel == finalLevel {
		outputResultURI = getFinalPartialResultURI(request.ResultDir, request.QueryID, h.Origin)
	} else {
		outputResultURI = query.GetRequestPartialResultURI(h.SharedDir, request.QueryID, request.QueryLevel)
	}

	args := []string{
		"--partial_report_uri=" + partialReportURI,
		"--expand_parameters_uri=" + expandParamsURI,
		"--partial_histogram_uri=" + outputResultURI,
		"--decrypted_report_uri=" + outputDecryptedReportURI,
		"--epsilon=" + fmt.Sprintf("%f", request.TotalEpsilon*config.PrivacyBudgetPerPrefix[request.QueryLevel]),
		"--private_key_params_uri=" + h.ServerCfg.PrivateKeyParamsURI,
		"--key_bit_size=" + fmt.Sprint(request.KeyBitSize),
		"--runner=" + h.PipelineRunner,
	}

	if h.PipelineRunner == "dataflow" {
		args = append(args,
			"--project="+h.DataflowCfg.Project,
			"--region="+h.DataflowCfg.Region,
			"--temp_location="+h.DataflowCfg.TempLocation,
			"--staging_location="+h.DataflowCfg.StagingLocation,
			// set jobname to queryID-level-origin
			"--job_name="+fmt.Sprintf("%s-%v-%s", request.QueryID, request.QueryLevel, h.Origin),
			"--worker_binary="+h.ServerCfg.DpfAggregatePartialReportBinary,
		)
	}

	str := h.ServerCfg.DpfAggregatePartialReportBinary
	for _, s := range args {
		str = fmt.Sprintf("%s\n%s", str, s)
	}
	log.Infof("Running command\n%s", str)

	cmd := exec.CommandContext(ctx, h.ServerCfg.DpfAggregatePartialReportBinary, args...)
	var out bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &stderr
	err = cmd.Run()
	if err != nil {
		log.Errorf("%s: %s", err, stderr.String())
		return err
	}
	log.Infof("output of cmd: %s", out.String())

	if request.QueryLevel == finalLevel {
		log.Infof("query %q complete", request.QueryID)
		return nil
	}

	// If the hierarchical query is not finished yet, publish the requests for the next-level aggregation.
	request.QueryLevel++
	_, topic, err := utils.ParsePubSubResourceName(h.RequestPubSubTopic)
	if err != nil {
		return err
	}
	return utils.PublishRequest(ctx, h.PubSubTopicClient, topic, request)
}

// ReadHelperSharedInfo reads the helper shared info from a URL.
func ReadHelperSharedInfo(client *http.Client, url, token string) (*query.HelperSharedInfo, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	if token != "" {
		req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", token))
	}

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.Status != "200 OK" {
		body, _ := ioutil.ReadAll(resp.Body)
		log.Infof("%v: %s", resp.Status, string(body))
		return nil, fmt.Errorf("Error reading shared info from %s: %s", url, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	info := &query.HelperSharedInfo{}
	if err := json.Unmarshal([]byte(body), info); err != nil {
		return nil, err
	}
	return info, nil
}
<commit_msg>Increasing Aggregator PubSub MaxExtension timeout to 24 hours<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package aggregatorservice contains the functions needed for handling the aggregation requests.
package aggregatorservice

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os/exec"
	"strings"
	"time"

	log "github.com/golang/glog"
	"cloud.google.com/go/pubsub"
	"cloud.google.com/go/storage"
	"github.com/google/privacy-sandbox-aggregation-service/pipeline/ioutils"
	"github.com/google/privacy-sandbox-aggregation-service/service/query"
	"github.com/google/privacy-sandbox-aggregation-service/service/utils"
)

// DataflowCfg contains parameters necessary for running pipelines on Dataflow.
type DataflowCfg struct {
	Project         string
	Region          string
	TempLocation    string
	StagingLocation string
}

// ServerCfg contains file URIs necessary for the service.
type ServerCfg struct {
	PrivateKeyParamsURI             string
	DpfAggregatePartialReportBinary string
	WorkspaceURI                    string
}

// SharedInfoHandler handles HTTP requests for the information shared with other helpers.
type SharedInfoHandler struct {
	SharedInfo *query.HelperSharedInfo
}

func (h *SharedInfoHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	err := json.NewEncoder(w).Encode(h.SharedInfo)
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// QueryHandler handles the request in the pubsub messages.
type QueryHandler struct {
	ServerCfg                 ServerCfg
	PipelineRunner            string
	DataflowCfg               DataflowCfg
	Origin                    string
	SharedDir                 string
	RequestPubSubTopic        string
	RequestPubsubSubscription string

	PubSubTopicClient, PubSubSubscriptionClient *pubsub.Client
	GCSClient                                   *storage.Client
}

// Setup creates the cloud API clients.
func (h *QueryHandler) Setup(ctx context.Context) error {
	topicProject, _, err := utils.ParsePubSubResourceName(h.RequestPubSubTopic)
	if err != nil {
		return err
	}
	h.PubSubTopicClient, err = pubsub.NewClient(ctx, topicProject)
	if err != nil {
		return err
	}

	subscriptionProject, _, err := utils.ParsePubSubResourceName(h.RequestPubsubSubscription)
	if err != nil {
		return err
	}

	if subscriptionProject == topicProject {
		h.PubSubSubscriptionClient = h.PubSubTopicClient
	} else {
		h.PubSubSubscriptionClient, err = pubsub.NewClient(ctx, subscriptionProject)
		if err != nil {
			return err
		}

	}

	h.GCSClient, err = storage.NewClient(ctx)
	return err
}

// Close closes the cloud API clients.
func (h *QueryHandler) Close() {
	h.PubSubTopicClient.Close()
	h.PubSubSubscriptionClient.Close()
	h.GCSClient.Close()
}

// SetupPullRequests gets ready to pull requests contained in a PubSub message subscription, and handles the request.
func (h *QueryHandler) SetupPullRequests(ctx context.Context) error {
	_, subID, err := utils.ParsePubSubResourceName(h.RequestPubsubSubscription)
	if err != nil {
		return err
	}
	sub := h.PubSubSubscriptionClient.Subscription(subID)

	// Only allow pulling one message at a time to avoid overloading the memory.
	sub.ReceiveSettings.Synchronous = true
	sub.ReceiveSettings.MaxOutstandingMessages = 1
	sub.ReceiveSettings.MaxExtension = 24 * time.Hour // extending from 60min default to 1 day
	return sub.Receive(ctx, func(ctx context.Context, msg *pubsub.Message) {
		request := &query.AggregateRequest{}
		err := json.Unmarshal(msg.Data, request)
		if err != nil {
			log.Error(err)
			msg.Nack()
			return
		}
		if err := h.aggregatePartialReportHierarchy(ctx, request); err != nil {
			log.Error(err)
			msg.Nack()
			return
		}
		msg.Ack()
	})
}

func getFinalPartialResultURI(resultDir, queryID, origin string) string {
	return ioutils.JoinPath(resultDir, fmt.Sprintf("%s_%s", queryID, strings.ReplaceAll(origin, ".", "_")))
}

func (h *QueryHandler) aggregatePartialReportHierarchy(ctx context.Context, request *query.AggregateRequest) error {
	config, err := query.ReadExpansionConfigFile(ctx, request.ExpandConfigURI)
	if err != nil {
		return nil
	}

	finalLevel := int32(len(config.PrefixLengths)) - 1
	if request.QueryLevel > finalLevel {
		return fmt.Errorf("expect request level <= finalLevel %d, got %d", finalLevel, request.QueryLevel)
	}

	partialReportURI := request.PartialReportURI
	outputDecryptedReportURI := ""
	if request.QueryLevel > 0 {
		// If it is not the first-level aggregation, check if the result from the partner helper is ready for the previous level.
		exist, err := utils.IsGCSObjectExist(ctx, h.GCSClient,
			query.GetRequestPartialResultURI(request.PartnerSharedInfo.SharedDir, request.QueryID, request.QueryLevel-1),
		)
		if err != nil {
			return err
		}
		if !exist {
			// When the partial result from the partner helper is not ready, nack the message with an error.
			return fmt.Errorf("result from %s for level %d of query %s is not ready", request.PartnerSharedInfo.Origin, request.QueryLevel-1, request.QueryID)
		}

		// If it is not the first-level aggregation, the pipeline should read the decrypted reports instead of the original encrypted ones.
		partialReportURI = query.GetRequestDecryptedReportURI(h.ServerCfg.WorkspaceURI, request.QueryID)
	} else {
		outputDecryptedReportURI = query.GetRequestDecryptedReportURI(h.ServerCfg.WorkspaceURI, request.QueryID)
	}

	expandParamsURI, err := query.GetRequestExpandParamsURI(ctx, config, request,
		h.ServerCfg.WorkspaceURI,
		h.SharedDir,
		request.PartnerSharedInfo.SharedDir,
	)
	if err != nil {
		return err
	}

	var outputResultURI string
	// The final-level results are not supposed to be shared with the partner helpers.
	if request.QueryLevel == finalLevel {
		outputResultURI = getFinalPartialResultURI(request.ResultDir, request.QueryID, h.Origin)
	} else {
		outputResultURI = query.GetRequestPartialResultURI(h.SharedDir, request.QueryID, request.QueryLevel)
	}

	args := []string{
		"--partial_report_uri=" + partialReportURI,
		"--expand_parameters_uri=" + expandParamsURI,
		"--partial_histogram_uri=" + outputResultURI,
		"--decrypted_report_uri=" + outputDecryptedReportURI,
		"--epsilon=" + fmt.Sprintf("%f", request.TotalEpsilon*config.PrivacyBudgetPerPrefix[request.QueryLevel]),
		"--private_key_params_uri=" + h.ServerCfg.PrivateKeyParamsURI,
		"--key_bit_size=" + fmt.Sprint(request.KeyBitSize),
		"--runner=" + h.PipelineRunner,
	}

	if h.PipelineRunner == "dataflow" {
		args = append(args,
			"--project="+h.DataflowCfg.Project,
			"--region="+h.DataflowCfg.Region,
			"--temp_location="+h.DataflowCfg.TempLocation,
			"--staging_location="+h.DataflowCfg.StagingLocation,
			// set jobname to queryID-level-origin
			"--job_name="+fmt.Sprintf("%s-%v-%s", request.QueryID, request.QueryLevel, h.Origin),
			"--worker_binary="+h.ServerCfg.DpfAggregatePartialReportBinary,
		)
	}

	str := h.ServerCfg.DpfAggregatePartialReportBinary
	for _, s := range args {
		str = fmt.Sprintf("%s\n%s", str, s)
	}
	log.Infof("Running command\n%s", str)

	cmd := exec.CommandContext(ctx, h.ServerCfg.DpfAggregatePartialReportBinary, args...)
	var out bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &stderr
	err = cmd.Run()
	if err != nil {
		log.Errorf("%s: %s", err, stderr.String())
		return err
	}
	log.Infof("output of cmd: %s", out.String())

	if request.QueryLevel == finalLevel {
		log.Infof("query %q complete", request.QueryID)
		return nil
	}

	// If the hierarchical query is not finished yet, publish the requests for the next-level aggregation.
	request.QueryLevel++
	_, topic, err := utils.ParsePubSubResourceName(h.RequestPubSubTopic)
	if err != nil {
		return err
	}
	return utils.PublishRequest(ctx, h.PubSubTopicClient, topic, request)
}

// ReadHelperSharedInfo reads the helper shared info from a URL.
func ReadHelperSharedInfo(client *http.Client, url, token string) (*query.HelperSharedInfo, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	if token != "" {
		req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", token))
	}

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.Status != "200 OK" {
		body, _ := ioutil.ReadAll(resp.Body)
		log.Infof("%v: %s", resp.Status, string(body))
		return nil, fmt.Errorf("Error reading shared info from %s: %s", url, resp.Status)
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	info := &query.HelperSharedInfo{}
	if err := json.Unmarshal([]byte(body), info); err != nil {
		return nil, err
	}
	return info, nil
}
<|endoftext|>


<commit_before>package builtin

import (
	"neon/build"
	"reflect"
	"strings"
)

func init() {
	build.AddBuiltin(build.BuiltinDesc{
		Name: "joinpath",
		Func: joinPath,
		Help: `Join file paths.

Arguments:

- The paths to join as a list of strings.

Returns:

- Joined path as a string.

Examples:

    # join paths "/foo", "bar" and "spam.txt"
    joinpath("foo", "bar", "spam.txt")
    # returns: "foo/bar/spam.txt" on a Linux box and "foo\bar\spam.txt" on
    # Windows`,
	})
}

func joinPath(paths ...interface{}) string {
	s := make([]string, len(paths))
	for i, e := range paths {
		s[i] = reflect.ValueOf(e).String()
	}
	return strings.Join(s, "/")
}
<commit_msg>Fixed doc of joinpath<commit_after>package builtin

import (
	"neon/build"
	"reflect"
	"strings"
)

func init() {
	build.AddBuiltin(build.BuiltinDesc{
		Name: "joinpath",
		Func: joinPath,
		Help: `Join file paths.

Arguments:

- The paths to join as a list of strings.

Returns:

- Joined path as a string.

Examples:

    # join paths "foo", "bar" and "spam.txt"
    joinpath("foo", "bar", "spam.txt")
    # returns: "foo/bar/spam.txt" on a Linux box and "foo\bar\spam.txt" on
    # Windows`,
	})
}

func joinPath(paths ...interface{}) string {
	s := make([]string, len(paths))
	for i, e := range paths {
		s[i] = reflect.ValueOf(e).String()
	}
	return strings.Join(s, "/")
}
<|endoftext|>


<commit_before>package problem3
<commit_msg>add some more<commit_after>package main

import (
	"fmt"
	"strconv"
	"strings"
)

// opcode
const (
	add  = 1
	mult = 2
	halt = 99
)

func main() {
	program := make([]int, 0)
	for _, s := range strings.Split(input, ",") {
		if s == "" {
			continue
		}
		i, err := strconv.ParseInt(s, 10, 64)
		if err != nil {
			panic(err)
		}
		program = append(program, int(i))
	}

	i := 0
	l := len(program)
	fmt.Println(l)
	fmt.Println(program)
Loop:
	for {
		assert(l, i)
		opcode := program[i]
		if opcode == add {
			assert(l, i+3)
			v1 := program[i+1]
			v2 := program[i+2]
			o1 := program[i+3]
			assert(l, v1, v2, o1)

			i1 := program[v1]
			i2 := program[v2]
			program[o1] = i1 + i2
			i += 4
		} else if opcode == mult {
			assert(l, i+3)
			v1 := program[i+1]
			v2 := program[i+2]
			o1 := program[i+3]
			assert(l, v1, v2, o1)

			i1 := program[v1]
			i2 := program[v2]
			program[o1] = i1 * i2
			i += 4
		} else if opcode == halt {
			break Loop
		}
	}

	fmt.Println(program[0])
}

func assert(len int, indices ...int) {
	for _, index := range indices {
		if index >= len {
			panic(fmt.Sprintf("overflow len %v indices %v", len, indices))
		}
	}
}

var input = `1,12,2,3,1,1,2,3,1,3,4,3,1,5,0,3,2,13,1,19,1,10,19,23,1,6,23,27,1,5,27,31,1,10,31,35,2,10,35,39,1,39,5,43,2,43,6,47,2,9,47,51,1,51,5,55,1,5,55,59,2,10,59,63,1,5,63,67,1,67,10,71,2,6,71,75,2,6,75,79,1,5,79,83,2,6,83,87,2,13,87,91,1,91,6,95,2,13,95,99,1,99,5,103,2,103,10,107,1,9,107,111,1,111,6,115,1,115,2,119,1,119,10,0,99,2,14,0,0`
<|endoftext|>


<commit_before>package cli

import (
	"errors"
	"fmt"
	"os"
	"text/tabwriter"

	"github.com/flynn/flynn/host/volume"
	"github.com/flynn/flynn/pkg/cluster"
	"github.com/flynn/go-docopt"
)

func init() {
	Register("volume", runVolume, `
usage: flynn-host volume list
       flynn-host volume create [--provider=<provider>] <host>
       flynn-host volume delete ID...
       flynn-host volume gc

Commands:
    list    Display a list of all volumes of known Flynn hosts
    create  Creates a volume on a host
    delete  Deletes volumes, destroying any data stored on them
    gc      Garbage collect currently unused volumes

Examples:

    $ flynn-host volume list

    $ flynn-host volume create --provider default host0

    $ flynn-host volume destroy 102fad07-07a3-4841-bded-d9e8a3eedbd6

    $ flynn-host volume gc
`)
}

func runVolume(args *docopt.Args, client *cluster.Client) error {
	switch {
	case args.Bool["list"]:
		return runVolumeList(args, client)
	case args.Bool["delete"]:
		return runVolumeDelete(args, client)
	case args.Bool["create"]:
		return runVolumeCreate(args, client)
	case args.Bool["gc"]:
		return runVolumeGarbageCollection(args, client)
	}
	return nil
}

func runVolumeGarbageCollection(args *docopt.Args, client *cluster.Client) error {
	// collect list of all volume ids currently attached to jobs
	hosts, err := client.Hosts()
	if err != nil {
		return fmt.Errorf("could not list hosts: %s", err)
	}
	if len(hosts) == 0 {
		return errors.New("no hosts found")
	}

	attached := make(map[string]struct{})
	for _, h := range hosts {
		jobs, err := h.ListJobs()
		if err != nil {
			fmt.Printf("error listing jobs on host %s: %s\n", h.ID(), err)
			continue
		}
		for _, j := range jobs {
			for _, vb := range j.Job.Config.Volumes {
				attached[vb.VolumeID] = struct{}{}
			}
		}
	}

	volumes, err := clusterVolumes(hosts)
	if err != nil {
		return err
	}

	// iterate over list of all volumes, deleting any not found in the attached list
	success := true
outer:
	for _, v := range volumes {
		if _, ok := attached[v.Volume.ID]; ok {
			// volume is attached, continue to next volume
			continue outer
		}
		if err := v.Host.DestroyVolume(v.Volume.ID); err != nil {
			success = false
			fmt.Printf("could not delete volume %s: %s\n", v.Volume.ID, err)
			continue outer
		}
		fmt.Println(v.Volume.ID, "deleted")
	}
	if !success {
		return errors.New("could not garbage collect all volumes")
	}

	return nil
}

func runVolumeDelete(args *docopt.Args, client *cluster.Client) error {
	success := true
	hosts, err := client.Hosts()
	if err != nil {
		return fmt.Errorf("could not list hosts: %s", err)
	}
	if len(hosts) == 0 {
		return errors.New("no hosts found")
	}

	volumes, err := clusterVolumes(hosts)
	if err != nil {
		return err
	}

outer:
	for _, id := range args.All["ID"].([]string) {
		// find this volume in the list
		for _, v := range volumes {
			if v.Volume.ID == id {
				if err := v.Host.DestroyVolume(id); err != nil {
					success = false
					fmt.Printf("could not delete volume %s: %s\n", id, err)
					continue outer
				}
				// delete the volume
				fmt.Println(id, "deleted")
				continue outer
			}
		}
		success = false
		fmt.Printf("could not delete volume %s: volume not found\n", id)
	}
	if !success {
		return errors.New("could not delete all volumes")
	}
	return nil
}

func runVolumeCreate(args *docopt.Args, client *cluster.Client) error {
	hostId := args.String["<host>"]
	hostClient, err := client.Host(hostId)
	if err != nil {
		fmt.Println("could not connect to host", hostId)
	}
	provider := "default"
	if args.String["--provider"] != "" {
		provider = args.String["--provider"]
	}
	v, err := hostClient.CreateVolume(provider)
	if err != nil {
		fmt.Printf("could not create volume: %s\n", err)
		return err
	}
	fmt.Printf("created volume %s on %s\n", v.ID, hostId)
	return nil
}

type hostVolume struct {
	Host   *cluster.Host
	Volume *volume.Info
}

func clusterVolumes(hosts []*cluster.Host) ([]hostVolume, error) {
	var volumes []hostVolume
	for _, h := range hosts {
		hostVolumes, err := h.ListVolumes()
		if err != nil {
			return volumes, fmt.Errorf("could not get volumes for host %s: %s", h.ID(), err)
		}
		for _, v := range hostVolumes {
			volumes = append(volumes, hostVolume{Host: h, Volume: v})
		}
	}
	return volumes, nil
}

func runVolumeList(args *docopt.Args, client *cluster.Client) error {
	hosts, err := client.Hosts()
	if err != nil {
		return fmt.Errorf("could not list hosts: %s", err)
	}
	if len(hosts) == 0 {
		return errors.New("no hosts found")
	}

	volumes, err := clusterVolumes(hosts)
	if err != nil {
		return err
	}

	w := tabwriter.NewWriter(os.Stdout, 1, 2, 2, ' ', 0)
	defer w.Flush()
	listRec(w,
		"ID",
		"HOST",
	)

	for _, volume := range volumes {
		listRec(w,
			volume.Volume.ID,
			volume.Host.ID(),
		)
	}
	return nil
}
<commit_msg>host/cli: Fix volume garbage collection<commit_after>package cli

import (
	"errors"
	"fmt"
	"os"
	"text/tabwriter"

	"github.com/flynn/flynn/host/types"
	"github.com/flynn/flynn/host/volume"
	"github.com/flynn/flynn/pkg/cluster"
	"github.com/flynn/go-docopt"
)

func init() {
	Register("volume", runVolume, `
usage: flynn-host volume list
       flynn-host volume create [--provider=<provider>] <host>
       flynn-host volume delete ID...
       flynn-host volume gc

Commands:
    list    Display a list of all volumes of known Flynn hosts
    create  Creates a volume on a host
    delete  Deletes volumes, destroying any data stored on them
    gc      Garbage collect currently unused volumes

Examples:

    $ flynn-host volume list

    $ flynn-host volume create --provider default host0

    $ flynn-host volume destroy 102fad07-07a3-4841-bded-d9e8a3eedbd6

    $ flynn-host volume gc
`)
}

func runVolume(args *docopt.Args, client *cluster.Client) error {
	switch {
	case args.Bool["list"]:
		return runVolumeList(args, client)
	case args.Bool["delete"]:
		return runVolumeDelete(args, client)
	case args.Bool["create"]:
		return runVolumeCreate(args, client)
	case args.Bool["gc"]:
		return runVolumeGarbageCollection(args, client)
	}
	return nil
}

func runVolumeGarbageCollection(args *docopt.Args, client *cluster.Client) error {
	// collect list of all volume ids currently attached to jobs
	hosts, err := client.Hosts()
	if err != nil {
		return fmt.Errorf("could not list hosts: %s", err)
	}
	if len(hosts) == 0 {
		return errors.New("no hosts found")
	}

	attached := make(map[string]struct{})
	for _, h := range hosts {
		jobs, err := h.ListJobs()
		if err != nil {
			fmt.Printf("error listing jobs on host %s: %s\n", h.ID(), err)
			continue
		}
		for _, j := range jobs {
			for _, vb := range j.Job.Config.Volumes {
				if j.Status == host.StatusRunning || j.Status == host.StatusStarting {
					attached[vb.VolumeID] = struct{}{}
				}
			}
		}
	}

	volumes, err := clusterVolumes(hosts)
	if err != nil {
		return err
	}

	// iterate over list of all volumes, deleting any not found in the attached list
	success := true
outer:
	for _, v := range volumes {
		if _, ok := attached[v.Volume.ID]; ok {
			// volume is attached, continue to next volume
			continue outer
		}
		if err := v.Host.DestroyVolume(v.Volume.ID); err != nil {
			success = false
			fmt.Printf("could not delete volume %s: %s\n", v.Volume.ID, err)
			continue outer
		}
		fmt.Println(v.Volume.ID, "deleted")
	}
	if !success {
		return errors.New("could not garbage collect all volumes")
	}

	return nil
}

func runVolumeDelete(args *docopt.Args, client *cluster.Client) error {
	success := true
	hosts, err := client.Hosts()
	if err != nil {
		return fmt.Errorf("could not list hosts: %s", err)
	}
	if len(hosts) == 0 {
		return errors.New("no hosts found")
	}

	volumes, err := clusterVolumes(hosts)
	if err != nil {
		return err
	}

outer:
	for _, id := range args.All["ID"].([]string) {
		// find this volume in the list
		for _, v := range volumes {
			if v.Volume.ID == id {
				if err := v.Host.DestroyVolume(id); err != nil {
					success = false
					fmt.Printf("could not delete volume %s: %s\n", id, err)
					continue outer
				}
				// delete the volume
				fmt.Println(id, "deleted")
				continue outer
			}
		}
		success = false
		fmt.Printf("could not delete volume %s: volume not found\n", id)
	}
	if !success {
		return errors.New("could not delete all volumes")
	}
	return nil
}

func runVolumeCreate(args *docopt.Args, client *cluster.Client) error {
	hostId := args.String["<host>"]
	hostClient, err := client.Host(hostId)
	if err != nil {
		fmt.Println("could not connect to host", hostId)
	}
	provider := "default"
	if args.String["--provider"] != "" {
		provider = args.String["--provider"]
	}
	v, err := hostClient.CreateVolume(provider)
	if err != nil {
		fmt.Printf("could not create volume: %s\n", err)
		return err
	}
	fmt.Printf("created volume %s on %s\n", v.ID, hostId)
	return nil
}

type hostVolume struct {
	Host   *cluster.Host
	Volume *volume.Info
}

func clusterVolumes(hosts []*cluster.Host) ([]hostVolume, error) {
	var volumes []hostVolume
	for _, h := range hosts {
		hostVolumes, err := h.ListVolumes()
		if err != nil {
			return volumes, fmt.Errorf("could not get volumes for host %s: %s", h.ID(), err)
		}
		for _, v := range hostVolumes {
			volumes = append(volumes, hostVolume{Host: h, Volume: v})
		}
	}
	return volumes, nil
}

func runVolumeList(args *docopt.Args, client *cluster.Client) error {
	hosts, err := client.Hosts()
	if err != nil {
		return fmt.Errorf("could not list hosts: %s", err)
	}
	if len(hosts) == 0 {
		return errors.New("no hosts found")
	}

	volumes, err := clusterVolumes(hosts)
	if err != nil {
		return err
	}

	w := tabwriter.NewWriter(os.Stdout, 1, 2, 2, ' ', 0)
	defer w.Flush()
	listRec(w,
		"ID",
		"HOST",
	)

	for _, volume := range volumes {
		listRec(w,
			volume.Volume.ID,
			volume.Host.ID(),
		)
	}
	return nil
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"log"
	"net/url"
	"sync"
	"time"

	consul "github.com/hashicorp/consul/api"
	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	syncInterval = 5 * time.Second
)

type trackedService struct {
	allocId     string
	task        *structs.Task
	serviceHash string
	service     *structs.Service
}

type trackedTask struct {
	allocID string
	task    *structs.Task
}

func (t *trackedService) IsServiceValid() bool {
	for _, service := range t.task.Services {
		if service.Id == t.service.Id && service.Hash() == t.serviceHash {
			return true
		}
	}

	return false
}

type ConsulService struct {
	client     *consul.Client
	logger     *log.Logger
	shutdownCh chan struct{}

	trackedServices map[string]*trackedService                // Service ID to Tracked Service Map
	trackedChecks   map[string]*consul.AgentCheckRegistration // List of check ids that is being tracked
	trackedTasks    map[string]*trackedTask
	trackedSrvLock  sync.Mutex
	trackedChkLock  sync.Mutex
	trackedTskLock  sync.Mutex
}

func NewConsulService(logger *log.Logger, consulAddr string) (*ConsulService, error) {
	var err error
	var c *consul.Client
	cfg := consul.DefaultConfig()
	cfg.Address = consulAddr
	if c, err = consul.NewClient(cfg); err != nil {
		return nil, err
	}

	consulService := ConsulService{
		client:          c,
		logger:          logger,
		trackedServices: make(map[string]*trackedService),
		trackedTasks:    make(map[string]*trackedTask),
		trackedChecks:   make(map[string]*consul.AgentCheckRegistration),
		shutdownCh:      make(chan struct{}),
	}

	return &consulService, nil
}

func (c *ConsulService) Register(task *structs.Task, allocID string) error {
	var mErr multierror.Error
	c.trackedTskLock.Lock()
	tt := &trackedTask{allocID: allocID, task: task}
	c.trackedTasks[fmt.Sprintf("%s-%s", allocID, task.Name)] = tt
	c.trackedTskLock.Unlock()
	for _, service := range task.Services {
		c.logger.Printf("[INFO] consul: Registering service %s with Consul.", service.Name)
		if err := c.registerService(service, task, allocID); err != nil {
			mErr.Errors = append(mErr.Errors, err)
		}
	}

	return mErr.ErrorOrNil()
}

func (c *ConsulService) Deregister(task *structs.Task, allocID string) error {
	var mErr multierror.Error
	c.trackedTskLock.Lock()
	delete(c.trackedTasks, fmt.Sprintf("%s-%s", allocID, task.Name))
	c.trackedTskLock.Unlock()
	for _, service := range task.Services {
		if service.Id == "" {
			continue
		}
		c.logger.Printf("[INFO] consul: De-Registering service %v with Consul", service.Name)
		if err := c.deregisterService(service.Id); err != nil {
			c.logger.Printf("[DEBUG] consul: Error in de-registering service %v from Consul", service.Name)
			mErr.Errors = append(mErr.Errors, err)
		}
	}
	return mErr.ErrorOrNil()
}

func (c *ConsulService) ShutDown() {
	close(c.shutdownCh)
}

func (c *ConsulService) SyncWithConsul() {
	sync := time.After(syncInterval)
	agent := c.client.Agent()

	for {
		select {
		case <-sync:
			c.performSync(agent)
			sync = time.After(syncInterval)
		case <-c.shutdownCh:
			c.logger.Printf("[INFO] Shutting down Consul Client")
			return
		}
	}
}

func (c *ConsulService) performSync(agent *consul.Agent) {
	var consulServices map[string]*consul.AgentService
	var consulChecks map[string]*consul.AgentCheck
	var err error

	// Remove the tracked services which tasks no longer references
	for serviceId, ts := range c.trackedServices {
		if !ts.IsServiceValid() {
			c.logger.Printf("[INFO] consul: Removing service: %s since the task doesn't have it anymore", ts.service.Name)
			c.deregisterService(serviceId)
		}
	}

	// Add additional services that we might not have added from tasks
	for _, trackedTask := range c.trackedTasks {
		for _, service := range trackedTask.task.Services {
			if _, ok := c.trackedServices[service.Id]; !ok {
				c.registerService(service, trackedTask.task, trackedTask.allocID)
			}
		}
	}

	// Get the list of the services that Consul knows about
	if consulServices, err = agent.Services(); err != nil {
		return
	}

	// See if we have services that Consul doesn't know about yet.
	// Register with Consul the services which are not registered
	for serviceId := range c.trackedServices {
		if _, ok := consulServices[serviceId]; !ok {
			ts := c.trackedServices[serviceId]
			c.registerService(ts.service, ts.task, ts.allocId)
		}
	}

	// See if consul thinks we have some services which are not running
	// anymore on the node. We de-register those services
	for serviceId := range consulServices {
		if serviceId == "consul" {
			continue
		}
		if _, ok := c.trackedServices[serviceId]; !ok {
			if err := c.deregisterService(serviceId); err != nil {
				c.logger.Printf("[DEBUG] consul: Error while de-registering service with ID: %s", serviceId)
			}
		}
	}

	if consulChecks, err = agent.Checks(); err != nil {
		return
	}

	// Remove checks that Consul knows about but we don't
	for checkID := range consulChecks {
		if _, ok := c.trackedChecks[checkID]; !ok {
			c.deregisterCheck(checkID)
		}
	}

	// Add checks that might not be present
	for _, trackedService := range c.trackedServices {
		host, port := trackedService.task.FindHostAndPortFor(trackedService.service.PortLabel)
		if host == "" || port == 0 {
			continue
		}
		checks := c.makeChecks(trackedService.service, host, port)
		for _, check := range checks {
			if _, ok := consulChecks[check.ID]; !ok {
				c.registerCheck(check)
			}
		}

	}
}

func (c *ConsulService) registerService(service *structs.Service, task *structs.Task, allocID string) error {
	var mErr multierror.Error
	service.Id = fmt.Sprintf("%s-%s", allocID, service.Name)
	host, port := task.FindHostAndPortFor(service.PortLabel)
	if host == "" || port == 0 {
		return fmt.Errorf("consul: The port:%s marked for registration of service: %s couldn't be found", service.PortLabel, service.Name)
	}
	asr := &consul.AgentServiceRegistration{
		ID:      service.Id,
		Name:    service.Name,
		Tags:    service.Tags,
		Port:    port,
		Address: host,
	}
	ts := &trackedService{
		allocId:     allocID,
		task:        task,
		serviceHash: service.Hash(),
		service:     service,
	}
	c.trackedSrvLock.Lock()
	c.trackedServices[service.Id] = ts
	c.trackedSrvLock.Unlock()

	if err := c.client.Agent().ServiceRegister(asr); err != nil {
		c.logger.Printf("[DEBUG] consul: Error while registering service %v with Consul: %v", service.Name, err)
		mErr.Errors = append(mErr.Errors, err)
	}
	checks := c.makeChecks(service, host, port)
	for _, check := range checks {
		if err := c.registerCheck(check); err != nil {
			c.logger.Printf("[ERROR] consul: Error while registerting check %v with Consul: %v", check.Name, err)
			mErr.Errors = append(mErr.Errors, err)
		}
	}
	return mErr.ErrorOrNil()
}

func (c *ConsulService) registerCheck(check *consul.AgentCheckRegistration) error {
	c.logger.Printf("[DEBUG] Registering Check with ID: %v for Service: %v", check.ID, check.ServiceID)
	c.trackedChkLock.Lock()
	c.trackedChecks[check.ID] = check
	c.trackedChkLock.Unlock()
	return c.client.Agent().CheckRegister(check)
}

func (c *ConsulService) deregisterCheck(checkID string) error {
	c.logger.Printf("[DEBUG] Removing check with ID: %v", checkID)
	c.trackedChkLock.Lock()
	delete(c.trackedChecks, checkID)
	c.trackedChkLock.Unlock()
	return c.client.Agent().CheckDeregister(checkID)
}

func (c *ConsulService) deregisterService(serviceId string) error {
	c.trackedSrvLock.Lock()
	delete(c.trackedServices, serviceId)
	c.trackedSrvLock.Unlock()

	if err := c.client.Agent().ServiceDeregister(serviceId); err != nil {
		return err
	}
	return nil
}

func (c *ConsulService) makeChecks(service *structs.Service, ip string, port int) []*consul.AgentCheckRegistration {
	var checks []*consul.AgentCheckRegistration
	for _, check := range service.Checks {
		if check.Name == "" {
			check.Name = fmt.Sprintf("service: '%s' check", service.Name)
		}
		cr := &consul.AgentCheckRegistration{
			ID:        check.Hash(),
			Name:      check.Name,
			ServiceID: service.Id,
		}
		cr.Interval = check.Interval.String()
		cr.Timeout = check.Timeout.String()
		switch check.Type {
		case structs.ServiceCheckHTTP:
			if check.Protocol == "" {
				check.Protocol = "http"
			}
			url := url.URL{
				Scheme: check.Protocol,
				Host:   fmt.Sprintf("%s:%d", ip, port),
				Path:   check.Path,
			}
			cr.HTTP = url.String()
		case structs.ServiceCheckTCP:
			cr.TCP = fmt.Sprintf("%s:%d", ip, port)
		case structs.ServiceCheckScript:
			cr.Script = check.Script // TODO This needs to include the path of the alloc dir and based on driver types
		}

		checks = append(checks, cr)
	}
	return checks
}
<commit_msg>Saving the host and port in tracked service<commit_after>package client

import (
	"fmt"
	"log"
	"net/url"
	"sync"
	"time"

	consul "github.com/hashicorp/consul/api"
	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	syncInterval = 5 * time.Second
)

type trackedService struct {
	allocId     string
	task        *structs.Task
	serviceHash string
	service     *structs.Service
	host        string
	port        int
}

type trackedTask struct {
	allocID string
	task    *structs.Task
}

func (t *trackedService) IsServiceValid() bool {
	for _, service := range t.task.Services {
		if service.Id == t.service.Id && service.Hash() == t.serviceHash {
			return true
		}
	}

	return false
}

type ConsulService struct {
	client     *consul.Client
	logger     *log.Logger
	shutdownCh chan struct{}

	trackedServices map[string]*trackedService                // Service ID to Tracked Service Map
	trackedChecks   map[string]*consul.AgentCheckRegistration // List of check ids that is being tracked
	trackedTasks    map[string]*trackedTask
	trackedSrvLock  sync.Mutex
	trackedChkLock  sync.Mutex
	trackedTskLock  sync.Mutex
}

func NewConsulService(logger *log.Logger, consulAddr string) (*ConsulService, error) {
	var err error
	var c *consul.Client
	cfg := consul.DefaultConfig()
	cfg.Address = consulAddr
	if c, err = consul.NewClient(cfg); err != nil {
		return nil, err
	}

	consulService := ConsulService{
		client:          c,
		logger:          logger,
		trackedServices: make(map[string]*trackedService),
		trackedTasks:    make(map[string]*trackedTask),
		trackedChecks:   make(map[string]*consul.AgentCheckRegistration),
		shutdownCh:      make(chan struct{}),
	}

	return &consulService, nil
}

func (c *ConsulService) Register(task *structs.Task, allocID string) error {
	var mErr multierror.Error
	c.trackedTskLock.Lock()
	tt := &trackedTask{allocID: allocID, task: task}
	c.trackedTasks[fmt.Sprintf("%s-%s", allocID, task.Name)] = tt
	c.trackedTskLock.Unlock()
	for _, service := range task.Services {
		c.logger.Printf("[INFO] consul: Registering service %s with Consul.", service.Name)
		if err := c.registerService(service, task, allocID); err != nil {
			mErr.Errors = append(mErr.Errors, err)
		}
	}

	return mErr.ErrorOrNil()
}

func (c *ConsulService) Deregister(task *structs.Task, allocID string) error {
	var mErr multierror.Error
	c.trackedTskLock.Lock()
	delete(c.trackedTasks, fmt.Sprintf("%s-%s", allocID, task.Name))
	c.trackedTskLock.Unlock()
	for _, service := range task.Services {
		if service.Id == "" {
			continue
		}
		c.logger.Printf("[INFO] consul: De-Registering service %v with Consul", service.Name)
		if err := c.deregisterService(service.Id); err != nil {
			c.logger.Printf("[DEBUG] consul: Error in de-registering service %v from Consul", service.Name)
			mErr.Errors = append(mErr.Errors, err)
		}
	}
	return mErr.ErrorOrNil()
}

func (c *ConsulService) ShutDown() {
	close(c.shutdownCh)
}

func (c *ConsulService) SyncWithConsul() {
	sync := time.After(syncInterval)
	agent := c.client.Agent()

	for {
		select {
		case <-sync:
			c.performSync(agent)
			sync = time.After(syncInterval)
		case <-c.shutdownCh:
			c.logger.Printf("[INFO] Shutting down Consul Client")
			return
		}
	}
}

func (c *ConsulService) performSync(agent *consul.Agent) {
	var consulServices map[string]*consul.AgentService
	var consulChecks map[string]*consul.AgentCheck
	var err error

	// Remove the tracked services which tasks no longer references
	for serviceId, ts := range c.trackedServices {
		if !ts.IsServiceValid() {
			c.logger.Printf("[INFO] consul: Removing service: %s since the task doesn't have it anymore", ts.service.Name)
			c.deregisterService(serviceId)
		}
	}

	// Add additional services that we might not have added from tasks
	for _, trackedTask := range c.trackedTasks {
		for _, service := range trackedTask.task.Services {
			if _, ok := c.trackedServices[service.Id]; !ok {
				c.registerService(service, trackedTask.task, trackedTask.allocID)
			}
		}
	}

	// Get the list of the services that Consul knows about
	if consulServices, err = agent.Services(); err != nil {
		return
	}

	// See if we have services that Consul doesn't know about yet.
	// Register with Consul the services which are not registered
	for serviceId := range c.trackedServices {
		if _, ok := consulServices[serviceId]; !ok {
			ts := c.trackedServices[serviceId]
			c.registerService(ts.service, ts.task, ts.allocId)
		}
	}

	// See if consul thinks we have some services which are not running
	// anymore on the node. We de-register those services
	for serviceId := range consulServices {
		if serviceId == "consul" {
			continue
		}
		if _, ok := c.trackedServices[serviceId]; !ok {
			if err := c.deregisterService(serviceId); err != nil {
				c.logger.Printf("[DEBUG] consul: Error while de-registering service with ID: %s", serviceId)
			}
		}
	}

	if consulChecks, err = agent.Checks(); err != nil {
		return
	}

	// Remove checks that Consul knows about but we don't
	for checkID := range consulChecks {
		if _, ok := c.trackedChecks[checkID]; !ok {
			c.deregisterCheck(checkID)
		}
	}

	// Add checks that might not be present
	for _, ts := range c.trackedServices {
		checks := c.makeChecks(ts.service, ts.host, ts.port)
		for _, check := range checks {
			if _, ok := consulChecks[check.ID]; !ok {
				c.registerCheck(check)
			}
		}

	}
}

func (c *ConsulService) registerService(service *structs.Service, task *structs.Task, allocID string) error {
	var mErr multierror.Error
	service.Id = fmt.Sprintf("%s-%s", allocID, service.Name)
	host, port := task.FindHostAndPortFor(service.PortLabel)
	if host == "" || port == 0 {
		return fmt.Errorf("consul: The port:%s marked for registration of service: %s couldn't be found", service.PortLabel, service.Name)
	}
	ts := &trackedService{
		allocId:     allocID,
		task:        task,
		serviceHash: service.Hash(),
		service:     service,
		host:        host,
		port:        port,
	}
	c.trackedSrvLock.Lock()
	c.trackedServices[service.Id] = ts
	c.trackedSrvLock.Unlock()

	asr := &consul.AgentServiceRegistration{
		ID:      service.Id,
		Name:    service.Name,
		Tags:    service.Tags,
		Port:    port,
		Address: host,
	}

	if err := c.client.Agent().ServiceRegister(asr); err != nil {
		c.logger.Printf("[DEBUG] consul: Error while registering service %v with Consul: %v", service.Name, err)
		mErr.Errors = append(mErr.Errors, err)
	}
	checks := c.makeChecks(service, host, port)
	for _, check := range checks {
		if err := c.registerCheck(check); err != nil {
			c.logger.Printf("[ERROR] consul: Error while registerting check %v with Consul: %v", check.Name, err)
			mErr.Errors = append(mErr.Errors, err)
		}
	}
	return mErr.ErrorOrNil()
}

func (c *ConsulService) registerCheck(check *consul.AgentCheckRegistration) error {
	c.logger.Printf("[DEBUG] Registering Check with ID: %v for Service: %v", check.ID, check.ServiceID)
	c.trackedChkLock.Lock()
	c.trackedChecks[check.ID] = check
	c.trackedChkLock.Unlock()
	return c.client.Agent().CheckRegister(check)
}

func (c *ConsulService) deregisterCheck(checkID string) error {
	c.logger.Printf("[DEBUG] Removing check with ID: %v", checkID)
	c.trackedChkLock.Lock()
	delete(c.trackedChecks, checkID)
	c.trackedChkLock.Unlock()
	return c.client.Agent().CheckDeregister(checkID)
}

func (c *ConsulService) deregisterService(serviceId string) error {
	c.trackedSrvLock.Lock()
	delete(c.trackedServices, serviceId)
	c.trackedSrvLock.Unlock()

	if err := c.client.Agent().ServiceDeregister(serviceId); err != nil {
		return err
	}
	return nil
}

func (c *ConsulService) makeChecks(service *structs.Service, ip string, port int) []*consul.AgentCheckRegistration {
	var checks []*consul.AgentCheckRegistration
	for _, check := range service.Checks {
		if check.Name == "" {
			check.Name = fmt.Sprintf("service: '%s' check", service.Name)
		}
		cr := &consul.AgentCheckRegistration{
			ID:        check.Hash(),
			Name:      check.Name,
			ServiceID: service.Id,
		}
		cr.Interval = check.Interval.String()
		cr.Timeout = check.Timeout.String()
		switch check.Type {
		case structs.ServiceCheckHTTP:
			if check.Protocol == "" {
				check.Protocol = "http"
			}
			url := url.URL{
				Scheme: check.Protocol,
				Host:   fmt.Sprintf("%s:%d", ip, port),
				Path:   check.Path,
			}
			cr.HTTP = url.String()
		case structs.ServiceCheckTCP:
			cr.TCP = fmt.Sprintf("%s:%d", ip, port)
		case structs.ServiceCheckScript:
			cr.Script = check.Script // TODO This needs to include the path of the alloc dir and based on driver types
		}

		checks = append(checks, cr)
	}
	return checks
}
<|endoftext|>


<commit_before>// Package remctl provides access to the remctl client library.
// http://www.eyrie.org/~eagle/software/remctl/
//
// Copyright 2014 Thomas L. Kula <kula@tproa.net>
// All Rights Reserved
//
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE file.

/*
    Package remctl provides access to the remctl client library from
    http://www.eyrie.org/~eagle/software/remctl/

    There is a 'simple' interface which effectively implements the
    'remctl' and 'remctl_result_free' functions from the C api
    (in Go programs remctl_result_free is called for you).
*/
package remctl

/*
#cgo pkg-config: libremctl
#include <remctl.h>
#include <stdlib.h>
#include <sys/uio.h>

// Cheerfully stolen from https://groups.google.com/forum/#!topic/golang-nuts/pQueMFdY0mk
static char** makeCharArray( int size ) {
    // We need to hold 'size' character arrays, and
    // null terminate the last entry for the remctl library

    char** a = calloc( sizeof(char*), size + 1 );
    a[size] = NULL;

    return a;
}

static void setArrayString( char **a, char *s, int n ) {
    a[n] = s;
}

static void freeCharArray( char **a, int size ) {
    int i;
    for( i = 0; i < size; i++ ) {
	free( a[i] );
    }
    free(a);
}

*/
import "C"

import (
    "errors"
    "fmt"
    "unsafe"
)

// The "Simple" interface

// RemctlResult is what is returned by the remctl 'simple' interface.
// Stdout and Stderr are the returned stdout and stderr, respectively,
// while Status is the return code.
type RemctlResult struct {
    Stdout  string
    Stderr  string
    Status  int
}

// Simple is the 'simple' interface.
//  If port is 0, the default remctl port is used
//  If principal is "", the default principal of "host/<hostname>" is used
//
// For more control of how the call is made, use the more 'complex'
// interface.
func Simple( host string, port uint16, principal string, command []string ) ( *RemctlResult, error ) {
    var res *C.struct_remctl_result

    host_c := C.CString( host )
    defer C.free( unsafe.Pointer( host_c ))

    var principal_c *_Ctype_char
    if principal != "" {
	principal_c = C.CString( principal )
	defer C.free( unsafe.Pointer( principal_c ))
    }

    command_len := len( command )
    command_c := C.makeCharArray( C.int( command_len ))
    defer C.freeCharArray( command_c, C.int( command_len ))
    for i, s := range command {
	C.setArrayString( command_c, C.CString( s ), C.int( i ))
    }

    res, err := C.remctl( host_c, C.ushort( port ), principal_c, command_c )

    if res == nil {
	return nil, err
    } else {
	defer C.remctl_result_free((*C.struct_remctl_result)(unsafe.Pointer( res )))
	result := &RemctlResult{}
	stdout_len := (C.int)(res.stdout_len)
	stderr_len := (C.int)(res.stderr_len)
	result.Stdout = C.GoStringN( res.stdout_buf, stdout_len )
	result.Stderr = C.GoStringN( res.stderr_buf, stderr_len )
	result.Status = (int)(res.status)
	if res.error != nil {
	    return nil, errors.New( C.GoString( res.error ))
	}

	return result, nil
    }
}

// The 'complex' interface

type Command []string

type Output struct {
    Data string
    Stream int
}

type Error struct {
    Data string
    Code int
}

type Done bool

type Status int

const (
    REMCTL_OUT_OUTPUT = iota
    REMCTL_OUT_STATUS
    REMCTL_OUT_ERROR
    REMCTL_OUT_DONE
)

// The following are from the remctl protocol specification
const (
    ERROR_NOCODE = iota		// 0 isn't used
    ERROR_INTERNAL		// Internal server failure
    ERROR_BAD_TOKEN		// Invalid format in token
    ERROR_UNKNOWN		// Unknown message type
    ERROR_BAD_COMMAND		// Invalid command format in token
    ERROR_UNKNOWN_COMMAND	// Command not defined on remctld server
    ERROR_ACCESS		// Access denied
    ERROR_TOOMANY_ARGS		// Argument count exceeds server limit
    ERROR_TOOMUCH_DATA		// Argument size exceeds server limit
    ERROR_UNEXPECTED_MESSAGE	// Message type not valid now
)

type remctl struct {
    ctx	    *C.struct_remctl
    Output chan Output
    Error chan Error
    Status chan Status
    Done chan Done  // We send true over this when REMCTL_OUT_DONE is returned
    open bool
}

func (r *remctl) get_error() ( error ) {
    return errors.New( C.GoString( C.remctl_error( r.ctx )))
}

func (r *remctl) assert_init() ( error ) {
    if r.ctx == nil {
	return errors.New( "Connection not inited" )
    }

    return nil
}

func (r *remctl) isopen() ( bool ) {
    return r.open
}

// All of the remctl calls have the convention that if the principal is
// NULL, use the default. we use the convention that if the principal is
// an empty string, we use the default. Translate.
func get_principal ( principal string ) (*_Ctype_char) {
    var principal_c *_Ctype_char
    if principal != "" {
	principal_c = C.CString( principal )
    }

    return principal_c
}

// Free a principal, but only if it's not NULL
func free_principal ( principal_c *_Ctype_char ) {
    if principal_c != nil {
	C.free( unsafe.Pointer( principal_c ))
    }
}

// Create a new remctl connection
func New() (*remctl, error ) {
    r := &remctl{}
    r.Output = make( chan Output )
    r.Error = make( chan Error )
    r.Done = make( chan Done )
    r.Status = make( chan Status )

    ctx, err := C.remctl_new()
    r.ctx = ctx
    r.open = false
    return r, err
}

// Setccache allows you to define the credentials cache the underlying
// GSSAPI library will use to make the connection.
//
// Your underlying GSSAPI library may not tolerate that, so be prepared
// for errors
func (r *remctl) Setccache( ccache string ) ( error ) {
    ccache_c := C.CString( ccache )
    defer C.free( unsafe.Pointer( ccache_c ))

    if set := C.remctl_set_ccache( r.ctx, ccache_c ); set != 1 {
	return r.get_error()
    }

    return nil
}

// Close() the remctl connection after you are done with it
//
// After calling, the struct is useless and should be deleted
func (r *remctl) Close() (error) {
    if open := r.assert_init(); open != nil {
	return open
    }

    C.remctl_close( r.ctx ) // The remctl library frees the memory pointed to by r.ctx
    r.ctx = nil
    return nil
}

// Open() a connection
//
// Open a connection to `host' on port `port'. If port is 0, use the default
// remctl port. You may specify a principal to use in `principal', if it is
// a blank string the remctl will use the default principal.
func (r *remctl) Open( host string, port uint16, principal string ) ( error ) {
    if r.isopen() {
	return errors.New( "Already open" )
    }

    host_c := C.CString( host )
    defer C.free( unsafe.Pointer( host_c ))
    port_c := C.ushort( port )

    principal_c := get_principal( principal )
    if principal != "" {
	// If principal is empty, principal_c is NULL, don't free
	defer C.free( unsafe.Pointer( principal_c ))
    }

    if opened := C.remctl_open( r.ctx, host_c, port_c, principal_c ); opened != 1 {
	return r.get_error()
    }

    return nil
}

// Execute() a command
//
// Executes a command whose arguments are a slice of strings. This starts a
// goroutine which will send Output over the remctl.Output channel.  It will
// continue to do this until there is no more output, after which it will send
// Status over the remctl.Status channel. If at any point there is an error
// message, it will send an Error over the remctl.Error channel. If the remctl
// server things the connection is done, `true' will be sent over the Done
// channel. In any case, after a single Status, Error or Done has been
// returned, the command is done and the remctl connection is ready for the
// next Execute() or Close()
//
// If this function returns an error, nothing except Close() is safe
func (r *remctl) Execute( cmd Command ) (error) {

    if !r.isopen() {
	return errors.New( "Not open" )
    }

    // Idea cheerfully stolen from go_fuse's syscall_linux.go
    iov := make([]_Ctype_struct_iovec, len(cmd))
    for n, v := range cmd {
	cmd_c := C.CString(v)
	defer C.free(unsafe.Pointer(cmd_c))
	iov[n].iov_base = unsafe.Pointer(cmd_c)
	iov[n].iov_len = C.size_t(len(v))
    }

    if sent := C.remctl_commandv( r.ctx, &iov[0], C.size_t(len(cmd))); sent == 1 {
	// It failed, return an error
	return r.get_error()
    }

    // Now we enter a goroutine that pulls various bits of data out of
    // the remctl connection and sends it down channels until we're
    // done
    go func(r *remctl) {
	var output *C.struct_remctl_output
	for {
	    output = C.remctl_output(r.ctx)
	    if output == nil {
		error_msg := Error{}
		err_txt := fmt.Sprintf( "%s", r.get_error())
		error_msg.Data = err_txt
		error_msg.Code = ERROR_NOCODE   // We fake this here
		r.Error <- error_msg
		// And we're done
		return
	    }

	    switch output_type := C.int( output._type ); output_type {
	    case REMCTL_OUT_OUTPUT:
		output_msg := Output{}
		data := C.GoStringN( output.data, C.int(output.length))
		output_msg.Data = data
		output_msg.Stream = int( output.stream )
		r.Output <- output_msg
	    case REMCTL_OUT_STATUS:
		r.Status <- Status( output.status )
		return
	    case REMCTL_OUT_ERROR:
		error_msg := Error{}
		data := C.GoStringN( output.data, C.int(output.length ))
		error_msg.Data = data
		error_msg.Code = int( output.error )
		r.Error <- error_msg
		return
	    case REMCTL_OUT_DONE:
		r.Done <- true
		return
	    }
	}
    }(r) // End of goroutine

    return nil
}
<commit_msg>Set open after Open is called<commit_after>// Package remctl provides access to the remctl client library.
// http://www.eyrie.org/~eagle/software/remctl/
//
// Copyright 2014 Thomas L. Kula <kula@tproa.net>
// All Rights Reserved
//
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE file.

/*
    Package remctl provides access to the remctl client library from
    http://www.eyrie.org/~eagle/software/remctl/

    There is a 'simple' interface which effectively implements the
    'remctl' and 'remctl_result_free' functions from the C api
    (in Go programs remctl_result_free is called for you).
*/
package remctl

/*
#cgo pkg-config: libremctl
#include <remctl.h>
#include <stdlib.h>
#include <sys/uio.h>

// Cheerfully stolen from https://groups.google.com/forum/#!topic/golang-nuts/pQueMFdY0mk
static char** makeCharArray( int size ) {
    // We need to hold 'size' character arrays, and
    // null terminate the last entry for the remctl library

    char** a = calloc( sizeof(char*), size + 1 );
    a[size] = NULL;

    return a;
}

static void setArrayString( char **a, char *s, int n ) {
    a[n] = s;
}

static void freeCharArray( char **a, int size ) {
    int i;
    for( i = 0; i < size; i++ ) {
	free( a[i] );
    }
    free(a);
}

*/
import "C"

import (
    "errors"
    "fmt"
    "unsafe"
)

// The "Simple" interface

// RemctlResult is what is returned by the remctl 'simple' interface.
// Stdout and Stderr are the returned stdout and stderr, respectively,
// while Status is the return code.
type RemctlResult struct {
    Stdout  string
    Stderr  string
    Status  int
}

// Simple is the 'simple' interface.
//  If port is 0, the default remctl port is used
//  If principal is "", the default principal of "host/<hostname>" is used
//
// For more control of how the call is made, use the more 'complex'
// interface.
func Simple( host string, port uint16, principal string, command []string ) ( *RemctlResult, error ) {
    var res *C.struct_remctl_result

    host_c := C.CString( host )
    defer C.free( unsafe.Pointer( host_c ))

    var principal_c *_Ctype_char
    if principal != "" {
	principal_c = C.CString( principal )
	defer C.free( unsafe.Pointer( principal_c ))
    }

    command_len := len( command )
    command_c := C.makeCharArray( C.int( command_len ))
    defer C.freeCharArray( command_c, C.int( command_len ))
    for i, s := range command {
	C.setArrayString( command_c, C.CString( s ), C.int( i ))
    }

    res, err := C.remctl( host_c, C.ushort( port ), principal_c, command_c )

    if res == nil {
	return nil, err
    } else {
	defer C.remctl_result_free((*C.struct_remctl_result)(unsafe.Pointer( res )))
	result := &RemctlResult{}
	stdout_len := (C.int)(res.stdout_len)
	stderr_len := (C.int)(res.stderr_len)
	result.Stdout = C.GoStringN( res.stdout_buf, stdout_len )
	result.Stderr = C.GoStringN( res.stderr_buf, stderr_len )
	result.Status = (int)(res.status)
	if res.error != nil {
	    return nil, errors.New( C.GoString( res.error ))
	}

	return result, nil
    }
}

// The 'complex' interface

type Command []string

type Output struct {
    Data string
    Stream int
}

type Error struct {
    Data string
    Code int
}

type Done bool

type Status int

const (
    REMCTL_OUT_OUTPUT = iota
    REMCTL_OUT_STATUS
    REMCTL_OUT_ERROR
    REMCTL_OUT_DONE
)

// The following are from the remctl protocol specification
const (
    ERROR_NOCODE = iota		// 0 isn't used
    ERROR_INTERNAL		// Internal server failure
    ERROR_BAD_TOKEN		// Invalid format in token
    ERROR_UNKNOWN		// Unknown message type
    ERROR_BAD_COMMAND		// Invalid command format in token
    ERROR_UNKNOWN_COMMAND	// Command not defined on remctld server
    ERROR_ACCESS		// Access denied
    ERROR_TOOMANY_ARGS		// Argument count exceeds server limit
    ERROR_TOOMUCH_DATA		// Argument size exceeds server limit
    ERROR_UNEXPECTED_MESSAGE	// Message type not valid now
)

type remctl struct {
    ctx	    *C.struct_remctl
    Output chan Output
    Error chan Error
    Status chan Status
    Done chan Done  // We send true over this when REMCTL_OUT_DONE is returned
    open bool
}

func (r *remctl) get_error() ( error ) {
    return errors.New( C.GoString( C.remctl_error( r.ctx )))
}

func (r *remctl) assert_init() ( error ) {
    if r.ctx == nil {
	return errors.New( "Connection not inited" )
    }

    return nil
}

func (r *remctl) isopen() ( bool ) {
    return r.open
}

// All of the remctl calls have the convention that if the principal is
// NULL, use the default. we use the convention that if the principal is
// an empty string, we use the default. Translate.
func get_principal ( principal string ) (*_Ctype_char) {
    var principal_c *_Ctype_char
    if principal != "" {
	principal_c = C.CString( principal )
    }

    return principal_c
}

// Free a principal, but only if it's not NULL
func free_principal ( principal_c *_Ctype_char ) {
    if principal_c != nil {
	C.free( unsafe.Pointer( principal_c ))
    }
}

// Create a new remctl connection
func New() (*remctl, error ) {
    r := &remctl{}
    r.Output = make( chan Output )
    r.Error = make( chan Error )
    r.Done = make( chan Done )
    r.Status = make( chan Status )

    ctx, err := C.remctl_new()
    r.ctx = ctx
    r.open = false
    return r, err
}

// Setccache allows you to define the credentials cache the underlying
// GSSAPI library will use to make the connection.
//
// Your underlying GSSAPI library may not tolerate that, so be prepared
// for errors
func (r *remctl) Setccache( ccache string ) ( error ) {
    ccache_c := C.CString( ccache )
    defer C.free( unsafe.Pointer( ccache_c ))

    if set := C.remctl_set_ccache( r.ctx, ccache_c ); set != 1 {
	return r.get_error()
    }

    return nil
}

// Close() the remctl connection after you are done with it
//
// After calling, the struct is useless and should be deleted
func (r *remctl) Close() (error) {
    if open := r.assert_init(); open != nil {
	return open
    }

    C.remctl_close( r.ctx ) // The remctl library frees the memory pointed to by r.ctx
    r.ctx = nil
    return nil
}

// Open() a connection
//
// Open a connection to `host' on port `port'. If port is 0, use the default
// remctl port. You may specify a principal to use in `principal', if it is
// a blank string the remctl will use the default principal.
func (r *remctl) Open( host string, port uint16, principal string ) ( error ) {
    if r.isopen() {
	return errors.New( "Already open" )
    }

    host_c := C.CString( host )
    defer C.free( unsafe.Pointer( host_c ))
    port_c := C.ushort( port )

    principal_c := get_principal( principal )
    if principal != "" {
	// If principal is empty, principal_c is NULL, don't free
	defer C.free( unsafe.Pointer( principal_c ))
    }

    if opened := C.remctl_open( r.ctx, host_c, port_c, principal_c ); opened != 1 {
	return r.get_error()
    }

    r.open = true

    return nil
}

// Execute() a command
//
// Executes a command whose arguments are a slice of strings. This starts a
// goroutine which will send Output over the remctl.Output channel.  It will
// continue to do this until there is no more output, after which it will send
// Status over the remctl.Status channel. If at any point there is an error
// message, it will send an Error over the remctl.Error channel. If the remctl
// server things the connection is done, `true' will be sent over the Done
// channel. In any case, after a single Status, Error or Done has been
// returned, the command is done and the remctl connection is ready for the
// next Execute() or Close()
//
// If this function returns an error, nothing except Close() is safe
func (r *remctl) Execute( cmd Command ) (error) {

    if !r.isopen() {
	return errors.New( "Not open" )
    }

    // Idea cheerfully stolen from go_fuse's syscall_linux.go
    iov := make([]_Ctype_struct_iovec, len(cmd))
    for n, v := range cmd {
	cmd_c := C.CString(v)
	defer C.free(unsafe.Pointer(cmd_c))
	iov[n].iov_base = unsafe.Pointer(cmd_c)
	iov[n].iov_len = C.size_t(len(v))
    }

    if sent := C.remctl_commandv( r.ctx, &iov[0], C.size_t(len(cmd))); sent == 1 {
	// It failed, return an error
	return r.get_error()
    }

    // Now we enter a goroutine that pulls various bits of data out of
    // the remctl connection and sends it down channels until we're
    // done
    go func(r *remctl) {
	var output *C.struct_remctl_output
	for {
	    output = C.remctl_output(r.ctx)
	    if output == nil {
		error_msg := Error{}
		err_txt := fmt.Sprintf( "%s", r.get_error())
		error_msg.Data = err_txt
		error_msg.Code = ERROR_NOCODE   // We fake this here
		r.Error <- error_msg
		// And we're done
		return
	    }

	    switch output_type := C.int( output._type ); output_type {
	    case REMCTL_OUT_OUTPUT:
		output_msg := Output{}
		data := C.GoStringN( output.data, C.int(output.length))
		output_msg.Data = data
		output_msg.Stream = int( output.stream )
		r.Output <- output_msg
	    case REMCTL_OUT_STATUS:
		r.Status <- Status( output.status )
		return
	    case REMCTL_OUT_ERROR:
		error_msg := Error{}
		data := C.GoStringN( output.data, C.int(output.length ))
		error_msg.Data = data
		error_msg.Code = int( output.error )
		r.Error <- error_msg
		return
	    case REMCTL_OUT_DONE:
		r.Done <- true
		return
	    }
	}
    }(r) // End of goroutine

    return nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/spf13/cobra"
)

var rootCmd = &cobra.Command{
	Use:              "apex",
	PersistentPreRun: pv.PreRun,
}

func init() {
	pf := rootCmd.PersistentFlags()

	pf.StringVarP(&pv.Chdir, "chdir", "C", "", "Working directory")
	pf.BoolVarP(&pv.DryRun, "dry-run", "D", false, "Perform a dry-run")
	pf.StringSliceVarP(&pv.Env, "env", "e", nil, "Environment variable")
	pf.StringVarP(&pv.LogLevel, "log-level", "l", "info", "Log severity level")
	pf.BoolVarP(&pv.Verbose, "verbose", "v", false, "Output verbose logs")
	pf.BoolVarP(&pv.Version, "version", "V", false, "Output version")
	pf.BoolVarP(&pv.Yes, "yes", "y", false, "Automatic yes to prompts")

	rootCmd.AddCommand(deployCmd)
	rootCmd.AddCommand(deleteCmd)
	rootCmd.AddCommand(invokeCmd)
	rootCmd.AddCommand(rollbackCmd)
	rootCmd.AddCommand(logsCmd)
	rootCmd.AddCommand(buildCmd)
	rootCmd.AddCommand(listCmd)
	rootCmd.AddCommand(wikiCmd)

	// Add a newline at the top of each help message
	rootCmd.SetHelpTemplate("\n" + rootCmd.HelpTemplate())
}
<commit_msg>Enabled SilenceErrors option for root command<commit_after>package main

import (
	"github.com/spf13/cobra"
)

var rootCmd = &cobra.Command{
	Use:              "apex",
	PersistentPreRun: pv.PreRun,
	SilenceErrors:    true,
}

func init() {
	pf := rootCmd.PersistentFlags()

	pf.StringVarP(&pv.Chdir, "chdir", "C", "", "Working directory")
	pf.BoolVarP(&pv.DryRun, "dry-run", "D", false, "Perform a dry-run")
	pf.StringSliceVarP(&pv.Env, "env", "e", nil, "Environment variable")
	pf.StringVarP(&pv.LogLevel, "log-level", "l", "info", "Log severity level")
	pf.BoolVarP(&pv.Verbose, "verbose", "v", false, "Output verbose logs")
	pf.BoolVarP(&pv.Version, "version", "V", false, "Output version")
	pf.BoolVarP(&pv.Yes, "yes", "y", false, "Automatic yes to prompts")

	rootCmd.AddCommand(deployCmd)
	rootCmd.AddCommand(deleteCmd)
	rootCmd.AddCommand(invokeCmd)
	rootCmd.AddCommand(rollbackCmd)
	rootCmd.AddCommand(logsCmd)
	rootCmd.AddCommand(buildCmd)
	rootCmd.AddCommand(listCmd)
	rootCmd.AddCommand(wikiCmd)

	// Add a newline at the top of each help message
	rootCmd.SetHelpTemplate("\n" + rootCmd.HelpTemplate())
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"time"

	_ "github.com/go-sql-driver/mysql"
	config "github.com/vardius/go-api-boilerplate/cmd/auth/internal/application/config"
	eventhandler "github.com/vardius/go-api-boilerplate/cmd/auth/internal/application/eventhandler"
	oauth2 "github.com/vardius/go-api-boilerplate/cmd/auth/internal/application/oauth2"
	client "github.com/vardius/go-api-boilerplate/cmd/auth/internal/domain/client"
	token "github.com/vardius/go-api-boilerplate/cmd/auth/internal/domain/token"
	persistence "github.com/vardius/go-api-boilerplate/cmd/auth/internal/infrastructure/persistence/mysql"
	repository "github.com/vardius/go-api-boilerplate/cmd/auth/internal/infrastructure/repository"
	auth_grpc "github.com/vardius/go-api-boilerplate/cmd/auth/internal/interfaces/grpc"
	auth_http "github.com/vardius/go-api-boilerplate/cmd/auth/internal/interfaces/http"
	application "github.com/vardius/go-api-boilerplate/internal/application"
	buildinfo "github.com/vardius/go-api-boilerplate/internal/buildinfo"
	commandbus "github.com/vardius/go-api-boilerplate/internal/commandbus"
	eventbus "github.com/vardius/go-api-boilerplate/internal/eventbus"
	eventstore "github.com/vardius/go-api-boilerplate/internal/eventstore/memory"
	grpc_utils "github.com/vardius/go-api-boilerplate/internal/grpc"
	log "github.com/vardius/go-api-boilerplate/internal/log"
	mysql "github.com/vardius/go-api-boilerplate/internal/mysql"
	pubsub_proto "github.com/vardius/pubsub/proto"
	"google.golang.org/grpc"
	grpc_health "google.golang.org/grpc/health"
	oauth2_models "gopkg.in/oauth2.v3/models"
)

func main() {
	buildinfo.PrintVersionOrContinue()

	ctx := context.Background()

	logger := log.New(config.Env.APP.Environment)
	eventStore := eventstore.New()
	grpcServer := grpc_utils.NewServer(
		grpc_utils.ServerConfig{
			ServerMinTime: config.Env.GRPC.ServerMinTime,
			ServerTime:    config.Env.GRPC.ServerTime,
			ServerTimeout: config.Env.GRPC.ServerTimeout,
		},
		logger,
	)
	commandBus := commandbus.New(config.Env.CommandBus.QueueSize, logger)

	mysqlConnection := mysql.NewConnection(
		ctx,
		mysql.ConnectionConfig{
			Host:            config.Env.DB.Host,
			Port:            config.Env.DB.Port,
			User:            config.Env.DB.User,
			Pass:            config.Env.DB.Pass,
			Database:        config.Env.DB.Database,
			ConnMaxLifetime: config.Env.DB.ConnMaxLifetime,
			MaxIdleConns:    config.Env.DB.MaxIdleConns,
			MaxOpenConns:    config.Env.DB.MaxOpenConns,
		},
		logger,
	)
	defer mysqlConnection.Close()
	grpcPubsubConn := grpc_utils.NewConnection(
		ctx,
		config.Env.PubSub.Host,
		config.Env.GRPC.Port,
		grpc_utils.ConnectionConfig{
			ConnTime:    config.Env.GRPC.ConnTime,
			ConnTimeout: config.Env.GRPC.ConnTimeout,
		},
		logger,
	)
	defer grpcPubsubConn.Close()
	grpcAuthConn := grpc_utils.NewConnection(
		ctx,
		config.Env.GRPC.Host,
		config.Env.GRPC.Port,
		grpc_utils.ConnectionConfig{
			ConnTime:    config.Env.GRPC.ConnTime,
			ConnTimeout: config.Env.GRPC.ConnTimeout,
		},
		logger,
	)
	defer grpcAuthConn.Close()

	grpPubsubClient := pubsub_proto.NewMessageBusClient(grpcPubsubConn)
	eventBus := eventbus.New(grpPubsubClient, logger)
	tokenRepository := repository.NewTokenRepository(eventStore, eventBus)
	clientRepository := repository.NewClientRepository(eventStore, eventBus)
	tokenPersistenceRepository := persistence.NewTokenRepository(mysqlConnection)
	clientPersistenceRepository := persistence.NewClientRepository(mysqlConnection)
	tokenStore := oauth2.NewTokenStore(tokenPersistenceRepository, commandBus)
	clientStore := oauth2.NewClientStore(clientPersistenceRepository)
	manager := oauth2.NewManager(tokenStore, clientStore, []byte(config.Env.APP.Secret))
	oauth2Server := oauth2.InitServer(manager, mysqlConnection, logger, config.Env.APP.Secret)
	grpcHealthServer := grpc_health.NewServer()
	grpcAuthServer := auth_grpc.NewServer(oauth2Server, logger, config.Env.APP.Secret)
	router := auth_http.NewRouter(
		logger,
		oauth2Server,
		mysqlConnection,
		map[string]*grpc.ClientConn{
			"auth":   grpcAuthConn,
			"pubsub": grpcPubsubConn,
		},
	)
	app := application.New(logger)

	// store our internal user service client
	clientStore.SetInternal(config.Env.User.ClientID, &oauth2_models.Client{
		ID:     config.Env.User.ClientID,
		Secret: config.Env.User.ClientSecret,
		Domain: fmt.Sprintf("http://%s:%d", config.Env.User.Host, config.Env.HTTP.Port),
	})

	commandBus.Subscribe((token.Create{}).GetName(), token.OnCreate(tokenRepository, mysqlConnection))
	commandBus.Subscribe((token.Remove{}).GetName(), token.OnRemove(tokenRepository, mysqlConnection))
	commandBus.Subscribe((client.Create{}).GetName(), client.OnCreate(clientRepository, mysqlConnection))
	commandBus.Subscribe((client.Remove{}).GetName(), client.OnRemove(clientRepository, mysqlConnection))

	go func() {
		go func() {
			eventhandler.Register(
				grpcPubsubConn,
				eventBus,
				map[string]eventbus.EventHandler{
					(token.WasCreated{}).GetType():  eventhandler.WhenTokenWasCreated(mysqlConnection, tokenPersistenceRepository),
					(token.WasRemoved{}).GetType():  eventhandler.WhenTokenWasRemoved(mysqlConnection, tokenPersistenceRepository),
					(client.WasCreated{}).GetType(): eventhandler.WhenClientWasCreated(mysqlConnection, clientPersistenceRepository),
					(client.WasRemoved{}).GetType(): eventhandler.WhenClientWasRemoved(mysqlConnection, clientPersistenceRepository),
				},
				5*time.Minute,
			)
		}()
	}()

	app.AddAdapters(
		auth_http.NewAdapter(
			fmt.Sprintf("%s:%d", config.Env.HTTP.Host, config.Env.HTTP.Port),
			router,
		),
		auth_grpc.NewAdapter(
			fmt.Sprintf("%s:%d", config.Env.GRPC.Host, config.Env.GRPC.Port),
			grpcServer,
			grpcHealthServer,
			grpcAuthServer,
		),
	)

	if config.Env.APP.Environment == "development" {
		app.AddAdapters(
			application.NewDebugAdapter(
				fmt.Sprintf("%s:%d", config.Env.Debug.Host, config.Env.Debug.Port),
			),
		)
	}

	app.WithShutdownTimeout(config.Env.APP.ShutdownTimeout)
	app.Run(ctx)
}
<commit_msg>Handle client store error<commit_after>package main

import (
	"context"
	"fmt"
	"time"

	_ "github.com/go-sql-driver/mysql"
	config "github.com/vardius/go-api-boilerplate/cmd/auth/internal/application/config"
	eventhandler "github.com/vardius/go-api-boilerplate/cmd/auth/internal/application/eventhandler"
	oauth2 "github.com/vardius/go-api-boilerplate/cmd/auth/internal/application/oauth2"
	client "github.com/vardius/go-api-boilerplate/cmd/auth/internal/domain/client"
	token "github.com/vardius/go-api-boilerplate/cmd/auth/internal/domain/token"
	persistence "github.com/vardius/go-api-boilerplate/cmd/auth/internal/infrastructure/persistence/mysql"
	repository "github.com/vardius/go-api-boilerplate/cmd/auth/internal/infrastructure/repository"
	auth_grpc "github.com/vardius/go-api-boilerplate/cmd/auth/internal/interfaces/grpc"
	auth_http "github.com/vardius/go-api-boilerplate/cmd/auth/internal/interfaces/http"
	application "github.com/vardius/go-api-boilerplate/internal/application"
	buildinfo "github.com/vardius/go-api-boilerplate/internal/buildinfo"
	commandbus "github.com/vardius/go-api-boilerplate/internal/commandbus"
	eventbus "github.com/vardius/go-api-boilerplate/internal/eventbus"
	eventstore "github.com/vardius/go-api-boilerplate/internal/eventstore/memory"
	grpc_utils "github.com/vardius/go-api-boilerplate/internal/grpc"
	log "github.com/vardius/go-api-boilerplate/internal/log"
	mysql "github.com/vardius/go-api-boilerplate/internal/mysql"
	pubsub_proto "github.com/vardius/pubsub/proto"
	"google.golang.org/grpc"
	grpc_health "google.golang.org/grpc/health"
	oauth2_models "gopkg.in/oauth2.v3/models"
)

func main() {
	buildinfo.PrintVersionOrContinue()

	ctx := context.Background()

	logger := log.New(config.Env.APP.Environment)
	eventStore := eventstore.New()
	grpcServer := grpc_utils.NewServer(
		grpc_utils.ServerConfig{
			ServerMinTime: config.Env.GRPC.ServerMinTime,
			ServerTime:    config.Env.GRPC.ServerTime,
			ServerTimeout: config.Env.GRPC.ServerTimeout,
		},
		logger,
	)
	commandBus := commandbus.New(config.Env.CommandBus.QueueSize, logger)

	mysqlConnection := mysql.NewConnection(
		ctx,
		mysql.ConnectionConfig{
			Host:            config.Env.DB.Host,
			Port:            config.Env.DB.Port,
			User:            config.Env.DB.User,
			Pass:            config.Env.DB.Pass,
			Database:        config.Env.DB.Database,
			ConnMaxLifetime: config.Env.DB.ConnMaxLifetime,
			MaxIdleConns:    config.Env.DB.MaxIdleConns,
			MaxOpenConns:    config.Env.DB.MaxOpenConns,
		},
		logger,
	)
	defer mysqlConnection.Close()
	grpcPubsubConn := grpc_utils.NewConnection(
		ctx,
		config.Env.PubSub.Host,
		config.Env.GRPC.Port,
		grpc_utils.ConnectionConfig{
			ConnTime:    config.Env.GRPC.ConnTime,
			ConnTimeout: config.Env.GRPC.ConnTimeout,
		},
		logger,
	)
	defer grpcPubsubConn.Close()
	grpcAuthConn := grpc_utils.NewConnection(
		ctx,
		config.Env.GRPC.Host,
		config.Env.GRPC.Port,
		grpc_utils.ConnectionConfig{
			ConnTime:    config.Env.GRPC.ConnTime,
			ConnTimeout: config.Env.GRPC.ConnTimeout,
		},
		logger,
	)
	defer grpcAuthConn.Close()

	grpPubsubClient := pubsub_proto.NewMessageBusClient(grpcPubsubConn)
	eventBus := eventbus.New(grpPubsubClient, logger)
	tokenRepository := repository.NewTokenRepository(eventStore, eventBus)
	clientRepository := repository.NewClientRepository(eventStore, eventBus)
	tokenPersistenceRepository := persistence.NewTokenRepository(mysqlConnection)
	clientPersistenceRepository := persistence.NewClientRepository(mysqlConnection)
	tokenStore := oauth2.NewTokenStore(tokenPersistenceRepository, commandBus)
	clientStore := oauth2.NewClientStore(clientPersistenceRepository)
	manager := oauth2.NewManager(tokenStore, clientStore, []byte(config.Env.APP.Secret))
	oauth2Server := oauth2.InitServer(manager, mysqlConnection, logger, config.Env.APP.Secret)
	grpcHealthServer := grpc_health.NewServer()
	grpcAuthServer := auth_grpc.NewServer(oauth2Server, logger, config.Env.APP.Secret)
	router := auth_http.NewRouter(
		logger,
		oauth2Server,
		mysqlConnection,
		map[string]*grpc.ClientConn{
			"auth":   grpcAuthConn,
			"pubsub": grpcPubsubConn,
		},
	)
	app := application.New(logger)

	// store our internal user service client
	if err := clientStore.SetInternal(config.Env.User.ClientID, &oauth2_models.Client{
		ID:     config.Env.User.ClientID,
		Secret: config.Env.User.ClientSecret,
		Domain: fmt.Sprintf("http://%s:%d", config.Env.User.Host, config.Env.HTTP.Port),
	}); err != nil {
		panic(err)
	}

	commandBus.Subscribe((token.Create{}).GetName(), token.OnCreate(tokenRepository, mysqlConnection))
	commandBus.Subscribe((token.Remove{}).GetName(), token.OnRemove(tokenRepository, mysqlConnection))
	commandBus.Subscribe((client.Create{}).GetName(), client.OnCreate(clientRepository, mysqlConnection))
	commandBus.Subscribe((client.Remove{}).GetName(), client.OnRemove(clientRepository, mysqlConnection))

	go func() {
		go func() {
			eventhandler.Register(
				grpcPubsubConn,
				eventBus,
				map[string]eventbus.EventHandler{
					(token.WasCreated{}).GetType():  eventhandler.WhenTokenWasCreated(mysqlConnection, tokenPersistenceRepository),
					(token.WasRemoved{}).GetType():  eventhandler.WhenTokenWasRemoved(mysqlConnection, tokenPersistenceRepository),
					(client.WasCreated{}).GetType(): eventhandler.WhenClientWasCreated(mysqlConnection, clientPersistenceRepository),
					(client.WasRemoved{}).GetType(): eventhandler.WhenClientWasRemoved(mysqlConnection, clientPersistenceRepository),
				},
				5*time.Minute,
			)
		}()
	}()

	app.AddAdapters(
		auth_http.NewAdapter(
			fmt.Sprintf("%s:%d", config.Env.HTTP.Host, config.Env.HTTP.Port),
			router,
		),
		auth_grpc.NewAdapter(
			fmt.Sprintf("%s:%d", config.Env.GRPC.Host, config.Env.GRPC.Port),
			grpcServer,
			grpcHealthServer,
			grpcAuthServer,
		),
	)

	if config.Env.APP.Environment == "development" {
		app.AddAdapters(
			application.NewDebugAdapter(
				fmt.Sprintf("%s:%d", config.Env.Debug.Host, config.Env.Debug.Port),
			),
		)
	}

	app.WithShutdownTimeout(config.Env.APP.ShutdownTimeout)
	app.Run(ctx)
}
<|endoftext|>


<commit_before>// Copyright © 2017 Compose, an IBM Company
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"
	"log"

	"github.com/spf13/cobra"
)

// backupgetCmd represents the backups get command
var backupgetCmd = &cobra.Command{
	Use:   "get [deployment id] [backup id/name]",
	Short: "Show Backup details for deployment",
	Long:  `Show the backup details for a deployment's backup`,
	Run: func(cmd *cobra.Command, args []string) {
		c := getComposeAPI()
		if len(args) != 2 {
			log.Fatal("Need a deployment id/name and a backup id")
		}
		depid, err := resolveDepID(c, args[0])
		if err != nil {
			log.Fatal(err)
		}
		backupid := args[1]
		if outputRaw {
			text, errs := c.GetBackupDetailsForDeploymentJSON(depid, backupid)
			bailOnErrs(errs)
			fmt.Println(text)
		} else {
			backup, errs := c.GetBackupDetailsForDeployment(depid, backupid)
			bailOnErrs(errs)

			if !outputJSON {
				fmt.Printf("%15s: %s\n", "Backup ID", backup.ID)
				fmt.Printf("%15s: %s\n", "Deployment ID", backup.Deploymentid)
				fmt.Printf("%15s: %s\n", "Backup Name", backup.Name)
				fmt.Printf("%15s: %s\n", "Type", backup.Type)
				fmt.Printf("%15s: %s\n", "Status", backup.Status)
				fmt.Printf("%15s: %s\n", "Download Link", backup.DownloadLink)
			} else {
				printAsJSON(backup)
			}
		}
	},
}

func init() {
	backupsCmd.AddCommand(backupgetCmd)
}
<commit_msg>Added download backup support<commit_after>// Copyright © 2017 Compose, an IBM Company
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"
	"log"
	"net/url"
	"path/filepath"

	"github.com/spf13/cobra"
)

var downloadBackupNow bool

// backupgetCmd represents the backups get command
var backupgetCmd = &cobra.Command{
	Use:   "get [deployment id] [backup id/name]",
	Short: "Show Backup details for deployment",
	Long:  `Show the backup details for a deployment's backup`,
	Run: func(cmd *cobra.Command, args []string) {
		c := getComposeAPI()
		if len(args) != 2 {
			log.Fatal("Need a deployment id/name and a backup id")
		}
		depid, err := resolveDepID(c, args[0])
		if err != nil {
			log.Fatal(err)
		}
		backupid := args[1]
		if outputRaw {
			text, errs := c.GetBackupDetailsForDeploymentJSON(depid, backupid)
			bailOnErrs(errs)
			fmt.Println(text)
		} else {
			backup, errs := c.GetBackupDetailsForDeployment(depid, backupid)
			bailOnErrs(errs)

			if !outputJSON {
				fmt.Printf("%15s: %s\n", "Backup ID", backup.ID)
				fmt.Printf("%15s: %s\n", "Deployment ID", backup.Deploymentid)
				fmt.Printf("%15s: %s\n", "Backup Name", backup.Name)
				fmt.Printf("%15s: %s\n", "Type", backup.Type)
				fmt.Printf("%15s: %s\n", "Status", backup.Status)
				fmt.Printf("%15s: %s\n", "Download Link", backup.DownloadLink)

				if downloadBackupNow && backup.IsDownloadable {
					u, err := url.Parse(backup.DownloadLink)
					err = DownloadFile(filepath.Base(u.Path), backup.DownloadLink)
					if err != nil {
						fmt.Printf("Error downloading: %s\n", err)
					} else {
						fmt.Printf("Downloaded: %s\n", backup.Name)
					}
				} else if downloadBackupNow && !backup.IsDownloadable {
					fmt.Println("Backup is not downloadable - Download skipped")
				}
			} else {
				printAsJSON(backup)
			}
		}
	},
}

func init() {
	backupsCmd.AddCommand(backupgetCmd)
	backupgetCmd.Flags().BoolVarP(&downloadBackupNow, "download", "d", false, "Download the selected backup")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/vially/gois"
	"log"
	"os"
)

func main() {
	if len(os.Args) == 1 {
		fmt.Println("Usage: gois DOMAIN")
		os.Exit(1)
	}

	record, err := gois.Whois(os.Args[1])
	if err != nil {
		log.Fatalln(err)
	}

	fmt.Println(record)
}
<commit_msg>Update gois command line client to take accept multiple domains<commit_after>package main

import (
	"fmt"
	"github.com/vially/gois"
	"os"
)

func main() {
	if len(os.Args) == 1 {
		fmt.Println("Usage: gois DOMAIN [DOMAIN...]")
		os.Exit(1)
	}

	records := gois.WhoisBulk(os.Args[1:])
	for domain, record := range records {
		fmt.Printf("%s: %v\n", domain, record)
	}
}
<|endoftext|>