// Copyright 2014 The Android Open Source Project
//
// 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.

#include "iptables.h"

#include <gtest/gtest.h>

#include "mock_iptables.h"

namespace firewalld {

using testing::_;
using testing::Return;

class IpTablesTest : public testing::Test {
 public:
  IpTablesTest() = default;
  ~IpTablesTest() override = default;

 protected:
  void SetMockExpectations(MockIpTables* iptables, bool success) {
    // Empty criterion matches all commands.
    iptables->SetRunInMinijailFailCriterion(std::vector<std::string>(),
                                            true, success);
  }

  void SetMockExpectationsPerExecutable(MockIpTables* iptables,
                                        bool ip4_success,
                                        bool ip6_success) {
    if (!ip4_success)
      iptables->SetRunInMinijailFailCriterion(
          std::vector<std::string>({kIpTablesPath}), true /* repeat */,
                                   false /* omit_failure */);
    if (!ip6_success)
      iptables->SetRunInMinijailFailCriterion(
          std::vector<std::string>({kIp6TablesPath}), true, false);
  }

 private:
  DISALLOW_COPY_AND_ASSIGN(IpTablesTest);
};

TEST_F(IpTablesTest, Port0Fails) {
  MockIpTables mock_iptables;
  // Don't fail on anything.
  int id = mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>(), true, true);

  // Try to punch hole for TCP port 0, port 0 is not a valid port.
  EXPECT_FALSE(mock_iptables.PunchTcpHole(0, "iface"));
  // Try to punch hole for UDP port 0, port 0 is not a valid port.
  EXPECT_FALSE(mock_iptables.PunchUdpHole(0, "iface"));

  // We should not be adding any rules for port 0.
  EXPECT_TRUE(mock_iptables.GetRunInMinijailCriterionMatchCount(id) == 0);
}

TEST_F(IpTablesTest, ValidInterfaceName) {
  MockIpTables mock_iptables;
  SetMockExpectations(&mock_iptables, true /* success */);

  EXPECT_TRUE(mock_iptables.PunchTcpHole(80, "shortname"));
  EXPECT_TRUE(mock_iptables.PunchUdpHole(53, "shortname"));
  EXPECT_TRUE(mock_iptables.PunchTcpHole(80, "middle-dash"));
  EXPECT_TRUE(mock_iptables.PunchUdpHole(53, "middle-dash"));
  EXPECT_TRUE(mock_iptables.PunchTcpHole(80, "middle.dot"));
  EXPECT_TRUE(mock_iptables.PunchUdpHole(53, "middle.dot"));
}

TEST_F(IpTablesTest, InvalidInterfaceName) {
  MockIpTables mock_iptables;
  int id = mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>(), true, true);

  EXPECT_FALSE(mock_iptables.PunchTcpHole(80, "reallylonginterfacename"));
  EXPECT_FALSE(mock_iptables.PunchTcpHole(80, "with spaces"));
  EXPECT_FALSE(mock_iptables.PunchTcpHole(80, "with$ymbols"));
  EXPECT_FALSE(mock_iptables.PunchTcpHole(80, "-startdash"));
  EXPECT_FALSE(mock_iptables.PunchTcpHole(80, "enddash-"));
  EXPECT_FALSE(mock_iptables.PunchTcpHole(80, ".startdot"));
  EXPECT_FALSE(mock_iptables.PunchTcpHole(80, "enddot."));

  EXPECT_FALSE(mock_iptables.PunchUdpHole(53, "reallylonginterfacename"));
  EXPECT_FALSE(mock_iptables.PunchUdpHole(53, "with spaces"));
  EXPECT_FALSE(mock_iptables.PunchUdpHole(53, "with$ymbols"));
  EXPECT_FALSE(mock_iptables.PunchUdpHole(53, "-startdash"));
  EXPECT_FALSE(mock_iptables.PunchUdpHole(53, "enddash-"));
  EXPECT_FALSE(mock_iptables.PunchUdpHole(53, ".startdot"));
  EXPECT_FALSE(mock_iptables.PunchUdpHole(53, "enddot."));

  // We should not be adding any rules for invalid interface names.
  EXPECT_TRUE(mock_iptables.GetRunInMinijailCriterionMatchCount(id) == 0);
}

