// getConn is used to get a connection from the pool
func (n *NetworkTransport) getConn(target net.Addr) (*netConn, error) {
// Check for a pooled conn
if conn := n.getPooledConn(target); conn != nil {
return conn, nil
}
// Dial a new connection
conn, err := n.stream.Dial(target.String(), n.timeout)
if err != nil {
return nil, err
}
// Wrap the conn
netConn := &netConn{
target: target,
conn: conn,
r: bufio.NewReader(conn),
w: bufio.NewWriter(conn),
}
// Setup encoder/decoders
netConn.dec = codec.NewDecoder(netConn.r, &codec.MsgpackHandle{})
netConn.enc = codec.NewEncoder(netConn.w, &codec.MsgpackHandle{})
// Done
return netConn, nil
}
func MockDecoder(buf []byte) interface{} {
out := new(MockData)
err := codec.NewDecoder(bytes.NewReader(buf), msgpackHandle).Decode(out)
if err != nil {
panic(err)
}
return out
}
func (m *MockFSM) Restore(inp io.ReadCloser) error {
m.Lock()
defer m.Unlock()
defer inp.Close()
hd := codec.MsgpackHandle{}
dec := codec.NewDecoder(inp, &hd)
m.logs = nil
return dec.Decode(&m.logs)
}
// decodeTags is used to decode a tag map
func (s *Serf) decodeTags(buf []byte) map[string]string {
tags := make(map[string]string)
// Backwards compatibility mode
if len(buf) == 0 || buf[0] != tagMagicByte {
tags["role"] = string(buf)
return tags
}
// Decode the tags
r := bytes.NewReader(buf[1:])
dec := codec.NewDecoder(r, &codec.MsgpackHandle{})
if err := dec.Decode(&tags); err != nil {
s.logger.Printf("[ERR] serf: Failed to decode tags: %v", err)
}
return tags
}
// handleConn is used to handle an inbound connection for its lifespan
func (n *NetworkTransport) handleConn(conn net.Conn) {
defer conn.Close()
r := bufio.NewReader(conn)
w := bufio.NewWriter(conn)
dec := codec.NewDecoder(r, &codec.MsgpackHandle{})
enc := codec.NewEncoder(w, &codec.MsgpackHandle{})
for {
if err := n.handleCommand(r, dec, enc); err != nil {
if err != io.EOF {
n.logger.Printf("[ERR] raft-net: Failed to decode incoming command: %v", err)
}
return
}
if err := w.Flush(); err != nil {
n.logger.Printf("[ERR] raft-net: Failed to flush response: %v", err)
return
}
}
}
// listen is a long running routine that listens for new clients
func (i *AgentRPC) listen() {
for {
conn, err := i.listener.Accept()
if err != nil {
if i.stop {
return
}
i.logger.Printf("[ERR] agent.rpc: Failed to accept client: %v", err)
continue
}
i.logger.Printf("[INFO] agent.rpc: Accepted client: %v", conn.RemoteAddr())
// Wrap the connection in a client
client := &rpcClient{
name: conn.RemoteAddr().String(),
conn: conn,
reader: bufio.NewReader(conn),
writer: bufio.NewWriter(conn),
}
client.dec = codec.NewDecoder(client.reader, msgpackHandle)
client.enc = codec.NewEncoder(client.writer, msgpackHandle)
if err != nil {
i.logger.Printf("[ERR] agent.rpc: Failed to create decoder: %v", err)
conn.Close()
continue
}
// Register the client
i.Lock()
if !i.stop {
i.clients[client.name] = client
go i.handleClient(client)
} else {
conn.Close()
}
i.Unlock()
}
}
// NewRPCClient is used to create a new RPC client given the address.
// This will properly dial, handshake, and start listening
func NewRPCClient(addr string) (*RPCClient, error) {
var conn net.Conn
var err error
if envAddr := os.Getenv("CONSUL_RPC_ADDR"); envAddr != "" {
addr = envAddr
}
// Try to dial to agent
mode := "tcp"
if strings.HasPrefix(addr, "/") {
mode = "unix"
}
if conn, err = net.Dial(mode, addr); err != nil {
return nil, err
}
// Create the client
client := &RPCClient{
seq: 0,
conn: conn,
reader: bufio.NewReader(conn),
writer: bufio.NewWriter(conn),
dispatch: make(map[uint64]seqHandler),
shutdownCh: make(chan struct{}),
}
client.dec = codec.NewDecoder(client.reader, msgpackHandle)
client.enc = codec.NewEncoder(client.writer, msgpackHandle)
go client.listen()
// Do the initial handshake
if err := client.handshake(); err != nil {
client.Close()
return nil, err
}
return client, err
}
func (c *consulFSM) Restore(old io.ReadCloser) error {
defer old.Close()
// Create a temporary path for the state store
tmpPath, err := ioutil.TempDir(c.path, "state")
if err != nil {
return err
}
// Create a new state store
state, err := NewStateStorePath(c.gc, tmpPath, c.logOutput)
if err != nil {
return err
}
c.state.Close()
c.state = state
// Create a decoder
dec := codec.NewDecoder(old, msgpackHandle)
// Read in the header
var header snapshotHeader
if err := dec.Decode(&header); err != nil {
return err
}
// Populate the new state
msgType := make([]byte, 1)
for {
// Read the message type
_, err := old.Read(msgType)
if err == io.EOF {
break
} else if err != nil {
return err
}
// Decode
switch structs.MessageType(msgType[0]) {
case structs.RegisterRequestType:
var req structs.RegisterRequest
if err := dec.Decode(&req); err != nil {
return err
}
c.applyRegister(&req, header.LastIndex)
case structs.KVSRequestType:
var req structs.DirEntry
if err := dec.Decode(&req); err != nil {
return err
}
if err := c.state.KVSRestore(&req); err != nil {
return err
}
case structs.SessionRequestType:
var req structs.Session
if err := dec.Decode(&req); err != nil {
return err
}
if err := c.state.SessionRestore(&req); err != nil {
return err
}
case structs.ACLRequestType:
var req structs.ACL
if err := dec.Decode(&req); err != nil {
return err
}
if err := c.state.ACLRestore(&req); err != nil {
return err
}
case structs.TombstoneRequestType:
var req structs.DirEntry
if err := dec.Decode(&req); err != nil {
return err
}
if err := c.state.TombstoneRestore(&req); err != nil {
return err
}
default:
return fmt.Errorf("Unrecognized msg type: %v", msgType)
}
}
return nil
}
func decodeMessage(buf []byte, out interface{}) error {
var handle codec.MsgpackHandle
return codec.NewDecoder(bytes.NewReader(buf), &handle).Decode(out)
}
// Decode reverses the encode operation on a byte slice input
func decode(buf []byte, out interface{}) error {
r := bytes.NewReader(buf)
hd := codec.MsgpackHandle{}
dec := codec.NewDecoder(r, &hd)
return dec.Decode(out)
}
// Decode is used to decode a MsgPack encoded object
func Decode(buf []byte, out interface{}) error {
return codec.NewDecoder(bytes.NewReader(buf), msgpackHandle).Decode(out)
}
// recvRemoteState is used to read the remote state from a connection
func (m *Memberlist) readRemoteState(conn net.Conn) (bool, []pushNodeState, []byte, error) {
// Setup a deadline
conn.SetDeadline(time.Now().Add(m.config.TCPTimeout))
// Created a buffered reader
var bufConn io.Reader = bufio.NewReader(conn)
// Read the message type
buf := [1]byte{0}
if _, err := bufConn.Read(buf[:]); err != nil {
return false, nil, nil, err
}
msgType := messageType(buf[0])
// Check if the message is encrypted
if msgType == encryptMsg {
if !m.config.EncryptionEnabled() {
return false, nil, nil,
fmt.Errorf("Remote state is encrypted and encryption is not configured")
}
plain, err := m.decryptRemoteState(bufConn)
if err != nil {
return false, nil, nil, err
}
// Reset message type and bufConn
msgType = messageType(plain[0])
bufConn = bytes.NewReader(plain[1:])
} else if m.config.EncryptionEnabled() {
return false, nil, nil,
fmt.Errorf("Encryption is configured but remote state is not encrypted")
}
// Get the msgPack decoders
hd := codec.MsgpackHandle{}
dec := codec.NewDecoder(bufConn, &hd)
// Check if we have a compressed message
if msgType == compressMsg {
var c compress
if err := dec.Decode(&c); err != nil {
return false, nil, nil, err
}
decomp, err := decompressBuffer(&c)
if err != nil {
return false, nil, nil, err
}
// Reset the message type
msgType = messageType(decomp[0])
// Create a new bufConn
bufConn = bytes.NewReader(decomp[1:])
// Create a new decoder
dec = codec.NewDecoder(bufConn, &hd)
}
// Quit if not push/pull
if msgType != pushPullMsg {
err := fmt.Errorf("received invalid msgType (%d)", msgType)
return false, nil, nil, err
}
// Read the push/pull header
var header pushPullHeader
if err := dec.Decode(&header); err != nil {
return false, nil, nil, err
}
// Allocate space for the transfer
remoteNodes := make([]pushNodeState, header.Nodes)
// Try to decode all the states
for i := 0; i < header.Nodes; i++ {
if err := dec.Decode(&remoteNodes[i]); err != nil {
return false, remoteNodes, nil, err
}
}
// Read the remote user state into a buffer
var userBuf []byte
if header.UserStateLen > 0 {
userBuf = make([]byte, header.UserStateLen)
bytes, err := io.ReadAtLeast(bufConn, userBuf, header.UserStateLen)
if err == nil && bytes != header.UserStateLen {
err = fmt.Errorf(
"Failed to read full user state (%d / %d)",
bytes, header.UserStateLen)
}
if err != nil {
return false, remoteNodes, nil, err
}
}
// For proto versions < 2, there is no port provided. Mask old
// behavior by using the configured port
for idx := range remoteNodes {
if m.ProtocolVersion() < 2 || remoteNodes[idx].Port == 0 {
remoteNodes[idx].Port = uint16(m.config.BindPort)
}
}
return header.Join, remoteNodes, userBuf, nil
}
func TestTCPPushPull(t *testing.T) {
m := GetMemberlist(t)
defer m.Shutdown()
m.nodes = append(m.nodes, &nodeState{
Node: Node{
Name: "Test 0",
Addr: net.ParseIP(m.config.BindAddr),
Port: uint16(m.config.BindPort),
},
Incarnation: 0,
State: stateSuspect,
StateChange: time.Now().Add(-1 * time.Second),
})
addr := fmt.Sprintf("%s:%d", m.config.BindAddr, m.config.BindPort)
conn, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("unexpected err %s", err)
}
defer conn.Close()
localNodes := make([]pushNodeState, 3)
localNodes[0].Name = "Test 0"
localNodes[0].Addr = net.ParseIP(m.config.BindAddr)
localNodes[0].Port = uint16(m.config.BindPort)
localNodes[0].Incarnation = 1
localNodes[0].State = stateAlive
localNodes[1].Name = "Test 1"
localNodes[1].Addr = net.ParseIP(m.config.BindAddr)
localNodes[1].Port = uint16(m.config.BindPort)
localNodes[1].Incarnation = 1
localNodes[1].State = stateAlive
localNodes[2].Name = "Test 2"
localNodes[2].Addr = net.ParseIP(m.config.BindAddr)
localNodes[2].Port = uint16(m.config.BindPort)
localNodes[2].Incarnation = 1
localNodes[2].State = stateAlive
// Send our node state
header := pushPullHeader{Nodes: 3}
hd := codec.MsgpackHandle{}
enc := codec.NewEncoder(conn, &hd)
// Send the push/pull indicator
conn.Write([]byte{byte(pushPullMsg)})
if err := enc.Encode(&header); err != nil {
t.Fatalf("unexpected err %s", err)
}
for i := 0; i < header.Nodes; i++ {
if err := enc.Encode(&localNodes[i]); err != nil {
t.Fatalf("unexpected err %s", err)
}
}
// Read the message type
var msgType messageType
if err := binary.Read(conn, binary.BigEndian, &msgType); err != nil {
t.Fatalf("unexpected err %s", err)
}
var bufConn io.Reader = conn
msghd := codec.MsgpackHandle{}
dec := codec.NewDecoder(bufConn, &msghd)
// Check if we have a compressed message
if msgType == compressMsg {
var c compress
if err := dec.Decode(&c); err != nil {
t.Fatalf("unexpected err %s", err)
}
decomp, err := decompressBuffer(&c)
if err != nil {
t.Fatalf("unexpected err %s", err)
}
// Reset the message type
msgType = messageType(decomp[0])
// Create a new bufConn
bufConn = bytes.NewReader(decomp[1:])
// Create a new decoder
dec = codec.NewDecoder(bufConn, &hd)
}
// Quit if not push/pull
if msgType != pushPullMsg {
t.Fatalf("bad message type")
}
if err := dec.Decode(&header); err != nil {
t.Fatalf("unexpected err %s", err)
}
// Allocate space for the transfer
remoteNodes := make([]pushNodeState, header.Nodes)
// Try to decode all the states
for i := 0; i < header.Nodes; i++ {
if err := dec.Decode(&remoteNodes[i]); err != nil {
t.Fatalf("unexpected err %s", err)
}
}
if len(remoteNodes) != 1 {
t.Fatalf("bad response")
}
n := &remoteNodes[0]
if n.Name != "Test 0" {
t.Fatalf("bad name")
}
if bytes.Compare(n.Addr, net.ParseIP(m.config.BindAddr)) != 0 {
t.Fatal("bad addr")
}
if n.Incarnation != 0 {
t.Fatal("bad incarnation")
}
if n.State != stateSuspect {
t.Fatal("bad state")
}
}
