github.phpd.cn/hashicorp/consul@v1.4.5/agent/consul/raft_rpc.go (about)

     1  package consul
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/hashicorp/consul/agent/pool"
    10  	"github.com/hashicorp/consul/tlsutil"
    11  	"github.com/hashicorp/raft"
    12  )
    13  
    14  // RaftLayer implements the raft.StreamLayer interface,
    15  // so that we can use a single RPC layer for Raft and Consul
    16  type RaftLayer struct {
    17  	// src is the address for outgoing connections.
    18  	src net.Addr
    19  
    20  	// addr is the listener address to return.
    21  	addr net.Addr
    22  
    23  	// connCh is used to accept connections.
    24  	connCh chan net.Conn
    25  
    26  	// TLS wrapper
    27  	tlsWrap tlsutil.Wrapper
    28  
    29  	// Tracks if we are closed
    30  	closed    bool
    31  	closeCh   chan struct{}
    32  	closeLock sync.Mutex
    33  
    34  	// tlsFunc is a callback to determine whether to use TLS for connecting to
    35  	// a given Raft server
    36  	tlsFunc func(raft.ServerAddress) bool
    37  }
    38  
    39  // NewRaftLayer is used to initialize a new RaftLayer which can
    40  // be used as a StreamLayer for Raft. If a tlsConfig is provided,
    41  // then the connection will use TLS.
    42  func NewRaftLayer(src, addr net.Addr, tlsWrap tlsutil.Wrapper, tlsFunc func(raft.ServerAddress) bool) *RaftLayer {
    43  	layer := &RaftLayer{
    44  		src:     src,
    45  		addr:    addr,
    46  		connCh:  make(chan net.Conn),
    47  		tlsWrap: tlsWrap,
    48  		closeCh: make(chan struct{}),
    49  		tlsFunc: tlsFunc,
    50  	}
    51  	return layer
    52  }
    53  
    54  // Handoff is used to hand off a connection to the
    55  // RaftLayer. This allows it to be Accept()'ed
    56  func (l *RaftLayer) Handoff(c net.Conn) error {
    57  	select {
    58  	case l.connCh <- c:
    59  		return nil
    60  	case <-l.closeCh:
    61  		return fmt.Errorf("Raft RPC layer closed")
    62  	}
    63  }
    64  
    65  // Accept is used to return connection which are
    66  // dialed to be used with the Raft layer
    67  func (l *RaftLayer) Accept() (net.Conn, error) {
    68  	select {
    69  	case conn := <-l.connCh:
    70  		return conn, nil
    71  	case <-l.closeCh:
    72  		return nil, fmt.Errorf("Raft RPC layer closed")
    73  	}
    74  }
    75  
    76  // Close is used to stop listening for Raft connections
    77  func (l *RaftLayer) Close() error {
    78  	l.closeLock.Lock()
    79  	defer l.closeLock.Unlock()
    80  
    81  	if !l.closed {
    82  		l.closed = true
    83  		close(l.closeCh)
    84  	}
    85  	return nil
    86  }
    87  
    88  // Addr is used to return the address of the listener
    89  func (l *RaftLayer) Addr() net.Addr {
    90  	return l.addr
    91  }
    92  
    93  // Dial is used to create a new outgoing connection
    94  func (l *RaftLayer) Dial(address raft.ServerAddress, timeout time.Duration) (net.Conn, error) {
    95  	d := &net.Dialer{LocalAddr: l.src, Timeout: timeout}
    96  	conn, err := d.Dial("tcp", string(address))
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  
   101  	// Check for tls mode
   102  	if l.tlsFunc(address) && l.tlsWrap != nil {
   103  		// Switch the connection into TLS mode
   104  		if _, err := conn.Write([]byte{byte(pool.RPCTLS)}); err != nil {
   105  			conn.Close()
   106  			return nil, err
   107  		}
   108  
   109  		// Wrap the connection in a TLS client
   110  		conn, err = l.tlsWrap(conn)
   111  		if err != nil {
   112  			return nil, err
   113  		}
   114  	}
   115  
   116  	// Write the Raft byte to set the mode
   117  	_, err = conn.Write([]byte{byte(pool.RPCRaft)})
   118  	if err != nil {
   119  		conn.Close()
   120  		return nil, err
   121  	}
   122  	return conn, err
   123  }