From 268e4b044df3d6e2598a219153665ab61dec4f6b Mon Sep 17 00:00:00 2001
From: wangsiyuan <2392948297@qq.com>
Date: Thu, 3 Aug 2023 00:21:48 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20main.go?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.go | 92 +++++++++++++++++++++++++++++++--------------------------
 1 file changed, 50 insertions(+), 42 deletions(-)

diff --git a/main.go b/main.go
index 84f1637..3ffc766 100644
--- a/main.go
+++ b/main.go
@@ -4,53 +4,61 @@ import (
 	"checkIP/modules"
 	"encoding/json"
 	"fmt"
-	_ "github.com/go-sql-driver/mysql"
 	"log"
 	"net/http"
 	"strings"
 )
 
-func main() {
-	//modules.ImportDatabaseToHex()
-	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		//测试用ip地址
-		//ip := "192.168.2.134"
-		//获取访问浏览器ip
-		ip := strings.Split(r.RemoteAddr, ":")[0]
-		//fmt.Println("客户端IP地址:", ip)
-		if ip == "" {
-			w.Write([]byte("Invalid IP"))
-			return
-		}
-		if ip != "" {
-			IPinfo := modules.GetIPInformation(ip)
-			fmt.Println("IP", IPinfo.IP)
-			fmt.Println("IsPrivate", IPinfo.IsPrivate)
-			fmt.Println("PrivateCIDR", IPinfo.PrivateCIDR)
-			if IPinfo.IsPrivate == true {
-				w.Write([]byte("Invalid IP"))
-				return
-			}
-			if IPinfo.IsPrivate != true {
-				//使用mmdb查询ip
-				IpIformation := modules.MmdbQueryIP(IPinfo.IP)
-				//使用MySql查询ip
-				//dsn := modules.Readconfig()
-				//IpIformation := modules.MysqlQuery(ip, dsn)
+func getIPInfoAndRespond(w http.ResponseWriter, r *http.Request) {
+	// 获取访问者IP地址
+	ip := ""
+	xffHeader := r.Header.Get("X-Forwarded-For")
+	if xffHeader != "" {
+		// 如果X-Forwarded-For头部不为空，则使用其中的第一个IP地址
+		ip = strings.Split(xffHeader, ", ")[0]
+	} else {
+		// 否则，使用RemoteAddr提取IP地址（可能包含端口信息）
+		ip = strings.Split(r.RemoteAddr, ":")[0]
+	}
 
-				log.Println("tag:main", IpIformation)
-				jsonData, err := json.Marshal(IpIformation)
-				if err != nil {
-					fmt.Println("JSON 序列化失败:", err)
-					return
-				}
-				// 设置响应的 Content-Type 为 application/json
-				w.Header().Set("Content-Type", "application/json")
-				// 返回 JSON 数据
-				//log.Println(jsonData)
-				w.Write(jsonData)
-			}
-		}
-	})
+	// 检查IP地址是否为空
+	if ip == "" {
+		w.Write([]byte("Invalid IP"))
+		return
+	}
+
+	// 获取IP信息
+	IPinfo := modules.GetIPInformation(ip)
+	fmt.Println("IP", IPinfo.IP)
+	fmt.Println("IsPrivate", IPinfo.IsPrivate)
+	fmt.Println("PrivateCIDR", IPinfo.PrivateCIDR)
+
+	// 检查是否为私有IP
+	if IPinfo.IsPrivate == true {
+		w.Write([]byte("Invalid IP"))
+		return
+	}
+
+	// 使用mmdb查询ip
+	IpIformation := modules.MmdbQueryIP(IPinfo.IP)
+	// 使用MySql查询ip
+	//dsn := modules.Readconfig()
+	//IpIformation := modules.MysqlQuery(ip, dsn)
+
+	log.Println("tag:main", IpIformation)
+	jsonData, err := json.Marshal(IpIformation)
+	if err != nil {
+		fmt.Println("JSON 序列化失败:", err)
+		return
+	}
+
+	// 设置响应的 Content-Type 为 application/json
+	w.Header().Set("Content-Type", "application/json")
+	// 返回 JSON 数据
+	w.Write(jsonData)
+}
+
+func main() {
+	http.HandleFunc("/", getIPInfoAndRespond)
 	http.ListenAndServe(":8341", nil)
 }