package lb

import (
	"fmt"
	"net"
	"strings"
	"testing"
	"time"

	coreapi "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

// the S1 fixtures: cp, w1 (holds 86.107.168.15), w2

func egressFor(name string, status EgressServiceStatus) *EgressService {
	return &EgressService{ObjectMeta: metav1.ObjectMeta{Namespace: "gitea", Name: name},
		Spec: EgressServiceSpec{Emit: EmitHolder}, Status: status}
}

func mailSvc() *coreapi.Service {
	return svc("mail", "86.107.168.15", func(s *coreapi.Service) {
		s.Spec.Selector = map[string]string{"app": "mail"}
	})
}

func pod(name, node, ip string, labels map[string]string, mutate ...func(*coreapi.Pod)) *coreapi.Pod {
	p := &coreapi.Pod{ObjectMeta: metav1.ObjectMeta{Namespace: "gitea", Name: name, Labels: labels,
		CreationTimestamp: metav1.NewTime(time.Unix(1000, 0))},
		Spec: coreapi.PodSpec{NodeName: node}, Status: coreapi.PodStatus{PodIP: ip, Phase: coreapi.PodRunning}}
	for _, m := range mutate {
		m(p)
	}
	return p
}

var mailLabels = map[string]string{"app": "mail"}

// the status the holder (w1) writes for three nodes, blocks in name order
func written() EgressServiceStatus {
	return EgressServiceStatus{Host: "w1", Address: "86.107.168.15", Blocks: []PortBlock{
		{Node: "cp", From: 61000, To: 61511}, {Node: "w1", From: 61512, To: 62023}, {Node: "w2", From: 62024, To: 62535}}}
}

func TestBlocksAreStableAndCapped(t *testing.T) {
	got := AllocateBlocks(nil, []string{"w2", "cp", "w1"})
	if !sameStatus(EgressServiceStatus{Blocks: got}, EgressServiceStatus{Blocks: written().Blocks}) {
		t.Fatalf("fresh blocks: %+v", got)
	}
	// a node leaves, another joins: survivors keep their blocks, the newcomer takes the freed lowest one
	got = AllocateBlocks(got, []string{"w1", "w2", "w3"})
	want := []PortBlock{{"w1", 61512, 62023}, {"w2", 62024, 62535}, {"w3", 61000, 61511}}
	if !sameStatus(EgressServiceStatus{Blocks: got}, EgressServiceStatus{Blocks: want}) {
		t.Fatalf("after churn: %+v", got)
	}
	// misaligned or duplicate blocks are not kept
	got = AllocateBlocks([]PortBlock{{"w1", 61001, 61512}, {"w2", 61000, 61511}, {"w3", 61000, 61511}}, []string{"w1", "w2", "w3"})
	for _, b := range got {
		if (b.From-BlockBase)%BlockSize != 0 || b.To != b.From+BlockSize-1 {
			t.Fatalf("invalid block kept: %+v", got)
		}
	}
	if got[1].Node != "w2" || got[1].From != 61000 {
		t.Fatalf("first valid claim wins: %+v", got)
	}
	var many []string
	for i := 0; i < BlockCount+2; i++ {
		many = append(many, string(rune('a'+i)))
	}
	if n := len(AllocateBlocks(nil, many)); n != BlockCount {
		t.Fatalf("more blocks than exist: %d", n)
	}
	for _, b := range AllocateBlocks(nil, many) {
		if b.To > 65535 || b.From < 61000 {
			t.Fatalf("block outside the range: %+v", b)
		}
	}
}

func TestHolderWritesStatusAndMembersWaitForIt(t *testing.T) {
	e := egressFor("mail", EgressServiceStatus{})
	pods := []*coreapi.Pod{pod("mail-0", "w2", "10.44.0.9", mailLabels)}
	d := ComputeEgress("w1", nodes, []*coreapi.Service{mailSvc()}, []*EgressService{e}, nil, nil)
	if len(d.Status) != 1 || !sameStatus(d.Status[0].Status, written()) {
		t.Fatalf("holder's status: %+v", d.Status)
	}
	// until the status is in the API, a member's outside traffic goes nowhere (fail closed)
	d = ComputeEgress("w2", nodes, []*coreapi.Service{mailSvc()}, []*EgressService{e}, pods, nil)
	if len(d.Status) != 0 || len(d.Members) != 1 || !d.Members[0].Unreachable {
		t.Fatalf("member before status: %+v", d)
	}
	if d.Tunnelled("w2") {
		t.Fatalf("no tunnel for an unreachable member")
	}
}

func TestMemberTunnelsToHolderWhichRoutesItsBlockBack(t *testing.T) {
	e := egressFor("mail", written())
	svcs := []*coreapi.Service{mailSvc()}
	pods := []*coreapi.Pod{
		pod("mail-0", "w2", "10.44.0.9", mailLabels),
		pod("web", "w2", "10.44.0.10", map[string]string{"app": "web"}), // not a member (A4)
	}
	d := ComputeEgress("w2", nodes, svcs, []*EgressService{e}, pods, nil)
	if len(d.Members) != 1 || d.Members[0].Pod != "gitea/mail-0" {
		t.Fatalf("members: %+v", d.Members)
	}
	m := d.Members[0]
	if m.Unreachable || !m.Address.Equal(net.ParseIP("86.107.168.15")) || m.From != 62024 || m.To != 62535 || m.Holder.Name != "w1" {
		t.Fatalf("member path: %+v", m)
	}
	if !d.Tunnelled("w2") || len(d.Peers("w2")) != 1 || d.Peers("w2")[0].Name != "w1" {
		t.Fatalf("w2 tunnels to w1 only: %+v", d.Peers("w2"))
	}
	if len(d.Status) != 0 || len(d.Routes) != 0 {
		t.Fatalf("only the holder writes and routes: %+v", d)
	}
	// the holder: blocks of the other nodes route back to them; its own block is plain conntrack
	h := ComputeEgress("w1", nodes, svcs, []*EgressService{e}, nil, nil)
	if len(h.Status) != 0 {
		t.Fatalf("status already right: %+v", h.Status)
	}
	if len(h.Routes) != 2 || h.Routes[0].Peer.Name != "cp" || h.Routes[1].Peer.Name != "w2" || h.Routes[1].From != 62024 {
		t.Fatalf("holder routes: %+v", h.Routes)
	}
	// a member on the holder leaves from there
	hm := ComputeEgress("w1", nodes, svcs, []*EgressService{e}, []*coreapi.Pod{pod("mail-1", "w1", "10.40.0.4", mailLabels)}, nil)
	if len(hm.Members) != 1 || hm.Members[0].Holder.Name != "w1" || hm.Members[0].From != 61512 {
		t.Fatalf("member on the holder: %+v", hm.Members)
	}
}

func TestEgressFailsClosed(t *testing.T) {
	pods := []*coreapi.Pod{pod("mail-0", "w2", "10.44.0.9", mailLabels)}
	// no node holds the address: status cleared (by the first node by name), member dropped
	unheld := svc("mail", "86.107.168.99", func(s *coreapi.Service) { s.Spec.Selector = mailLabels })
	for _, self := range []string{"cp", "w2"} {
		d := ComputeEgress(self, nodes, []*coreapi.Service{unheld}, []*EgressService{egressFor("mail", written())}, pods, nil)
		clears := len(d.Status) == 1 && d.Status[0].Status.Host == ""
		if clears != (self == "cp") {
			t.Fatalf("%s clears=%v: %+v", self, clears, d.Status)
		}
		if self == "w2" && (len(d.Members) != 1 || !d.Members[0].Unreachable) {
			t.Fatalf("member of an unheld address: %+v", d.Members)
		}
	}
	// two EgressServices select the pod: dropped, reported
	other := svc("mail2", "86.107.168.15", func(s *coreapi.Service) { s.Spec.Selector = mailLabels })
	d := ComputeEgress("w2", nodes, []*coreapi.Service{mailSvc(), other},
		[]*EgressService{egressFor("mail", written()), egressFor("mail2", written())}, pods, nil)
	if len(d.Members) != 1 || !d.Members[0].Unreachable || !strings.Contains(strings.Join(d.Problems, "\n"), "several EgressServices") {
		t.Fatalf("double member: %+v", d)
	}
	// a spec this agent does not implement is refused, not guessed at
	bad := egressFor("mail", written())
	bad.Spec.Emit = "Direct"
	d = ComputeEgress("w2", nodes, []*coreapi.Service{mailSvc()}, []*EgressService{bad}, pods, nil)
	if len(d.Members) != 1 || !d.Members[0].Unreachable {
		t.Fatalf("unknown emit mode: %+v", d)
	}
	// hostNetwork pods are never members (their traffic is the node's)
	hn := pod("mail-h", "w2", "86.107.168.88", mailLabels, func(p *coreapi.Pod) { p.Spec.HostNetwork = true })
	d = ComputeEgress("w2", nodes, []*coreapi.Service{mailSvc()}, []*EgressService{egressFor("mail", written())}, []*coreapi.Pod{hn}, nil)
	if len(d.Members) != 0 {
		t.Fatalf("hostNetwork pod became a member: %+v", d.Members)
	}
}

func TestAdmissionDecidesWhoOwnsAnIP(t *testing.T) {
	e := egressFor("mail", written())
	svcs := []*coreapi.Service{mailSvc()}
	old := pod("mail-0", "w2", "10.44.0.9", mailLabels, func(p *coreapi.Pod) {
		now := metav1.Now()
		p.DeletionTimestamp = &now
	})
	newer := pod("web", "w2", "", map[string]string{"app": "web"}) // API has no IP for it yet
	// the CNI plugin just gave 10.44.0.9 to web: the terminating member must not keep it (A4)
	d := ComputeEgress("w2", nodes, svcs, []*EgressService{e}, []*coreapi.Pod{old, newer},
		map[string]Admission{"10.44.0.9": {Namespace: "gitea", Name: "web"}})
	if len(d.Members) != 0 {
		t.Fatalf("reused IP kept the old member's rules: %+v", d.Members)
	}
	// and the other way round: a member admitted before the API shows its IP is protected at once (A1)
	fresh := pod("mail-1", "w2", "", mailLabels)
	d = ComputeEgress("w2", nodes, svcs, []*EgressService{e}, []*coreapi.Pod{fresh},
		map[string]Admission{"10.44.0.11": {Namespace: "gitea", Name: "mail-1"}})
	if len(d.Members) != 1 || d.Members[0].IP.String() != "10.44.0.11" || d.Members[0].Unreachable {
		t.Fatalf("admitted member: %+v", d.Members)
	}
	// without an admission, a live pod wins a shared IP over a terminating one
	live := pod("web", "w2", "10.44.0.9", map[string]string{"app": "web"})
	d = ComputeEgress("w2", nodes, svcs, []*EgressService{e}, []*coreapi.Pod{old, live}, nil)
	if len(d.Members) != 0 {
		t.Fatalf("terminating member won the IP: %+v", d.Members)
	}
}

func TestRenderedEgressRules(t *testing.T) {
	_, podNet, _ := net.ParseCIDR("10.32.0.0/12")
	e := egressFor("mail", written())
	svcs := []*coreapi.Service{mailSvc()}
	toW1, toW2, toCP := TunnelName("w1"), TunnelName("w2"), TunnelName("cp")
	member := ComputeEgress("w2", nodes, svcs, []*EgressService{e}, []*coreapi.Pod{pod("mail-0", "w2", "10.44.0.9", mailLabels)}, nil)
	r := Render(Desired{}, member, "w2", podNet, nil)
	for _, want := range []string{
		"-A REWEAVE-LB-CLASS -d 10.32.0.0/12 -j RETURN\n",
		"-A REWEAVE-LB-CLASS -d 86.107.168.24/32 -j RETURN\n", // a node is not outside
		"-A REWEAVE-LB-CLASS -s 10.44.0.9/32 -m comment --comment \"gitea/mail-0\" -j CONNMARK --set-xmark 0x100000/0x100000\n",
		"-A REWEAVE-LB-CLASS -s 10.44.0.9/32 -j MARK --set-xmark 0x100000/0x100000\n",
		// port protocols within this node's block, into the holder's own tunnel
		"-A REWEAVE-LB-POST -s 10.44.0.9/32 -o " + toW1 + " -m mark --mark 0x100000/0x100000 -p tcp -j SNAT --to-source 86.107.168.15:62024-62535\n",
		"-A REWEAVE-LB-POST -s 10.44.0.9/32 -o " + toW1 + " -m mark --mark 0x100000/0x100000 -p udp -j SNAT --to-source 86.107.168.15:62024-62535\n",
		"-A REWEAVE-LB-POST -s 10.44.0.9/32 -o " + toW1 + " -m mark --mark 0x100000/0x100000 -p sctp -j SNAT --to-source 86.107.168.15:62024-62535\n",
		// everything else (ping, GRE, ...) from the address alone
		"-A REWEAVE-LB-POST -s 10.44.0.9/32 -o " + toW1 + " -m mark --mark 0x100000/0x100000 -j SNAT --to-source 86.107.168.15\n",
	} {
		if !strings.Contains(r.Nat, want) {
			t.Fatalf("member nat lacks %q:\n%s", want, r.Nat)
		}
	}
	// the catch-all SNAT comes after the port-ranged ones (or ports would not stay in the block)
	if strings.Index(r.Nat, "-j SNAT --to-source 86.107.168.15\n") < strings.Index(r.Nat, "-p sctp -j SNAT") {
		t.Fatalf("SNAT order:\n%s", r.Nat)
	}
	for _, want := range []string{
		"-A REWEAVE-LB-FWD -s 10.44.0.9/32 -o " + toW1 + " -m mark --mark 0x100000/0x100000 -j RETURN\n",
		"-A REWEAVE-LB-FWD -s 10.44.0.9/32 -m mark --mark 0x100000/0x100000 -m comment --comment \"gitea/mail-0\" -j DROP\n",
		"-A REWEAVE-LB-FWD -i rweg+ -d 86.107.168.15/32 -j DROP\n",
		"-A REWEAVE-LB-FWD -i rweg+ -m conntrack --ctstate ESTABLISHED,RELATED -j RETURN\n",
		"-A REWEAVE-LB-FWD -i rweg+ -j DROP\n",
		"-A REWEAVE-LB-IN -p 4 -s 10.40.0.0/32 -m mac --mac-source aa:aa:aa:aa:aa:02 -j RETURN\n",
		"-A REWEAVE-LB-IN -p 4 -j DROP\n",
	} {
		if !strings.Contains(r.Filter, want) {
			t.Fatalf("member filter lacks %q:\n%s", want, r.Filter)
		}
	}
	for _, want := range []string{
		"-A REWEAVE-LB-MARK -m conntrack --ctdir ORIGINAL -m connmark --mark 0x100000/0x100000 -j MARK --set-xmark 0x100000/0x100000\n",
		"-A REWEAVE-LB-MSS -o rweg+ -p tcp -m tcp --tcp-flags SYN,RST SYN -j TCPMSS --clamp-mss-to-pmtu\n",
	} {
		if !strings.Contains(r.Mangle, want) {
			t.Fatalf("member mangle lacks %q:\n%s", want, r.Mangle)
		}
	}
	if strings.Contains(r.Mangle, "--set-xmark 0x") && strings.Contains(r.Mangle, "-i "+toW1) {
		t.Fatalf("a member's node does not tag what its holder sends back:\n%s", r.Mangle)
	}
	// the guard's DROP comes after the member's own RETURN: nothing else of its leaves
	if strings.Index(r.Filter, toW1+" -m mark --mark 0x100000/0x100000 -j RETURN") > strings.Index(r.Filter, "\"gitea/mail-0\" -j DROP") {
		t.Fatalf("guard order:\n%s", r.Filter)
	}

	// the holder: tags each tunnel's connections with its node, accepts each node's block, nothing foreign
	h := ComputeEgress("w1", nodes, svcs, []*EgressService{e}, []*coreapi.Pod{pod("mail-1", "w1", "10.40.0.4", mailLabels)}, nil)
	r = Render(Desired{}, h, "w1", podNet, nil)
	idx := peerIndex(h, "w1")
	tagW2 := fmt.Sprintf("0x%x", peerTag(idx["w2"]))
	for _, want := range []string{
		"-A REWEAVE-LB-MARK -m conntrack --ctdir REPLY -m connmark --mark 0x200000/0x200000 -j CONNMARK --restore-mark --nfmask 0xff200000 --ctmask 0xff200000\n",
		"-A REWEAVE-LB-MARK -i " + toW2 + " -m conntrack --ctstate NEW -m connmark --mark 0x0/0xff200000 -m comment --comment \"from w2\" -j CONNMARK --set-xmark " + tagW2 + "/0xff200000\n",
	} {
		if !strings.Contains(r.Mangle, want) {
			t.Fatalf("holder mangle lacks %q:\n%s", want, r.Mangle)
		}
	}
	for _, want := range []string{
		"-A REWEAVE-LB-FWD -i " + toW2 + " -m conntrack --ctdir ORIGINAL -m connmark ! --mark " + tagW2 + "/0xff200000 -j DROP\n",
		"-A REWEAVE-LB-FWD -i " + toW2 + " -m conntrack --ctdir ORIGINAL -s 86.107.168.15/32 -p tcp -m tcp --sport 62024:62535 -m comment --comment \"block of w2\" -j RETURN\n",
		"-A REWEAVE-LB-FWD -i " + toCP + " -m conntrack --ctdir ORIGINAL -s 86.107.168.15/32 -p sctp -m sctp --sport 61000:61511 -m comment --comment \"block of cp\" -j RETURN\n",
		"-A REWEAVE-LB-FWD -i " + toW2 + " -m conntrack --ctdir ORIGINAL -s 86.107.168.15/32 -p udp -j DROP\n",
		"-A REWEAVE-LB-FWD -i " + toW2 + " -m conntrack --ctdir ORIGINAL -s 86.107.168.15/32 -j RETURN\n",
		"-A REWEAVE-LB-FWD -s 10.40.0.4/32 -o weave -m mark --mark 0x100000/0x100000 -j DROP\n",
	} {
		if !strings.Contains(r.Filter, want) {
			t.Fatalf("holder filter lacks %q:\n%s", want, r.Filter)
		}
	}
	// foreign tags are refused before anything of that tunnel is accepted; a port outside the block before the
	// catch-all accept
	if strings.Index(r.Filter, "! --mark "+tagW2) > strings.Index(r.Filter, "\"block of w2\"") ||
		strings.Index(r.Filter, "-i "+toW2+" -m conntrack --ctdir ORIGINAL -s 86.107.168.15/32 -p udp -j DROP") >
			strings.Index(r.Filter, "-i "+toW2+" -m conntrack --ctdir ORIGINAL -s 86.107.168.15/32 -j RETURN") {
		t.Fatalf("holder filter order:\n%s", r.Filter)
	}
	if !strings.Contains(r.Nat, "-A REWEAVE-LB-POST -s 10.40.0.4/32 ! -o weave -m mark --mark 0x100000/0x100000 -p tcp -j SNAT --to-source 86.107.168.15:61512-62023\n") {
		t.Fatalf("holder's own member:\n%s", r.Nat)
	}
	for _, line := range strings.Split(r.Nat+r.Filter, "\n") {
		if strings.Contains(line, "10.44.0.10") {
			t.Fatalf("a rule names a non-member: %s", line)
		}
	}
}

func TestTunnelNames(t *testing.T) {
	a, b := TunnelName("reweave-worker-fe19b040"), TunnelName("reweave-cp-94a8db43")
	if a == b || len(a) > 15 || !strings.HasPrefix(a, TunnelPrefix) || a != TunnelName("reweave-worker-fe19b040") {
		t.Fatalf("tunnel names must be distinct, stable, short: %q %q", a, b)
	}
}

func TestPolicyRules(t *testing.T) {
	e := egressFor("mail", written())
	svcs := []*coreapi.Service{mailSvc()}
	h := ComputeEgress("w1", nodes, svcs, []*EgressService{e}, nil, nil)
	rules := wantRules(h, "w1")
	tables := peerTables(h, "w1")
	idx := peerIndex(h, "w1")
	ports, tags := 0, 0
	for _, r := range rules {
		if r.priority != HolderRulePriority || r.src != "" {
			t.Fatalf("holder rule: %+v", r)
		}
		switch {
		case r.ipproto != 0:
			ports++
			if r.dst != "86.107.168.15/32" || r.mark != 0 {
				t.Fatalf("port rule: %+v", r)
			}
			if (r.from == 62024) != (r.table == tables["w2"]) || (r.from == 61000) != (r.table == tables["cp"]) {
				t.Fatalf("block routed to the wrong node: %+v %v", r, tables)
			}
		default:
			tags++
			if r.mask != peerMask || r.dst != "" {
				t.Fatalf("tag rule: %+v", r)
			}
			for name, i := range idx {
				if r.mark == peerTag(i) && r.table != tables[name] {
					t.Fatalf("tag of %s routed to table %d, want %d", name, r.table, tables[name])
				}
			}
		}
	}
	if ports != 2*len(portProtocols) || tags != 2 {
		t.Fatalf("2 peers: %d port rules (tcp, udp, sctp each), %d tag rules: %+v", ports, tags, rules)
	}
	if HolderRulePriority >= LocalRulePriority {
		t.Fatalf("holder rules must precede the local table")
	}
	m := ComputeEgress("w2", nodes, svcs, []*EgressService{e}, []*coreapi.Pod{pod("mail-0", "w2", "10.44.0.9", mailLabels)}, nil)
	rules = wantRules(m, "w2")
	if len(rules) != 1 || rules[0].src != "10.44.0.9/32" || rules[0].mark != EgressMark || rules[0].mask != EgressMark ||
		rules[0].table != peerTables(m, "w2")["w1"] {
		t.Fatalf("member rule: %+v", rules)
	}
	u := ComputeEgress("w2", nodes, svcs, []*EgressService{egressFor("mail", EgressServiceStatus{})}, []*coreapi.Pod{pod("mail-0", "w2", "10.44.0.9", mailLabels)}, nil)
	rules = wantRules(u, "w2")
	if len(rules) != 1 || rules[0].table != EgressTableBase {
		t.Fatalf("unreachable member must hit the unreachable table: %+v", rules)
	}
}