TEST_F(IpTablesTest, PunchTcpHoleSucceeds) {
  MockIpTables mock_iptables;
  SetMockExpectations(&mock_iptables, true /* success */);

  // Punch hole for TCP port 80, should succeed.
  EXPECT_TRUE(mock_iptables.PunchTcpHole(80, "iface"));
  // Punch again, should still succeed.
  EXPECT_TRUE(mock_iptables.PunchTcpHole(80, "iface"));
  // Plug the hole, should succeed.
  EXPECT_TRUE(mock_iptables.PlugTcpHole(80, "iface"));
}

TEST_F(IpTablesTest, PlugTcpHoleSucceeds) {
  MockIpTables mock_iptables;
  SetMockExpectations(&mock_iptables, true /* success */);

  // Punch hole for TCP port 80, should succeed.
  EXPECT_TRUE(mock_iptables.PunchTcpHole(80, "iface"));
  // Plug the hole, should succeed.
  EXPECT_TRUE(mock_iptables.PlugTcpHole(80, "iface"));
  // Plug again, should fail.
  EXPECT_FALSE(mock_iptables.PlugTcpHole(80, "iface"));
}

TEST_F(IpTablesTest, PunchUdpHoleSucceeds) {
  MockIpTables mock_iptables;
  SetMockExpectations(&mock_iptables, true /* success */);

  // Punch hole for UDP port 53, should succeed.
  EXPECT_TRUE(mock_iptables.PunchUdpHole(53, "iface"));
  // Punch again, should still succeed.
  EXPECT_TRUE(mock_iptables.PunchUdpHole(53, "iface"));
  // Plug the hole, should succeed.
  EXPECT_TRUE(mock_iptables.PlugUdpHole(53, "iface"));
}

TEST_F(IpTablesTest, PlugUdpHoleSucceeds) {
  MockIpTables mock_iptables;
  SetMockExpectations(&mock_iptables, true /* success */);

  // Punch hole for UDP port 53, should succeed.
  EXPECT_TRUE(mock_iptables.PunchUdpHole(53, "iface"));
  // Plug the hole, should succeed.
  EXPECT_TRUE(mock_iptables.PlugUdpHole(53, "iface"));
  // Plug again, should fail.
  EXPECT_FALSE(mock_iptables.PlugUdpHole(53, "iface"));
}

TEST_F(IpTablesTest, PunchTcpHoleFails) {
  MockIpTables mock_iptables;
  SetMockExpectations(&mock_iptables, false /* success */);
  // Punch hole for TCP port 80, should fail.
  ASSERT_FALSE(mock_iptables.PunchTcpHole(80, "iface"));
}

TEST_F(IpTablesTest, PunchUdpHoleFails) {
  MockIpTables mock_iptables;
  SetMockExpectations(&mock_iptables, false /* success */);
  // Punch hole for UDP port 53, should fail.
  ASSERT_FALSE(mock_iptables.PunchUdpHole(53, "iface"));
}

TEST_F(IpTablesTest, PunchTcpHoleIpv6Fails) {
  MockIpTables mock_iptables;
  SetMockExpectationsPerExecutable(
      &mock_iptables, true /* ip4_success */, false /* ip6_success */);
  // Punch hole for TCP port 80, should fail because 'ip6tables' fails.
  ASSERT_FALSE(mock_iptables.PunchTcpHole(80, "iface"));
}

TEST_F(IpTablesTest, PunchUdpHoleIpv6Fails) {
  MockIpTables mock_iptables;
  SetMockExpectationsPerExecutable(
      &mock_iptables, true /* ip4_success */, false /* ip6_success */);
  // Punch hole for UDP port 53, should fail because 'ip6tables' fails.
  ASSERT_FALSE(mock_iptables.PunchUdpHole(53, "iface"));
}

