use my_WinHttpCrackUrl from segatools

hooklib/dns.c

@@ -23,6 +23,8 @@
    pointers for those args. So are most of the fields in this structure, for
    that matter. */
 
+wchar_t url_combined[MAX_PATH];
+
 typedef struct POLYFILL_DNS_QUERY_REQUEST {
     ULONG Version;
     PCWSTR QueryName;
@@ -492,29 +494,116 @@ end:
     return result;
 }
 
+bool match_domain(const wchar_t* target, const wchar_t* pattern) {
+    assert(target != NULL);
+    assert(pattern != NULL);
+    bool need_free = false;
+    if (_wcsicmp(pattern, target) == 0) {
+        return true;
+    }
+
+    int pattern_ptr_index = 0;
+    int target_ptr_index = 0;
+    wchar_t* proto = wcsstr(target, L"://");
+    
+    if (proto != NULL) {
+        proto += 3;
+    } else {
+        need_free = true;
+        size_t target_size_in_bytes = wcslen(target) * sizeof(wchar_t);
+        proto = malloc(target_size_in_bytes);
+        memcpy_s(proto, target_size_in_bytes, target, target_size_in_bytes);
+    }
+
+    while (pattern[pattern_ptr_index] != '\0' && proto[target_ptr_index] != '\0') {
+        if (pattern[pattern_ptr_index] == '*') {
+            pattern_ptr_index++; // Check next character for wildcard match.
+
+            while (pattern[pattern_ptr_index] != proto[target_ptr_index]) {
+                target_ptr_index++;
+
+                if (proto[target_ptr_index] == '\0') { 
+                    if (need_free) { free(proto); }
+                    return false;
+                }
+            }
+        }
+        else if (pattern[pattern_ptr_index] != proto[target_ptr_index]) {
+            if (need_free) { free(proto); }
+            return false;
+        }
+        else {
+            pattern_ptr_index++;
+            target_ptr_index++;
+        }
+    }
+
+    if (need_free) { free(proto); }
+    return pattern[pattern_ptr_index] == '\0' && (proto[target_ptr_index] == '\0' || proto[target_ptr_index] == '/');
+}
+
 static BOOL my_WinHttpCrackUrl(
         LPCWSTR pwszUrl,
         DWORD dwUrlLength,
         DWORD dwFlags,
         LPURL_COMPONENTS lpUrlComponents)
 {
-    if (!next_WinHttpCrackUrl(pwszUrl, dwUrlLength, dwFlags, lpUrlComponents)) {
-        dprintf("DNS: next_WinHttpCrackUrl FAIL %08X\n", (int)GetLastError());
-        return false;
-    }
+    const struct dns_hook_entry *pos;
+    size_t i;
 
     EnterCriticalSection(&dns_hook_lock);
 
-    for (size_t i = 0 ; i < dns_hook_nentries ; i++) {
-        struct dns_hook_entry* pos = &dns_hook_entries[i];
-        if (_wcsnicmp(lpUrlComponents->lpszHostName, pos->from, wcslen(pos->from)) == 0) {
-            dprintf("DNS: Replace url %S -> %S\n", lpUrlComponents->lpszHostName, pos->to);
-            lpUrlComponents->lpszHostName = pos->to;
-            lpUrlComponents->dwHostNameLength = wcslen(pos->to);
-            break;
+    for (i = 0 ; i < dns_hook_nentries ; i++) {
+        pos = &dns_hook_entries[i];
+
+        if (match_domain(pwszUrl, pos->from)) {
+            wchar_t* toAddr = pos->to;
+            wchar_t* proto = wcsstr(pwszUrl, L"://");
+            wchar_t* path = wcsstr(pwszUrl, L"/");
+            memset(url_combined, 0, sizeof(url_combined));
+
+            dprintf("DNS: Replace URL %S with %S\n", pwszUrl, toAddr);
+            
+            if (proto != NULL) {
+                // I'm sure there is a way to only concat a subset of a string
+                // but I don't feel like figuring it out
+                // Store the first character after the slash and then null it out in the string
+                wchar_t tmp = proto[3];
+                memset(&proto[3], 0, sizeof(wchar_t));
+                
+                // Concat the protocol onto the string and restore the temp value after
+                wcscat_s(url_combined, MAX_PATH, pwszUrl);
+                proto[3] = tmp;
+                
+                // Push the proto up by 3 to point to the first hostname character and 
+                // re-check for the forward slash indicating the path
+                proto += 3;
+                path = wcsstr(proto, L"/");
+            }
+
+            wcscat_s(url_combined, MAX_PATH, toAddr);
+
+            if (path != NULL) {
+                wcscat_s(url_combined, MAX_PATH, path);
+            }
+            
+            bool result = next_WinHttpCrackUrl(
+                url_combined,
+                wcslen(url_combined),
+                dwFlags,
+                lpUrlComponents
+            );
+            LeaveCriticalSection(&dns_hook_lock);
+            return result;
         }
     }
 
     LeaveCriticalSection(&dns_hook_lock);
-    return true;
+    return next_WinHttpCrackUrl(
+        pwszUrl,
+        dwUrlLength,
+        dwFlags,
+        lpUrlComponents
+    );
+
 }