ReadDataToolTests.cs

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license.

using Microsoft.Extensions.Logging;
using Moq;
using Mssql.McpServer;

namespace MssqlMcp.Tests
{
    [Collection("Database Tests")]
    public sealed class MssqlMcpReadDataToolTests : IDisposable
    {
        private readonly string _tableName;
        private readonly Tools _tools;

        public MssqlMcpReadDataToolTests()
        {
            _tableName = $"ReadDataTest_{Guid.NewGuid():N}";
            var connectionFactory = new SqlConnectionFactory();
            var loggerMock = new Mock<ILogger<Tools>>();
            _tools = new Tools(connectionFactory, loggerMock.Object);
        }

        public void Dispose()
        {
            // Clean up test table if it exists
            var _ = _tools.DropTable($"DROP TABLE IF EXISTS {_tableName}").GetAwaiter().GetResult();
        }


        [Fact]
        public async Task ReadData_ReturnsData_WhenSqlIsValid()
        {
            // Set up test table with data
            var createResult = await _tools.CreateTable($"CREATE TABLE {_tableName} (Id INT PRIMARY KEY)") as DbOperationResult;
            Assert.NotNull(createResult);
            Assert.True(createResult.Success);
            var insertResult = await _tools.InsertData($"INSERT INTO {_tableName} (Id) VALUES (1)") as DbOperationResult;
            Assert.NotNull(insertResult);
            Assert.True(insertResult.Success);

            var sql = $"SELECT * FROM {_tableName}";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.True(result.Success);
            Assert.NotNull(result.Data);
        }

        [Fact]
        public async Task ReadData_ReturnsError_WhenSqlIsInvalid()
        {
            var sql = "SELECT FROM";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Database query execution failed", result.Error ?? string.Empty, StringComparison.OrdinalIgnoreCase);
        }


