<commit_before>// Copyright 2017 Google 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
//
//     https://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 https provides an client.Communicator which connects to the
// Fleetspeak server using https.
package https

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"errors"
	"fmt"
	"io"
	"math/big"
	mrand "math/rand"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/golang/protobuf/proto"
	"github.com/google/fleetspeak/fleetspeak/src/client/comms"
	"github.com/google/fleetspeak/fleetspeak/src/common"

	clpb "github.com/google/fleetspeak/fleetspeak/src/client/proto/fleetspeak_client"
	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
)

const (
	sendBytesThreshold = 15 * 1024 * 1024
	sendCountThreshold = 100
	closeWaitThreshold = 30 * time.Second // Matches IdleTimeout in server/https.
)

// Communicator implements comms.Communicator and communicates with a Fleetspeak
// server using https.
type Communicator struct {
	cctx        comms.Context
	conf        *clpb.CommunicatorConfig
	id          common.ClientID
	hc          *http.Client
	ctx         context.Context
	done        context.CancelFunc
	working     sync.WaitGroup
	hosts       []string
	hostLock    sync.RWMutex                                                      // Protects hosts
	DialContext func(ctx context.Context, network, addr string) (net.Conn, error) // If set, will be used to initiate network connections to the server.
}

func (c *Communicator) Setup(cl comms.Context) error {
	c.cctx = cl
	return c.configure()
}

func (c *Communicator) configure() error {
	id, err := c.cctx.CurrentIdentity()
	if err != nil {
		return err
	}

	c.id = id.ID

	tmpl := x509.Certificate{
		Issuer:       pkix.Name{Organization: []string{"GRR Client"}},
		Subject:      pkix.Name{Organization: []string{id.ID.String()}},
		SerialNumber: big.NewInt(1),
	}
	certBytes, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, id.Public, id.Private)
	if err != nil {
		return fmt.Errorf("unable to configure communicator, could not create client cert: %v", err)
	}
	certPair := tls.Certificate{
		Certificate: [][]byte{certBytes},
		PrivateKey:  id.Private,
	}

	si, err := c.cctx.ServerInfo()
	if err != nil {
		return fmt.Errorf("unable to configure communicator, could not get server information: %v", err)
	}
	c.hosts = append([]string(nil), si.Servers...)
	if len(c.hosts) == 0 {
		return errors.New("no server_addresses in client configuration")
	}

	c.conf = c.cctx.CommunicatorConfig()
	if c.conf == nil {
		return errors.New("no communicator_config in client configuration")
	}
	if c.conf.MaxPollDelaySeconds == 0 {
		c.conf.MaxPollDelaySeconds = 60 * 5
	}
	if c.conf.MaxBufferDelaySeconds == 0 {
		c.conf.MaxBufferDelaySeconds = 5
	}
	if c.conf.MinFailureDelaySeconds == 0 {
		c.conf.MinFailureDelaySeconds = 60 * 5
	}
	if c.conf.FailureSuicideTimeSeconds == 0 {
		c.conf.FailureSuicideTimeSeconds = 60 * 60 * 24 * 7
	}

	if c.DialContext == nil {
		c.DialContext = (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
			DualStack: true,
		}).DialContext
	}

	c.hc = &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			TLSClientConfig: &tls.Config{
				RootCAs:      si.TrustedCerts,
				Certificates: []tls.Certificate{certPair},
				CipherSuites: []uint16{
					// We implement both endpoints, so we might as well require long keys and
					// perfect forward secrecy. Note that TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256
					// is required by the https library.
					tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
					tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
					tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
					tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256},
				VerifyPeerCertificate: c.peerCertVerifier,
			},
			DialContext:           c.DialContext,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		},
		Timeout: 5 * time.Minute,
	}
	c.ctx, c.done = context.WithCancel(context.Background())
	return nil
}

func (c *Communicator) Start() error {
	c.working.Add(1)
	go c.processingLoop()
	return nil
}

func (c *Communicator) Stop() {
	c.done()
	c.working.Wait()
}

// jitter adds up to 50% random jitter, and converts to time.Duration.
func jitter(seconds int32) time.Duration {
	return time.Duration((1.0 + 0.5*mrand.Float32()) * float32(seconds) * float32(time.Second))
}

func (c *Communicator) peerCertVerifier(_ [][]byte, chains [][]*x509.Certificate) error {
	for _, chain := range chains {
		if !c.cctx.ChainRevoked(chain) {
			return nil
		}
	}
	return errors.New("certificate revoked")
}

// processingLoop polls the server according to the configured policies while
// the communicator is active.
//
// It is run on an internal goroutine by c.Start, stops when c.ctx is canceled
// and indicates completion by decrementing c.working.
func (c *Communicator) processingLoop() {
	defer c.working.Done()

	// These are the variables that we need to keep tract of between polls.

	// Message we are trying to deliver. Nack anything leftover during shutdown.
	var toSend []comms.MessageInfo
	defer func() {
		for _, m := range toSend {
			m.Nack()
		}
	}()

	var toSendSize int // approximate size of toSend in bytes
	var lastPoll, oldestUnsent, lastActive time.Time

	// poll performs a poll (actually implemented by c.poll), records any errors
	// and updates the variables defined above. In case of failure it also sleeps
	// for the MinFailureDelay.
	poll := func() {
		if c.cctx.CurrentID() != c.id {
			c.configure()
		}
		active, err := c.poll(toSend)
		if err != nil {
			log.Warningf("Failure during polling: %v", err)
			for _, m := range toSend {
				m.Nack()
			}
			toSend = nil
			toSendSize = 0

			if (lastPoll != time.Time{}) && (time.Since(lastPoll) > time.Duration(c.conf.FailureSuicideTimeSeconds)*time.Second) {
				// Die in the hopes that our replacement will be better configured, or otherwise have better luck.
				log.Fatalf("Too Lonely! Failed to contact server in %v.", time.Since(lastPoll))
			}

			t := time.NewTimer(jitter(c.conf.MinFailureDelaySeconds))
			select {
			case <-t.C:
			case <-c.ctx.Done():
				t.Stop()
			}
			return
		}
		for _, m := range toSend {
			m.Ack()
		}
		toSend = nil
		toSendSize = 0
		oldestUnsent = time.Time{}
		lastPoll = time.Now()
		if active {
			lastActive = time.Now()
		}
	}

	for {
		// Stop if we are shutting down, e.g. during the wait after a previous poll
		// failure.
		if c.ctx.Err() != nil {
			return
		}

		// Compute the time that we should next send (assuming we don't hit a send
		// threshold). This could be MaxPollDelay after the last successful send.
		deadline := lastPoll.Add(jitter(c.conf.MaxPollDelaySeconds))

		// If we received something recently, we reduce it to 200ms + 1/10 of the
		// time since we last received a message. (Instructions often lead to more
		// instructions.)
		if !lastActive.IsZero() {
			fpd := lastPoll.Add(200*time.Millisecond + time.Since(lastActive)/10)
			if fpd.Before(deadline) {
				deadline = fpd
			}
		}

		// If we already have something, we should wait at most MaxBufferDelay from
		// receipt of it before sending.
		if !oldestUnsent.IsZero() {
			bd := oldestUnsent.Add(jitter(c.conf.MaxBufferDelaySeconds))
			if bd.Before(deadline) {
				deadline = bd
			}
		}

		now := time.Now()
		if now.After(deadline) || toSendSize > sendBytesThreshold || len(toSend) >= sendCountThreshold {
			poll()
			if c.ctx.Err() != nil {
				return
			}
			continue
		}

		delay := deadline.Sub(now)
		t := time.NewTimer(delay)

		if delay > closeWaitThreshold {
			// Our planned sleep is longer than the idle timeout, so go ahead and kill
			// any idle connection now.
			c.hc.Transport.(*http.Transport).CloseIdleConnections()
		}

		select {
		case <-c.ctx.Done():
			t.Stop()
			if toSendSize > 0 {
				poll()
			}
			return
		case _ = <-t.C:
			poll()
		case m := <-c.cctx.Outbox():
			t.Stop()
			toSend = append(toSend, m)
			toSendSize += 2 + proto.Size(m.M)
			if toSendSize >= sendBytesThreshold ||
				len(toSend) >= sendCountThreshold {
				poll()
			} else {
				if oldestUnsent.IsZero() {
					oldestUnsent = time.Now()
				}
			}
		}
	}
}

func (c *Communicator) poll(toSend []comms.MessageInfo) (bool, error) {
	var sent bool // records whether an interesting (non-LOW) priority message was sent.
	msgs := make([]*fspb.Message, 0, len(toSend))
	for _, m := range toSend {
		msgs = append(msgs, m.M)
		if m.M.Priority != fspb.Message_LOW {
			sent = true
		}
	}
	d, err := c.cctx.MakeContactData(msgs)
	if err != nil {
		return false, err
	}
	data, err := proto.Marshal(d)

	if err != nil {
		return false, fmt.Errorf("unable to marshal outgoing messages: %v", err)
	}

	for i, h := range c.hosts {
		u := url.URL{Scheme: "https", Host: h, Path: "/message"}

		resp, err := c.hc.Post(u.String(), "", bytes.NewReader(data))
		if err != nil {
			log.Warningf("POST to %v failed with error: %v", u, err)
			continue
		}

		if resp.StatusCode != 200 {
			log.Warningf("POST to %v failed with status: %v", u, resp.StatusCode)
			continue
		}

		var b bytes.Buffer
		if _, err := b.ReadFrom(resp.Body); err != nil {
			resp.Body.Close()
			log.Warning("Unable to read response body.")
			continue
		}
		resp.Body.Close()

		var r fspb.ContactData
		if err := proto.Unmarshal(b.Bytes(), &r); err != nil {
			log.Warningf("Unable to parse ContactData from server: %v", err)
			continue
		}

		if err := c.cctx.ProcessContactData(&r); err != nil {
			log.Warningf("Error processing ContactData from server: %v", err)
			continue
		}

		if i != 0 {
			c.hostLock.Lock()
			// Swap, so we check this host first next time.
			c.hosts[0], c.hosts[i] = c.hosts[i], c.hosts[0]
			c.hostLock.Unlock()
		}
		return sent || (len(r.Messages) != 0), nil
	}
	return false, errors.New("unable to contact any server")
}