TEST_F(IpTablesTest, ApplyVpnSetupAdd_Success) {
  const std::vector<std::string> usernames = {"testuser0", "testuser1"};
  const std::string interface = "ifc0";
  const bool add = true;

  MockIpTables mock_iptables;
  ASSERT_TRUE(
      mock_iptables.ApplyVpnSetup(usernames, interface, add));
}

TEST_F(IpTablesTest, ApplyVpnSetupAdd_FailureInUsername) {
  const std::vector<std::string> usernames = {"testuser0", "testuser1"};
  const std::string interface = "ifc0";
  const bool add = true;

  MockIpTables mock_iptables;
  // Fail on all commands issued for testuser1.
  int id = mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>({usernames[1], "-A"}), true, false);
  // Fail on all the calls to RunInMinijail issued by the
  // attempt to apply marks for the last user.
  ASSERT_FALSE(
      mock_iptables.ApplyVpnSetup(usernames, interface, add));
  EXPECT_TRUE(mock_iptables.CheckCommandsUndone());

  // Should have failed only on the first testuser1 command.
  EXPECT_TRUE(mock_iptables.GetRunInMinijailCriterionMatchCount(id) == 1);
}

TEST_F(IpTablesTest, ApplyVpnSetupAdd_FailureInMasquerade) {
  const std::vector<std::string> usernames = {"testuser0", "testuser1"};
  // First two calls are issued by ApplyRuleForUserTraffic, next two
  // are issued by ApplyMasquerade. We make both of those fail.
  const std::string interface = "ifc0";
  const bool add = true;

  MockIpTables mock_iptables;
  // Fail on calls adding MASQUERADE.
  int id = mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>({"MASQUERADE", "-A"}), true, false);
  ASSERT_FALSE(
      mock_iptables.ApplyVpnSetup(usernames, interface, add));
  EXPECT_TRUE(mock_iptables.CheckCommandsUndone());
  EXPECT_TRUE(mock_iptables.GetRunInMinijailCriterionMatchCount(id) == 1);
}

TEST_F(IpTablesTest, ApplyVpnSetupAdd_FailureInRuleForUserTraffic) {
  const std::vector<std::string> usernames = {"testuser0", "testuser1"};
  const std::string interface = "ifc0";
  const bool add = true;

  MockIpTables mock_iptables;
  int id = mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>({"rule", "add", "fwmark"}), true, false);
  ASSERT_FALSE(mock_iptables.ApplyVpnSetup(usernames, interface, add));
  EXPECT_TRUE(mock_iptables.CheckCommandsUndone());
  EXPECT_TRUE(mock_iptables.GetRunInMinijailCriterionMatchCount(id) == 1);
}

TEST_F(IpTablesTest, ApplyVpnSetupRemove_Success) {
  const std::vector<std::string> usernames = {"testuser0", "testuser1"};
  const std::string interface = "ifc0";
  const bool remove = false;

  MockIpTables mock_iptables;
  // The fail criterions ensure that the ApplyVpnSetup function will delete
  // the rules that are added. In addition to this, test that when all of
  // these commands succeed, ApplyVpnSetup returns true.
  mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>({"MASQUERADE", "-D"}), true, true);
  mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>({"--uid-owner", "-D"}), true, true);
  mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>({"rule", "delete", "fwmark"}), true, true);

  ASSERT_TRUE(mock_iptables.ApplyVpnSetup(usernames, interface, remove));
}

TEST_F(IpTablesTest, ApplyVpnSetupRemove_Failure) {
  const std::vector<std::string> usernames = {"testuser0", "testuser1"};
  const std::string interface = "ifc0";
  const bool remove = false;

  MockIpTables mock_iptables;
  // Make all removing commands fail.
  mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>({"delete"}), true, false);
  mock_iptables.SetRunInMinijailFailCriterion(
      std::vector<std::string>({"-D"}), true, false);
  ASSERT_FALSE(mock_iptables.ApplyVpnSetup(usernames, interface, remove));
}

}  // namespace firewalld