        [Fact]
        public async Task ReadData_Security_RejectsDeleteStatement()
        {
            var sql = "DELETE FROM users WHERE id = 1";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Query must start with SELECT", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsDropStatement()
        {
            var sql = "DROP TABLE users";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Query must start with SELECT", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsUpdateStatement()
        {
            var sql = "UPDATE users SET admin = 1";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Query must start with SELECT", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsInsertStatement()
        {
            var sql = "INSERT INTO users VALUES ('hacker', 'password')";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            // INSERT gets caught either by "must start with SELECT" or keyword detection
            Assert.True(
                result.Error?.Contains("Query must start with SELECT") == true ||
                result.Error?.Contains("Dangerous keyword 'INSERT'") == true
            );
        }

        [Theory]
        [InlineData("DELETE", "DELETE FROM users")]
        [InlineData("DROP", "SELECT * FROM users WHERE 1=1 OR DROP TABLE accounts")]
        [InlineData("TRUNCATE", "SELECT * FROM users; TRUNCATE TABLE logs")]
        [InlineData("EXEC", "SELECT * FROM users EXEC sp_help")]
        [InlineData("EXECUTE", "SELECT * FROM users EXECUTE xp_cmdshell")]
        [InlineData("ALTER", "SELECT * FROM users; ALTER TABLE users ADD admin BIT")]
        [InlineData("CREATE", "SELECT * FROM users; CREATE TABLE hacked (id INT)")]
        [InlineData("GRANT", "SELECT * FROM users; GRANT ALL TO hacker")]
        [InlineData("REVOKE", "SELECT * FROM users; REVOKE SELECT ON users FROM public")]
        [InlineData("BACKUP", "SELECT * FROM users; BACKUP DATABASE test TO DISK='hack.bak'")]
        [InlineData("RESTORE", "SELECT * FROM users; RESTORE DATABASE test FROM DISK='hack.bak'")]
        [InlineData("SHUTDOWN", "SELECT * FROM users; SHUTDOWN")]
        public async Task ReadData_Security_RejectsAllDangerousKeywords(string keyword, string sql)
        {
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            // Different keywords might trigger different validation rules - some caught by "must start with SELECT", others by keyword detection
            Assert.True(
                result.Error?.Contains($"Dangerous keyword '{keyword}'", StringComparison.OrdinalIgnoreCase) == true ||
                result.Error?.Contains("Query must start with SELECT") == true ||
                result.Error?.Contains("malicious SQL pattern") == true ||
                result.Error?.Contains("Multiple SQL statements") == true
            );
        }


        [Fact]
        public async Task ReadData_Security_RejectsSemicolonInjection()
        {
            var sql = "SELECT * FROM users; DROP TABLE accounts--";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            // Should catch either multiple statements or the DROP keyword
            Assert.True(
                result.Error?.Contains("Multiple SQL statements") == true ||
                result.Error?.Contains("Dangerous keyword 'DROP'") == true
            );
        }

        [Fact]
        public async Task ReadData_Security_RejectsUnionWithDangerousKeyword()
        {
            var sql = "SELECT id FROM users UNION SELECT * FROM passwords; DELETE FROM logs";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Dangerous keyword 'DELETE' detected", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsStoredProcedureExecution()
        {
            var sql = "SELECT * FROM users WHERE id = 1 EXEC xp_cmdshell 'dir'";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Dangerous keyword 'EXEC' detected", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsCommentInjection()
        {
            // The validation strips comments first, so DELETE in comments should be caught
            var sql = "SELECT * FROM users /* injected DELETE FROM accounts */";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Dangerous keyword 'DELETE' detected", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsLineCommentInjection()
        {
            var sql = "SELECT * FROM users -- DELETE FROM accounts";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            // After stripping comments, should be valid unless DELETE is in the actual query
            // This one should actually pass since DELETE is only in the comment
            // Let's test one where it matters
            var sql2 = "SELECT * FROM users WHERE id = 1 OR 1=1 -- UNION DELETE";
            var result2 = await _tools.ReadData(sql2) as DbOperationResult;
            // This should pass as DELETE is in comment
        }

        [Fact]
        public async Task ReadData_Security_RejectsWaitforDelay()
        {
            var sql = "SELECT * FROM users WHERE id = 1 WAITFOR DELAY '00:00:05'";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Dangerous keyword 'WAITFOR' detected", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsCharObfuscation()
        {
            var sql = "SELECT * FROM users WHERE name = 'test' + CHAR(59) + 'DROP TABLE users'";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Dangerous keyword 'DROP' detected", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsNCharObfuscation()
        {
            var sql = "SELECT * FROM users WHERE name = NCHAR(0x44) + NCHAR(0x52) + NCHAR(0x4F) + NCHAR(0x50)";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Potentially malicious SQL pattern detected.", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsBulkOperations()
        {
            var sql = "SELECT * FROM OPENROWSET('SQLNCLI', 'Server=hack;Trusted_Connection=yes;', 'SELECT * FROM users')";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            // Should be caught by OPENROWSET keyword or pattern
            Assert.True(
                result.Error?.Contains("Dangerous keyword 'OPENROWSET'") == true ||
                result.Error?.Contains("malicious SQL pattern") == true
            );
        }


        [Fact]
        public async Task ReadData_Security_RejectsNonSelectStatement()
        {
            var sql = "INSERT INTO users VALUES ('hacker', 'password')";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Query must start with SELECT", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsEmptyQuery()
        {
            var sql = "   ";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Query must be a non-empty string", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsNullQuery()
        {
            string sql = null!;
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Query must be a non-empty string", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsVeryLongQuery()
        {
            // Build a query that exceeds the 10,000 character limit
            var sql = "SELECT " + new string('a', 10001) + " FROM users";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Query is too long", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsMultipleStatements()
        {
            var sql = "SELECT * FROM users; SELECT * FROM passwords";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.False(result.Success);
            Assert.Contains("Security validation failed: Potentially malicious SQL pattern detected. Only simple SELECT queries are allowed.", result.Error ?? string.Empty);
        }

        [Fact]
        public async Task ReadData_Security_RejectsCaseVariationsOfDangerousKeywords()
        {
            var queries = new[]
            {
                "DeLeTe FROM users",
                "dRoP TABLE users",
                "UpDaTe users SET admin = 1",
                "iNsErT INTO users VALUES (1)",
                "tRuNcAtE TABLE logs"
            };

            foreach (var sql in queries)
            {
                var result = await _tools.ReadData(sql) as DbOperationResult;
                Assert.NotNull(result);
                Assert.False(result.Success);
                // Should be rejected either for not starting with SELECT or dangerous keyword
                Assert.True(
                    result.Error?.Contains("Query must start with SELECT") == true ||
                    result.Error?.Contains("Dangerous keyword") == true
                );
            }
        }


        [Theory]
        [InlineData("SELECT * FROM users")]
        [InlineData("SELECT id, name FROM customers WHERE active = 1")]
        [InlineData("SELECT COUNT(*) FROM orders")]
        [InlineData("SELECT updated_at, created_at FROM logs")] // Tests that 'UPDATE' in column name is OK
        [InlineData("SELECT * FROM user_updates")] // Tests that 'UPDATE' in table name is OK
        [InlineData("select * from users")] // Lowercase should work
        [InlineData("SeLeCt * FrOm users")] // Mixed case should work
        [InlineData("SELECT TOP 10 * FROM users ORDER BY created_at DESC")] // TOP and ORDER BY
        [InlineData("SELECT u.*, o.order_date FROM users u JOIN orders o ON u.id = o.user_id")] // JOIN syntax
        public async Task ReadData_Security_AllowsValidSelectQueries(string sql)
        {
            // Create test table to actually run these queries
            var testTableName = $"ReadDataTest_{Guid.NewGuid():N}";
            await _tools.CreateTable($"CREATE TABLE {testTableName} (id INT, name VARCHAR(50), active BIT, updated_at DATETIME, created_at DATETIME, user_id INT, order_date DATETIME)");
            
            // Replace placeholder table names with our test table
            sql = sql.Replace("users", testTableName)
                     .Replace("customers", testTableName)
                     .Replace("orders", testTableName)
                     .Replace("logs", testTableName)
                     .Replace("user_updates", testTableName);
            
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.True(result.Success);
            
            // Clean up
            await _tools.DropTable($"DROP TABLE IF EXISTS {testTableName}");
        }

        [Fact]
        public async Task ReadData_Security_AllowsSelectWithUpdatedAtColumn()
        {
            // This specifically tests that we don't false-positive on column names containing keywords
            var testTableName = $"ReadDataTest_{Guid.NewGuid():N}";
            await _tools.CreateTable($"CREATE TABLE {testTableName} (id INT, updated_at DATETIME, deleted_flag BIT, created_by VARCHAR(50))");
            
            var sql = $"SELECT id, updated_at, deleted_flag, created_by FROM {testTableName}";
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.True(result.Success);
            
            // Clean up
            await _tools.DropTable($"DROP TABLE IF EXISTS {testTableName}");
        }

        [Fact]
        public async Task ReadData_Security_AllowsComplexValidQuery()
        {
            // Test a complex but valid SELECT query
            var testTableName = $"ReadDataTest_{Guid.NewGuid():N}";
            await _tools.CreateTable($"CREATE TABLE {testTableName} (id INT, category VARCHAR(50), amount DECIMAL(10,2), created_at DATETIME)");
            
            var sql = $@"SELECT 
                            category,
                            COUNT(*) as count,
                            SUM(amount) as total,
                            AVG(amount) as average,
                            MIN(created_at) as first_created,
                            MAX(created_at) as last_created
                        FROM {testTableName}
                        WHERE amount > 0
                        GROUP BY category
                        HAVING COUNT(*) > 1
                        ORDER BY total DESC";
            
            var result = await _tools.ReadData(sql) as DbOperationResult;
            Assert.NotNull(result);
            Assert.True(result.Success);
            
            // Clean up
            await _tools.DropTable($"DROP TABLE IF EXISTS {testTableName}");
        }


        [Fact]
        public async Task ReadData_SanitizeResult_RemovesSuspiciousCharactersFromColumnNames()
        {
            // Create a table with suspicious characters in column names
            var testTableName = $"ReadDataTest_{Guid.NewGuid():N}";
            await _tools.CreateTable($@"CREATE TABLE {testTableName} (
                [normal_id] INT,
                [bad<script>] VARCHAR(50),
                [evil&injection] VARCHAR(50),
                [ok_name] VARCHAR(50)
            )");
            
            // Insert test data
            await _tools.InsertData($"INSERT INTO {testTableName} VALUES (1, 'test', 'data', 'ok')");
            
            // Query the data - this will go through SanitizeResult
            var result = await _tools.ReadData($"SELECT * FROM {testTableName}") as DbOperationResult;
            
            Assert.NotNull(result);
            Assert.True(result.Success);
            Assert.NotNull(result.Data);
            
            // Check that suspicious characters were removed from column names
            var firstRow = ((List<Dictionary<string, object?>>)result.Data).First();
            Assert.Contains("normal_id", firstRow.Keys);
            Assert.Contains("badscript", firstRow.Keys); // <script> should be removed
            Assert.Contains("evilinjection", firstRow.Keys); // &injection should be sanitized
            Assert.Contains("ok_name", firstRow.Keys);
            
            // Cleanup
            await _tools.DropTable($"DROP TABLE IF EXISTS {testTableName}");
        }
    }
}