func (c *Communicator) GetFileIfModified(ctx context.Context, service, name string, modSince time.Time) (io.ReadCloser, time.Time, error) {
	var lastErr error
	c.hostLock.RLock()
	defer c.hostLock.RUnlock()
	for _, h := range c.hosts {
		u := url.URL{Scheme: "https", Host: h,
			Path: "/files/" + url.PathEscape(service) + "/" + url.PathEscape(name)}

		req, err := http.NewRequest("GET", u.String(), nil)
		if err != nil {
			lastErr = err
			continue
		}
		req = req.WithContext(ctx)

		if (modSince != time.Time{}) {
			req.Header.Set("If-Modified-Since", modSince.Format(http.TimeFormat))
		}

		resp, err := c.hc.Do(req)
		if err != nil {
			lastErr = err
			if ctx.Err() != nil {
				break
			}
			continue
		}

		switch resp.StatusCode {
		case http.StatusOK:
			return resp.Body, time.Time{}, nil
		case http.StatusNotModified:
			resp.Body.Close()
			return nil, time.Time{}, nil
		default:
			resp.Body.Close()
			lastErr = fmt.Errorf("failed with http response code: %v", resp.StatusCode)
			continue
		}
	}

	return nil, time.Time{}, fmt.Errorf("unable to retrieve file, last attempt failed with: %v", lastErr)
}
<commit_msg>Add verbose logging messages useful for debugging fastpoll. (#31)<commit_after>// Copyright 2017 Google 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
//
//     https://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 https provides an client.Communicator which connects to the
// Fleetspeak server using https.
package https

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"errors"
	"fmt"
	"io"
	"math/big"
	mrand "math/rand"
	"net"
	"net/http"
	"net/url"
	"sync"
	"time"

	log "github.com/golang/glog"
	"github.com/golang/protobuf/proto"
	"github.com/google/fleetspeak/fleetspeak/src/client/comms"
	"github.com/google/fleetspeak/fleetspeak/src/common"

	clpb "github.com/google/fleetspeak/fleetspeak/src/client/proto/fleetspeak_client"
	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
)

const (
	sendBytesThreshold = 15 * 1024 * 1024
	sendCountThreshold = 100
	closeWaitThreshold = 30 * time.Second // Matches IdleTimeout in server/https.
)

// Communicator implements comms.Communicator and communicates with a Fleetspeak
// server using https.
type Communicator struct {
	cctx        comms.Context
	conf        *clpb.CommunicatorConfig
	id          common.ClientID
	hc          *http.Client
	ctx         context.Context
	done        context.CancelFunc
	working     sync.WaitGroup
	hosts       []string
	hostLock    sync.RWMutex                                                      // Protects hosts
	DialContext func(ctx context.Context, network, addr string) (net.Conn, error) // If set, will be used to initiate network connections to the server.
}

func (c *Communicator) Setup(cl comms.Context) error {
	c.cctx = cl
	return c.configure()
}

func (c *Communicator) configure() error {
	id, err := c.cctx.CurrentIdentity()
	if err != nil {
		return err
	}

	c.id = id.ID

	tmpl := x509.Certificate{
		Issuer:       pkix.Name{Organization: []string{"GRR Client"}},
		Subject:      pkix.Name{Organization: []string{id.ID.String()}},
		SerialNumber: big.NewInt(1),
	}
	certBytes, err := x509.CreateCertificate(rand.Reader, &tmpl, &tmpl, id.Public, id.Private)
	if err != nil {
		return fmt.Errorf("unable to configure communicator, could not create client cert: %v", err)
	}
	certPair := tls.Certificate{
		Certificate: [][]byte{certBytes},
		PrivateKey:  id.Private,
	}

	si, err := c.cctx.ServerInfo()
	if err != nil {
		return fmt.Errorf("unable to configure communicator, could not get server information: %v", err)
	}
	c.hosts = append([]string(nil), si.Servers...)
	if len(c.hosts) == 0 {
		return errors.New("no server_addresses in client configuration")
	}

	c.conf = c.cctx.CommunicatorConfig()
	if c.conf == nil {
		return errors.New("no communicator_config in client configuration")
	}
	if c.conf.MaxPollDelaySeconds == 0 {
		c.conf.MaxPollDelaySeconds = 60 * 5
	}
	if c.conf.MaxBufferDelaySeconds == 0 {
		c.conf.MaxBufferDelaySeconds = 5
	}
	if c.conf.MinFailureDelaySeconds == 0 {
		c.conf.MinFailureDelaySeconds = 60 * 5
	}
	if c.conf.FailureSuicideTimeSeconds == 0 {
		c.conf.FailureSuicideTimeSeconds = 60 * 60 * 24 * 7
	}

	if c.DialContext == nil {
		c.DialContext = (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
			DualStack: true,
		}).DialContext
	}

	c.hc = &http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			TLSClientConfig: &tls.Config{
				RootCAs:      si.TrustedCerts,
				Certificates: []tls.Certificate{certPair},
				CipherSuites: []uint16{
					// We implement both endpoints, so we might as well require long keys and
					// perfect forward secrecy. Note that TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256
					// is required by the https library.
					tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
					tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
					tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
					tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256},
				VerifyPeerCertificate: c.peerCertVerifier,
			},
			DialContext:           c.DialContext,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		},
		Timeout: 5 * time.Minute,
	}
	c.ctx, c.done = context.WithCancel(context.Background())
	return nil
}

func (c *Communicator) Start() error {
	c.working.Add(1)
	go c.processingLoop()
	return nil
}

func (c *Communicator) Stop() {
	c.done()
	c.working.Wait()
}

// jitter adds up to 50% random jitter, and converts to time.Duration.
func jitter(seconds int32) time.Duration {
	return time.Duration((1.0 + 0.5*mrand.Float32()) * float32(seconds) * float32(time.Second))
}

func (c *Communicator) peerCertVerifier(_ [][]byte, chains [][]*x509.Certificate) error {
	for _, chain := range chains {
		if !c.cctx.ChainRevoked(chain) {
			return nil
		}
	}
	return errors.New("certificate revoked")
}

// processingLoop polls the server according to the configured policies while
// the communicator is active.
//
// It is run on an internal goroutine by c.Start, stops when c.ctx is canceled
// and indicates completion by decrementing c.working.
func (c *Communicator) processingLoop() {
	defer c.working.Done()

	// These are the variables that we need to keep tract of between polls.

	// Message we are trying to deliver. Nack anything leftover during shutdown.
	var toSend []comms.MessageInfo
	defer func() {
		for _, m := range toSend {
			m.Nack()
		}
	}()

	var toSendSize int // approximate size of toSend in bytes
	var lastPoll, oldestUnsent, lastActive time.Time

	// poll performs a poll (actually implemented by c.poll), records any errors
	// and updates the variables defined above. In case of failure it also sleeps
	// for the MinFailureDelay.
	poll := func() {
		if c.cctx.CurrentID() != c.id {
			c.configure()
		}
		active, err := c.poll(toSend)
		if err != nil {
			log.Warningf("Failure during polling: %v", err)
			for _, m := range toSend {
				m.Nack()
			}
			toSend = nil
			toSendSize = 0

			if (lastPoll != time.Time{}) && (time.Since(lastPoll) > time.Duration(c.conf.FailureSuicideTimeSeconds)*time.Second) {
				// Die in the hopes that our replacement will be better configured, or otherwise have better luck.
				log.Fatalf("Too Lonely! Failed to contact server in %v.", time.Since(lastPoll))
			}

			t := time.NewTimer(jitter(c.conf.MinFailureDelaySeconds))
			select {
			case <-t.C:
			case <-c.ctx.Done():
				t.Stop()
			}
			return
		}
		for _, m := range toSend {
			m.Ack()
		}
		toSend = nil
		toSendSize = 0
		oldestUnsent = time.Time{}
		lastPoll = time.Now()
		if active {
			lastActive = time.Now()
		}
	}

	for {
		// Stop if we are shutting down, e.g. during the wait after a previous poll
		// failure.
		if c.ctx.Err() != nil {
			return
		}

		// Compute the time that we should next send (assuming we don't hit a send
		// threshold). This could be MaxPollDelay after the last successful send.
		deadline := lastPoll.Add(jitter(c.conf.MaxPollDelaySeconds))
		if log.V(2) {
			log.Infof("Base wait of %v", deadline.Sub(time.Now()))
		}

		// If we received something recently, we reduce it to 200ms + 1/10 of the
		// time since we last received a message. (Instructions often lead to more
		// instructions.)
		if !lastActive.IsZero() {
			fpd := lastPoll.Add(200*time.Millisecond + time.Since(lastActive)/10)
			if fpd.Before(deadline) {
				deadline = fpd
				if log.V(2) {
					log.Infof("Last active %v ago, reduced wait to %v.", time.Since(lastActive), deadline.Sub(time.Now()))
				}
			}
		}

		// If we already have something, we should wait at most MaxBufferDelay from
		// receipt of it before sending.
		if !oldestUnsent.IsZero() {
			bd := oldestUnsent.Add(jitter(c.conf.MaxBufferDelaySeconds))
			if bd.Before(deadline) {
				deadline = bd
			}
		}

		now := time.Now()
		if now.After(deadline) || toSendSize > sendBytesThreshold || len(toSend) >= sendCountThreshold {
			log.V(1).Info("Polling without delay.")
			poll()
			if c.ctx.Err() != nil {
				return
			}
			continue
		}

		delay := deadline.Sub(now)
		t := time.NewTimer(delay)

		if delay > closeWaitThreshold {
			// Our planned sleep is longer than the idle timeout, so go ahead and kill
			// any idle connection now.
			c.hc.Transport.(*http.Transport).CloseIdleConnections()
		}
		log.V(1).Infof("Waiting %v for next poll.", delay)

		select {
		case <-c.ctx.Done():
			t.Stop()
			if toSendSize > 0 {
				poll()
			}
			return
		case _ = <-t.C:
			poll()
		case m := <-c.cctx.Outbox():
			t.Stop()
			toSend = append(toSend, m)
			toSendSize += 2 + proto.Size(m.M)
			if toSendSize >= sendBytesThreshold ||
				len(toSend) >= sendCountThreshold {
				poll()
			} else {
				if oldestUnsent.IsZero() {
					oldestUnsent = time.Now()
				}
			}
		}
	}
}

func (c *Communicator) poll(toSend []comms.MessageInfo) (bool, error) {
	var sent bool // records whether an interesting (non-LOW) priority message was sent.
	msgs := make([]*fspb.Message, 0, len(toSend))
	for _, m := range toSend {
		msgs = append(msgs, m.M)
		if m.M.Priority != fspb.Message_LOW {
			if !sent && bool(log.V(2)) {
				log.Infof("Activity: %s - %s", m.M.Destination.ServiceName, m.M.MessageType)
			}
			sent = true
		}
	}
	d, err := c.cctx.MakeContactData(msgs)
	if err != nil {
		return false, err
	}
	data, err := proto.Marshal(d)

	if err != nil {
		return false, fmt.Errorf("unable to marshal outgoing messages: %v", err)
	}

	for i, h := range c.hosts {
		u := url.URL{Scheme: "https", Host: h, Path: "/message"}

		resp, err := c.hc.Post(u.String(), "", bytes.NewReader(data))
		if err != nil {
			log.Warningf("POST to %v failed with error: %v", u, err)
			continue
		}

		if resp.StatusCode != 200 {
			log.Warningf("POST to %v failed with status: %v", u, resp.StatusCode)
			continue
		}

		var b bytes.Buffer
		if _, err := b.ReadFrom(resp.Body); err != nil {
			resp.Body.Close()
			log.Warning("Unable to read response body.")
			continue
		}
		resp.Body.Close()

		var r fspb.ContactData
		if err := proto.Unmarshal(b.Bytes(), &r); err != nil {
			log.Warningf("Unable to parse ContactData from server: %v", err)
			continue
		}

		if err := c.cctx.ProcessContactData(&r); err != nil {
			log.Warningf("Error processing ContactData from server: %v", err)
			continue
		}

		if i != 0 {
			c.hostLock.Lock()
			// Swap, so we check this host first next time.
			c.hosts[0], c.hosts[i] = c.hosts[i], c.hosts[0]
			c.hostLock.Unlock()
		}
		return sent || (len(r.Messages) != 0), nil
	}
	return false, errors.New("unable to contact any server")
}

