// Copyright 2023 The gVisor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package netfilter_test

import (
	_ "embed"
	"testing"

	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
	"gvisor.dev/gvisor/pkg/sentry/socket/netfilter"
	"gvisor.dev/gvisor/pkg/tcpip/header"
	"gvisor.dev/gvisor/pkg/tcpip/stack"
)

// FakeIDMapper implements IDMapper.
type FakeIDMapper struct{}

// MapToKUID implements IDMapper.
func (*FakeIDMapper) MapToKUID(auth.UID) auth.KUID {
	return 0
}

// MapToKGID implements IDMapper.
func (*FakeIDMapper) MapToKGID(auth.GID) auth.KGID {
	return 0
}

// istioBlob is an iptables NAT ruleset generated by Istio. It is already in
// the format ready to by passed to IPT_SO_SET_REPLACE.
//
// Updating this requires running Istio, calling IPT_SO_GET_INFO and
// IPT_SO_GET_ENTRIES, then stitching the structs up. So be careful when
// updating!
//
//go:embed istio_blob
var istioBlob []byte

// TestIstioBlob tests that we support the iptables ruleset generated by Istio,
// i.e. that we can parse the rules and set them in netstack.
func TestIstioBlob(t *testing.T) {
	mapper := FakeIDMapper{}
	stk := stack.New(stack.Options{})
	if err := netfilter.SetEntries(&mapper, stk, istioBlob, false); err != nil {
		t.Fatalf("failed to set Istio rules, try setting enableLogging and running again: %v", err)
	}
}

// acceptBlob is an iptables NAT ruleset that instructs iptables to ACCEPT all
// and modify nothing. This is the Linux default: a newly booted machine just
// doesn't NAT anything until someone sets the rules.
//
// Updating this requires running Istio, calling IPT_SO_GET_INFO and
// IPT_SO_GET_ENTRIES, then stitching the structs up. So be careful when
// updating!
//
//go:embed accept_blob
var acceptBlob []byte

// TestAcceptBlob tests that updating the default (all ACCEPT) NAT table with
// an identical table doesn't actually set the modified flag. In doing so, we
// can preserve our performance optimization in which iptables is off until
// non-default rules are set.
func TestAcceptBlob(t *testing.T) {
	mapper := FakeIDMapper{}
	stk := stack.New(stack.Options{
		DefaultIPTables: netfilter.DefaultLinuxTables,
	})
	if err := netfilter.SetEntries(&mapper, stk, acceptBlob, false); err != nil {
		t.Fatalf("failed to set Istio rules, try setting enableLogging and running again: %v", err)
	}
	if stk.IPTables().Modified() {
		t.Fatalf("ACCEPT rules shouldn't cause iptables modifications, but did")
	}
}

// TestMarkMatchingCheckOutput verifies that iptables Output filtering
// correctly matches and drops/accepts packets based on their mark.
func TestMarkMatchingCheckOutput(t *testing.T) {
	stk := stack.New(stack.Options{
		DefaultIPTables: netfilter.DefaultLinuxTables,
	})

	table := stack.Table{
		Rules: []stack.Rule{
			{
				Target: &stack.AcceptTarget{},
			},
			{
				Target: &stack.AcceptTarget{},
			},
			{
				// Output hook
				Matchers: []stack.Matcher{netfilter.NewMarkMatcher(0x1234, 0xffffffff, false)},
				Target:   &stack.DropTarget{},
			},
			{
				Target: &stack.AcceptTarget{},
			},
			{
				Target: &stack.ErrorTarget{},
			},
		},
		BuiltinChains: [stack.NumHooks]int{
			stack.Prerouting:  stack.HookUnset,
			stack.Input:       0,
			stack.Forward:     1,
			stack.Output:      2,
			stack.Postrouting: stack.HookUnset,
		},
		Underflows: [stack.NumHooks]int{
			stack.Prerouting:  stack.HookUnset,
			stack.Input:       0,
			stack.Forward:     1,
			stack.Output:      2,
			stack.Postrouting: stack.HookUnset,
		},
	}
	stk.IPTables().ForceReplaceTable(stack.FilterID, table, false /* ipv4 */)

	pkt1 := stack.NewPacketBuffer(stack.PacketBufferOptions{
		ReserveHeaderBytes: header.IPv4MinimumSize,
	})
	defer pkt1.DecRef()
	pkt1.NetworkProtocolNumber = header.IPv4ProtocolNumber
	ip1 := header.IPv4(pkt1.NetworkHeader().Push(header.IPv4MinimumSize))
	ip1.Encode(&header.IPv4Fields{
		TotalLength: header.IPv4MinimumSize,
		Protocol:    99,
	})
	pkt1.Mark = 0x5678 // Does not match 0x1234
	if ok := stk.IPTables().CheckOutput(pkt1, nil, ""); !ok {
		t.Errorf("CheckOutput(Mark=0x5678) = false (dropped), want true (accepted)")
	}

	pkt2 := stack.NewPacketBuffer(stack.PacketBufferOptions{
		ReserveHeaderBytes: header.IPv4MinimumSize,
	})
	defer pkt2.DecRef()
	pkt2.NetworkProtocolNumber = header.IPv4ProtocolNumber
	ip2 := header.IPv4(pkt2.NetworkHeader().Push(header.IPv4MinimumSize))
	ip2.Encode(&header.IPv4Fields{
		TotalLength: header.IPv4MinimumSize,
		Protocol:    99,
	})
	pkt2.Mark = 0x1234 // Matches 0x1234
	if ok := stk.IPTables().CheckOutput(pkt2, nil, ""); ok {
		t.Errorf("CheckOutput(Mark=0x1234) = true (accepted), want false (dropped)")
	}
}
