<commit_before>package mocknet

import (
	"bytes"
	"errors"
	"io"
	"net"
	"time"

	inet "github.com/libp2p/go-libp2p-net"
	protocol "github.com/libp2p/go-libp2p-protocol"
)

// stream implements inet.Stream
type stream struct {
	write     *io.PipeWriter
	read      *io.PipeReader
	conn      *conn
	toDeliver chan *transportObject

	reset  chan struct{}
	close  chan struct{}
	closed chan struct{}

	writeErr error

	protocol protocol.ID
	stat     inet.Stat
}

var ErrReset error = errors.New("stream reset")
var ErrClosed error = errors.New("stream closed")

type transportObject struct {
	msg         []byte
	arrivalTime time.Time
}

func NewStream(w *io.PipeWriter, r *io.PipeReader, dir inet.Direction) *stream {
	s := &stream{
		read:      r,
		write:     w,
		reset:     make(chan struct{}, 1),
		close:     make(chan struct{}, 1),
		closed:    make(chan struct{}),
		toDeliver: make(chan *transportObject),
		stat:      inet.Stat{Direction: dir},
	}

	go s.transport()
	return s
}

//  How to handle errors with writes?
func (s *stream) Write(p []byte) (n int, err error) {
	l := s.conn.link
	delay := l.GetLatency() + l.RateLimit(len(p))
	t := time.Now().Add(delay)
	select {
	case <-s.closed: // bail out if we're closing.
		return 0, s.writeErr
	case s.toDeliver <- &transportObject{msg: p, arrivalTime: t}:
	}
	return len(p), nil
}

func (s *stream) Protocol() protocol.ID {
	return s.protocol
}

func (s *stream) Stat() inet.Stat {
	return s.stat
}

func (s *stream) SetProtocol(proto protocol.ID) {
	s.protocol = proto
}

func (s *stream) Close() error {
	select {
	case s.close <- struct{}{}:
	default:
	}
	<-s.closed
	if s.writeErr != ErrClosed {
		return s.writeErr
	}
	return nil
}

func (s *stream) Reset() error {
	// Cancel any pending reads/writes with an error.
	s.write.CloseWithError(ErrReset)
	s.read.CloseWithError(ErrReset)

	select {
	case s.reset <- struct{}{}:
	default:
	}
	<-s.closed

	// No meaningful error case here.
	return nil
}

func (s *stream) teardown() {
	// at this point, no streams are writing.
	s.conn.removeStream(s)

	// Mark as closed.
	close(s.closed)

	s.conn.net.notifyAll(func(n inet.Notifiee) {
		n.ClosedStream(s.conn.net, s)
	})
}

func (s *stream) Conn() inet.Conn {
	return s.conn
}

func (s *stream) SetDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "pipe", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (s *stream) SetReadDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "pipe", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (s *stream) SetWriteDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "pipe", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (s *stream) Read(b []byte) (int, error) {
	return s.read.Read(b)
}

// transport will grab message arrival times, wait until that time, and
// then write the message out when it is scheduled to arrive
func (s *stream) transport() {
	defer s.teardown()

	bufsize := 256
	buf := new(bytes.Buffer)
	timer := time.NewTimer(0)
	if !timer.Stop() {
		select {
		case <-timer.C:
		default:
		}
	}

	// cleanup
	defer timer.Stop()

	// writeBuf writes the contents of buf through to the s.Writer.
	// done only when arrival time makes sense.
	drainBuf := func() error {
		if buf.Len() > 0 {
			_, err := s.write.Write(buf.Bytes())
			if err != nil {
				return err
			}
			buf.Reset()
		}
		return nil
	}

	// deliverOrWait is a helper func that processes
	// an incoming packet. it waits until the arrival time,
	// and then writes things out.
	deliverOrWait := func(o *transportObject) error {
		buffered := len(o.msg) + buf.Len()

		// Yes, we can end up extending a timer multiple times if we
		// keep on making small writes but that shouldn't be too much of an
		// issue. Fixing that would be painful.
		if !timer.Stop() {
			// FIXME: So, we *shouldn't* need to do this but we hang
			// here if we don't... Go bug?
			select {
			case <-timer.C:
			default:
			}
		}
		delay := o.arrivalTime.Sub(time.Now())
		if delay >= 0 {
			timer.Reset(delay)
		} else {
			timer.Reset(0)
		}

		if buffered >= bufsize {
			select {
			case <-timer.C:
			case <-s.reset:
				select {
				case s.reset <- struct{}{}:
				default:
				}
				return ErrReset
			}
			if err := drainBuf(); err != nil {
				return err
			}
			// write this message.
			_, err := s.write.Write(o.msg)
			if err != nil {
				return err
			}
		} else {
			buf.Write(o.msg)
		}
		return nil
	}

	for {
		// Reset takes precedent.
		select {
		case <-s.reset:
			s.writeErr = ErrReset
			return
		default:
		}

		select {
		case <-s.reset:
			s.writeErr = ErrReset
			return
		case <-s.close:
			if err := drainBuf(); err != nil {
				s.resetWith(err)
				return
			}
			s.writeErr = s.write.Close()
			if s.writeErr == nil {
				s.writeErr = ErrClosed
			}
			return
		case o := <-s.toDeliver:
			if err := deliverOrWait(o); err != nil {
				s.resetWith(err)
				return
			}
		case <-timer.C: // ok, due to write it out.
			if err := drainBuf(); err != nil {
				s.resetWith(err)
				return
			}
		}
	}
}

func (s *stream) resetWith(err error) {
	s.write.CloseWithError(err)
	s.read.CloseWithError(err)
	s.writeErr = err
}
<commit_msg>fix a data-race in the mock net<commit_after>package mocknet

import (
	"bytes"
	"errors"
	"io"
	"net"
	"time"

	inet "github.com/libp2p/go-libp2p-net"
	protocol "github.com/libp2p/go-libp2p-protocol"
)

// stream implements inet.Stream
type stream struct {
	write     *io.PipeWriter
	read      *io.PipeReader
	conn      *conn
	toDeliver chan *transportObject

	reset  chan struct{}
	close  chan struct{}
	closed chan struct{}

	writeErr error

	protocol protocol.ID
	stat     inet.Stat
}

var ErrReset error = errors.New("stream reset")
var ErrClosed error = errors.New("stream closed")

type transportObject struct {
	msg         []byte
	arrivalTime time.Time
}

func NewStream(w *io.PipeWriter, r *io.PipeReader, dir inet.Direction) *stream {
	s := &stream{
		read:      r,
		write:     w,
		reset:     make(chan struct{}, 1),
		close:     make(chan struct{}, 1),
		closed:    make(chan struct{}),
		toDeliver: make(chan *transportObject),
		stat:      inet.Stat{Direction: dir},
	}

	go s.transport()
	return s
}

//  How to handle errors with writes?
func (s *stream) Write(p []byte) (n int, err error) {
	l := s.conn.link
	delay := l.GetLatency() + l.RateLimit(len(p))
	t := time.Now().Add(delay)

	// Copy it.
	cpy := make([]byte, len(p))
	copy(cpy, p)

	select {
	case <-s.closed: // bail out if we're closing.
		return 0, s.writeErr
	case s.toDeliver <- &transportObject{msg: p, arrivalTime: t}:
	}
	return len(p), nil
}

func (s *stream) Protocol() protocol.ID {
	return s.protocol
}

func (s *stream) Stat() inet.Stat {
	return s.stat
}

func (s *stream) SetProtocol(proto protocol.ID) {
	s.protocol = proto
}

func (s *stream) Close() error {
	select {
	case s.close <- struct{}{}:
	default:
	}
	<-s.closed
	if s.writeErr != ErrClosed {
		return s.writeErr
	}
	return nil
}

func (s *stream) Reset() error {
	// Cancel any pending reads/writes with an error.
	s.write.CloseWithError(ErrReset)
	s.read.CloseWithError(ErrReset)

	select {
	case s.reset <- struct{}{}:
	default:
	}
	<-s.closed

	// No meaningful error case here.
	return nil
}

func (s *stream) teardown() {
	// at this point, no streams are writing.
	s.conn.removeStream(s)

	// Mark as closed.
	close(s.closed)

	s.conn.net.notifyAll(func(n inet.Notifiee) {
		n.ClosedStream(s.conn.net, s)
	})
}

func (s *stream) Conn() inet.Conn {
	return s.conn
}