func (c *Communicator) GetFileIfModified(ctx context.Context, service, name string, modSince time.Time) (io.ReadCloser, time.Time, error) {
	var lastErr error
	c.hostLock.RLock()
	defer c.hostLock.RUnlock()
	for _, h := range c.hosts {
		u := url.URL{Scheme: "https", Host: h,
			Path: "/files/" + url.PathEscape(service) + "/" + url.PathEscape(name)}

		req, err := http.NewRequest("GET", u.String(), nil)
		if err != nil {
			lastErr = err
			continue
		}
		req = req.WithContext(ctx)

		if (modSince != time.Time{}) {
			req.Header.Set("If-Modified-Since", modSince.Format(http.TimeFormat))
		}

		resp, err := c.hc.Do(req)
		if err != nil {
			lastErr = err
			if ctx.Err() != nil {
				break
			}
			continue
		}

		switch resp.StatusCode {
		case http.StatusOK:
			return resp.Body, time.Time{}, nil
		case http.StatusNotModified:
			resp.Body.Close()
			return nil, time.Time{}, nil
		default:
			resp.Body.Close()
			lastErr = fmt.Errorf("failed with http response code: %v", resp.StatusCode)
			continue
		}
	}

	return nil, time.Time{}, fmt.Errorf("unable to retrieve file, last attempt failed with: %v", lastErr)
}
<|endoftext|>


<commit_before>// Copyright 2017 Google 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
//
//     https://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 admin defines an administrative interface into the fleetspeak system.
package admin

import (
	"bytes"
	"errors"
	"fmt"

	"context"

	"github.com/google/fleetspeak/fleetspeak/src/common"
	"github.com/google/fleetspeak/fleetspeak/src/server/db"

	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
	sgrpc "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
	spb "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
)

// NewServer returns an admin_grpc.AdminServer which performs operations using
// the provided db.Store.
func NewServer(s db.Store) sgrpc.AdminServer {
	return adminServer{s}
}

type adminServer struct {
	store db.Store
}

func (s adminServer) CreateBroadcast(ctx context.Context, req *spb.CreateBroadcastRequest) (*fspb.EmptyMessage, error) {
	if err := s.store.CreateBroadcast(ctx, req.Broadcast, req.Limit); err != nil {
		return nil, err
	}
	return &fspb.EmptyMessage{}, nil
}

func (s adminServer) ListActiveBroadcasts(ctx context.Context, req *spb.ListActiveBroadcastsRequest) (*spb.ListActiveBroadcastsResponse, error) {
	var ret spb.ListActiveBroadcastsResponse
	bis, err := s.store.ListActiveBroadcasts(ctx)
	if err != nil {
		return nil, err
	}
	for _, bi := range bis {
		if req.ServiceName != "" && req.ServiceName != bi.Broadcast.Source.ServiceName {
			continue
		}
		ret.Broadcasts = append(ret.Broadcasts, bi.Broadcast)
	}
	return &ret, nil
}

func (s adminServer) GetMessageStatus(ctx context.Context, req *spb.GetMessageStatusRequest) (*spb.GetMessageStatusResponse, error) {
	mid, err := common.BytesToMessageID(req.MessageId)
	if err != nil {
		return nil, err
	}

	msgs, err := s.store.GetMessages(ctx, []common.MessageID{mid}, false)
	if err != nil {
		if s.store.IsNotFound(err) {
			return &spb.GetMessageStatusResponse{}, nil
		}
		return nil, err
	}
	if len(msgs) != 1 {
		return nil, fmt.Errorf("Internal error, expected 1 message, got %d", len(msgs))
	}

	return &spb.GetMessageStatusResponse{
			CreationTime: msgs[0].CreationTime,
			Result:       msgs[0].Result},
		nil
}

func (s adminServer) ListClients(ctx context.Context, req *spb.ListClientsRequest) (*spb.ListClientsResponse, error) {
	ids := make([]common.ClientID, 0, len(req.ClientIds))
	for i, b := range req.ClientIds {
		id, err := common.BytesToClientID(b)
		if err != nil {
			return nil, fmt.Errorf("unable to parse id [%d]: %v", i, err)
		}
		ids = append(ids, id)
	}
	clients, err := s.store.ListClients(ctx, ids)
	if err != nil {
		return nil, err
	}

	return &spb.ListClientsResponse{
		Clients: clients,
	}, nil
}

func (s adminServer) ListClientContacts(ctx context.Context, req *spb.ListClientContactsRequest) (*spb.ListClientContactsResponse, error) {
	id, err := common.BytesToClientID(req.ClientId)
	if err != nil {
		return nil, fmt.Errorf("unable to parse id [%d]: %v", req.ClientId, err)
	}

	contacts, err := s.store.ListClientContacts(ctx, id)
	if err != nil {
		return nil, err
	}
	return &spb.ListClientContactsResponse{
		Contacts: contacts,
	}, nil
}

// InsertMessage implements sgrpc.AdminServer.
func (s adminServer) InsertMessage(ctx context.Context, m *fspb.Message) (*fspb.EmptyMessage, error) {
	// At this point, we mostly trust the message we get, but do some basic
	// sanity checks and generate missing metadata.
	if m.Destination == nil || m.Destination.ServiceName == "" {
		return nil, errors.New("message must have Destination")
	}
	if m.Source == nil || m.Source.ServiceName == "" {
		return nil, errors.New("message must have Source")
	}
	if len(m.MessageId) == 0 {
		id, err := common.RandomMessageID()
		if err != nil {
			return nil, fmt.Errorf("unable to create random MessageID: %v", err)
		}
		m.MessageId = id.Bytes()
	}
	if m.CreationTime == nil {
		m.CreationTime = db.NowProto()
	}
	if err := s.store.StoreMessages(ctx, []*fspb.Message{m}, ""); err != nil {
		return nil, err
	}
	return &fspb.EmptyMessage{}, nil
}

// StoreFile implements sgrpc.AdminServer.
func (s adminServer) StoreFile(ctx context.Context, req *spb.StoreFileRequest) (*fspb.EmptyMessage, error) {
	if req.ServiceName == "" || req.FileName == "" {
		return nil, errors.New("file must have service_name and file_name")
	}
	if err := s.store.StoreFile(ctx, req.ServiceName, req.FileName, bytes.NewReader(req.Data)); err != nil {
		return nil, err
	}
	return &fspb.EmptyMessage{}, nil
}

// KeepAlive implements sgrpc.AdminServer.
func (s adminServer) KeepAlive(ctx context.Context, _ *fspb.EmptyMessage) (*fspb.EmptyMessage, error) {
	return &fspb.EmptyMessage{}, nil
}
<commit_msg>Add BlacklistClient to admin interface implementation.<commit_after>// Copyright 2017 Google 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
//
//     https://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 admin defines an administrative interface into the fleetspeak system.
package admin

import (
	"bytes"
	"errors"
	"fmt"

	"context"

	"github.com/google/fleetspeak/fleetspeak/src/common"
	"github.com/google/fleetspeak/fleetspeak/src/server/db"

	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
	sgrpc "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
	spb "github.com/google/fleetspeak/fleetspeak/src/server/proto/fleetspeak_server"
)

// NewServer returns an admin_grpc.AdminServer which performs operations using
// the provided db.Store.
func NewServer(s db.Store) sgrpc.AdminServer {
	return adminServer{s}
}

type adminServer struct {
	store db.Store
}

func (s adminServer) CreateBroadcast(ctx context.Context, req *spb.CreateBroadcastRequest) (*fspb.EmptyMessage, error) {
	if err := s.store.CreateBroadcast(ctx, req.Broadcast, req.Limit); err != nil {
		return nil, err
	}
	return &fspb.EmptyMessage{}, nil
}

func (s adminServer) ListActiveBroadcasts(ctx context.Context, req *spb.ListActiveBroadcastsRequest) (*spb.ListActiveBroadcastsResponse, error) {
	var ret spb.ListActiveBroadcastsResponse
	bis, err := s.store.ListActiveBroadcasts(ctx)
	if err != nil {
		return nil, err
	}
	for _, bi := range bis {
		if req.ServiceName != "" && req.ServiceName != bi.Broadcast.Source.ServiceName {
			continue
		}
		ret.Broadcasts = append(ret.Broadcasts, bi.Broadcast)
	}
	return &ret, nil
}

func (s adminServer) GetMessageStatus(ctx context.Context, req *spb.GetMessageStatusRequest) (*spb.GetMessageStatusResponse, error) {
	mid, err := common.BytesToMessageID(req.MessageId)
	if err != nil {
		return nil, err
	}

	msgs, err := s.store.GetMessages(ctx, []common.MessageID{mid}, false)
	if err != nil {
		if s.store.IsNotFound(err) {
			return &spb.GetMessageStatusResponse{}, nil
		}
		return nil, err
	}
	if len(msgs) != 1 {
		return nil, fmt.Errorf("Internal error, expected 1 message, got %d", len(msgs))
	}

	return &spb.GetMessageStatusResponse{
			CreationTime: msgs[0].CreationTime,
			Result:       msgs[0].Result},
		nil
}

func (s adminServer) ListClients(ctx context.Context, req *spb.ListClientsRequest) (*spb.ListClientsResponse, error) {
	ids := make([]common.ClientID, 0, len(req.ClientIds))
	for i, b := range req.ClientIds {
		id, err := common.BytesToClientID(b)
		if err != nil {
			return nil, fmt.Errorf("unable to parse id [%d]: %v", i, err)
		}
		ids = append(ids, id)
	}
	clients, err := s.store.ListClients(ctx, ids)
	if err != nil {
		return nil, err
	}

	return &spb.ListClientsResponse{
		Clients: clients,
	}, nil
}

func (s adminServer) ListClientContacts(ctx context.Context, req *spb.ListClientContactsRequest) (*spb.ListClientContactsResponse, error) {
	id, err := common.BytesToClientID(req.ClientId)
	if err != nil {
		return nil, fmt.Errorf("unable to parse id [%d]: %v", req.ClientId, err)
	}

	contacts, err := s.store.ListClientContacts(ctx, id)
	if err != nil {
		return nil, err
	}
	return &spb.ListClientContactsResponse{
		Contacts: contacts,
	}, nil
}

// InsertMessage implements sgrpc.AdminServer.
func (s adminServer) InsertMessage(ctx context.Context, m *fspb.Message) (*fspb.EmptyMessage, error) {
	// At this point, we mostly trust the message we get, but do some basic
	// sanity checks and generate missing metadata.
	if m.Destination == nil || m.Destination.ServiceName == "" {
		return nil, errors.New("message must have Destination")
	}
	if m.Source == nil || m.Source.ServiceName == "" {
		return nil, errors.New("message must have Source")
	}
	if len(m.MessageId) == 0 {
		id, err := common.RandomMessageID()
		if err != nil {
			return nil, fmt.Errorf("unable to create random MessageID: %v", err)
		}
		m.MessageId = id.Bytes()
	}
	if m.CreationTime == nil {
		m.CreationTime = db.NowProto()
	}
	if err := s.store.StoreMessages(ctx, []*fspb.Message{m}, ""); err != nil {
		return nil, err
	}
	return &fspb.EmptyMessage{}, nil
}

