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 }