package memberlist import ( "bytes" "fmt" "io" "net" "strconv" "time" ) // MockNetwork is used as a factory that produces MockTransport instances which // are uniquely addressed and wired up to talk to each other. type MockNetwork struct { transportsByAddr map[string]*MockTransport transportsByName map[string]*MockTransport port int } // NewTransport returns a new MockTransport with a unique address, wired up to // talk to the other transports in the MockNetwork. func (n *MockNetwork) NewTransport(name string) *MockTransport { n.port += 1 addr := fmt.Sprintf("127.0.0.1:%d", n.port) transport := &MockTransport{ net: n, addr: &MockAddress{addr, name}, packetCh: make(chan *Packet), streamCh: make(chan net.Conn), } if n.transportsByAddr == nil { n.transportsByAddr = make(map[string]*MockTransport) } n.transportsByAddr[addr] = transport if n.transportsByName == nil { n.transportsByName = make(map[string]*MockTransport) } n.transportsByName[name] = transport return transport } // MockAddress is a wrapper which adds the net.Addr interface to our mock // address scheme. type MockAddress struct { addr string name string } // See net.Addr. func (a *MockAddress) Network() string { return "mock" } // See net.Addr. func (a *MockAddress) String() string { return a.addr } // MockTransport directly plumbs messages to other transports its MockNetwork. type MockTransport struct { net *MockNetwork addr *MockAddress packetCh chan *Packet streamCh chan net.Conn } var _ NodeAwareTransport = (*MockTransport)(nil) // See Transport. func (t *MockTransport) FinalAdvertiseAddr(string, int) (net.IP, int, error) { host, portStr, err := net.SplitHostPort(t.addr.String()) if err != nil { return nil, 0, err } ip := net.ParseIP(host) if ip == nil { return nil, 0, fmt.Errorf("Failed to parse IP %q", host) } port, err := strconv.ParseInt(portStr, 10, 16) if err != nil { return nil, 0, err } return ip, int(port), nil } // See Transport. func (t *MockTransport) WriteTo(b []byte, addr string) (time.Time, error) { a := Address{Addr: addr, Name: ""} return t.WriteToAddress(b, a) } // See NodeAwareTransport. func (t *MockTransport) WriteToAddress(b []byte, a Address) (time.Time, error) { dest, err := t.getPeer(a) if err != nil { return time.Time{}, err } now := time.Now() dest.packetCh <- &Packet{ Buf: b, From: t.addr, Timestamp: now, } return now, nil } // See Transport. func (t *MockTransport) PacketCh() <-chan *Packet { return t.packetCh } // See NodeAwareTransport. func (t *MockTransport) IngestPacket(conn net.Conn, addr net.Addr, now time.Time, shouldClose bool) error { if shouldClose { defer conn.Close() } // Copy everything from the stream into packet buffer. var buf bytes.Buffer if _, err := io.Copy(&buf, conn); err != nil { return fmt.Errorf("failed to read packet: %v", err) } // Check the length - it needs to have at least one byte to be a proper // message. This is checked elsewhere for writes coming in directly from // the UDP socket. if n := buf.Len(); n < 1 { return fmt.Errorf("packet too short (%d bytes) %s", n, LogAddress(addr)) } // Inject the packet. t.packetCh <- &Packet{ Buf: buf.Bytes(), From: addr, Timestamp: now, } return nil } // See Transport. func (t *MockTransport) DialTimeout(addr string, timeout time.Duration) (net.Conn, error) { a := Address{Addr: addr, Name: ""} return t.DialAddressTimeout(a, timeout) } // See NodeAwareTransport. func (t *MockTransport) DialAddressTimeout(a Address, timeout time.Duration) (net.Conn, error) { dest, err := t.getPeer(a) if err != nil { return nil, err } p1, p2 := net.Pipe() dest.streamCh <- p1 return p2, nil } // See Transport. func (t *MockTransport) StreamCh() <-chan net.Conn { return t.streamCh } // See NodeAwareTransport. func (t *MockTransport) IngestStream(conn net.Conn) error { t.streamCh <- conn return nil } // See Transport. func (t *MockTransport) Shutdown() error { return nil } func (t *MockTransport) getPeer(a Address) (*MockTransport, error) { var ( dest *MockTransport ok bool ) if a.Name != "" { dest, ok = t.net.transportsByName[a.Name] } else { dest, ok = t.net.transportsByAddr[a.Addr] } if !ok { return nil, fmt.Errorf("No route to %s", a) } return dest, nil }