// StoreFile implements sgrpc.AdminServer.
func (s adminServer) StoreFile(ctx context.Context, req *spb.StoreFileRequest) (*fspb.EmptyMessage, error) {
	if req.ServiceName == "" || req.FileName == "" {
		return nil, errors.New("file must have service_name and file_name")
	}
	if err := s.store.StoreFile(ctx, req.ServiceName, req.FileName, bytes.NewReader(req.Data)); err != nil {
		return nil, err
	}
	return &fspb.EmptyMessage{}, nil
}

// KeepAlive implements sgrpc.AdminServer.
func (s adminServer) KeepAlive(ctx context.Context, _ *fspb.EmptyMessage) (*fspb.EmptyMessage, error) {
	return &fspb.EmptyMessage{}, nil
}

// BlacklistClient implements sgrpc.BlacklistClient.
func (s adminServer) BlacklistClient(ctx context.Context, req *spb.BlacklistClientRequest) (*fspb.EmptyMessage, error) {
	id, err := common.BytesToClientID(req.ClientId)
	if err != nil {
		return nil, fmt.Errorf("unable to parse id [%d]: %v", req.ClientId, err)
	}
	if err := s.store.BlacklistClient(ctx, id); err != nil {
		return nil, err
	}
	return &fspb.EmptyMessage{}, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, Inc.
// Copyright 2015 Rancher Labs, 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 cloudinitsave

import (
	"errors"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	yaml "github.com/cloudfoundry-incubator/candiedyaml"

	"github.com/coreos/coreos-cloudinit/config"
	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/datasource/configdrive"
	"github.com/coreos/coreos-cloudinit/datasource/file"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/digitalocean"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/ec2"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/packet"
	"github.com/coreos/coreos-cloudinit/datasource/proc_cmdline"
	"github.com/coreos/coreos-cloudinit/datasource/url"
	"github.com/coreos/coreos-cloudinit/pkg"
	"github.com/docker/docker/pkg/mount"
	"github.com/rancher/os/cmd/cloudinitsave/gce"
	"github.com/rancher/os/cmd/network"
	rancherConfig "github.com/rancher/os/config"
	"github.com/rancher/os/log"
	"github.com/rancher/os/netconf"
	"github.com/rancher/os/util"
)

const (
	datasourceInterval    = 100 * time.Millisecond
	datasourceMaxInterval = 30 * time.Second
	datasourceTimeout     = 5 * time.Minute
	configDevName         = "config-2"
	configDev             = "LABEL=" + configDevName
	configDevMountPoint   = "/media/config-2"
)

func Main() {
	log.InitLogger()
	log.Info("Running cloud-init-save")

	cfg := rancherConfig.LoadConfig()
	network.ApplyNetworkConfig(cfg)

	if err := SaveCloudConfig(true); err != nil {
		log.Errorf("Failed to save cloud-config: %v", err)
	}
}

func MountConfigDrive() error {
	if err := os.MkdirAll(configDevMountPoint, 644); err != nil {
		return err
	}

	configDev := util.ResolveDevice(configDev)

	if configDev == "" {
		return mount.Mount(configDevName, configDevMountPoint, "9p", "trans=virtio,version=9p2000.L")
	}

	return mount.Mount(configDev, configDevMountPoint, "iso9660,vfat", "")
}

func UnmountConfigDrive() error {
	return syscall.Unmount(configDevMountPoint, 0)
}

func SaveCloudConfig(network bool) error {
	userDataBytes, metadata, err := fetchUserData(network)
	if err != nil {
		return err
	}

	userData := string(userDataBytes)
	scriptBytes := []byte{}

	if config.IsScript(userData) {
		scriptBytes = userDataBytes
		userDataBytes = []byte{}
	} else if isCompose(userData) {
		if userDataBytes, err = composeToCloudConfig(userDataBytes); err != nil {
			log.Errorf("Failed to convert compose to cloud-config syntax: %v", err)
			return err
		}
	} else if config.IsCloudConfig(userData) {
		if _, err := rancherConfig.ReadConfig(userDataBytes, false); err != nil {
			log.WithFields(log.Fields{"cloud-config": userData, "err": err}).Warn("Failed to parse cloud-config, not saving.")
			userDataBytes = []byte{}
		}
	} else {
		log.Errorf("Unrecognized user-data\n%s", userData)
		userDataBytes = []byte{}
	}

	if _, err := rancherConfig.ReadConfig(userDataBytes, false); err != nil {
		log.WithFields(log.Fields{"cloud-config": userData, "err": err}).Warn("Failed to parse cloud-config")
		return errors.New("Failed to parse cloud-config")
	}

	return saveFiles(userDataBytes, scriptBytes, metadata)
}

func RequiresNetwork(datasource string) bool {
	parts := strings.SplitN(datasource, ":", 2)
	requiresNetwork, ok := map[string]bool{
		"ec2":          true,
		"file":         false,
		"url":          true,
		"cmdline":      true,
		"configdrive":  false,
		"digitalocean": true,
		"gce":          true,
		"packet":       true,
	}[parts[0]]
	return ok && requiresNetwork
}

func saveFiles(cloudConfigBytes, scriptBytes []byte, metadata datasource.Metadata) error {
	os.MkdirAll(rancherConfig.CloudConfigDir, os.ModeDir|0600)

	if len(scriptBytes) > 0 {
		log.Infof("Writing to %s", rancherConfig.CloudConfigScriptFile)
		if err := util.WriteFileAtomic(rancherConfig.CloudConfigScriptFile, scriptBytes, 500); err != nil {
			log.Errorf("Error while writing file %s: %v", rancherConfig.CloudConfigScriptFile, err)
			return err
		}
	}

	if len(cloudConfigBytes) > 0 {
		if err := util.WriteFileAtomic(rancherConfig.CloudConfigBootFile, cloudConfigBytes, 400); err != nil {
			return err
		}
		log.Infof("Written to %s:\n%s", rancherConfig.CloudConfigBootFile, string(cloudConfigBytes))
	}

	metaDataBytes, err := yaml.Marshal(metadata)
	if err != nil {
		return err
	}

	if err = util.WriteFileAtomic(rancherConfig.MetaDataFile, metaDataBytes, 400); err != nil {
		return err
	}
	log.Infof("Written to %s:\n%s", rancherConfig.MetaDataFile, string(metaDataBytes))

	return nil
}

func currentDatasource(network bool) (datasource.Datasource, error) {
	cfg := rancherConfig.LoadConfig()

	dss := getDatasources(cfg, network)
	if len(dss) == 0 {
		return nil, nil
	}

	ds := selectDatasource(dss)
	return ds, nil
}

func fetchUserData(network bool) ([]byte, datasource.Metadata, error) {
	var metadata datasource.Metadata
	ds, err := currentDatasource(network)
	if err != nil || ds == nil {
		log.Errorf("Failed to select datasource: %v", err)
		return nil, metadata, err
	}
	log.Infof("Fetching user-data from datasource %v", ds.Type())
	userDataBytes, err := ds.FetchUserdata()
	if err != nil {
		log.Errorf("Failed fetching user-data from datasource: %v", err)
		return nil, metadata, err
	}
	log.Infof("Fetching meta-data from datasource of type %v", ds.Type())
	metadata, err = ds.FetchMetadata()
	if err != nil {
		log.Errorf("Failed fetching meta-data from datasource: %v", err)
		return nil, metadata, err
	}
	return userDataBytes, metadata, nil
}

// getDatasources creates a slice of possible Datasources for cloudinit based
// on the different source command-line flags.
func getDatasources(cfg *rancherConfig.CloudConfig, network bool) []datasource.Datasource {
	dss := make([]datasource.Datasource, 0, 5)

	for _, ds := range cfg.Rancher.CloudInit.Datasources {
		parts := strings.SplitN(ds, ":", 2)

		switch parts[0] {
		case "ec2":
			if network {
				if len(parts) == 1 {
					dss = append(dss, ec2.NewDatasource(ec2.DefaultAddress))
				} else {
					dss = append(dss, ec2.NewDatasource(parts[1]))
				}
			}
		case "file":
			if len(parts) == 2 {
				dss = append(dss, file.NewDatasource(parts[1]))
			}
		case "url":
			if network {
				if len(parts) == 2 {
					dss = append(dss, url.NewDatasource(parts[1]))
				}
			}
		case "cmdline":
			if network {
				if len(parts) == 1 {
					dss = append(dss, proc_cmdline.NewDatasource())
				}
			}
		case "configdrive":
			if len(parts) == 2 {
				dss = append(dss, configdrive.NewDatasource(parts[1]))
			}
		case "digitalocean":
			if network {
				if len(parts) == 1 {
					dss = append(dss, digitalocean.NewDatasource(digitalocean.DefaultAddress))
				} else {
					dss = append(dss, digitalocean.NewDatasource(parts[1]))
				}
			} else {
				enableDoLinkLocal()
			}
		case "gce":
			if network {
				dss = append(dss, gce.NewDatasource("http://metadata.google.internal/"))
			}
		case "packet":
			if !network {
				enablePacketNetwork(&cfg.Rancher)
			}
			dss = append(dss, packet.NewDatasource("https://metadata.packet.net/"))
		}
	}

	return dss
}

func enableDoLinkLocal() {
	err := netconf.ApplyNetworkConfigs(&rancherConfig.NetworkConfig{
		Interfaces: map[string]rancherConfig.InterfaceConfig{
			"eth0": {
				IPV4LL: true,
			},
		},
	})
	if err != nil {
		log.Errorf("Failed to apply link local on eth0: %v", err)
	}
}

// selectDatasource attempts to choose a valid Datasource to use based on its
// current availability. The first Datasource to report to be available is
// returned. Datasources will be retried if possible if they are not
// immediately available. If all Datasources are permanently unavailable or
// datasourceTimeout is reached before one becomes available, nil is returned.
func selectDatasource(sources []datasource.Datasource) datasource.Datasource {
	ds := make(chan datasource.Datasource)
	stop := make(chan struct{})
	var wg sync.WaitGroup

	for _, s := range sources {
		wg.Add(1)
		go func(s datasource.Datasource) {
			defer wg.Done()

			duration := datasourceInterval
			for {
				log.Infof("Checking availability of %q\n", s.Type())
				if s.IsAvailable() {
					ds <- s
					return
				} else if !s.AvailabilityChanges() {
					return
				}
				select {
				case <-stop:
					return
				case <-time.After(duration):
					duration = pkg.ExpBackoff(duration, datasourceMaxInterval)
				}
			}
		}(s)
	}

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	var s datasource.Datasource
	select {
	case s = <-ds:
	case <-done:
	case <-time.After(datasourceTimeout):
	}

	close(stop)
	return s
}

func isCompose(content string) bool {
	return strings.HasPrefix(content, "#compose\n")
}

func composeToCloudConfig(bytes []byte) ([]byte, error) {
	compose := make(map[interface{}]interface{})
	err := yaml.Unmarshal(bytes, &compose)
	if err != nil {
		return nil, err
	}

	return yaml.Marshal(map[interface{}]interface{}{
		"rancher": map[interface{}]interface{}{
			"services": compose,
		},
	})
}
<commit_msg>Run udev before early cloud-init<commit_after>// Copyright 2015 CoreOS, Inc.
// Copyright 2015 Rancher Labs, 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 cloudinitsave

import (
	"errors"
	"os"
	"strings"
	"sync"
	"syscall"
	"time"

	yaml "github.com/cloudfoundry-incubator/candiedyaml"

	"github.com/coreos/coreos-cloudinit/config"
	"github.com/coreos/coreos-cloudinit/datasource"
	"github.com/coreos/coreos-cloudinit/datasource/configdrive"
	"github.com/coreos/coreos-cloudinit/datasource/file"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/digitalocean"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/ec2"
	"github.com/coreos/coreos-cloudinit/datasource/metadata/packet"
	"github.com/coreos/coreos-cloudinit/datasource/proc_cmdline"
	"github.com/coreos/coreos-cloudinit/datasource/url"
	"github.com/coreos/coreos-cloudinit/pkg"
	"github.com/docker/docker/pkg/mount"
	"github.com/rancher/os/cmd/cloudinitsave/gce"
	"github.com/rancher/os/cmd/control"
	"github.com/rancher/os/cmd/network"
	rancherConfig "github.com/rancher/os/config"
	"github.com/rancher/os/log"
	"github.com/rancher/os/netconf"
	"github.com/rancher/os/util"
)

const (
	datasourceInterval    = 100 * time.Millisecond
	datasourceMaxInterval = 30 * time.Second
	datasourceTimeout     = 5 * time.Minute
	configDevName         = "config-2"
	configDev             = "LABEL=" + configDevName
	configDevMountPoint   = "/media/config-2"
)

func Main() {
	log.InitLogger()
	log.Info("Running cloud-init-save")

	if err := control.UdevSettle(); err != nil {
		log.Errorf("Failed to run udev settle: %v", err)
	}

	cfg := rancherConfig.LoadConfig()
	network.ApplyNetworkConfig(cfg)

	if err := SaveCloudConfig(true); err != nil {
		log.Errorf("Failed to save cloud-config: %v", err)
	}
}

func MountConfigDrive() error {
	if err := os.MkdirAll(configDevMountPoint, 644); err != nil {
		return err
	}

	configDev := util.ResolveDevice(configDev)

	if configDev == "" {
		return mount.Mount(configDevName, configDevMountPoint, "9p", "trans=virtio,version=9p2000.L")
	}

	return mount.Mount(configDev, configDevMountPoint, "iso9660,vfat", "")
}

func UnmountConfigDrive() error {
	return syscall.Unmount(configDevMountPoint, 0)
}

func SaveCloudConfig(network bool) error {
	userDataBytes, metadata, err := fetchUserData(network)
	if err != nil {
		return err
	}

	userData := string(userDataBytes)
	scriptBytes := []byte{}

	if config.IsScript(userData) {
		scriptBytes = userDataBytes
		userDataBytes = []byte{}
	} else if isCompose(userData) {
		if userDataBytes, err = composeToCloudConfig(userDataBytes); err != nil {
			log.Errorf("Failed to convert compose to cloud-config syntax: %v", err)
			return err
		}
	} else if config.IsCloudConfig(userData) {
		if _, err := rancherConfig.ReadConfig(userDataBytes, false); err != nil {
			log.WithFields(log.Fields{"cloud-config": userData, "err": err}).Warn("Failed to parse cloud-config, not saving.")
			userDataBytes = []byte{}
		}
	} else {
		log.Errorf("Unrecognized user-data\n%s", userData)
		userDataBytes = []byte{}
	}

	if _, err := rancherConfig.ReadConfig(userDataBytes, false); err != nil {
		log.WithFields(log.Fields{"cloud-config": userData, "err": err}).Warn("Failed to parse cloud-config")
		return errors.New("Failed to parse cloud-config")
	}

	return saveFiles(userDataBytes, scriptBytes, metadata)
}

func RequiresNetwork(datasource string) bool {
	parts := strings.SplitN(datasource, ":", 2)
	requiresNetwork, ok := map[string]bool{
		"ec2":          true,
		"file":         false,
		"url":          true,
		"cmdline":      true,
		"configdrive":  false,
		"digitalocean": true,
		"gce":          true,
		"packet":       true,
	}[parts[0]]
	return ok && requiresNetwork
}

func saveFiles(cloudConfigBytes, scriptBytes []byte, metadata datasource.Metadata) error {
	os.MkdirAll(rancherConfig.CloudConfigDir, os.ModeDir|0600)

	if len(scriptBytes) > 0 {
		log.Infof("Writing to %s", rancherConfig.CloudConfigScriptFile)
		if err := util.WriteFileAtomic(rancherConfig.CloudConfigScriptFile, scriptBytes, 500); err != nil {
			log.Errorf("Error while writing file %s: %v", rancherConfig.CloudConfigScriptFile, err)
			return err
		}
	}

	if len(cloudConfigBytes) > 0 {
		if err := util.WriteFileAtomic(rancherConfig.CloudConfigBootFile, cloudConfigBytes, 400); err != nil {
			return err
		}
		log.Infof("Written to %s:\n%s", rancherConfig.CloudConfigBootFile, string(cloudConfigBytes))
	}

	metaDataBytes, err := yaml.Marshal(metadata)
	if err != nil {
		return err
	}

	if err = util.WriteFileAtomic(rancherConfig.MetaDataFile, metaDataBytes, 400); err != nil {
		return err
	}
	log.Infof("Written to %s:\n%s", rancherConfig.MetaDataFile, string(metaDataBytes))

	return nil
}

func currentDatasource(network bool) (datasource.Datasource, error) {
	cfg := rancherConfig.LoadConfig()

	dss := getDatasources(cfg, network)
	if len(dss) == 0 {
		return nil, nil
	}

	ds := selectDatasource(dss)
	return ds, nil
}

func fetchUserData(network bool) ([]byte, datasource.Metadata, error) {
	var metadata datasource.Metadata
	ds, err := currentDatasource(network)
	if err != nil || ds == nil {
		log.Errorf("Failed to select datasource: %v", err)
		return nil, metadata, err
	}
	log.Infof("Fetching user-data from datasource %v", ds.Type())
	userDataBytes, err := ds.FetchUserdata()
	if err != nil {
		log.Errorf("Failed fetching user-data from datasource: %v", err)
		return nil, metadata, err
	}
	log.Infof("Fetching meta-data from datasource of type %v", ds.Type())
	metadata, err = ds.FetchMetadata()
	if err != nil {
		log.Errorf("Failed fetching meta-data from datasource: %v", err)
		return nil, metadata, err
	}
	return userDataBytes, metadata, nil
}

// getDatasources creates a slice of possible Datasources for cloudinit based
// on the different source command-line flags.
func getDatasources(cfg *rancherConfig.CloudConfig, network bool) []datasource.Datasource {
	dss := make([]datasource.Datasource, 0, 5)

	for _, ds := range cfg.Rancher.CloudInit.Datasources {
		parts := strings.SplitN(ds, ":", 2)

		switch parts[0] {
		case "ec2":
			if network {
				if len(parts) == 1 {
					dss = append(dss, ec2.NewDatasource(ec2.DefaultAddress))
				} else {
					dss = append(dss, ec2.NewDatasource(parts[1]))
				}
			}
		case "file":
			if len(parts) == 2 {
				dss = append(dss, file.NewDatasource(parts[1]))
			}
		case "url":
			if network {
				if len(parts) == 2 {
					dss = append(dss, url.NewDatasource(parts[1]))
				}
			}
		case "cmdline":
			if network {
				if len(parts) == 1 {
					dss = append(dss, proc_cmdline.NewDatasource())
				}
			}
		case "configdrive":
			if len(parts) == 2 {
				dss = append(dss, configdrive.NewDatasource(parts[1]))
			}
		case "digitalocean":
			if network {
				if len(parts) == 1 {
					dss = append(dss, digitalocean.NewDatasource(digitalocean.DefaultAddress))
				} else {
					dss = append(dss, digitalocean.NewDatasource(parts[1]))
				}
			} else {
				enableDoLinkLocal()
			}
		case "gce":
			if network {
				dss = append(dss, gce.NewDatasource("http://metadata.google.internal/"))
			}
		case "packet":
			if !network {
				enablePacketNetwork(&cfg.Rancher)
			}
			dss = append(dss, packet.NewDatasource("https://metadata.packet.net/"))
		}
	}

	return dss
}

func enableDoLinkLocal() {
	err := netconf.ApplyNetworkConfigs(&rancherConfig.NetworkConfig{
		Interfaces: map[string]rancherConfig.InterfaceConfig{
			"eth0": {
				IPV4LL: true,
			},
		},
	})
	if err != nil {
		log.Errorf("Failed to apply link local on eth0: %v", err)
	}
}

// selectDatasource attempts to choose a valid Datasource to use based on its
// current availability. The first Datasource to report to be available is
// returned. Datasources will be retried if possible if they are not
// immediately available. If all Datasources are permanently unavailable or
// datasourceTimeout is reached before one becomes available, nil is returned.
func selectDatasource(sources []datasource.Datasource) datasource.Datasource {
	ds := make(chan datasource.Datasource)
	stop := make(chan struct{})
	var wg sync.WaitGroup

	for _, s := range sources {
		wg.Add(1)
		go func(s datasource.Datasource) {
			defer wg.Done()

			duration := datasourceInterval
			for {
				log.Infof("Checking availability of %q\n", s.Type())
				if s.IsAvailable() {
					ds <- s
					return
				} else if !s.AvailabilityChanges() {
					return
				}
				select {
				case <-stop:
					return
				case <-time.After(duration):
					duration = pkg.ExpBackoff(duration, datasourceMaxInterval)
				}
			}
		}(s)
	}

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()

	var s datasource.Datasource
	select {
	case s = <-ds:
	case <-done:
	case <-time.After(datasourceTimeout):
	}

	close(stop)
	return s
}

func isCompose(content string) bool {
	return strings.HasPrefix(content, "#compose\n")
}

func composeToCloudConfig(bytes []byte) ([]byte, error) {
	compose := make(map[interface{}]interface{})
	err := yaml.Unmarshal(bytes, &compose)
	if err != nil {
		return nil, err
	}

	return yaml.Marshal(map[interface{}]interface{}{
		"rancher": map[interface{}]interface{}{
			"services": compose,
		},
	})
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 config_test

import (
	"github.com/jacobsa/comeback/config"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestParse(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type ParseTest struct {
	data string
	cfg *config.Config
	err error
}

func init() { RegisterTestSuite(&ParseTest{}) }

func (t *ParseTest) parse() {
	t.cfg, t.err = config.Parse([]byte(t.data))
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ParseTest) TotalJunk() {
	t.data = "sdhjklfghdskjghdjkfgj"
	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("invalid")))
}

func (t *ParseTest) Null() {
	t.data = `null`
	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("null")))
}

