From bdccdd5bc9e9b290915f02d332ef384570bd5f93 Mon Sep 17 00:00:00 2001
From: lqqyt2423 <974923609@qq.com>
Date: Thu, 16 Feb 2023 11:18:21 +0800
Subject: [PATCH] Add support for wildcard domain matching in matchHost
 function

---
 cmd/go-mitmproxy/main.go      | 12 ++++++-
 cmd/go-mitmproxy/main_test.go | 59 +++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/cmd/go-mitmproxy/main.go b/cmd/go-mitmproxy/main.go
index 57d5f41..c7de19c 100644
--- a/cmd/go-mitmproxy/main.go
+++ b/cmd/go-mitmproxy/main.go
@@ -97,13 +97,23 @@ func matchHost(address string, hosts []string) bool {
 	hostname, port := splitHostPort(address)
 	for _, host := range hosts {
 		h, p := splitHostPort(host)
-		if h == hostname && (p == "" || p == port) {
+		if matchHostname(hostname, h) && (p == "" || p == port) {
 			return true
 		}
 	}
 	return false
 }
 
+func matchHostname(hostname string, h string) bool {
+	if h == "*" {
+		return true
+	}
+	if strings.HasPrefix(h, "*.") {
+		return hostname == h[2:] || strings.HasSuffix(hostname, h[1:])
+	}
+	return h == hostname
+}
+
 func splitHostPort(address string) (string, string) {
 	index := strings.LastIndex(address, ":")
 	if index == -1 {
diff --git a/cmd/go-mitmproxy/main_test.go b/cmd/go-mitmproxy/main_test.go
index 56963e2..87def52 100644
--- a/cmd/go-mitmproxy/main_test.go
+++ b/cmd/go-mitmproxy/main_test.go
@@ -5,6 +5,7 @@ import (
 )
 
 func TestMatchHost(t *testing.T) {
+	// Test case 1: Exact match
 	address := "www.baidu.com:443"
 	hosts := []string{
 		"www.baidu.com:443",
@@ -17,6 +18,7 @@ func TestMatchHost(t *testing.T) {
 		t.Errorf("Expected %t but got %t", expected, result)
 	}
 
+	// Test case 2: Exact match with port
 	address = "www.google.com:80"
 	hosts = []string{
 		"www.baidu.com:443",
@@ -29,6 +31,7 @@ func TestMatchHost(t *testing.T) {
 		t.Errorf("Expected %t but got %t", expected, result)
 	}
 
+	// Test case 3: No match
 	address = "www.test.com:80"
 	hosts = []string{
 		"www.baidu.com:443",
@@ -40,4 +43,60 @@ func TestMatchHost(t *testing.T) {
 	if result != expected {
 		t.Errorf("Expected %t but got %t", expected, result)
 	}
+
+	// Test case 4: Wildcard match
+	address = "test.baidu.com:443"
+	hosts = []string{
+		"*.baidu.com",
+		"www.baidu.com:443",
+		"www.baidu.com",
+		"www.google.com",
+	}
+	expected = true
+	result = matchHost(address, hosts)
+	if result != expected {
+		t.Errorf("Expected %t but got %t", expected, result)
+	}
+
+	// Test case 5: Wildcard match with port
+	address = "test.baidu.com:443"
+	hosts = []string{
+		"*.baidu.com:443",
+		"www.baidu.com:443",
+		"www.baidu.com",
+		"www.google.com",
+	}
+	expected = true
+	result = matchHost(address, hosts)
+	if result != expected {
+		t.Errorf("Expected %t but got %t", expected, result)
+	}
+
+	// Test case 6: Wildcard mismatch
+	address = "test.baidu.com:80"
+	hosts = []string{
+		"*.baidu.com:443",
+		"www.baidu.com:443",
+		"www.baidu.com",
+		"www.google.com",
+	}
+	expected = false
+	result = matchHost(address, hosts)
+	if result != expected {
+		t.Errorf("Expected %t but got %t", expected, result)
+	}
+
+	// Test case 7: Wildcard mismatch
+	address = "test.google.com:80"
+	hosts = []string{
+		"*.baidu.com",
+		"www.baidu.com:443",
+		"www.baidu.com",
+		"www.google.com",
+	}
+	expected = false
+	result = matchHost(address, hosts)
+	if result != expected {
+		t.Errorf("Expected %t but got %t", expected, result)
+	}
 }