func (s *stream) SetDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "pipe", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (s *stream) SetReadDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "pipe", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (s *stream) SetWriteDeadline(t time.Time) error {
	return &net.OpError{Op: "set", Net: "pipe", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
}

func (s *stream) Read(b []byte) (int, error) {
	return s.read.Read(b)
}

// transport will grab message arrival times, wait until that time, and
// then write the message out when it is scheduled to arrive
func (s *stream) transport() {
	defer s.teardown()

	bufsize := 256
	buf := new(bytes.Buffer)
	timer := time.NewTimer(0)
	if !timer.Stop() {
		select {
		case <-timer.C:
		default:
		}
	}

	// cleanup
	defer timer.Stop()

	// writeBuf writes the contents of buf through to the s.Writer.
	// done only when arrival time makes sense.
	drainBuf := func() error {
		if buf.Len() > 0 {
			_, err := s.write.Write(buf.Bytes())
			if err != nil {
				return err
			}
			buf.Reset()
		}
		return nil
	}

	// deliverOrWait is a helper func that processes
	// an incoming packet. it waits until the arrival time,
	// and then writes things out.
	deliverOrWait := func(o *transportObject) error {
		buffered := len(o.msg) + buf.Len()

		// Yes, we can end up extending a timer multiple times if we
		// keep on making small writes but that shouldn't be too much of an
		// issue. Fixing that would be painful.
		if !timer.Stop() {
			// FIXME: So, we *shouldn't* need to do this but we hang
			// here if we don't... Go bug?
			select {
			case <-timer.C:
			default:
			}
		}
		delay := o.arrivalTime.Sub(time.Now())
		if delay >= 0 {
			timer.Reset(delay)
		} else {
			timer.Reset(0)
		}

		if buffered >= bufsize {
			select {
			case <-timer.C:
			case <-s.reset:
				select {
				case s.reset <- struct{}{}:
				default:
				}
				return ErrReset
			}
			if err := drainBuf(); err != nil {
				return err
			}
			// write this message.
			_, err := s.write.Write(o.msg)
			if err != nil {
				return err
			}
		} else {
			buf.Write(o.msg)
		}
		return nil
	}

	for {
		// Reset takes precedent.
		select {
		case <-s.reset:
			s.writeErr = ErrReset
			return
		default:
		}

		select {
		case <-s.reset:
			s.writeErr = ErrReset
			return
		case <-s.close:
			if err := drainBuf(); err != nil {
				s.resetWith(err)
				return
			}
			s.writeErr = s.write.Close()
			if s.writeErr == nil {
				s.writeErr = ErrClosed
			}
			return
		case o := <-s.toDeliver:
			if err := deliverOrWait(o); err != nil {
				s.resetWith(err)
				return
			}
		case <-timer.C: // ok, due to write it out.
			if err := drainBuf(); err != nil {
				s.resetWith(err)
				return
			}
		}
	}
}

func (s *stream) resetWith(err error) {
	s.write.CloseWithError(err)
	s.read.CloseWithError(err)
	s.writeErr = err
}
<|endoftext|>


<commit_before>/*
 *
 * Copyright 2014 gRPC authors.
 *
 * 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 credentials

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net"

	"google.golang.org/grpc/credentials/internal"
)

// TLSInfo contains the auth information for a TLS authenticated connection.
// It implements the AuthInfo interface.
type TLSInfo struct {
	State tls.ConnectionState
	CommonAuthInfo
}

// AuthType returns the type of TLSInfo as a string.
func (t TLSInfo) AuthType() string {
	return "tls"
}

// GetSecurityValue returns security info requested by channelz.
func (t TLSInfo) GetSecurityValue() ChannelzSecurityValue {
	v := &TLSChannelzSecurityValue{
		StandardName: cipherSuiteLookup[t.State.CipherSuite],
	}
	// Currently there's no way to get LocalCertificate info from tls package.
	if len(t.State.PeerCertificates) > 0 {
		v.RemoteCertificate = t.State.PeerCertificates[0].Raw
	}
	return v
}

// tlsCreds is the credentials required for authenticating a connection using TLS.
type tlsCreds struct {
	// TLS configuration
	config *tls.Config
}

func (c tlsCreds) Info() ProtocolInfo {
	return ProtocolInfo{
		SecurityProtocol: "tls",
		SecurityVersion:  "1.2",
		ServerName:       c.config.ServerName,
	}
}

func (c *tlsCreds) ClientHandshake(ctx context.Context, authority string, rawConn net.Conn) (_ net.Conn, _ AuthInfo, err error) {
	// use local cfg to avoid clobbering ServerName if using multiple endpoints
	cfg := cloneTLSConfig(c.config)
	if cfg.ServerName == "" {
		serverName, _, err := net.SplitHostPort(authority)
		if err != nil {
			// If the authority had no host port or if the authority cannot be parsed, use it as-is.
			serverName = authority
		}
		cfg.ServerName = serverName
	}
	conn := tls.Client(rawConn, cfg)
	errChannel := make(chan error, 1)
	go func() {
		errChannel <- conn.Handshake()
		close(errChannel)
	}()
	select {
	case err := <-errChannel:
		if err != nil {
			conn.Close()
			return nil, nil, err
		}
	case <-ctx.Done():
		conn.Close()
		return nil, nil, ctx.Err()
	}
	return internal.WrapSyscallConn(rawConn, conn), TLSInfo{conn.ConnectionState(), CommonAuthInfo{PrivacyAndIntegrity}}, nil
}

func (c *tlsCreds) ServerHandshake(rawConn net.Conn) (net.Conn, AuthInfo, error) {
	conn := tls.Server(rawConn, c.config)
	if err := conn.Handshake(); err != nil {
		conn.Close()
		return nil, nil, err
	}
	return internal.WrapSyscallConn(rawConn, conn), TLSInfo{conn.ConnectionState(), CommonAuthInfo{PrivacyAndIntegrity}}, nil
}

func (c *tlsCreds) Clone() TransportCredentials {
	return NewTLS(c.config)
}

func (c *tlsCreds) OverrideServerName(serverNameOverride string) error {
	c.config.ServerName = serverNameOverride
	return nil
}

const alpnProtoStrH2 = "h2"

func appendH2ToNextProtos(ps []string) []string {
	for _, p := range ps {
		if p == alpnProtoStrH2 {
			return ps
		}
	}
	ret := make([]string, 0, len(ps)+1)
	ret = append(ret, ps...)
	return append(ret, alpnProtoStrH2)
}

// NewTLS uses c to construct a TransportCredentials based on TLS.
func NewTLS(c *tls.Config) TransportCredentials {
	tc := &tlsCreds{cloneTLSConfig(c)}
	tc.config.NextProtos = appendH2ToNextProtos(tc.config.NextProtos)
	return tc
}

// NewClientTLSFromCert constructs TLS credentials from the input certificate for client.
// serverNameOverride is for testing only. If set to a non empty string,
// it will override the virtual host name of authority (e.g. :authority header field) in requests.
func NewClientTLSFromCert(cp *x509.CertPool, serverNameOverride string) TransportCredentials {
	return NewTLS(&tls.Config{ServerName: serverNameOverride, RootCAs: cp})
}

// NewClientTLSFromFile constructs TLS credentials from the input certificate file for client.
// serverNameOverride is for testing only. If set to a non empty string,
// it will override the virtual host name of authority (e.g. :authority header field) in requests.
func NewClientTLSFromFile(certFile, serverNameOverride string) (TransportCredentials, error) {
	b, err := ioutil.ReadFile(certFile)
	if err != nil {
		return nil, err
	}
	cp := x509.NewCertPool()
	if !cp.AppendCertsFromPEM(b) {
		return nil, fmt.Errorf("credentials: failed to append certificates")
	}
	return NewTLS(&tls.Config{ServerName: serverNameOverride, RootCAs: cp}), nil
}

// NewServerTLSFromCert constructs TLS credentials from the input certificate for server.
func NewServerTLSFromCert(cert *tls.Certificate) TransportCredentials {
	return NewTLS(&tls.Config{Certificates: []tls.Certificate{*cert}})
}

// NewServerTLSFromFile constructs TLS credentials from the input certificate file and key
// file for server.
func NewServerTLSFromFile(certFile, keyFile string) (TransportCredentials, error) {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return nil, err
	}
	return NewTLS(&tls.Config{Certificates: []tls.Certificate{cert}}), nil
}

// TLSChannelzSecurityValue defines the struct that TLS protocol should return
// from GetSecurityValue(), containing security info like cipher and certificate used.
//
// This API is EXPERIMENTAL.
type TLSChannelzSecurityValue struct {
	ChannelzSecurityValue
	StandardName      string
	LocalCertificate  []byte
	RemoteCertificate []byte
}

var cipherSuiteLookup = map[uint16]string{
	tls.TLS_RSA_WITH_RC4_128_SHA:                "TLS_RSA_WITH_RC4_128_SHA",
	tls.TLS_RSA_WITH_3DES_EDE_CBC_SHA:           "TLS_RSA_WITH_3DES_EDE_CBC_SHA",
	tls.TLS_RSA_WITH_AES_128_CBC_SHA:            "TLS_RSA_WITH_AES_128_CBC_SHA",
	tls.TLS_RSA_WITH_AES_256_CBC_SHA:            "TLS_RSA_WITH_AES_256_CBC_SHA",
	tls.TLS_RSA_WITH_AES_128_GCM_SHA256:         "TLS_RSA_WITH_AES_128_GCM_SHA256",
	tls.TLS_RSA_WITH_AES_256_GCM_SHA384:         "TLS_RSA_WITH_AES_256_GCM_SHA384",
	tls.TLS_ECDHE_ECDSA_WITH_RC4_128_SHA:        "TLS_ECDHE_ECDSA_WITH_RC4_128_SHA",
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA:    "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA",
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA:    "TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA",
	tls.TLS_ECDHE_RSA_WITH_RC4_128_SHA:          "TLS_ECDHE_RSA_WITH_RC4_128_SHA",
	tls.TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA:     "TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA",
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA:      "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA",
	tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA:      "TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA",
	tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256:   "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256: "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256",
	tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384:   "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384: "TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384",
	tls.TLS_FALLBACK_SCSV:                       "TLS_FALLBACK_SCSV",
	tls.TLS_RSA_WITH_AES_128_CBC_SHA256:         "TLS_RSA_WITH_AES_128_CBC_SHA256",
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256: "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256",
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256:   "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256",
	tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305:    "TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305",
	tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305:  "TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305",
}

// cloneTLSConfig returns a shallow clone of the exported
// fields of cfg, ignoring the unexported sync.Once, which
// contains a mutex and must not be copied.
//
// If cfg is nil, a new zero tls.Config is returned.
//
// TODO: inline this function if possible.
func cloneTLSConfig(cfg *tls.Config) *tls.Config {
	if cfg == nil {
		return &tls.Config{}
	}

	return cfg.Clone()
}
<commit_msg>credentials: Update doc strings for NewClientTLSFromCert et. al. (#3508)<commit_after>/*
 *
 * Copyright 2014 gRPC authors.
 *
 * 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 credentials

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net"

	"google.golang.org/grpc/credentials/internal"
)

// TLSInfo contains the auth information for a TLS authenticated connection.
// It implements the AuthInfo interface.
type TLSInfo struct {
	State tls.ConnectionState
	CommonAuthInfo
}

// AuthType returns the type of TLSInfo as a string.
func (t TLSInfo) AuthType() string {
	return "tls"
}

// GetSecurityValue returns security info requested by channelz.
func (t TLSInfo) GetSecurityValue() ChannelzSecurityValue {
	v := &TLSChannelzSecurityValue{
		StandardName: cipherSuiteLookup[t.State.CipherSuite],
	}
	// Currently there's no way to get LocalCertificate info from tls package.
	if len(t.State.PeerCertificates) > 0 {
		v.RemoteCertificate = t.State.PeerCertificates[0].Raw
	}
	return v
}

// tlsCreds is the credentials required for authenticating a connection using TLS.
type tlsCreds struct {
	// TLS configuration
	config *tls.Config
}

func (c tlsCreds) Info() ProtocolInfo {
	return ProtocolInfo{
		SecurityProtocol: "tls",
		SecurityVersion:  "1.2",
		ServerName:       c.config.ServerName,
	}
}

func (c *tlsCreds) ClientHandshake(ctx context.Context, authority string, rawConn net.Conn) (_ net.Conn, _ AuthInfo, err error) {
	// use local cfg to avoid clobbering ServerName if using multiple endpoints
	cfg := cloneTLSConfig(c.config)
	if cfg.ServerName == "" {
		serverName, _, err := net.SplitHostPort(authority)
		if err != nil {
			// If the authority had no host port or if the authority cannot be parsed, use it as-is.
			serverName = authority
		}
		cfg.ServerName = serverName
	}
	conn := tls.Client(rawConn, cfg)
	errChannel := make(chan error, 1)
	go func() {
		errChannel <- conn.Handshake()
		close(errChannel)
	}()
	select {
	case err := <-errChannel:
		if err != nil {
			conn.Close()
			return nil, nil, err
		}
	case <-ctx.Done():
		conn.Close()
		return nil, nil, ctx.Err()
	}
	return internal.WrapSyscallConn(rawConn, conn), TLSInfo{conn.ConnectionState(), CommonAuthInfo{PrivacyAndIntegrity}}, nil
}

func (c *tlsCreds) ServerHandshake(rawConn net.Conn) (net.Conn, AuthInfo, error) {
	conn := tls.Server(rawConn, c.config)
	if err := conn.Handshake(); err != nil {
		conn.Close()
		return nil, nil, err
	}
	return internal.WrapSyscallConn(rawConn, conn), TLSInfo{conn.ConnectionState(), CommonAuthInfo{PrivacyAndIntegrity}}, nil
}

func (c *tlsCreds) Clone() TransportCredentials {
	return NewTLS(c.config)
}

func (c *tlsCreds) OverrideServerName(serverNameOverride string) error {
	c.config.ServerName = serverNameOverride
	return nil
}

const alpnProtoStrH2 = "h2"

func appendH2ToNextProtos(ps []string) []string {
	for _, p := range ps {
		if p == alpnProtoStrH2 {
			return ps
		}
	}
	ret := make([]string, 0, len(ps)+1)
	ret = append(ret, ps...)
	return append(ret, alpnProtoStrH2)
}

// NewTLS uses c to construct a TransportCredentials based on TLS.
func NewTLS(c *tls.Config) TransportCredentials {
	tc := &tlsCreds{cloneTLSConfig(c)}
	tc.config.NextProtos = appendH2ToNextProtos(tc.config.NextProtos)
	return tc
}

// NewClientTLSFromCert constructs TLS credentials from the provided root
// certificate authority certificate(s) to validate server connections. If
// certificates to establish the identity of the client need to be included in
// the credentials (eg: for mTLS), use NewTLS instead, where a complete
// tls.Config can be specified.
// serverNameOverride is for testing only. If set to a non empty string,
// it will override the virtual host name of authority (e.g. :authority header
// field) in requests.
func NewClientTLSFromCert(cp *x509.CertPool, serverNameOverride string) TransportCredentials {
	return NewTLS(&tls.Config{ServerName: serverNameOverride, RootCAs: cp})
}

// NewClientTLSFromFile constructs TLS credentials from the provided root
// certificate authority certificate file(s) to validate server connections. If
// certificates to establish the identity of the client need to be included in
// the credentials (eg: for mTLS), use NewTLS instead, where a complete
// tls.Config can be specified.
// serverNameOverride is for testing only. If set to a non empty string,
// it will override the virtual host name of authority (e.g. :authority header
// field) in requests.
func NewClientTLSFromFile(certFile, serverNameOverride string) (TransportCredentials, error) {
	b, err := ioutil.ReadFile(certFile)
	if err != nil {
		return nil, err
	}
	cp := x509.NewCertPool()
	if !cp.AppendCertsFromPEM(b) {
		return nil, fmt.Errorf("credentials: failed to append certificates")
	}
	return NewTLS(&tls.Config{ServerName: serverNameOverride, RootCAs: cp}), nil
}

// NewServerTLSFromCert constructs TLS credentials from the input certificate for server.
func NewServerTLSFromCert(cert *tls.Certificate) TransportCredentials {
	return NewTLS(&tls.Config{Certificates: []tls.Certificate{*cert}})
}

// NewServerTLSFromFile constructs TLS credentials from the input certificate file and key
// file for server.
func NewServerTLSFromFile(certFile, keyFile string) (TransportCredentials, error) {
	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return nil, err
	}
	return NewTLS(&tls.Config{Certificates: []tls.Certificate{cert}}), nil
}

// TLSChannelzSecurityValue defines the struct that TLS protocol should return
// from GetSecurityValue(), containing security info like cipher and certificate used.
//
// This API is EXPERIMENTAL.
type TLSChannelzSecurityValue struct {
	ChannelzSecurityValue
	StandardName      string
	LocalCertificate  []byte
	RemoteCertificate []byte
}

var cipherSuiteLookup = map[uint16]string{
	tls.TLS_RSA_WITH_RC4_128_SHA:                "TLS_RSA_WITH_RC4_128_SHA",
	tls.TLS_RSA_WITH_3DES_EDE_CBC_SHA:           "TLS_RSA_WITH_3DES_EDE_CBC_SHA",
	tls.TLS_RSA_WITH_AES_128_CBC_SHA:            "TLS_RSA_WITH_AES_128_CBC_SHA",
	tls.TLS_RSA_WITH_AES_256_CBC_SHA:            "TLS_RSA_WITH_AES_256_CBC_SHA",
	tls.TLS_RSA_WITH_AES_128_GCM_SHA256:         "TLS_RSA_WITH_AES_128_GCM_SHA256",
	tls.TLS_RSA_WITH_AES_256_GCM_SHA384:         "TLS_RSA_WITH_AES_256_GCM_SHA384",
	tls.TLS_ECDHE_ECDSA_WITH_RC4_128_SHA:        "TLS_ECDHE_ECDSA_WITH_RC4_128_SHA",
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA:    "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA",
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA:    "TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA",
	tls.TLS_ECDHE_RSA_WITH_RC4_128_SHA:          "TLS_ECDHE_RSA_WITH_RC4_128_SHA",
	tls.TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA:     "TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA",
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA:      "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA",
	tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA:      "TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA",
	tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256:   "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256",
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256: "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256",
	tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384:   "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384",
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384: "TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384",
	tls.TLS_FALLBACK_SCSV:                       "TLS_FALLBACK_SCSV",
	tls.TLS_RSA_WITH_AES_128_CBC_SHA256:         "TLS_RSA_WITH_AES_128_CBC_SHA256",
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256: "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256",
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256:   "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256",
	tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305:    "TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305",
	tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305:  "TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305",
}

// cloneTLSConfig returns a shallow clone of the exported
// fields of cfg, ignoring the unexported sync.Once, which
// contains a mutex and must not be copied.
//
// If cfg is nil, a new zero tls.Config is returned.
//
// TODO: inline this function if possible.
func cloneTLSConfig(cfg *tls.Config) *tls.Config {
	if cfg == nil {
		return &tls.Config{}
	}

	return cfg.Clone()
}
<|endoftext|>


<commit_before>package dockergoal

import (
	"errors"
	"regexp"
	"testing"

	"github.com/facebookgo/ensure"
	"github.com/facebookgo/stackerr"
	"github.com/samalba/dockerclient"
)

func TestNewContainerError(t *testing.T) {
	givenErr := errors.New("")
	c, err := NewContainer(func(*Container) error {
		return givenErr
	})
	ensure.True(t, err == givenErr)
	ensure.True(t, c == nil)
}

func TestContainerName(t *testing.T) {
	const name = "foo"
	c, err := NewContainer(ContainerName(name))
	ensure.Nil(t, err)
	ensure.DeepEqual(t, c.name, name)
}

func TestContainerNameMissing(t *testing.T) {
	c, err := NewContainer()
	ensure.DeepEqual(t, err, errNameMissing)
	ensure.True(t, c == nil)
}

func TestContainerRemoveExisting(t *testing.T) {
	c, err := NewContainer(
		ContainerName("x"),
		ContainerRemoveExisting(),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.removeExisting)
}

func TestContainerForceRemoveExisting(t *testing.T) {
	c, err := NewContainer(
		ContainerName("x"),
		ContainerForceRemoveExisting(),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.forceRemoveExisting)
}

func TestContainerConfig(t *testing.T) {
	config := &dockerclient.ContainerConfig{}
	c, err := NewContainer(
		ContainerName("x"),
		ContainerConfig(config),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.containerConfig == config)
}

func TestContainerHostConfig(t *testing.T) {
	config := &dockerclient.HostConfig{}
	c, err := NewContainer(
		ContainerName("x"),
		ContainerHostConfig(config),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.hostConfig == config)
}

func TestContainerAuthConfig(t *testing.T) {
	config := &dockerclient.AuthConfig{}
	c, err := NewContainer(
		ContainerName("x"),
		ContainerAuthConfig(config),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.authConfig == config)
}

func TestApplyMakesNew(t *testing.T) {
	const givenName = "x"
	const givenID = "y"
	givenContainerConfig := &dockerclient.ContainerConfig{Image: "foo"}
	givenHostConfig := &dockerclient.HostConfig{}
	var inspectCalls, createCalls, startCalls int
	container, err := NewContainer(
		ContainerName(givenName),
		ContainerConfig(givenContainerConfig),
		ContainerHostConfig(givenHostConfig),
	)
	ensure.Nil(t, err)
	client := &mockClient{
		inspectContainer: func(name string) (*dockerclient.ContainerInfo, error) {
			inspectCalls++
			switch inspectCalls {
			case 1:
				ensure.DeepEqual(t, name, givenName)
				return nil, dockerclient.ErrNotFound
			case 2:
				return &dockerclient.ContainerInfo{Id: givenID}, nil
			}
			panic("not reached")
		},
		createContainer: func(config *dockerclient.ContainerConfig, name string) (string, error) {
			createCalls++
			ensure.True(t, config == givenContainerConfig)
			ensure.DeepEqual(t, name, givenName)
			return "", nil
		},
		startContainer: func(id string, config *dockerclient.HostConfig) error {
			startCalls++
			ensure.DeepEqual(t, id, givenID)
			ensure.True(t, config == givenHostConfig)
			return nil
		},
	}
	ensure.Nil(t, container.Apply(client))
	ensure.DeepEqual(t, inspectCalls, 2)
	ensure.DeepEqual(t, createCalls, 1)
	ensure.DeepEqual(t, startCalls, 1)
}

func TestApplyForceRemoveExisting(t *testing.T) {
	const removeID = "y"
	const newID = "z"
	givenContainerConfig := &dockerclient.ContainerConfig{Image: "foo"}
	var removeCalls, inspectCalls, createCalls, startCalls int
	container, err := NewContainer(
		ContainerName("x"),
		ContainerConfig(givenContainerConfig),
		ContainerForceRemoveExisting(),
	)
	ensure.Nil(t, err)
	client := &mockClient{
		removeContainer: func(id string, force, volumes bool) error {
			removeCalls++
			ensure.DeepEqual(t, id, removeID)
			ensure.True(t, force)
			ensure.False(t, volumes)
			return nil
		},
		inspectContainer: func(name string) (*dockerclient.ContainerInfo, error) {
			inspectCalls++
			switch inspectCalls {
			case 1:
				return &dockerclient.ContainerInfo{Id: removeID}, nil
			case 2:
				return &dockerclient.ContainerInfo{Id: newID}, nil
			}
			panic("not reached")
		},
		createContainer: func(config *dockerclient.ContainerConfig, name string) (string, error) {
			createCalls++
			return "", nil
		},
		startContainer: func(id string, config *dockerclient.HostConfig) error {
			startCalls++
			ensure.DeepEqual(t, id, newID)
			return nil
		},
	}
	ensure.Nil(t, container.Apply(client))
	ensure.DeepEqual(t, removeCalls, 1)
	ensure.DeepEqual(t, inspectCalls, 2)
	ensure.DeepEqual(t, createCalls, 1)
	ensure.DeepEqual(t, startCalls, 1)
}

func TestApplyForceRemoveExistingError(t *testing.T) {
	container, err := NewContainer(
		ContainerName("x"),
		ContainerConfig(&dockerclient.ContainerConfig{Image: "foo"}),
		ContainerForceRemoveExisting(),
	)
	ensure.Nil(t, err)
	givenErr := errors.New("")
	client := &mockClient{
		removeContainer: func(id string, force, volumes bool) error {
			return givenErr
		},
		inspectContainer: func(name string) (*dockerclient.ContainerInfo, error) {
			return &dockerclient.ContainerInfo{Id: "x"}, nil
		},
	}
	err = container.Apply(client)
	ensure.True(t, stackerr.HasUnderlying(err, stackerr.Equals(givenErr)))
}

func TestApplyForceRemoveExistingWhenNotFound(t *testing.T) {
	var inspectCalls int
	container, err := NewContainer(
		ContainerName("x"),
		ContainerConfig(&dockerclient.ContainerConfig{Image: "foo"}),
		ContainerForceRemoveExisting(),
	)
	ensure.Nil(t, err)
	client := &mockClient{
		inspectContainer: func(name string) (*dockerclient.ContainerInfo, error) {
			inspectCalls++
			switch inspectCalls {
			case 1:
				return nil, dockerclient.ErrNotFound
			case 2:
				return &dockerclient.ContainerInfo{Id: "y"}, nil
			}
			panic("not reached")
		},
		createContainer: func(config *dockerclient.ContainerConfig, name string) (string, error) {
			return "", nil
		},
		startContainer: func(id string, config *dockerclient.HostConfig) error {
			return nil
		},
	}
	ensure.Nil(t, container.Apply(client))
}

func TestCheckRunningWithDesiredImage(t *testing.T) {
	const image = "x"
	const id = "y"
	container := &Container{
		containerConfig: &dockerclient.ContainerConfig{
			Image: image,
		},
	}
	client := &mockClient{
		listImages: func() ([]*dockerclient.Image, error) {
			return []*dockerclient.Image{
				{
					RepoTags: []string{image},
					Id:       id,
				},
			}, nil
		},
	}
	ok, err := container.checkRunning(client, &dockerclient.ContainerInfo{Image: id})
	ensure.Nil(t, err)
	ensure.True(t, ok)
}

func TestCheckRunningWithoutDesiredImageAndNoRemoveExisting(t *testing.T) {
	const image = "x"
	container := &Container{
		containerConfig: &dockerclient.ContainerConfig{
			Image: image,
		},
	}
	client := &mockClient{
		listImages: func() ([]*dockerclient.Image, error) {
			return []*dockerclient.Image{
				{
					RepoTags: []string{image},
					Id:       "y",
				},
			}, nil
		},
	}
	ok, err := container.checkRunning(client, &dockerclient.ContainerInfo{Image: "z"})
	ensure.Err(t, err, regexp.MustCompile("but desired image is"))
	ensure.False(t, ok)
}

func TestCheckRunningImageIdentifyError(t *testing.T) {
	givenErr := errors.New("")
	const image = "x"
	container := &Container{
		containerConfig: &dockerclient.ContainerConfig{
			Image: image,
		},
	}
	client := &mockClient{
		listImages: func() ([]*dockerclient.Image, error) {
			return nil, givenErr
		},
	}
	ok, err := container.checkRunning(client, &dockerclient.ContainerInfo{Image: "z"})
	ensure.True(t, stackerr.HasUnderlying(err, stackerr.Equals(givenErr)))
	ensure.False(t, ok)
}

func TestCheckRunningWithoutDesiredImageWithRemoveExisting(t *testing.T) {
	const image = "x"
	const removeID = "y"
	var removeCalls int
	container := &Container{
		containerConfig: &dockerclient.ContainerConfig{
			Image: image,
		},
		removeExisting: true,
	}
	client := &mockClient{
		listImages: func() ([]*dockerclient.Image, error) {
			return []*dockerclient.Image{
				{
					RepoTags: []string{image},
					Id:       "y",
				},
			}, nil
		},
		removeContainer: func(id string, force, volumes bool) error {
			removeCalls++
			ensure.DeepEqual(t, id, removeID)
			return nil
		},
	}
	ci := &dockerclient.ContainerInfo{
		Image: "z",
		Id:    removeID,
	}
	ok, err := container.checkRunning(client, ci)
	ensure.Nil(t, err)
	ensure.False(t, ok)
	ensure.DeepEqual(t, removeCalls, 1)
}
<commit_msg>TestCheckRunningRemoveError<commit_after>package dockergoal

import (
	"errors"
	"regexp"
	"testing"

	"github.com/facebookgo/ensure"
	"github.com/facebookgo/stackerr"
	"github.com/samalba/dockerclient"
)

func TestNewContainerError(t *testing.T) {
	givenErr := errors.New("")
	c, err := NewContainer(func(*Container) error {
		return givenErr
	})
	ensure.True(t, err == givenErr)
	ensure.True(t, c == nil)
}

func TestContainerName(t *testing.T) {
	const name = "foo"
	c, err := NewContainer(ContainerName(name))
	ensure.Nil(t, err)
	ensure.DeepEqual(t, c.name, name)
}

func TestContainerNameMissing(t *testing.T) {
	c, err := NewContainer()
	ensure.DeepEqual(t, err, errNameMissing)
	ensure.True(t, c == nil)
}

func TestContainerRemoveExisting(t *testing.T) {
	c, err := NewContainer(
		ContainerName("x"),
		ContainerRemoveExisting(),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.removeExisting)
}

func TestContainerForceRemoveExisting(t *testing.T) {
	c, err := NewContainer(
		ContainerName("x"),
		ContainerForceRemoveExisting(),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.forceRemoveExisting)
}

func TestContainerConfig(t *testing.T) {
	config := &dockerclient.ContainerConfig{}
	c, err := NewContainer(
		ContainerName("x"),
		ContainerConfig(config),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.containerConfig == config)
}

func TestContainerHostConfig(t *testing.T) {
	config := &dockerclient.HostConfig{}
	c, err := NewContainer(
		ContainerName("x"),
		ContainerHostConfig(config),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.hostConfig == config)
}

func TestContainerAuthConfig(t *testing.T) {
	config := &dockerclient.AuthConfig{}
	c, err := NewContainer(
		ContainerName("x"),
		ContainerAuthConfig(config),
	)
	ensure.Nil(t, err)
	ensure.True(t, c.authConfig == config)
}

func TestApplyMakesNew(t *testing.T) {
	const givenName = "x"
	const givenID = "y"
	givenContainerConfig := &dockerclient.ContainerConfig{Image: "foo"}
	givenHostConfig := &dockerclient.HostConfig{}
	var inspectCalls, createCalls, startCalls int
	container, err := NewContainer(
		ContainerName(givenName),
		ContainerConfig(givenContainerConfig),
		ContainerHostConfig(givenHostConfig),
	)
	ensure.Nil(t, err)
	client := &mockClient{
		inspectContainer: func(name string) (*dockerclient.ContainerInfo, error) {
			inspectCalls++
			switch inspectCalls {
			case 1:
				ensure.DeepEqual(t, name, givenName)
				return nil, dockerclient.ErrNotFound
			case 2:
				return &dockerclient.ContainerInfo{Id: givenID}, nil
			}
			panic("not reached")
		},
		createContainer: func(config *dockerclient.ContainerConfig, name string) (string, error) {
			createCalls++
			ensure.True(t, config == givenContainerConfig)
			ensure.DeepEqual(t, name, givenName)
			return "", nil
		},
		startContainer: func(id string, config *dockerclient.HostConfig) error {
			startCalls++
			ensure.DeepEqual(t, id, givenID)
			ensure.True(t, config == givenHostConfig)
			return nil
		},
	}
	ensure.Nil(t, container.Apply(client))
	ensure.DeepEqual(t, inspectCalls, 2)
	ensure.DeepEqual(t, createCalls, 1)
	ensure.DeepEqual(t, startCalls, 1)
}

func TestApplyForceRemoveExisting(t *testing.T) {
	const removeID = "y"
	const newID = "z"
	givenContainerConfig := &dockerclient.ContainerConfig{Image: "foo"}
	var removeCalls, inspectCalls, createCalls, startCalls int
	container, err := NewContainer(
		ContainerName("x"),
		ContainerConfig(givenContainerConfig),
		ContainerForceRemoveExisting(),
	)
	ensure.Nil(t, err)
	client := &mockClient{
		removeContainer: func(id string, force, volumes bool) error {
			removeCalls++
			ensure.DeepEqual(t, id, removeID)
			ensure.True(t, force)
			ensure.False(t, volumes)
			return nil
		},
		inspectContainer: func(name string) (*dockerclient.ContainerInfo, error) {
			inspectCalls++
			switch inspectCalls {
			case 1:
				return &dockerclient.ContainerInfo{Id: removeID}, nil
			case 2:
				return &dockerclient.ContainerInfo{Id: newID}, nil
			}
			panic("not reached")
		},
		createContainer: func(config *dockerclient.ContainerConfig, name string) (string, error) {
			createCalls++
			return "", nil
		},
		startContainer: func(id string, config *dockerclient.HostConfig) error {
			startCalls++
			ensure.DeepEqual(t, id, newID)
			return nil
		},
	}
	ensure.Nil(t, container.Apply(client))
	ensure.DeepEqual(t, removeCalls, 1)
	ensure.DeepEqual(t, inspectCalls, 2)
	ensure.DeepEqual(t, createCalls, 1)
	ensure.DeepEqual(t, startCalls, 1)
}

func TestApplyForceRemoveExistingError(t *testing.T) {
	container, err := NewContainer(
		ContainerName("x"),
		ContainerConfig(&dockerclient.ContainerConfig{Image: "foo"}),
		ContainerForceRemoveExisting(),
	)
	ensure.Nil(t, err)
	givenErr := errors.New("")
	client := &mockClient{
		removeContainer: func(id string, force, volumes bool) error {
			return givenErr
		},
		inspectContainer: func(name string) (*dockerclient.ContainerInfo, error) {
			return &dockerclient.ContainerInfo{Id: "x"}, nil
		},
	}
	err = container.Apply(client)
	ensure.True(t, stackerr.HasUnderlying(err, stackerr.Equals(givenErr)))
}

func TestApplyForceRemoveExistingWhenNotFound(t *testing.T) {
	var inspectCalls int
	container, err := NewContainer(
		ContainerName("x"),
		ContainerConfig(&dockerclient.ContainerConfig{Image: "foo"}),
		ContainerForceRemoveExisting(),
	)
	ensure.Nil(t, err)
	client := &mockClient{
		inspectContainer: func(name string) (*dockerclient.ContainerInfo, error) {
			inspectCalls++
			switch inspectCalls {
			case 1:
				return nil, dockerclient.ErrNotFound
			case 2:
				return &dockerclient.ContainerInfo{Id: "y"}, nil
			}
			panic("not reached")
		},
		createContainer: func(config *dockerclient.ContainerConfig, name string) (string, error) {
			return "", nil
		},
		startContainer: func(id string, config *dockerclient.HostConfig) error {
			return nil
		},
	}
	ensure.Nil(t, container.Apply(client))
}

func TestCheckRunningWithDesiredImage(t *testing.T) {
	const image = "x"
	const id = "y"
	container := &Container{
		containerConfig: &dockerclient.ContainerConfig{
			Image: image,
		},
	}
	client := &mockClient{
		listImages: func() ([]*dockerclient.Image, error) {
			return []*dockerclient.Image{
				{
					RepoTags: []string{image},
					Id:       id,
				},
			}, nil
		},
	}
	ok, err := container.checkRunning(client, &dockerclient.ContainerInfo{Image: id})
	ensure.Nil(t, err)
	ensure.True(t, ok)
}

func TestCheckRunningWithoutDesiredImageAndNoRemoveExisting(t *testing.T) {
	const image = "x"
	container := &Container{
		containerConfig: &dockerclient.ContainerConfig{
			Image: image,
		},
	}
	client := &mockClient{
		listImages: func() ([]*dockerclient.Image, error) {
			return []*dockerclient.Image{
				{
					RepoTags: []string{image},
					Id:       "y",
				},
			}, nil
		},
	}
	ok, err := container.checkRunning(client, &dockerclient.ContainerInfo{Image: "z"})
	ensure.Err(t, err, regexp.MustCompile("but desired image is"))
	ensure.False(t, ok)
}

func TestCheckRunningImageIdentifyError(t *testing.T) {
	givenErr := errors.New("")
	const image = "x"
	container := &Container{
		containerConfig: &dockerclient.ContainerConfig{
			Image: image,
		},
	}
	client := &mockClient{
		listImages: func() ([]*dockerclient.Image, error) {
			return nil, givenErr
		},
	}
	ok, err := container.checkRunning(client, &dockerclient.ContainerInfo{Image: "z"})
	ensure.True(t, stackerr.HasUnderlying(err, stackerr.Equals(givenErr)))
	ensure.False(t, ok)
}

func TestCheckRunningWithoutDesiredImageWithRemoveExisting(t *testing.T) {
	const image = "x"
	const removeID = "y"
	var removeCalls int
	container := &Container{
		containerConfig: &dockerclient.ContainerConfig{
			Image: image,
		},
		removeExisting: true,
	}
	client := &mockClient{
		listImages: func() ([]*dockerclient.Image, error) {
			return []*dockerclient.Image{
				{
					RepoTags: []string{image},
					Id:       "y",
				},
			}, nil
		},
		removeContainer: func(id string, force, volumes bool) error {
			removeCalls++
			ensure.DeepEqual(t, id, removeID)
			return nil
		},
	}
	ci := &dockerclient.ContainerInfo{
		Image: "z",
		Id:    removeID,
	}
	ok, err := container.checkRunning(client, ci)
	ensure.Nil(t, err)
	ensure.False(t, ok)
	ensure.DeepEqual(t, removeCalls, 1)
}

func TestCheckRunningRemovError(t *testing.T) {
	const image = "x"
	givenErr := errors.New("")
	container := &Container{
		containerConfig: &dockerclient.ContainerConfig{
			Image: image,
		},
		removeExisting: true,
	}
	client := &mockClient{
		listImages: func() ([]*dockerclient.Image, error) {
			return []*dockerclient.Image{
				{
					RepoTags: []string{image},
					Id:       "y",
				},
			}, nil
		},
		removeContainer: func(id string, force, volumes bool) error {
			return givenErr
		},
	}
	ci := &dockerclient.ContainerInfo{
		Id:    "y",
		Image: "z",
	}
	ok, err := container.checkRunning(client, ci)
	ensure.True(t, stackerr.HasUnderlying(err, stackerr.Equals(givenErr)))
	ensure.False(t, ok)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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 provider

import (
	"context"
	"fmt"
	"os"
	"sort"
	"strconv"
	"strings"

	cloudflare "github.com/cloudflare/cloudflare-go"
	log "github.com/sirupsen/logrus"

	"github.com/kubernetes-incubator/external-dns/endpoint"
	"github.com/kubernetes-incubator/external-dns/plan"
	"github.com/kubernetes-incubator/external-dns/source"
)

const (
	// cloudFlareCreate is a ChangeAction enum value
	cloudFlareCreate = "CREATE"
	// cloudFlareDelete is a ChangeAction enum value
	cloudFlareDelete = "DELETE"
	// cloudFlareUpdate is a ChangeAction enum value
	cloudFlareUpdate = "UPDATE"
	// defaultCloudFlareRecordTTL 1 = automatic
	defaultCloudFlareRecordTTL = 1
)

var cloudFlareTypeNotSupported = map[string]bool{
	"LOC": true,
	"MX":  true,
	"NS":  true,
	"SPF": true,
	"TXT": true,
	"SRV": true,
}

// cloudFlareDNS is the subset of the CloudFlare API that we actually use.  Add methods as required. Signatures must match exactly.
type cloudFlareDNS interface {
	UserDetails() (cloudflare.User, error)
	ZoneIDByName(zoneName string) (string, error)
	ListZones(zoneID ...string) ([]cloudflare.Zone, error)
	ListZonesContext(ctx context.Context, opts ...cloudflare.ReqOption) (cloudflare.ZonesResponse, error)
	DNSRecords(zoneID string, rr cloudflare.DNSRecord) ([]cloudflare.DNSRecord, error)
	CreateDNSRecord(zoneID string, rr cloudflare.DNSRecord) (*cloudflare.DNSRecordResponse, error)
	DeleteDNSRecord(zoneID, recordID string) error
	UpdateDNSRecord(zoneID, recordID string, rr cloudflare.DNSRecord) error
}

type zoneService struct {
	service *cloudflare.API
}

func (z zoneService) UserDetails() (cloudflare.User, error) {
	return z.service.UserDetails()
}

func (z zoneService) ListZones(zoneID ...string) ([]cloudflare.Zone, error) {
	return z.service.ListZones(zoneID...)
}

func (z zoneService) ZoneIDByName(zoneName string) (string, error) {
	return z.service.ZoneIDByName(zoneName)
}

func (z zoneService) CreateDNSRecord(zoneID string, rr cloudflare.DNSRecord) (*cloudflare.DNSRecordResponse, error) {
	return z.service.CreateDNSRecord(zoneID, rr)
}

func (z zoneService) DNSRecords(zoneID string, rr cloudflare.DNSRecord) ([]cloudflare.DNSRecord, error) {
	return z.service.DNSRecords(zoneID, rr)
}
func (z zoneService) UpdateDNSRecord(zoneID, recordID string, rr cloudflare.DNSRecord) error {
	return z.service.UpdateDNSRecord(zoneID, recordID, rr)
}
func (z zoneService) DeleteDNSRecord(zoneID, recordID string) error {
	return z.service.DeleteDNSRecord(zoneID, recordID)
}

func (z zoneService) ListZonesContext(ctx context.Context, opts ...cloudflare.ReqOption) (cloudflare.ZonesResponse, error) {
	return z.service.ListZonesContext(ctx, opts...)
}

// CloudFlareProvider is an implementation of Provider for CloudFlare DNS.
type CloudFlareProvider struct {
	Client cloudFlareDNS
	// only consider hosted zones managing domains ending in this suffix
	domainFilter      DomainFilter
	zoneIDFilter      ZoneIDFilter
	proxiedByDefault  bool
	DryRun            bool
	PaginationOptions cloudflare.PaginationOptions
}

// cloudFlareChange differentiates between ChangActions
type cloudFlareChange struct {
	Action            string
	ResourceRecordSet []cloudflare.DNSRecord
}

// NewCloudFlareProvider initializes a new CloudFlare DNS based Provider.
func NewCloudFlareProvider(domainFilter DomainFilter, zoneIDFilter ZoneIDFilter, zonesPerPage int, proxiedByDefault bool, dryRun bool) (*CloudFlareProvider, error) {
	// initialize via chosen auth method and returns new API object
	var (
    config *cloudflare.API
	  err error
  )
	if os.Getenv("CF_API_TOKEN") != "" {
		config, err = cloudflare.NewWithAPIToken(os.Getenv("CF_API_TOKEN"))
	} else {
		config, err = cloudflare.New(os.Getenv("CF_API_KEY"), os.Getenv("CF_API_EMAIL"))
	}
	if err != nil {
		return nil, fmt.Errorf("failed to initialize cloudflare provider: %v", err)
	}
	provider := &CloudFlareProvider{
		//Client: config,
		Client:           zoneService{config},
		domainFilter:     domainFilter,
		zoneIDFilter:     zoneIDFilter,
		proxiedByDefault: proxiedByDefault,
		DryRun:           dryRun,
		PaginationOptions: cloudflare.PaginationOptions{
			PerPage: zonesPerPage,
			Page:    1,
		},
	}
	return provider, nil
}

// Zones returns the list of hosted zones.
func (p *CloudFlareProvider) Zones() ([]cloudflare.Zone, error) {
	result := []cloudflare.Zone{}
	ctx := context.TODO()
	p.PaginationOptions.Page = 1

	for {
		zonesResponse, err := p.Client.ListZonesContext(ctx, cloudflare.WithPagination(p.PaginationOptions))
		if err != nil {
			return nil, err
		}

		for _, zone := range zonesResponse.Result {
			if !p.domainFilter.Match(zone.Name) {
				continue
			}

			if !p.zoneIDFilter.Match(zone.ID) {
				continue
			}
			result = append(result, zone)
		}
		if p.PaginationOptions.Page == zonesResponse.ResultInfo.TotalPages {
			break
		}
		p.PaginationOptions.Page++
	}
	return result, nil
}

// Records returns the list of records.
func (p *CloudFlareProvider) Records() ([]*endpoint.Endpoint, error) {
	zones, err := p.Zones()
	if err != nil {
		return nil, err
	}

	endpoints := []*endpoint.Endpoint{}
	for _, zone := range zones {
		records, err := p.Client.DNSRecords(zone.ID, cloudflare.DNSRecord{})
		if err != nil {
			return nil, err
		}

		// As CloudFlare does not support "sets" of targets, but instead returns
		// a single entry for each name/type/target, we have to group by name
		// and record to allow the planner to calculate the correct plan. See #992.
		endpoints = append(endpoints, groupByNameAndType(records)...)
	}

	return endpoints, nil
}

// ApplyChanges applies a given set of changes in a given zone.
func (p *CloudFlareProvider) ApplyChanges(ctx context.Context, changes *plan.Changes) error {
	proxiedByDefault := p.proxiedByDefault

	combinedChanges := make([]*cloudFlareChange, 0, len(changes.Create)+len(changes.UpdateNew)+len(changes.Delete))

	combinedChanges = append(combinedChanges, newCloudFlareChanges(cloudFlareCreate, changes.Create, proxiedByDefault)...)
	combinedChanges = append(combinedChanges, newCloudFlareChanges(cloudFlareUpdate, changes.UpdateNew, proxiedByDefault)...)
	combinedChanges = append(combinedChanges, newCloudFlareChanges(cloudFlareDelete, changes.Delete, proxiedByDefault)...)

	return p.submitChanges(combinedChanges)
}

// submitChanges takes a zone and a collection of Changes and sends them as a single transaction.
func (p *CloudFlareProvider) submitChanges(changes []*cloudFlareChange) error {
	// return early if there is nothing to change
	if len(changes) == 0 {
		return nil
	}

	zones, err := p.Zones()
	if err != nil {
		return err
	}
	// separate into per-zone change sets to be passed to the API.
	changesByZone := p.changesByZone(zones, changes)

	for zoneID, changes := range changesByZone {
		records, err := p.Client.DNSRecords(zoneID, cloudflare.DNSRecord{})
		if err != nil {
			return fmt.Errorf("could not fetch records from zone, %v", err)
		}
		for _, change := range changes {
			logFields := log.Fields{
				"record":  change.ResourceRecordSet[0].Name,
				"type":    change.ResourceRecordSet[0].Type,
				"ttl":     change.ResourceRecordSet[0].TTL,
				"targets": len(change.ResourceRecordSet),
				"action":  change.Action,
				"zone":    zoneID,
			}

			log.WithFields(logFields).Info("Changing record.")

			if p.DryRun {
				continue
			}

			recordIDs := p.getRecordIDs(records, change.ResourceRecordSet[0])

			// to simplify bookkeeping for multiple records, an update is executed as delete+create
			if change.Action == cloudFlareDelete || change.Action == cloudFlareUpdate {
				for _, recordID := range recordIDs {
					err := p.Client.DeleteDNSRecord(zoneID, recordID)
					if err != nil {
						log.WithFields(logFields).Errorf("failed to delete record: %v", err)
					}
				}
			}

			if change.Action == cloudFlareCreate || change.Action == cloudFlareUpdate {
				for _, record := range change.ResourceRecordSet {
					_, err := p.Client.CreateDNSRecord(zoneID, record)
					if err != nil {
						log.WithFields(logFields).Errorf("failed to create record: %v", err)
					}
				}
			}
		}
	}
	return nil
}

// changesByZone separates a multi-zone change into a single change per zone.
func (p *CloudFlareProvider) changesByZone(zones []cloudflare.Zone, changeSet []*cloudFlareChange) map[string][]*cloudFlareChange {
	changes := make(map[string][]*cloudFlareChange)
	zoneNameIDMapper := zoneIDName{}

	for _, z := range zones {
		zoneNameIDMapper.Add(z.ID, z.Name)
		changes[z.ID] = []*cloudFlareChange{}
	}

	for _, c := range changeSet {
		zoneID, _ := zoneNameIDMapper.FindZone(c.ResourceRecordSet[0].Name)
		if zoneID == "" {
			log.Debugf("Skipping record %s because no hosted zone matching record DNS Name was detected ", c.ResourceRecordSet[0].Name)
			continue
		}
		changes[zoneID] = append(changes[zoneID], c)
	}

	return changes
}

func (p *CloudFlareProvider) getRecordIDs(records []cloudflare.DNSRecord, record cloudflare.DNSRecord) []string {
	recordIDs := make([]string, 0)
	for _, zoneRecord := range records {
		if zoneRecord.Name == record.Name && zoneRecord.Type == record.Type {
			recordIDs = append(recordIDs, zoneRecord.ID)
		}
	}
	sort.Strings(recordIDs)
	return recordIDs
}

// newCloudFlareChanges returns a collection of Changes based on the given records and action.
func newCloudFlareChanges(action string, endpoints []*endpoint.Endpoint, proxiedByDefault bool) []*cloudFlareChange {
	changes := make([]*cloudFlareChange, 0, len(endpoints))

	for _, endpoint := range endpoints {
		changes = append(changes, newCloudFlareChange(action, endpoint, proxiedByDefault))
	}

	return changes
}

func newCloudFlareChange(action string, endpoint *endpoint.Endpoint, proxiedByDefault bool) *cloudFlareChange {
	ttl := defaultCloudFlareRecordTTL
	proxied := shouldBeProxied(endpoint, proxiedByDefault)

	if endpoint.RecordTTL.IsConfigured() {
		ttl = int(endpoint.RecordTTL)
	}

	resourceRecordSet := make([]cloudflare.DNSRecord, len(endpoint.Targets))

	for i := range endpoint.Targets {
		resourceRecordSet[i] = cloudflare.DNSRecord{
			Name:    endpoint.DNSName,
			TTL:     ttl,
			Proxied: proxied,
			Type:    endpoint.RecordType,
			Content: endpoint.Targets[i],
		}
	}

	return &cloudFlareChange{
		Action:            action,
		ResourceRecordSet: resourceRecordSet,
	}
}

func shouldBeProxied(endpoint *endpoint.Endpoint, proxiedByDefault bool) bool {
	proxied := proxiedByDefault

	for _, v := range endpoint.ProviderSpecific {
		if v.Name == source.CloudflareProxiedKey {
			b, err := strconv.ParseBool(v.Value)
			if err != nil {
				log.Errorf("Failed to parse annotation [%s]: %v", source.CloudflareProxiedKey, err)
			} else {
				proxied = b
			}
			break
		}
	}

	if cloudFlareTypeNotSupported[endpoint.RecordType] || strings.Contains(endpoint.DNSName, "*") {
		proxied = false
	}
	return proxied
}

func groupByNameAndType(records []cloudflare.DNSRecord) []*endpoint.Endpoint {
	endpoints := []*endpoint.Endpoint{}

	// group supported records by name and type
	groups := map[string][]cloudflare.DNSRecord{}

	for _, r := range records {
		if !supportedRecordType(r.Type) {
			continue
		}

		groupBy := r.Name + r.Type
		if _, ok := groups[groupBy]; !ok {
			groups[groupBy] = []cloudflare.DNSRecord{}
		}

		groups[groupBy] = append(groups[groupBy], r)
	}

	// create single endpoint with all the targets for each name/type
	for _, records := range groups {
		targets := make([]string, len(records))
		for i, record := range records {
			targets[i] = record.Content
		}
		endpoints = append(endpoints,
			endpoint.NewEndpointWithTTL(
				records[0].Name,
				records[0].Type,
				endpoint.TTL(records[0].TTL),
				targets...).
				WithProviderSpecific(source.CloudflareProxiedKey, strconv.FormatBool(records[0].Proxied)))
	}

	return endpoints
}
<commit_msg>Formatting<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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 provider

import (
	"context"
	"fmt"
	"os"
	"sort"
	"strconv"
	"strings"

	cloudflare "github.com/cloudflare/cloudflare-go"
	log "github.com/sirupsen/logrus"

	"github.com/kubernetes-incubator/external-dns/endpoint"
	"github.com/kubernetes-incubator/external-dns/plan"
	"github.com/kubernetes-incubator/external-dns/source"
)

const (
	// cloudFlareCreate is a ChangeAction enum value
	cloudFlareCreate = "CREATE"
	// cloudFlareDelete is a ChangeAction enum value
	cloudFlareDelete = "DELETE"
	// cloudFlareUpdate is a ChangeAction enum value
	cloudFlareUpdate = "UPDATE"
	// defaultCloudFlareRecordTTL 1 = automatic
	defaultCloudFlareRecordTTL = 1
)

var cloudFlareTypeNotSupported = map[string]bool{
	"LOC": true,
	"MX":  true,
	"NS":  true,
	"SPF": true,
	"TXT": true,
	"SRV": true,
}

// cloudFlareDNS is the subset of the CloudFlare API that we actually use.  Add methods as required. Signatures must match exactly.
type cloudFlareDNS interface {
	UserDetails() (cloudflare.User, error)
	ZoneIDByName(zoneName string) (string, error)
	ListZones(zoneID ...string) ([]cloudflare.Zone, error)
	ListZonesContext(ctx context.Context, opts ...cloudflare.ReqOption) (cloudflare.ZonesResponse, error)
	DNSRecords(zoneID string, rr cloudflare.DNSRecord) ([]cloudflare.DNSRecord, error)
	CreateDNSRecord(zoneID string, rr cloudflare.DNSRecord) (*cloudflare.DNSRecordResponse, error)
	DeleteDNSRecord(zoneID, recordID string) error
	UpdateDNSRecord(zoneID, recordID string, rr cloudflare.DNSRecord) error
}

type zoneService struct {
	service *cloudflare.API
}

func (z zoneService) UserDetails() (cloudflare.User, error) {
	return z.service.UserDetails()
}

func (z zoneService) ListZones(zoneID ...string) ([]cloudflare.Zone, error) {
	return z.service.ListZones(zoneID...)
}

func (z zoneService) ZoneIDByName(zoneName string) (string, error) {
	return z.service.ZoneIDByName(zoneName)
}

func (z zoneService) CreateDNSRecord(zoneID string, rr cloudflare.DNSRecord) (*cloudflare.DNSRecordResponse, error) {
	return z.service.CreateDNSRecord(zoneID, rr)
}

func (z zoneService) DNSRecords(zoneID string, rr cloudflare.DNSRecord) ([]cloudflare.DNSRecord, error) {
	return z.service.DNSRecords(zoneID, rr)
}
func (z zoneService) UpdateDNSRecord(zoneID, recordID string, rr cloudflare.DNSRecord) error {
	return z.service.UpdateDNSRecord(zoneID, recordID, rr)
}
func (z zoneService) DeleteDNSRecord(zoneID, recordID string) error {
	return z.service.DeleteDNSRecord(zoneID, recordID)
}

func (z zoneService) ListZonesContext(ctx context.Context, opts ...cloudflare.ReqOption) (cloudflare.ZonesResponse, error) {
	return z.service.ListZonesContext(ctx, opts...)
}

// CloudFlareProvider is an implementation of Provider for CloudFlare DNS.
type CloudFlareProvider struct {
	Client cloudFlareDNS
	// only consider hosted zones managing domains ending in this suffix
	domainFilter      DomainFilter
	zoneIDFilter      ZoneIDFilter
	proxiedByDefault  bool
	DryRun            bool
	PaginationOptions cloudflare.PaginationOptions
}

// cloudFlareChange differentiates between ChangActions
type cloudFlareChange struct {
	Action            string
	ResourceRecordSet []cloudflare.DNSRecord
}

// NewCloudFlareProvider initializes a new CloudFlare DNS based Provider.
func NewCloudFlareProvider(domainFilter DomainFilter, zoneIDFilter ZoneIDFilter, zonesPerPage int, proxiedByDefault bool, dryRun bool) (*CloudFlareProvider, error) {
	// initialize via chosen auth method and returns new API object
	var (
		config *cloudflare.API
		err    error
	)
	if os.Getenv("CF_API_TOKEN") != "" {
		config, err = cloudflare.NewWithAPIToken(os.Getenv("CF_API_TOKEN"))
	} else {
		config, err = cloudflare.New(os.Getenv("CF_API_KEY"), os.Getenv("CF_API_EMAIL"))
	}
	if err != nil {
		return nil, fmt.Errorf("failed to initialize cloudflare provider: %v", err)
	}
	provider := &CloudFlareProvider{
		//Client: config,
		Client:           zoneService{config},
		domainFilter:     domainFilter,
		zoneIDFilter:     zoneIDFilter,
		proxiedByDefault: proxiedByDefault,
		DryRun:           dryRun,
		PaginationOptions: cloudflare.PaginationOptions{
			PerPage: zonesPerPage,
			Page:    1,
		},
	}
	return provider, nil
}

// Zones returns the list of hosted zones.
func (p *CloudFlareProvider) Zones() ([]cloudflare.Zone, error) {
	result := []cloudflare.Zone{}
	ctx := context.TODO()
	p.PaginationOptions.Page = 1

	for {
		zonesResponse, err := p.Client.ListZonesContext(ctx, cloudflare.WithPagination(p.PaginationOptions))
		if err != nil {
			return nil, err
		}

		for _, zone := range zonesResponse.Result {
			if !p.domainFilter.Match(zone.Name) {
				continue
			}

			if !p.zoneIDFilter.Match(zone.ID) {
				continue
			}
			result = append(result, zone)
		}
		if p.PaginationOptions.Page == zonesResponse.ResultInfo.TotalPages {
			break
		}
		p.PaginationOptions.Page++
	}
	return result, nil
}

// Records returns the list of records.
func (p *CloudFlareProvider) Records() ([]*endpoint.Endpoint, error) {
	zones, err := p.Zones()
	if err != nil {
		return nil, err
	}

	endpoints := []*endpoint.Endpoint{}
	for _, zone := range zones {
		records, err := p.Client.DNSRecords(zone.ID, cloudflare.DNSRecord{})
		if err != nil {
			return nil, err
		}

		// As CloudFlare does not support "sets" of targets, but instead returns
		// a single entry for each name/type/target, we have to group by name
		// and record to allow the planner to calculate the correct plan. See #992.
		endpoints = append(endpoints, groupByNameAndType(records)...)
	}

	return endpoints, nil
}

// ApplyChanges applies a given set of changes in a given zone.
func (p *CloudFlareProvider) ApplyChanges(ctx context.Context, changes *plan.Changes) error {
	proxiedByDefault := p.proxiedByDefault

	combinedChanges := make([]*cloudFlareChange, 0, len(changes.Create)+len(changes.UpdateNew)+len(changes.Delete))

	combinedChanges = append(combinedChanges, newCloudFlareChanges(cloudFlareCreate, changes.Create, proxiedByDefault)...)
	combinedChanges = append(combinedChanges, newCloudFlareChanges(cloudFlareUpdate, changes.UpdateNew, proxiedByDefault)...)
	combinedChanges = append(combinedChanges, newCloudFlareChanges(cloudFlareDelete, changes.Delete, proxiedByDefault)...)

	return p.submitChanges(combinedChanges)
}

// submitChanges takes a zone and a collection of Changes and sends them as a single transaction.
func (p *CloudFlareProvider) submitChanges(changes []*cloudFlareChange) error {
	// return early if there is nothing to change
	if len(changes) == 0 {
		return nil
	}

	zones, err := p.Zones()
	if err != nil {
		return err
	}
	// separate into per-zone change sets to be passed to the API.
	changesByZone := p.changesByZone(zones, changes)

	for zoneID, changes := range changesByZone {
		records, err := p.Client.DNSRecords(zoneID, cloudflare.DNSRecord{})
		if err != nil {
			return fmt.Errorf("could not fetch records from zone, %v", err)
		}
		for _, change := range changes {
			logFields := log.Fields{
				"record":  change.ResourceRecordSet[0].Name,
				"type":    change.ResourceRecordSet[0].Type,
				"ttl":     change.ResourceRecordSet[0].TTL,
				"targets": len(change.ResourceRecordSet),
				"action":  change.Action,
				"zone":    zoneID,
			}

			log.WithFields(logFields).Info("Changing record.")

			if p.DryRun {
				continue
			}

			recordIDs := p.getRecordIDs(records, change.ResourceRecordSet[0])

			// to simplify bookkeeping for multiple records, an update is executed as delete+create
			if change.Action == cloudFlareDelete || change.Action == cloudFlareUpdate {
				for _, recordID := range recordIDs {
					err := p.Client.DeleteDNSRecord(zoneID, recordID)
					if err != nil {
						log.WithFields(logFields).Errorf("failed to delete record: %v", err)
					}
				}
			}

			if change.Action == cloudFlareCreate || change.Action == cloudFlareUpdate {
				for _, record := range change.ResourceRecordSet {
					_, err := p.Client.CreateDNSRecord(zoneID, record)
					if err != nil {
						log.WithFields(logFields).Errorf("failed to create record: %v", err)
					}
				}
			}
		}
	}
	return nil
}

// changesByZone separates a multi-zone change into a single change per zone.
func (p *CloudFlareProvider) changesByZone(zones []cloudflare.Zone, changeSet []*cloudFlareChange) map[string][]*cloudFlareChange {
	changes := make(map[string][]*cloudFlareChange)
	zoneNameIDMapper := zoneIDName{}

	for _, z := range zones {
		zoneNameIDMapper.Add(z.ID, z.Name)
		changes[z.ID] = []*cloudFlareChange{}
	}

	for _, c := range changeSet {
		zoneID, _ := zoneNameIDMapper.FindZone(c.ResourceRecordSet[0].Name)
		if zoneID == "" {
			log.Debugf("Skipping record %s because no hosted zone matching record DNS Name was detected ", c.ResourceRecordSet[0].Name)
			continue
		}
		changes[zoneID] = append(changes[zoneID], c)
	}

	return changes
}

func (p *CloudFlareProvider) getRecordIDs(records []cloudflare.DNSRecord, record cloudflare.DNSRecord) []string {
	recordIDs := make([]string, 0)
	for _, zoneRecord := range records {
		if zoneRecord.Name == record.Name && zoneRecord.Type == record.Type {
			recordIDs = append(recordIDs, zoneRecord.ID)
		}
	}
	sort.Strings(recordIDs)
	return recordIDs
}

// newCloudFlareChanges returns a collection of Changes based on the given records and action.
func newCloudFlareChanges(action string, endpoints []*endpoint.Endpoint, proxiedByDefault bool) []*cloudFlareChange {
	changes := make([]*cloudFlareChange, 0, len(endpoints))

	for _, endpoint := range endpoints {
		changes = append(changes, newCloudFlareChange(action, endpoint, proxiedByDefault))
	}

	return changes
}

func newCloudFlareChange(action string, endpoint *endpoint.Endpoint, proxiedByDefault bool) *cloudFlareChange {
	ttl := defaultCloudFlareRecordTTL
	proxied := shouldBeProxied(endpoint, proxiedByDefault)

	if endpoint.RecordTTL.IsConfigured() {
		ttl = int(endpoint.RecordTTL)
	}

	resourceRecordSet := make([]cloudflare.DNSRecord, len(endpoint.Targets))

	for i := range endpoint.Targets {
		resourceRecordSet[i] = cloudflare.DNSRecord{
			Name:    endpoint.DNSName,
			TTL:     ttl,
			Proxied: proxied,
			Type:    endpoint.RecordType,
			Content: endpoint.Targets[i],
		}
	}

	return &cloudFlareChange{
		Action:            action,
		ResourceRecordSet: resourceRecordSet,
	}
}

func shouldBeProxied(endpoint *endpoint.Endpoint, proxiedByDefault bool) bool {
	proxied := proxiedByDefault

	for _, v := range endpoint.ProviderSpecific {
		if v.Name == source.CloudflareProxiedKey {
			b, err := strconv.ParseBool(v.Value)
			if err != nil {
				log.Errorf("Failed to parse annotation [%s]: %v", source.CloudflareProxiedKey, err)
			} else {
				proxied = b
			}
			break
		}
	}

	if cloudFlareTypeNotSupported[endpoint.RecordType] || strings.Contains(endpoint.DNSName, "*") {
		proxied = false
	}
	return proxied
}

func groupByNameAndType(records []cloudflare.DNSRecord) []*endpoint.Endpoint {
	endpoints := []*endpoint.Endpoint{}

	// group supported records by name and type
	groups := map[string][]cloudflare.DNSRecord{}

	for _, r := range records {
		if !supportedRecordType(r.Type) {
			continue
		}

		groupBy := r.Name + r.Type
		if _, ok := groups[groupBy]; !ok {
			groups[groupBy] = []cloudflare.DNSRecord{}
		}

		groups[groupBy] = append(groups[groupBy], r)
	}

	// create single endpoint with all the targets for each name/type
	for _, records := range groups {
		targets := make([]string, len(records))
		for i, record := range records {
			targets[i] = record.Content
		}
		endpoints = append(endpoints,
			endpoint.NewEndpointWithTTL(
				records[0].Name,
				records[0].Type,
				endpoint.TTL(records[0].TTL),
				targets...).
				WithProviderSpecific(source.CloudflareProxiedKey, strconv.FormatBool(records[0].Proxied)))
	}

	return endpoints
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"encoding/gob"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/raft"
)

var (
	raftDir = flag.String("raftdir", "/tmp/r", "")
	peers   = flag.String("peers", "", "comma-separated host:port tuples")
	listen  = flag.String("listen", ":8000", "")
	join    = flag.String("join", "", "raft master to join")
	node    *raft.Raft
)

// trivial log store, writing one entry into one file each.
// fulfills the raft.LogStore interface.
type fancyLogStore struct {
	l         sync.RWMutex
	lowIndex  uint64
	highIndex uint64
}

func (s *fancyLogStore) FirstIndex() (uint64, error) {
	s.l.RLock()
	defer s.l.RUnlock()
	return s.lowIndex, nil
}

func (s *fancyLogStore) LastIndex() (uint64, error) {
	s.l.RLock()
	defer s.l.RUnlock()
	return s.highIndex, nil
}

func (s *fancyLogStore) GetLog(index uint64, rlog *raft.Log) error {
	s.l.Lock()
	defer s.l.Unlock()
	f, err := os.Open(filepath.Join(*raftDir, fmt.Sprintf("fancylogs/entry.%d", index)))
	if err != nil {
		return err
	}
	defer f.Close()

	var elog raft.Log
	if err := gob.NewDecoder(f).Decode(&elog); err != nil {
		return err
	}
	*rlog = elog
	return nil
}

func (s *fancyLogStore) StoreLog(log *raft.Log) error {
	return s.StoreLogs([]*raft.Log{log})
}

func (s *fancyLogStore) StoreLogs(logs []*raft.Log) error {
	s.l.Lock()
	defer s.l.Unlock()

	for _, entry := range logs {
		log.Printf("writing index %d to file (%v)\n", entry.Index, entry)
		f, err := os.Create(filepath.Join(*raftDir, fmt.Sprintf("fancylogs/entry.%d", entry.Index)))
		if err != nil {
			return err
		}
		defer f.Close()
		if err := gob.NewEncoder(f).Encode(entry); err != nil {
			return err
		}
	}

	return nil
}

func (s *fancyLogStore) DeleteRange(min, max uint64) error {
	s.l.Lock()
	defer s.l.Unlock()
	for i := min; i <= max; i++ {
		log.Printf("deleting index %d\n", i)
		if err := os.Remove(filepath.Join(*raftDir, fmt.Sprintf("fancylogs/entry.%d", i))); err != nil {
			return err
		}
	}
	return nil
}

type fancyStableStore struct {
}

func (s *fancyStableStore) Set(key []byte, val []byte) error {
	return ioutil.WriteFile(filepath.Join(*raftDir, "fancystable", string(key)), val, 0600)
}

func (s *fancyStableStore) Get(key []byte) ([]byte, error) {
	b, err := ioutil.ReadFile(filepath.Join(*raftDir, "fancystable", string(key)))
	if err != nil && os.IsNotExist(err) {
		return []byte{}, fmt.Errorf("not found")
	}
	return b, err
}

func (s *fancyStableStore) SetUint64(key []byte, val uint64) error {
	return s.Set(key, []byte(fmt.Sprintf("%d", val)))
}

func (s *fancyStableStore) GetUint64(key []byte) (uint64, error) {
	b, err := s.Get(key)
	if err != nil {
		return 0, err
	}
	i, err := strconv.ParseInt(string(b), 0, 64)
	if err != nil {
		return 0, err
	}
	return uint64(i), nil
}

// Snapshot holds the data needed to serialize storage
type Snapshot struct {
	uuids   [][]byte
	entries []string
}

func uint32ToBytes(u uint32) []byte {
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, u)
	return buf
}

// Converts bytes to an integer
func bytesToUint32(b []byte) uint32 {
	return binary.BigEndian.Uint32(b)
}

// Persist writes a snapshot to a file. We just serialize all active entries.
func (s *Snapshot) Persist(sink raft.SnapshotSink) error {
	_, err := sink.Write([]byte{0x0})
	if err != nil {
		sink.Cancel()
		return err
	}

	for i, e := range s.entries {
		_, err = sink.Write(s.uuids[i])
		if err != nil {
			sink.Cancel()
			return err
		}

		b := []byte(e)
		_, err = sink.Write(uint32ToBytes(uint32(len(b))))
		if err != nil {
			sink.Cancel()
			return err
		}

		_, err = sink.Write(b)
		if err != nil {
			sink.Cancel()
			return err
		}
	}

	return sink.Close()
}

// Release cleans up a snapshot. We don't need to do anything.
func (s *Snapshot) Release() {
}

type FSM struct {
	store *Storage
}

func (fsm *FSM) Apply(l *raft.Log) interface{} {
	log.Printf("TODO: apply %v\n", l)
	return nil
}

func (fsm *FSM) Snapshot() (raft.FSMSnapshot, error) {
	uuids, entries, err := fsm.store.GetAll()
	snapshot := &Snapshot{uuids, entries}
	return snapshot, err
}

func (fsm *FSM) Restore(snap io.ReadCloser) error {
	defer snap.Close()

	s, err := NewStorage()
	if err != nil {
		return err
	}

	// swap in the restored storage, with time emission channels.
	s.c = fsm.store.c
	s.C = fsm.store.C
	fsm.store.Close()
	fsm.store = s

	b := make([]byte, 1)
	_, err = snap.Read(b)
	if b[0] != byte(0x00) {
		msg := "Unknown snapshot schema version"
		log.Printf(msg)
		return fmt.Errorf(msg)
	}

	uuid := make([]byte, 16)
	size := make([]byte, 4)
	count := 0
	for {
		_, err = snap.Read(uuid)
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}

		_, err = snap.Read(size)
		if err != nil {
			return err
		}

		eb := make([]byte, bytesToUint32(size))
		_, err = snap.Read(eb)
		if err != nil {
			return err
		}
		e := string(eb)
		err = s.Add(uuid, e)
		if err != nil {
			return err
		}

		count++
	}

	log.Printf("Restored snapshot. entries=%d", count)
	return nil
}

func handleRequest(res http.ResponseWriter, req *http.Request) {
	log.Printf("reading body…\n")
	data, err := ioutil.ReadAll(req.Body)
	if err != nil {
		log.Println("error reading request:", err)
		return
	}
	log.Printf("proposing()\n")
	node.Apply(data, 10*time.Second)
	log.Println("proposed")
}

func handleJoin(res http.ResponseWriter, req *http.Request) {
	log.Println("Join request from", req.RemoteAddr)
	if node.State() != raft.Leader {
		log.Println("rejecting join request, I am not the leader")
		http.Redirect(res, req, fmt.Sprintf("http://%s/join", node.Leader()), 307)
		return
	}

	data, err := ioutil.ReadAll(req.Body)
	if err != nil {
		log.Println("Could not read body:", err)
		return
	}

	type joinRequest struct {
		Addr string
	}
	var r joinRequest

	if err = json.Unmarshal(data, &r); err != nil {
		log.Println("Could not decode request:", err)
		http.Error(res, fmt.Sprintf("Could not decode your request"), 400)
		return
	}

	log.Printf("Joining peer: %q\n", r.Addr)

	a, err := net.ResolveTCPAddr("tcp", r.Addr)
	if err != nil {
		log.Printf("Could not resolve addr %q: %v\n", r.Addr, err)
		http.Error(res, "Could not resolve your address", 400)
		return
	}

	if err = node.AddPeer(a).Error(); err != nil {
		log.Println("Could not add peer:", err)
		http.Error(res, "Could not add peer", 500)
		return
	}
	res.WriteHeader(200)
}

func joinMaster(addr string, peerStore *raft.JSONPeers) []net.Addr {
	master, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		log.Fatalf("Could not resolve %q: %v", addr, err)
	}

	type joinRequest struct {
		Addr string
	}
	var buf *bytes.Buffer
	if data, err := json.Marshal(joinRequest{*listen}); err != nil {
		log.Fatal("Could not marshal join request:", err)
	} else {
		buf = bytes.NewBuffer(data)
	}
	if res, err := http.Post(fmt.Sprintf("http://%s/join", addr), "application/json", buf); err != nil {
		log.Fatal("Could not send join request:", err)
	} else if res.StatusCode > 399 {
		data, _ := ioutil.ReadAll(res.Body)
		log.Fatal("Join request failed:", string(data))
	} else if res.StatusCode > 299 {
		loc := res.Header.Get("Location")
		if loc == "" {
			log.Fatal("Redirect has no Location header")
		}
		u, err := url.Parse(loc)
		if err != nil {
			log.Fatalf("Could not parse redirection %q: %v", loc, err)
		}

		return joinMaster(u.Host, peerStore)
	}

	log.Printf("Adding master %v as peer\n", master)
	p, err := peerStore.Peers()
	if err != nil {
		log.Fatal("Could not read peers:", err)
	}
	p = raft.AddUniquePeer(p, master)
	peerStore.SetPeers(p)
	return p
}

func main() {
	flag.Parse()
	log.Printf("hey\n")

	a, err := net.ResolveTCPAddr("tcp", *listen)
	if err != nil {
		log.Fatal(err)
	}

	transport := NewTransport(a)
	http.Handle("/raft/", transport)

	peerStore := raft.NewJSONPeers(*raftDir, transport)

	var p []net.Addr

	config := raft.DefaultConfig()
	if *peers == "" {
		config.EnableSingleNode = true
	} else {
		p, err = peerStore.Peers()
		if err != nil {
			log.Fatal(err)
		}
		for _, addr := range strings.Split(*peers, ",") {
			peer, err := net.ResolveTCPAddr("tcp", addr)
			if err != nil {
				log.Fatal(err)
			}
			log.Printf("adding %v\n", peer)
			p = raft.AddUniquePeer(p, peer)
			peerStore.SetPeers(p)
		}
	}

	if *join != "" {
		p = joinMaster(*join, peerStore)
	}

	fss, err := raft.NewFileSnapshotStore(*raftDir, 1, nil)
	if err != nil {
		log.Fatal(err)
	}

	storage, err := NewStorage()
	if err != nil {
		log.Fatal(err)
	}

	fsm := &FSM{storage}

	if err := os.MkdirAll(filepath.Join(*raftDir, "fancylogs"), 0755); err != nil {
		log.Fatal(err)
	}

	if err := os.MkdirAll(filepath.Join(*raftDir, "fancystable"), 0755); err != nil {
		log.Fatal(err)
	}

	ls := &fancyLogStore{}
	stablestore := &fancyStableStore{}

	// NewRaft(*Config, FSM, LogStore, StableStore, SnapshotStore, PeerStore, Transport)
	node, err = raft.NewRaft(config, fsm, ls, stablestore, fss, peerStore, transport)
	if err != nil {
		log.Fatal(err)
	}

	// TODO: observeleaderchanges?
	// TODO: observenexttime?

	//http.HandleFunc("/msg", handleMessages)
	http.HandleFunc("/put", handleRequest)
	http.HandleFunc("/join", handleJoin)
	go func() {
		err := http.ListenAndServe(*listen, nil)
		if err != nil {
			log.Fatal(err)
		}
	}()

	node.SetPeers(p)

	for {

		time.Sleep(1000 * time.Millisecond)

		// err = s.raft.Add(b, raftMaxTime)
		// if err != nil {
		// 	eWrapper.Done(false)
		// }

		// eWrapper.Done(true)

	}
}
<commit_msg>Status page on /<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"encoding/gob"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/raft"
)

var (
	raftDir   = flag.String("raftdir", "/tmp/r", "")
	peers     = flag.String("peers", "", "comma-separated host:port tuples")
	listen    = flag.String("listen", ":8000", "")
	join      = flag.String("join", "", "raft master to join")
	node      *raft.Raft
	peerStore *raft.JSONPeers
	logStore  *fancyLogStore
)

// trivial log store, writing one entry into one file each.
// fulfills the raft.LogStore interface.
type fancyLogStore struct {
	l         sync.RWMutex
	lowIndex  uint64
	highIndex uint64
}

func (s *fancyLogStore) FirstIndex() (uint64, error) {
	s.l.RLock()
	defer s.l.RUnlock()
	return s.lowIndex, nil
}

func (s *fancyLogStore) LastIndex() (uint64, error) {
	s.l.RLock()
	defer s.l.RUnlock()
	return s.highIndex, nil
}

func (s *fancyLogStore) GetLog(index uint64, rlog *raft.Log) error {
	s.l.Lock()
	defer s.l.Unlock()
	f, err := os.Open(filepath.Join(*raftDir, fmt.Sprintf("fancylogs/entry.%d", index)))
	if err != nil {
		return err
	}
	defer f.Close()

	var elog raft.Log
	if err := gob.NewDecoder(f).Decode(&elog); err != nil {
		return err
	}
	*rlog = elog
	return nil
}

func (s *fancyLogStore) StoreLog(log *raft.Log) error {
	return s.StoreLogs([]*raft.Log{log})
}

func (s *fancyLogStore) StoreLogs(logs []*raft.Log) error {
	s.l.Lock()
	defer s.l.Unlock()

	for _, entry := range logs {
		log.Printf("writing index %d to file (%v)\n", entry.Index, entry)
		f, err := os.Create(filepath.Join(*raftDir, fmt.Sprintf("fancylogs/entry.%d", entry.Index)))
		if err != nil {
			return err
		}
		defer f.Close()
		if err := gob.NewEncoder(f).Encode(entry); err != nil {
			return err
		}
		if s.lowIndex == 0 {
			s.lowIndex = entry.Index
		}
		if entry.Index > s.highIndex {
			s.highIndex = entry.Index
		}
	}

	return nil
}

func (s *fancyLogStore) DeleteRange(min, max uint64) error {
	s.l.Lock()
	defer s.l.Unlock()
	for i := min; i <= max; i++ {
		log.Printf("deleting index %d\n", i)
		if err := os.Remove(filepath.Join(*raftDir, fmt.Sprintf("fancylogs/entry.%d", i))); err != nil {
			return err
		}
	}
	return nil
}

type fancyStableStore struct {
}

func (s *fancyStableStore) Set(key []byte, val []byte) error {
	return ioutil.WriteFile(filepath.Join(*raftDir, "fancystable", string(key)), val, 0600)
}

func (s *fancyStableStore) Get(key []byte) ([]byte, error) {
	b, err := ioutil.ReadFile(filepath.Join(*raftDir, "fancystable", string(key)))
	if err != nil && os.IsNotExist(err) {
		return []byte{}, fmt.Errorf("not found")
	}
	return b, err
}

func (s *fancyStableStore) SetUint64(key []byte, val uint64) error {
	return s.Set(key, []byte(fmt.Sprintf("%d", val)))
}

func (s *fancyStableStore) GetUint64(key []byte) (uint64, error) {
	b, err := s.Get(key)
	if err != nil {
		return 0, err
	}
	i, err := strconv.ParseInt(string(b), 0, 64)
	if err != nil {
		return 0, err
	}
	return uint64(i), nil
}

// Snapshot holds the data needed to serialize storage
type Snapshot struct {
	uuids   [][]byte
	entries []string
}

func uint32ToBytes(u uint32) []byte {
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, u)
	return buf
}

// Converts bytes to an integer
func bytesToUint32(b []byte) uint32 {
	return binary.BigEndian.Uint32(b)
}

// Persist writes a snapshot to a file. We just serialize all active entries.
func (s *Snapshot) Persist(sink raft.SnapshotSink) error {
	_, err := sink.Write([]byte{0x0})
	if err != nil {
		sink.Cancel()
		return err
	}

	for i, e := range s.entries {
		_, err = sink.Write(s.uuids[i])
		if err != nil {
			sink.Cancel()
			return err
		}

		b := []byte(e)
		_, err = sink.Write(uint32ToBytes(uint32(len(b))))
		if err != nil {
			sink.Cancel()
			return err
		}

		_, err = sink.Write(b)
		if err != nil {
			sink.Cancel()
			return err
		}
	}

	return sink.Close()
}

// Release cleans up a snapshot. We don't need to do anything.
func (s *Snapshot) Release() {
}

type FSM struct {
	store *Storage
}

func (fsm *FSM) Apply(l *raft.Log) interface{} {
	log.Printf("TODO: apply %v\n", l)
	return nil
}

func (fsm *FSM) Snapshot() (raft.FSMSnapshot, error) {
	uuids, entries, err := fsm.store.GetAll()
	snapshot := &Snapshot{uuids, entries}
	return snapshot, err
}

func (fsm *FSM) Restore(snap io.ReadCloser) error {
	defer snap.Close()

	s, err := NewStorage()
	if err != nil {
		return err
	}

	// swap in the restored storage, with time emission channels.
	s.c = fsm.store.c
	s.C = fsm.store.C
	fsm.store.Close()
	fsm.store = s

	b := make([]byte, 1)
	_, err = snap.Read(b)
	if b[0] != byte(0x00) {
		msg := "Unknown snapshot schema version"
		log.Printf(msg)
		return fmt.Errorf(msg)
	}

	uuid := make([]byte, 16)
	size := make([]byte, 4)
	count := 0
	for {
		_, err = snap.Read(uuid)
		if err == io.EOF {
			break
		} else if err != nil {
			return err
		}

		_, err = snap.Read(size)
		if err != nil {
			return err
		}

		eb := make([]byte, bytesToUint32(size))
		_, err = snap.Read(eb)
		if err != nil {
			return err
		}
		e := string(eb)
		err = s.Add(uuid, e)
		if err != nil {
			return err
		}

		count++
	}

	log.Printf("Restored snapshot. entries=%d", count)
	return nil
}

func handleRequest(res http.ResponseWriter, req *http.Request) {
	log.Printf("reading body…\n")
	data, err := ioutil.ReadAll(req.Body)
	if err != nil {
		log.Println("error reading request:", err)
		return
	}
	log.Printf("proposing()\n")
	node.Apply(data, 10*time.Second)
	log.Println("proposed")
}

func handleJoin(res http.ResponseWriter, req *http.Request) {
	log.Println("Join request from", req.RemoteAddr)
	if node.State() != raft.Leader {
		log.Println("rejecting join request, I am not the leader")
		http.Redirect(res, req, fmt.Sprintf("http://%s/join", node.Leader()), 307)
		return
	}

	data, err := ioutil.ReadAll(req.Body)
	if err != nil {
		log.Println("Could not read body:", err)
		return
	}

	type joinRequest struct {
		Addr string
	}
	var r joinRequest

	if err = json.Unmarshal(data, &r); err != nil {
		log.Println("Could not decode request:", err)
		http.Error(res, fmt.Sprintf("Could not decode your request"), 400)
		return
	}

	log.Printf("Joining peer: %q\n", r.Addr)

	a, err := net.ResolveTCPAddr("tcp", r.Addr)
	if err != nil {
		log.Printf("Could not resolve addr %q: %v\n", r.Addr, err)
		http.Error(res, "Could not resolve your address", 400)
		return
	}

	if err = node.AddPeer(a).Error(); err != nil {
		log.Println("Could not add peer:", err)
		http.Error(res, "Could not add peer", 500)
		return
	}
	res.WriteHeader(200)
}

var statusTpl = template.Must(template.New("index").Parse(`<!DOCTYPE html>

		<title>Status of fancyirc node {{ .Addr }}</title>

Status of fancyirc node {{ .Addr }}

Status

{{ .State }}

Current Leader

{{ .Leader }}

Peers

{{ .Peers }}

Log

Entries {{ .First }} through {{ .Last }}

		{{ range .Entries }}

{{ . }}

		{{ end }}

Stats

		{{ range $key, $val := .Stats }}

{{ $key }}

{{ $val }}

		{{ end }}

`))

func handleStatus(res http.ResponseWriter, req *http.Request) {
	p, _ := peerStore.Peers()

	lo, err := logStore.FirstIndex()
	if err != nil {
		log.Printf("Could not get first index: %v", err)
		http.Error(res, "internal error", 500)
		return
	}
	hi, err := logStore.LastIndex()
	if err != nil {
		log.Printf("Could not get last index: %v", err)
		http.Error(res, "internal error", 500)
		return
	}

	var entries []*raft.Log
	if lo != 0 && hi != 0 {
		for i := lo; i <= hi; i++ {
			l := new(raft.Log)

			if err := logStore.GetLog(i, l); err != nil {
				log.Printf("Could not get entry %d: %v", i, err)
				http.Error(res, "internal error", 500)
				return
			}
			entries = append(entries, l)
		}
	}

	args := struct {
		Addr    string
		State   raft.RaftState
		Leader  net.Addr
		Peers   []net.Addr
		First   uint64
		Last    uint64
		Entries []*raft.Log
		Stats   map[string]string
	}{
		*listen,
		node.State(),
		node.Leader(),
		p,
		lo,
		hi,
		entries,
		node.Stats(),
	}

	statusTpl.Execute(res, args)
}

func joinMaster(addr string, peerStore *raft.JSONPeers) []net.Addr {
	master, err := net.ResolveTCPAddr("tcp", addr)
	if err != nil {
		log.Fatalf("Could not resolve %q: %v", addr, err)
	}

	type joinRequest struct {
		Addr string
	}
	var buf *bytes.Buffer
	if data, err := json.Marshal(joinRequest{*listen}); err != nil {
		log.Fatal("Could not marshal join request:", err)
	} else {
		buf = bytes.NewBuffer(data)
	}
	if res, err := http.Post(fmt.Sprintf("http://%s/join", addr), "application/json", buf); err != nil {
		log.Fatal("Could not send join request:", err)
	} else if res.StatusCode > 399 {
		data, _ := ioutil.ReadAll(res.Body)
		log.Fatal("Join request failed:", string(data))
	} else if res.StatusCode > 299 {
		loc := res.Header.Get("Location")
		if loc == "" {
			log.Fatal("Redirect has no Location header")
		}
		u, err := url.Parse(loc)
		if err != nil {
			log.Fatalf("Could not parse redirection %q: %v", loc, err)
		}

		return joinMaster(u.Host, peerStore)
	}

	log.Printf("Adding master %v as peer\n", master)
	p, err := peerStore.Peers()
	if err != nil {
		log.Fatal("Could not read peers:", err)
	}
	p = raft.AddUniquePeer(p, master)
	peerStore.SetPeers(p)
	return p
}

func main() {
	flag.Parse()
	log.Printf("hey\n")

	a, err := net.ResolveTCPAddr("tcp", *listen)
	if err != nil {
		log.Fatal(err)
	}

	transport := NewTransport(a)
	http.Handle("/raft/", transport)

	peerStore = raft.NewJSONPeers(*raftDir, transport)

	var p []net.Addr

	config := raft.DefaultConfig()
	if *peers == "" {
		config.EnableSingleNode = true
	} else {
		p, err = peerStore.Peers()
		if err != nil {
			log.Fatal(err)
		}
		for _, addr := range strings.Split(*peers, ",") {
			peer, err := net.ResolveTCPAddr("tcp", addr)
			if err != nil {
				log.Fatal(err)
			}
			log.Printf("adding %v\n", peer)
			p = raft.AddUniquePeer(p, peer)
			peerStore.SetPeers(p)
		}
	}

	if *join != "" {
		p = joinMaster(*join, peerStore)
	}

	fss, err := raft.NewFileSnapshotStore(*raftDir, 1, nil)
	if err != nil {
		log.Fatal(err)
	}

	storage, err := NewStorage()
	if err != nil {
		log.Fatal(err)
	}

	fsm := &FSM{storage}

	if err := os.MkdirAll(filepath.Join(*raftDir, "fancylogs"), 0755); err != nil {
		log.Fatal(err)
	}

	if err := os.MkdirAll(filepath.Join(*raftDir, "fancystable"), 0755); err != nil {
		log.Fatal(err)
	}

	logStore = &fancyLogStore{}
	stablestore := &fancyStableStore{}

	// NewRaft(*Config, FSM, LogStore, StableStore, SnapshotStore, PeerStore, Transport)
	node, err = raft.NewRaft(config, fsm, logStore, stablestore, fss, peerStore, transport)
	if err != nil {
		log.Fatal(err)
	}

	// TODO: observeleaderchanges?
	// TODO: observenexttime?

	//http.HandleFunc("/msg", handleMessages)
	http.HandleFunc("/", handleStatus)
	http.HandleFunc("/put", handleRequest)
	http.HandleFunc("/join", handleJoin)
	go func() {
		err := http.ListenAndServe(*listen, nil)
		if err != nil {
			log.Fatal(err)
		}
	}()

	node.SetPeers(p)

	for {

		time.Sleep(1000 * time.Millisecond)

		// err = s.raft.Add(b, raftMaxTime)
		// if err != nil {
		// 	eWrapper.Done(false)
		// }

		// eWrapper.Done(true)

	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Alea Soluciones SLL. All rights reserved.  Use of this
// source code is governed by a MIT-style license that can be found in the
// LICENSE file.

package timetoken_test

import (
	"testing"
	"time"

	. "github.com/aleasoluciones/goaleasoluciones/timetoken"

	"github.com/stretchr/testify/assert"
)

func tokenManager(periodicFuncTimesCalled *int) *TokenManager {
	return NewTokenManager(1*time.Millisecond, 3*time.Millisecond, func(id string) {
		*periodicFuncTimesCalled += 1
	})
}

func TestPeriodicFunctionCalledUntilTTL(t *testing.T) {
	t.Parallel()
	periodicFuncTimesCalled := 0
	tm := tokenManager(&periodicFuncTimesCalled)
	tm.Add("id", 3*time.Millisecond)
	time.Sleep(2 * time.Millisecond)
	assert.Equal(t, 2, periodicFuncTimesCalled)
}

func TestPeriodicFunctionCalledEveryPeriodeUntilTTL(t *testing.T) {
	t.Parallel()
	periodicFuncTimesCalled := 0
	tm := tokenManager(&periodicFuncTimesCalled)
	tm.Add("id", 3*time.Millisecond)

	time.Sleep(1 * time.Millisecond)
	assert.Equal(t, 1, periodicFuncTimesCalled)

	time.Sleep(1 * time.Millisecond)
	assert.Equal(t, 2, periodicFuncTimesCalled)
}

func TestPeriodicFunctionNotCalledAfterTTL(t *testing.T) {
	t.Parallel()
	periodicFuncTimesCalled := 0
	tm := tokenManager(&periodicFuncTimesCalled)
	tm.Add("id", 2*time.Millisecond)

	time.Sleep(4 * time.Millisecond)
	assert.Equal(t, 2, periodicFuncTimesCalled)
}
<commit_msg>Rewrite non deterministic test for tokenManager<commit_after>// Copyright 2014 Alea Soluciones SLL. All rights reserved.  Use of this
// source code is governed by a MIT-style license that can be found in the
// LICENSE file.

package timetoken_test

import (
	"testing"
	"time"

	. "github.com/aleasoluciones/goaleasoluciones/timetoken"

	"github.com/stretchr/testify/assert"
)

const (
	ZERO_TTL_FOR_TOKEN_MANAGER = 0*time.Millisecond
	ONE_CYCLE_FOR_TOKEN_MANAGER = 700 * time.Millisecond
)

func tokenManager(periodicFuncTimesCalled *int) *TokenManager {
	return NewTokenManager(ONE_CYCLE_FOR_TOKEN_MANAGER, ZERO_TTL_FOR_TOKEN_MANAGER, func(id string) {
		*periodicFuncTimesCalled += 1
	})
}

func TestPeriodicFunctionCalledUntilTokenTTL(t *testing.T) {
	t.Parallel()
	periodicFuncTimesCalled := 0
	tm := tokenManager(&periodicFuncTimesCalled)
	var tokenTTL = 2 * ONE_CYCLE_FOR_TOKEN_MANAGER
	tm.Add("id", tokenTTL)
	time.Sleep(tokenTTL + 50*time.Millisecond)
	assert.Equal(t, 2, periodicFuncTimesCalled)
}

func TestPeriodicFunctionCalledEveryPeriodeUntilTokenTTL(t *testing.T) {
	t.Parallel()
	periodicFuncTimesCalled := 0
	tm := tokenManager(&periodicFuncTimesCalled)
	var tokenTTL = 3 * ONE_CYCLE_FOR_TOKEN_MANAGER
	var lessThanOneCycle = ONE_CYCLE_FOR_TOKEN_MANAGER - 100 * time.Millisecond
	tm.Add("id", tokenTTL)
	time.Sleep(lessThanOneCycle)
	assert.Equal(t, 1, periodicFuncTimesCalled)
	time.Sleep(lessThanOneCycle)
	assert.Equal(t, 2, periodicFuncTimesCalled)
	time.Sleep(lessThanOneCycle)
	assert.Equal(t, 3, periodicFuncTimesCalled)
	time.Sleep(lessThanOneCycle)
	assert.Equal(t, 3, periodicFuncTimesCalled)
}

func TestPeriodicFunctionNotCalledAfterTokenTTL(t *testing.T) {
	t.Parallel()
	periodicFuncTimesCalled := 0
	tm := tokenManager(&periodicFuncTimesCalled)
	var tokenTTL = 1 * ONE_CYCLE_FOR_TOKEN_MANAGER
	tm.Add("id", tokenTTL)
	time.Sleep(2 * tokenTTL)
	assert.Equal(t, 1, periodicFuncTimesCalled)
}
<|endoftext|>