func (t *ParseTest) Array() {
	t.data = `[17, 19]`
	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("array")))
}

func (t *ParseTest) MissingTrailingBrace() {
	t.data = `
	{
		"jobs": {}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("unexpected end")))
}

func (t *ParseTest) BasePathIsNumber() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": 17
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("number")))
}

func (t *ParseTest) BasePathIsNull() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": null
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("null")))
}

func (t *ParseTest) BasePathIsObject() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": {}
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("object")))
}

func (t *ParseTest) OneExcludeDoesntCompile() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": "/foo",
				"excludes": ["a"],
			},
			"burrito": {
				"base_path": "/bar",
				"excludes": ["b", "(c"]
			},
			"enchilada": {
				"base_path": "/foo",
				"excludes": ["d"]
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("TODO")))
}

func (t *ParseTest) EmptyConfig() {
	t.data = `{}`
	t.parse()

	AssertEq(nil, t.err)
	ExpectNe(nil, t.cfg.Jobs)
	ExpectEq(0, len(t.cfg.Jobs))
}

func (t *ParseTest) MissingExcludesArray() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": "/foo"
			}
		}
	}
	`

	t.parse()

	AssertEq(nil, t.err)
	AssertEq(1, len(t.cfg.Jobs))

	AssertNe(nil, t.cfg.Jobs["taco"])
	ExpectThat(t.cfg.Jobs["taco"].Excludes, ElementsAre())
}

