C++: Improve the cpp/path-injection qhelp

cpp/ql/test/query-tests/Security/CWE/CWE-022/semmle/tests/test.c

@@ -2,7 +2,7 @@
 // Associated with CWE-022: Improper Limitation of a Pathname to a Restricted Directory. http://cwe.mitre.org/data/definitions/22.html
 
 #include "stdlib.h"
-
+#define PATH_MAX 4096
 ///// Test code /////
 
 int main(int argc, char** argv) {
@@ -56,6 +56,44 @@ int main(int argc, char** argv) {
     void read(const char *fileName);
     read(argv[1]); // BAD
   }
+
+  {
+    char *userAndFile = argv[2];
+    // Check for invalid sequences in the user input
+    if (strstr(userAndFile, "..") || strchr(userAndFile, '/') || strchr(userAndFile, '\\')) {
+        // printf("Invalid filename.\n");
+        return 1;
+    }
+
+    char fileBuffer[FILENAME_MAX] = "/home/user/files/";
+    // Ensure buffer overflow is prevented
+    strncat(fileBuffer, userAndFile, FILENAME_MAX - strlen(fileBuffer) - 1);
+    // GOOD: We know that the filename is safe and stays within the public folder. But we currently get an FP here.
+    FILE *file = fopen(fileBuffer, "wb+");
+  }
+
+  {
+    char *userAndFile = argv[2];
+    char baseDir[PATH_MAX] = "/home/user/public/";
+    char fullPath[PATH_MAX];
+    char resolvedPath[PATH_MAX];
+
+    // Attempt to concatenate the base directory and the user-supplied path
+    snprintf(fullPath, sizeof(fullPath), "%s%s", baseDir, userAndFile);
+
+    // Resolve the absolute path, normalizing any ".." or "."
+    if (realpath(fullPath, resolvedPath) == 0) {
+        return 1;
+    }
+
+    // Check if the resolved path starts with the base directory
+    if (strncmp(baseDir, resolvedPath, strlen(baseDir)) != 0) {
+        return 1;
+    }
+
+    // GOOD: Path is within the intended directory
+    FILE *file = fopen(resolvedPath, "wb+");
+  }
 }
 
 void read(char *fileName) {