// Copyright 2016 The Netstack Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ipv6

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/google/netstack/tcpip"
	"github.com/google/netstack/tcpip/buffer"
	"github.com/google/netstack/tcpip/header"
	"github.com/google/netstack/tcpip/link/channel"
	"github.com/google/netstack/tcpip/link/sniffer"
	"github.com/google/netstack/tcpip/stack"
)

const (
	linkAddr0 = tcpip.LinkAddress("\x01\x02\x03\x04\x05\x06")
	linkAddr1 = tcpip.LinkAddress("\x0a\x0b\x0c\x0d\x0e\x0f")
)

var (
	lladdr0 = LinkLocalAddr(linkAddr0)
	lladdr1 = LinkLocalAddr(linkAddr1)
)

type testContext struct {
	t  *testing.T
	s0 *stack.Stack
	s1 *stack.Stack

	linkEP0 *channel.Endpoint
	linkEP1 *channel.Endpoint

	icmpCh chan header.ICMPv6Type
}

func newTestContext(t *testing.T) *testContext {
	c := &testContext{
		t:      t,
		s0:     stack.New([]string{ProtocolName}, nil),
		s1:     stack.New([]string{ProtocolName}, nil),
		icmpCh: make(chan header.ICMPv6Type, 10),
	}

	const defaultMTU = 65536
	id0, linkEP0 := channel.New(256, defaultMTU, linkAddr0)
	c.linkEP0 = linkEP0
	if testing.Verbose() {
		id0 = sniffer.New(id0)
	}
	if err := c.s0.CreateNIC(1, id0); err != nil {
		t.Fatalf("CreateNIC s0: %v", err)
	}
	if err := c.s0.AddAddress(1, ProtocolNumber, lladdr0); err != nil {
		t.Fatalf("AddAddress lladdr0: %v", err)
	}
	if err := c.s0.AddAddress(1, ProtocolNumber, SolicitedNodeAddr(lladdr0)); err != nil {
		t.Fatalf("AddAddress sn lladdr0: %v", err)
	}

	id1, linkEP1 := channel.New(256, defaultMTU, linkAddr1)
	c.linkEP1 = linkEP1
	if err := c.s1.CreateNIC(1, id1); err != nil {
		t.Fatalf("CreateNIC failed: %v", err)
	}
	if err := c.s1.AddAddress(1, ProtocolNumber, lladdr1); err != nil {
		t.Fatalf("AddAddress lladdr1: %v", err)
	}
	if err := c.s1.AddAddress(1, ProtocolNumber, SolicitedNodeAddr(lladdr1)); err != nil {
		t.Fatalf("AddAddress sn lladdr1: %v", err)
	}

	routeTable := []tcpip.Route{{
		Destination: tcpip.Address(strings.Repeat("\x00", 16)),
		Mask:        tcpip.Address(strings.Repeat("\x00", 16)),
		NIC:         1,
	}}
	c.s0.SetRouteTable(routeTable)
	c.s1.SetRouteTable(routeTable)

	go c.routePackets(linkEP0.C, linkEP1)
	go c.routePackets(linkEP1.C, linkEP0)

	return c
}

func (c *testContext) countPacket(pkt channel.PacketInfo) {
	if pkt.Proto != header.IPv6ProtocolNumber {
		return
	}
	ipv6 := header.IPv6(pkt.Header)
	transProto := tcpip.TransportProtocolNumber(ipv6.NextHeader())
	if transProto != header.ICMPv6ProtocolNumber {
		return
	}
	b := pkt.Header[header.IPv6MinimumSize:]
	icmp := header.ICMPv6(b)
	c.icmpCh <- icmp.Type()
}

func (c *testContext) routePackets(ch <-chan channel.PacketInfo, ep *channel.Endpoint) {
	for pkt := range ch {
		c.countPacket(pkt)
		views := []buffer.View{pkt.Header, pkt.Payload}
		size := len(pkt.Header) + len(pkt.Payload)
		vv := buffer.NewVectorisedView(size, views)
		ep.InjectLinkAddr(pkt.Proto, ep.LinkAddress(), &vv)
	}
}

func (c *testContext) cleanup() {
	close(c.linkEP0.C)
	close(c.linkEP1.C)
}

func TestLinkResolution(t *testing.T) {
	c := newTestContext(t)
	defer c.cleanup()
	r, err := c.s0.FindRoute(1, lladdr0, lladdr1, header.IPv6ProtocolNumber)
	if err != nil {
		t.Fatal(err)
	}
	defer r.Release()

	hdr := buffer.NewPrependable(int(r.MaxHeaderLength()) + header.IPv6MinimumSize + header.ICMPv6EchoMinimumSize)
	pkt := header.ICMPv6(hdr.Prepend(header.ICMPv6EchoMinimumSize))
	pkt.SetType(header.ICMPv6EchoRequest)
	pkt.SetChecksum(icmpChecksum(pkt, r.LocalAddress, r.RemoteAddress, nil))
	if err := r.WritePacket(&hdr, nil, header.ICMPv6ProtocolNumber, header.IPv6DefaultHopLimit); err != nil {
		t.Fatal(err)
	}

	// This actually takes about 10 milliseconds, so no need to wait for
	// a multi-minute go test timeout if something is broken.
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	stats := make(map[header.ICMPv6Type]int)
	for {
		select {
		case <-ctx.Done():
			t.Errorf("timeout waiting for ICMP, got: %#+v", stats)
			return
		case typ := <-c.icmpCh:
			stats[typ]++

			if stats[header.ICMPv6NeighborSolicit] > 0 &&
				stats[header.ICMPv6NeighborAdvert] > 0 &&
				stats[header.ICMPv6EchoRequest] > 0 &&
				stats[header.ICMPv6EchoReply] > 0 {
				return
			}
		}
	}
}