func (t *ParseTest) DuplicateJobName() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": "/foo"
			},
			"burrito": {
				"base_path": "/bar"
			},
			"taco": {
				"base_path": "/enchilada"
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("TODO")))
}

func (t *ParseTest) MultipleValidJobs() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": "/foo"
				"excludes": ["a.b"],
			},
			"burrito": {
				"base_path": "/bar"
				"excludes": ["c", "d"],
			}
		}
	}
	`

	t.parse()

	AssertEq(nil, t.err)
	AssertEq(2, len(t.cfg.Jobs))

	AssertNe(nil, t.cfg.Jobs["taco"])
	ExpectEq("/foo", t.cfg.Jobs["taco"].BasePath)
	AssertThat(t.cfg.Jobs["taco"].Excludes, ElementsAre(Any()))
	ExpectEq("a.b", t.cfg.Jobs["taco"].Excludes[0])

	AssertNe(nil, t.cfg.Jobs["burrito"])
	ExpectEq("/bar", t.cfg.Jobs["burrito"].BasePath)
	AssertThat(t.cfg.Jobs["burrito"].Excludes, ElementsAre(Any(), Any()))
	ExpectEq("c", t.cfg.Jobs["burrito"].Excludes[0])
	ExpectEq("d", t.cfg.Jobs["burrito"].Excludes[1])
}
<commit_msg>Fixed some syntax errors.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 config_test

import (
	"github.com/jacobsa/comeback/config"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestParse(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type ParseTest struct {
	data string
	cfg *config.Config
	err error
}

func init() { RegisterTestSuite(&ParseTest{}) }

func (t *ParseTest) parse() {
	t.cfg, t.err = config.Parse([]byte(t.data))
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ParseTest) TotalJunk() {
	t.data = "sdhjklfghdskjghdjkfgj"
	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("invalid")))
}

func (t *ParseTest) Null() {
	t.data = `null`
	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("null")))
}

func (t *ParseTest) Array() {
	t.data = `[17, 19]`
	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("array")))
}

func (t *ParseTest) MissingTrailingBrace() {
	t.data = `
	{
		"jobs": {}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("unexpected end")))
}

func (t *ParseTest) BasePathIsNumber() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": 17
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("number")))
}

func (t *ParseTest) BasePathIsNull() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": null
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("null")))
}

func (t *ParseTest) BasePathIsObject() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": {}
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("JSON")))
	ExpectThat(t.err, Error(HasSubstr("object")))
}

func (t *ParseTest) OneExcludeDoesntCompile() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": "/foo",
				"excludes": ["a"]
			},
			"burrito": {
				"base_path": "/bar",
				"excludes": ["b", "(c"]
			},
			"enchilada": {
				"base_path": "/foo",
				"excludes": ["d"]
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("TODO")))
}

func (t *ParseTest) EmptyConfig() {
	t.data = `{}`
	t.parse()

	AssertEq(nil, t.err)
	ExpectNe(nil, t.cfg.Jobs)
	ExpectEq(0, len(t.cfg.Jobs))
}

func (t *ParseTest) MissingExcludesArray() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": "/foo"
			}
		}
	}
	`

	t.parse()

	AssertEq(nil, t.err)
	AssertEq(1, len(t.cfg.Jobs))

	AssertNe(nil, t.cfg.Jobs["taco"])
	ExpectThat(t.cfg.Jobs["taco"].Excludes, ElementsAre())
}

func (t *ParseTest) DuplicateJobName() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": "/foo"
			},
			"burrito": {
				"base_path": "/bar"
			},
			"taco": {
				"base_path": "/enchilada"
			}
		}
	}
	`

	t.parse()

	ExpectThat(t.err, Error(HasSubstr("TODO")))
}

func (t *ParseTest) MultipleValidJobs() {
	t.data = `
	{
		"jobs": {
			"taco": {
				"base_path": "/foo"
				"excludes": ["a.b"],
			},
			"burrito": {
				"base_path": "/bar",
				"excludes": ["c", "d"]
			}
		}
	}
	`

	t.parse()

	AssertEq(nil, t.err)
	AssertEq(2, len(t.cfg.Jobs))

	AssertNe(nil, t.cfg.Jobs["taco"])
	ExpectEq("/foo", t.cfg.Jobs["taco"].BasePath)
	AssertThat(t.cfg.Jobs["taco"].Excludes, ElementsAre(Any()))
	ExpectEq("a.b", t.cfg.Jobs["taco"].Excludes[0])

	AssertNe(nil, t.cfg.Jobs["burrito"])
	ExpectEq("/bar", t.cfg.Jobs["burrito"].BasePath)
	AssertThat(t.cfg.Jobs["burrito"].Excludes, ElementsAre(Any(), Any()))
	ExpectEq("c", t.cfg.Jobs["burrito"].Excludes[0])
	ExpectEq("d", t.cfg.Jobs["burrito"].Excludes[1])
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 config_test

import (
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestParse(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type ParseTest struct {
}

func init() { RegisterTestSuite(&ParseTest{}) }

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ParseTest) TotalJunk() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) MissingTrailingBrace() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) BasePathIsNumber() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) BasePathIsNull() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) BasePathIsObject() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) OneExcludeDoesntCompile() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) EmptyConfig() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) MissingExcludesArray() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) DuplicateJobName() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) StructurallyValid() {
	ExpectEq("TODO", "")
}
<commit_msg>Added some helpers.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 config_test

import (
	"github.com/jacobsa/comeback/config"
	. "github.com/jacobsa/ogletest"
	"testing"
)

func TestParse(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type ParseTest struct {
	data []byte
	cfg *config.Config
	err error
}

func init() { RegisterTestSuite(&ParseTest{}) }

func (t *ParseTest) parse() {
	t.cfg, t.err = config.Parse(t.data)
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *ParseTest) TotalJunk() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) MissingTrailingBrace() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) BasePathIsNumber() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) BasePathIsNull() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) BasePathIsObject() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) OneExcludeDoesntCompile() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) EmptyConfig() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) MissingExcludesArray() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) DuplicateJobName() {
	ExpectEq("TODO", "")
}

func (t *ParseTest) StructurallyValid() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>package config

import (
	"fmt"
	"net/url"
	"strings"
)

type URLConfig struct {
	git Environment
}

func NewURLConfig(git Environment) *URLConfig {
	return &URLConfig{
		git: git,
	}
}

// Get retrieves a `http.{url}.{key}` for the given key and urls, following the
// rules in https://git-scm.com/docs/git-config#git-config-httplturlgt.
// The value for `http.{key}` is returned as a fallback if no config keys are
// set for the given urls.
func (c *URLConfig) Get(prefix, key string, rawurl string) (string, bool) {
	key = strings.ToLower(key)
	prefix = strings.ToLower(prefix)
	if v, ok := c.get(key, rawurl); ok {
		return v, ok
	}
	return c.git.Get(strings.Join([]string{prefix, key}, "."))
}

func (c *URLConfig) get(key, rawurl string) (string, bool) {
	u, err := url.Parse(rawurl)
	if err != nil {
		return "", false
	}

	hosts := c.hosts(u)

	pLen := len(u.Path)
	if pLen > 2 {
		end := pLen
		if strings.HasSuffix(u.Path, "/") {
			end -= 1
		}

		paths := strings.Split(u.Path[1:end], "/")
		for i := len(paths); i > 0; i-- {
			for _, host := range hosts {
				path := strings.Join(paths[:i], "/")
				if v, ok := c.git.Get(fmt.Sprintf("http.%s/%s.%s", host, path, key)); ok {
					return v, ok
				}
				if v, ok := c.git.Get(fmt.Sprintf("http.%s/%s/.%s", host, path, key)); ok {
					return v, ok
				}
			}
		}
	}

	for _, host := range hosts {
		if v, ok := c.git.Get(fmt.Sprintf("http.%s.%s", host, key)); ok {
			return v, ok
		}
		if v, ok := c.git.Get(fmt.Sprintf("http.%s/.%s", host, key)); ok {
			return v, ok
		}
	}
	return "", false

}

func (c *URLConfig) hosts(u *url.URL) []string {
	hosts := make([]string, 0, 1)

	if u.User != nil {
		hosts = append(hosts, fmt.Sprintf("%s://%s@%s", u.Scheme, u.User.Username(), u.Host))
	}
	hosts = append(hosts, fmt.Sprintf("%s://%s", u.Scheme, u.Host))

	return hosts
}
<commit_msg>config/url_config: extract paths() function<commit_after>package config

import (
	"fmt"
	"net/url"
	"strings"
)

type URLConfig struct {
	git Environment
}

func NewURLConfig(git Environment) *URLConfig {
	return &URLConfig{
		git: git,
	}
}

// Get retrieves a `http.{url}.{key}` for the given key and urls, following the
// rules in https://git-scm.com/docs/git-config#git-config-httplturlgt.
// The value for `http.{key}` is returned as a fallback if no config keys are
// set for the given urls.
func (c *URLConfig) Get(prefix, key string, rawurl string) (string, bool) {
	key = strings.ToLower(key)
	prefix = strings.ToLower(prefix)
	if v, ok := c.get(key, rawurl); ok {
		return v, ok
	}
	return c.git.Get(strings.Join([]string{prefix, key}, "."))
}

func (c *URLConfig) get(key, rawurl string) (string, bool) {
	u, err := url.Parse(rawurl)
	if err != nil {
		return "", false
	}

	hosts := c.hosts(u)
	paths := c.paths(u.Path)

	for i := len(paths); i > 0; i-- {
		for _, host := range hosts {
			path := strings.Join(paths[:i], "/")
			if v, ok := c.git.Get(fmt.Sprintf("http.%s/%s.%s", host, path, key)); ok {
				return v, ok
			}
			if v, ok := c.git.Get(fmt.Sprintf("http.%s/%s/.%s", host, path, key)); ok {
				return v, ok
			}
		}
	}

	for _, host := range hosts {
		if v, ok := c.git.Get(fmt.Sprintf("http.%s.%s", host, key)); ok {
			return v, ok
		}
		if v, ok := c.git.Get(fmt.Sprintf("http.%s/.%s", host, key)); ok {
			return v, ok
		}
	}
	return "", false

}

func (c *URLConfig) hosts(u *url.URL) []string {
	hosts := make([]string, 0, 1)

	if u.User != nil {
		hosts = append(hosts, fmt.Sprintf("%s://%s@%s", u.Scheme, u.User.Username(), u.Host))
	}
	hosts = append(hosts, fmt.Sprintf("%s://%s", u.Scheme, u.Host))

	return hosts
}

func (c *URLConfig) paths(path string) []string {
	pLen := len(path)
	if pLen <= 2 {
		return nil
	}

	end := pLen
	if strings.HasSuffix(path, "/") {
		end -= 1
	}
	return strings.Split(path[1:end], "/")
}
<|endoftext|>


<commit_before>package goconfigparser

import (
	"io/ioutil"
	"sort"
	"strings"
	"testing"

	. "launchpad.net/gocheck"
)

// Hook up gocheck into the "go test" runner
func Test(t *testing.T) { TestingT(t) }

// partition specific testsuite
type ConfigParserTestSuite struct {
	cfg *ConfigParser
}

var _ = Suite(&ConfigParserTestSuite{})

const SAMPLE_INI = `
[service]
base: system-image.ubuntu.com
http_port: 80
https_port: 443
channel: ubuntu-core/devel-proposed
device: generic_amd64
build_number: 246
version_detail: ubuntu=20150121,raw-device=20150121,version=246

[foo]
bar: baz
yesbool: On
nobool: off
float: 3.14

[testOptions]
One: 1
Two: 2
`

func (s *ConfigParserTestSuite) SetUpTest(c *C) {
	s.cfg = New()
	c.Assert(s.cfg, NotNil)
	err := s.cfg.ReadString(SAMPLE_INI)
	c.Assert(err, IsNil)
}

func (s *ConfigParserTestSuite) TestSection(c *C) {
	sections := s.cfg.Sections()
	sort.Strings(sections)
	c.Assert(sections, DeepEquals, []string{"foo", "service", "testOptions"})
}

func (s *ConfigParserTestSuite) TestOptions(c *C) {
	options, err := s.cfg.Options("testOptions")
	c.Assert(err, IsNil)
	sort.Strings(options)
	c.Assert(options, DeepEquals, []string{"One", "Two"})
}

func (s *ConfigParserTestSuite) TestGet(c *C) {
	val, err := s.cfg.Get("service", "base")
	c.Assert(err, IsNil)
	c.Assert(val, Equals, "system-image.ubuntu.com")
}

func (s *ConfigParserTestSuite) TestGetint(c *C) {
	intval, err := s.cfg.Getint("service", "http_port")
	c.Assert(err, IsNil)
	c.Assert(intval, Equals, 80)
}

func (s *ConfigParserTestSuite) TestGetfloat(c *C) {
	intval, err := s.cfg.Getfloat("foo", "float")
	c.Assert(err, IsNil)
	c.Assert(intval, Equals, 3.14)
}

func (s *ConfigParserTestSuite) TestGetbool(c *C) {
	boolval, err := s.cfg.Getbool("foo", "yesbool")
	c.Assert(err, IsNil)
	c.Assert(boolval, Equals, true)

	boolval, err = s.cfg.Getbool("foo", "nobool")
	c.Assert(err, IsNil)
	c.Assert(boolval, Equals, false)

	boolval, err = s.cfg.Getbool("foo", "bar")
	c.Assert(err.Error(), Equals, "No boolean: baz")
}

func (s *ConfigParserTestSuite) TestErrors(c *C) {
	val, err := s.cfg.Get("foo", "bar")
	c.Assert(err, IsNil)
	c.Assert(val, Equals, "baz")

	val, err = s.cfg.Get("foo", "no-such-option")
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, "No option no-such-option in section foo")

	val, err = s.cfg.Get("no-such-section", "no-such-value")
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, "No section: no-such-section")
}

func (s *ConfigParserTestSuite) TestAllowNoSection(c *C) {
	s.cfg = New()
	s.cfg.AllowNoSectionHeader = true
	err := s.cfg.Read(strings.NewReader(`foo=bar`))
	c.Assert(err, IsNil)
	val, err := s.cfg.Get("", "foo")
	c.Assert(val, Equals, "bar")
}

func (s *ConfigParserTestSuite) TestReadFile(c *C) {
	tmp, err := ioutil.TempFile("", "")
	c.Assert(err, IsNil)
	tmp.Write([]byte(SAMPLE_INI))

	s.cfg = New()
	err = s.cfg.ReadFile(tmp.Name())
	c.Assert(err, IsNil)
	val, err := s.cfg.Get("foo", "bar")
	c.Assert(val, Equals, "baz")
}
<commit_msg>use gopkg.in/check.v1 for the tests<commit_after>package goconfigparser

import (
	"io/ioutil"
	"sort"
	"strings"
	"testing"

	. "gopkg.in/check.v1"
)

// Hook up gocheck into the "go test" runner
func Test(t *testing.T) { TestingT(t) }

// partition specific testsuite
type ConfigParserTestSuite struct {
	cfg *ConfigParser
}

var _ = Suite(&ConfigParserTestSuite{})

const SAMPLE_INI = `
[service]
base: system-image.ubuntu.com
http_port: 80
https_port: 443
channel: ubuntu-core/devel-proposed
device: generic_amd64
build_number: 246
version_detail: ubuntu=20150121,raw-device=20150121,version=246

[foo]
bar: baz
yesbool: On
nobool: off
float: 3.14

[testOptions]
One: 1
Two: 2
`

func (s *ConfigParserTestSuite) SetUpTest(c *C) {
	s.cfg = New()
	c.Assert(s.cfg, NotNil)
	err := s.cfg.ReadString(SAMPLE_INI)
	c.Assert(err, IsNil)
}

func (s *ConfigParserTestSuite) TestSection(c *C) {
	sections := s.cfg.Sections()
	sort.Strings(sections)
	c.Assert(sections, DeepEquals, []string{"foo", "service", "testOptions"})
}

func (s *ConfigParserTestSuite) TestOptions(c *C) {
	options, err := s.cfg.Options("testOptions")
	c.Assert(err, IsNil)
	sort.Strings(options)
	c.Assert(options, DeepEquals, []string{"One", "Two"})
}

func (s *ConfigParserTestSuite) TestGet(c *C) {
	val, err := s.cfg.Get("service", "base")
	c.Assert(err, IsNil)
	c.Assert(val, Equals, "system-image.ubuntu.com")
}

func (s *ConfigParserTestSuite) TestGetint(c *C) {
	intval, err := s.cfg.Getint("service", "http_port")
	c.Assert(err, IsNil)
	c.Assert(intval, Equals, 80)
}

func (s *ConfigParserTestSuite) TestGetfloat(c *C) {
	intval, err := s.cfg.Getfloat("foo", "float")
	c.Assert(err, IsNil)
	c.Assert(intval, Equals, 3.14)
}

func (s *ConfigParserTestSuite) TestGetbool(c *C) {
	boolval, err := s.cfg.Getbool("foo", "yesbool")
	c.Assert(err, IsNil)
	c.Assert(boolval, Equals, true)

	boolval, err = s.cfg.Getbool("foo", "nobool")
	c.Assert(err, IsNil)
	c.Assert(boolval, Equals, false)

	boolval, err = s.cfg.Getbool("foo", "bar")
	c.Assert(err.Error(), Equals, "No boolean: baz")
}

func (s *ConfigParserTestSuite) TestErrors(c *C) {
	val, err := s.cfg.Get("foo", "bar")
	c.Assert(err, IsNil)
	c.Assert(val, Equals, "baz")

	val, err = s.cfg.Get("foo", "no-such-option")
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, "No option no-such-option in section foo")

	val, err = s.cfg.Get("no-such-section", "no-such-value")
	c.Assert(err, NotNil)
	c.Assert(err.Error(), Equals, "No section: no-such-section")
}

func (s *ConfigParserTestSuite) TestAllowNoSection(c *C) {
	s.cfg = New()
	s.cfg.AllowNoSectionHeader = true
	err := s.cfg.Read(strings.NewReader(`foo=bar`))
	c.Assert(err, IsNil)
	val, err := s.cfg.Get("", "foo")
	c.Assert(val, Equals, "bar")
}

func (s *ConfigParserTestSuite) TestReadFile(c *C) {
	tmp, err := ioutil.TempFile("", "")
	c.Assert(err, IsNil)
	tmp.Write([]byte(SAMPLE_INI))

	s.cfg = New()
	err = s.cfg.ReadFile(tmp.Name())
	c.Assert(err, IsNil)
	val, err := s.cfg.Get("foo", "bar")
	c.Assert(val, Equals, "baz")
}
<|endoftext|>


<commit_before>package consumer

import (
	"encoding/json"
	"flume-log-sdk/consumer/client"
	"fmt"
	"github.com/blackbeans/redigo/redis"
	"log"
	"math/rand"
	_ "os"
	"strconv"
	"time"
)

// 用于向flume中作为sink 通过thrift客户端写入日志

type SinkServer struct {
	redisPool       map[string][]*redis.Pool
	flumeClientPool []*flumeClientPool
	isStop          bool
}

func NewSinkServer(option *Option) (server *SinkServer) {

	redisPool := make(map[string][]*redis.Pool, 0)

	//创建redis的消费连接
	for _, v := range option.queueHostPorts {

		pool := redis.NewPool(func() (conn redis.Conn, err error) {

			conn, err = redis.DialTimeout("tcp", v.Host+":"+strconv.Itoa(v.Port),
				time.Duration(v.Timeout)*time.Second,
				time.Duration(v.Timeout)*time.Second,
				time.Duration(v.Timeout)*time.Second)

			return
		}, time.Duration(v.Timeout*2)*time.Second, v.Maxconn/2, v.Maxconn)

		pools, ok := redisPool[v.QueueName]
		if !ok {
			pools = make([]*redis.Pool, 0)
			redisPool[v.QueueName] = pools
		}

		redisPool[v.QueueName] = append(pools, pool)

	}

	pools := make([]*flumeClientPool, 0)
	//创建flume的client
	for _, v := range option.flumeAgents {

		pool := newFlumeClientPool(20, 50, 100, 10*time.Second, func() *client.FlumeClient {
			flumeclient := client.NewFlumeClient(v.Host, v.Port)
			flumeclient.Connect()
			return flumeclient
		})
		pools = append(pools, pool)

		go monitorPool(v.Host+strconv.Itoa(v.Port), pool)
	}

	sinkserver := &SinkServer{redisPool: redisPool, flumeClientPool: pools}

	return sinkserver
}

func monitorPool(hostport string, pool *flumeClientPool) {
	for {
		time.Sleep(1 * time.Second)
		log.Printf("flume:%s|active:%d,core:%d,max:%d",
			hostport, pool.ActivePoolSize(), pool.CorePoolSize(), pool.maxPoolSize)
	}

}

//启动pop
func (self *SinkServer) Start() {

	self.isStop = false
	ch := make(chan int, 1)
	var count = 0
	for k, v := range self.redisPool {

		log.Println("start redis queueserver succ " + k)
		for _, pool := range v {
			count++
			defer pool.Close()
			go func(queuename string, pool *redis.Pool, end chan int) {
				conn := pool.Get()
				defer pool.Release(conn)
				for !self.isStop {

					// log.Println("pool active count :", strconv.Itoa(pool.ActiveCount()))
					reply, err := conn.Do("LPOP", queuename)
					if nil != err || nil == reply {
						if nil != err {
							log.Printf("LPOP|FAIL|%s", err)
							conn.Close()
							conn = pool.Get()
						} else {
							time.Sleep(100 * time.Millisecond)
						}

						continue
					}

					resp := reply.([]byte)
					var cmd command
					err = json.Unmarshal(resp, &cmd)

					if nil != err {
						log.Printf("command unmarshal fail ! %s | error:%s\n", resp, err.Error())
						continue
					} else if rand.Int()%10 == 0 {
						log.Println("trace|command|%s", cmd)
					}

					//
					momoid := cmd.Params["momoid"].(string)

					businessName := cmd.Params["businessName"].(string)

					action := cmd.Params["type"].(string)

					bodyContent := cmd.Params["body"]

					//将businessName 加入到body中
					bodyMap := bodyContent.(map[string]interface{})
					bodyMap["business_type"] = businessName

					body, err := json.Marshal(bodyContent)

					if nil != err {
						log.Printf("marshal log body fail %s", err.Error())
						continue
					}

					//这里需要优化一下body,需要采用其他的方式定义Body格式，写入

					log.Printf("%s,%s,%s,%s", momoid, businessName, action, string(body))

					//启动处理任务
					go self.innerSend(momoid, businessName, action, string(body))

				}
				end <- -1
			}(k, pool, ch)
		}
	}

	for {
		count += <-ch
		if count <= 0 {
			log.Printf("redis conn  close %d", count)
			break
		}
	}

}

func (self *SinkServer) innerSend(momoid, businessName, action string, body string) {

	for i := 0; i < 3; i++ {
		pool := self.getFlumeClientPool(businessName, action)
		flumeclient, err := pool.Get(5 * time.Second)
		if nil != err {
			continue
		}
		//拼装头部信息
		header := make(map[string]string, 1)
		header["businessName"] = businessName
		header["type"] = action

		//拼Body
		flumeBody := fmt.Sprintf("%s\t%s\t%s", momoid, action, body)
		err = flumeclient.Append(header, []byte(flumeBody))
		defer func() {
			if err := recover(); nil != err {
				//回收这个坏的连接
				pool.ReleaseBroken(flumeclient)
			}
		}()

		if nil != err {
			log.Printf("send 2 flume fail %s \t err:%s\n", body, err.Error())
		} else {
			log.Printf("send 2 flume succ %s\n", body)
			pool.Release(flumeclient)
			break
		}
	}
}

//仅供测试使用推送数据
func (self *SinkServer) testPushLog(queuename, logger string) {

	for _, v := range self.redisPool {
		for _, pool := range v {
			conn := pool.Get()
			defer pool.Release(conn)

			reply, err := conn.Do("RPUSH", queuename, logger)
			log.Printf("%s|err:%s", reply, err)
			break

		}
	}

}

func (self *SinkServer) Stop() {
	self.isStop = true
	for _, v := range self.flumeClientPool {
		v.Destroy()
	}

	for _, v := range self.redisPool {
		for _, p := range v {
			p.Close()
		}
	}
}

func (self *SinkServer) getFlumeClientPool(businessName, action string) *flumeClientPool {

	//使用随机算法直接获得

	idx := rand.Intn(len(self.flumeClientPool))
	return self.flumeClientPool[idx]

}
<commit_msg>修改日志格式 	modified:   consumer/log_sink.go<commit_after>package consumer

import (
	"encoding/json"
	"flume-log-sdk/consumer/client"
	"fmt"
	"github.com/blackbeans/redigo/redis"
	"log"
	"math/rand"
	_ "os"
	"strconv"
	"time"
)

// 用于向flume中作为sink 通过thrift客户端写入日志

type SinkServer struct {
	redisPool       map[string][]*redis.Pool
	flumeClientPool []*flumeClientPool
	isStop          bool
}

func NewSinkServer(option *Option) (server *SinkServer) {

	redisPool := make(map[string][]*redis.Pool, 0)

	//创建redis的消费连接
	for _, v := range option.queueHostPorts {

		pool := redis.NewPool(func() (conn redis.Conn, err error) {

			conn, err = redis.DialTimeout("tcp", v.Host+":"+strconv.Itoa(v.Port),
				time.Duration(v.Timeout)*time.Second,
				time.Duration(v.Timeout)*time.Second,
				time.Duration(v.Timeout)*time.Second)

			return
		}, time.Duration(v.Timeout*2)*time.Second, v.Maxconn/2, v.Maxconn)

		pools, ok := redisPool[v.QueueName]
		if !ok {
			pools = make([]*redis.Pool, 0)
			redisPool[v.QueueName] = pools
		}

		redisPool[v.QueueName] = append(pools, pool)

	}

	pools := make([]*flumeClientPool, 0)
	//创建flume的client
	for _, v := range option.flumeAgents {

		pool := newFlumeClientPool(20, 50, 100, 10*time.Second, func() *client.FlumeClient {
			flumeclient := client.NewFlumeClient(v.Host, v.Port)
			flumeclient.Connect()
			return flumeclient
		})
		pools = append(pools, pool)

		go monitorPool(v.Host+":"+strconv.Itoa(v.Port), pool)
	}

	sinkserver := &SinkServer{redisPool: redisPool, flumeClientPool: pools}

	return sinkserver
}

func monitorPool(hostport string, pool *flumeClientPool) {
	for {
		time.Sleep(1 * time.Second)
		log.Printf("flume:%s|active:%d,core:%d,max:%d",
			hostport, pool.ActivePoolSize(), pool.CorePoolSize(), pool.maxPoolSize)
	}

}

//启动pop
func (self *SinkServer) Start() {

	self.isStop = false
	ch := make(chan int, 1)
	var count = 0
	for k, v := range self.redisPool {

		log.Println("start redis queueserver succ " + k)
		for _, pool := range v {
			count++
			defer pool.Close()
			go func(queuename string, pool *redis.Pool, end chan int) {
				conn := pool.Get()
				defer pool.Release(conn)
				for !self.isStop {

					// log.Println("pool active count :", strconv.Itoa(pool.ActiveCount()))
					reply, err := conn.Do("LPOP", queuename)
					if nil != err || nil == reply {
						if nil != err {
							log.Printf("LPOP|FAIL|%s", err)
							conn.Close()
							conn = pool.Get()
						} else {
							time.Sleep(100 * time.Millisecond)
						}

						continue
					}

					resp := reply.([]byte)
					var cmd command
					err = json.Unmarshal(resp, &cmd)

					if nil != err {
						log.Printf("command unmarshal fail ! %s | error:%s\n", resp, err.Error())
						continue
					} else if rand.Int()%10 == 0 {
						log.Println("trace|command|%s", cmd)
					}

					//
					momoid := cmd.Params["momoid"].(string)

					businessName := cmd.Params["businessName"].(string)

					action := cmd.Params["type"].(string)

					bodyContent := cmd.Params["body"]

					//将businessName 加入到body中
					bodyMap := bodyContent.(map[string]interface{})
					bodyMap["business_type"] = businessName

					body, err := json.Marshal(bodyContent)

					if nil != err {
						log.Printf("marshal log body fail %s", err.Error())
						continue
					}

					//这里需要优化一下body,需要采用其他的方式定义Body格式，写入

					log.Printf("%s,%s,%s,%s", momoid, businessName, action, string(body))

					//启动处理任务
					go self.innerSend(momoid, businessName, action, string(body))

				}
				end <- -1
			}(k, pool, ch)
		}
	}

	for {
		count += <-ch
		if count <= 0 {
			log.Printf("redis conn  close %d", count)
			break
		}
	}

}

func (self *SinkServer) innerSend(momoid, businessName, action string, body string) {

	for i := 0; i < 3; i++ {
		pool := self.getFlumeClientPool(businessName, action)
		flumeclient, err := pool.Get(5 * time.Second)
		if nil != err {
			continue
		}
		//拼装头部信息
		header := make(map[string]string, 1)
		header["businessName"] = businessName
		header["type"] = action

		//拼Body
		flumeBody := fmt.Sprintf("%s\t%s\t%s", momoid, action, body)
		err = flumeclient.Append(header, []byte(flumeBody))
		defer func() {
			if err := recover(); nil != err {
				//回收这个坏的连接
				pool.ReleaseBroken(flumeclient)
			}
		}()

		if nil != err {
			log.Printf("send 2 flume fail %s \t err:%s\n", body, err.Error())
		} else {
			log.Printf("send 2 flume succ %s\n", body)
			pool.Release(flumeclient)
			break
		}
	}
}

//仅供测试使用推送数据
func (self *SinkServer) testPushLog(queuename, logger string) {

	for _, v := range self.redisPool {
		for _, pool := range v {
			conn := pool.Get()
			defer pool.Release(conn)

			reply, err := conn.Do("RPUSH", queuename, logger)
			log.Printf("%s|err:%s", reply, err)
			break

		}
	}

}

func (self *SinkServer) Stop() {
	self.isStop = true
	for _, v := range self.flumeClientPool {
		v.Destroy()
	}

	for _, v := range self.redisPool {
		for _, p := range v {
			p.Close()
		}
	}
}

func (self *SinkServer) getFlumeClientPool(businessName, action string) *flumeClientPool {

	//使用随机算法直接获得

	idx := rand.Intn(len(self.flumeClientPool))
	return self.flumeClientPool[idx]

}
<|endoftext|>


<commit_before>package configs

type Mount struct {
	// Source path for the mount.
	Source string `json:"source"`

	// Destination path for the mount inside the container.
	Destination string `json:"destination"`

	// Device the mount is for.
	Device string `json:"device"`

	// Mount flags.
	Flags int `json:"flags"`

	// Propagation Flags
	PropagationFlags []int `json:"propagation_flags"`

	// Mount data applied to the mount.
	Data string `json:"data"`

	// Relabel source if set, "z" indicates shared, "Z" indicates unshared.
	Relabel string `json:"relabel"`

	// Optional Command to be run before Source is mounted.
	PremountCmds []Command `json:"premount_cmds"`

	// Optional Command to be run after Source is mounted.
	PostmountCmds []Command `json:"postmount_cmds"`
}
<commit_msg>config: Add new Extensions flag to support custom mount options in runc<commit_after>package configs

const (
	// EXT_COPYUP is a directive to copy up the contents of a directory when
	// a tmpfs is mounted over it.
	EXT_COPYUP = 1 << iota
)

type Mount struct {
	// Source path for the mount.
	Source string `json:"source"`

	// Destination path for the mount inside the container.
	Destination string `json:"destination"`

	// Device the mount is for.
	Device string `json:"device"`

	// Mount flags.
	Flags int `json:"flags"`

	// Propagation Flags
	PropagationFlags []int `json:"propagation_flags"`

	// Mount data applied to the mount.
	Data string `json:"data"`

	// Relabel source if set, "z" indicates shared, "Z" indicates unshared.
	Relabel string `json:"relabel"`

	// Extensions are additional flags that are specific to runc.
	Extensions int `json:"extensions"`

	// Optional Command to be run before Source is mounted.
	PremountCmds []Command `json:"premount_cmds"`

	// Optional Command to be run after Source is mounted.
	PostmountCmds []Command `json:"postmount_cmds"`
}
<|endoftext|>