From 3819ae680883566995fc6b6a8a5b2af47d4f4fa1 Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Sat, 20 May 2023 17:53:43 +0800
Subject: [PATCH] Add the source field to the json output.

---
 internal/db/db.go              | 14 ++++++++++----
 pkg/cdn/cdn.go                 |  4 ++++
 pkg/dbif/db.go                 |  1 +
 pkg/geoip/geoip.go             |  4 ++++
 pkg/ip2location/ip2location.go |  8 ++++++--
 pkg/ip2region/ip2region.go     |  4 ++++
 pkg/ipip/ipip.go               |  4 ++++
 pkg/qqwry/qqwry.go             |  4 ++++
 pkg/zxipv6wry/zxipv6wry.go     |  4 ++++
 9 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/internal/db/db.go b/internal/db/db.go
index f8a4c2a..91d82e4 100644
--- a/internal/db/db.go
+++ b/internal/db/db.go
@@ -13,6 +13,11 @@ import (
 	"github.com/zu1k/nali/pkg/zxipv6wry"
 )
 
+type Result struct {
+	Source string `json:"source"`
+	common.Result
+}
+
 func GetDB(typ dbif.QueryType) (db dbif.DB) {
 	if db, found := dbTypeCache[typ]; found {
 		return db
@@ -69,15 +74,16 @@ func GetDB(typ dbif.QueryType) (db dbif.DB) {
 	return
 }
 
-func Find(typ dbif.QueryType, query string) common.Result {
+func Find(typ dbif.QueryType, query string) *Result {
 	if result, found := queryCache.Load(query); found {
-		return result.(common.Result)
+		return result.(*Result)
 	}
-	result, err := GetDB(typ).Find(query)
+	db := GetDB(typ)
+	result, err := db.Find(query)
 	if err != nil {
 		return nil
 	}
-	res := result
+	res := &Result{db.Name(), result}
 	queryCache.Store(query, res)
 	return res
 }
diff --git a/pkg/cdn/cdn.go b/pkg/cdn/cdn.go
index eba97c5..e39ccee 100644
--- a/pkg/cdn/cdn.go
+++ b/pkg/cdn/cdn.go
@@ -103,6 +103,10 @@ func (db CDN) Find(query string, params ...string) (result fmt.Stringer, err err
 	return nil, errors.New("not found")
 }
 
+func (db CDN) Name() string {
+	return "cdn"
+}
+
 func parseBaseCname(domain string) (result []string) {
 	parts := strings.Split(domain, ".")
 	size := len(parts)
diff --git a/pkg/dbif/db.go b/pkg/dbif/db.go
index 90c4bdd..ef55f5f 100644
--- a/pkg/dbif/db.go
+++ b/pkg/dbif/db.go
@@ -22,6 +22,7 @@ const (
 
 type DB interface {
 	Find(query string, params ...string) (result fmt.Stringer, err error)
+	Name() string
 }
 
 var (
diff --git a/pkg/geoip/geoip.go b/pkg/geoip/geoip.go
index e6f817d..3e16f76 100644
--- a/pkg/geoip/geoip.go
+++ b/pkg/geoip/geoip.go
@@ -54,6 +54,10 @@ func (g GeoIP) Find(query string, params ...string) (result fmt.Stringer, err er
 	return
 }
 
+func (db GeoIP) Name() string {
+	return "geoip"
+}
+
 type Result struct {
 	Country string `json:"country"`
 	Area    string `json:"area"`
diff --git a/pkg/ip2location/ip2location.go b/pkg/ip2location/ip2location.go
index 11f9f0b..7fb4eff 100644
--- a/pkg/ip2location/ip2location.go
+++ b/pkg/ip2location/ip2location.go
@@ -31,12 +31,12 @@ func NewIP2Location(filePath string) (*IP2Location, error) {
 	}
 }
 
-func (x IP2Location) Find(query string, params ...string) (result fmt.Stringer, err error) {
+func (db IP2Location) Find(query string, params ...string) (result fmt.Stringer, err error) {
 	ip := net.ParseIP(query)
 	if ip == nil {
 		return nil, errors.New("Query should be valid IP")
 	}
-	record, err := x.db.Get_all(ip.String())
+	record, err := db.db.Get_all(ip.String())
 
 	if err != nil {
 		return
@@ -50,6 +50,10 @@ func (x IP2Location) Find(query string, params ...string) (result fmt.Stringer,
 	return
 }
 
+func (db IP2Location) Name() string {
+	return "ip2location"
+}
+
 type Result struct {
 	Country string `json:"country"`
 	Region  string `json:"region"`
diff --git a/pkg/ip2region/ip2region.go b/pkg/ip2region/ip2region.go
index ae6dcc1..b00f0bf 100644
--- a/pkg/ip2region/ip2region.go
+++ b/pkg/ip2region/ip2region.go
@@ -67,3 +67,7 @@ func (db Ip2Region) Find(query string, params ...string) (result fmt.Stringer, e
 
 	return nil, errors.New("ip2region 未初始化")
 }
+
+func (db Ip2Region) Name() string {
+	return "ip2region"
+}
diff --git a/pkg/ipip/ipip.go b/pkg/ipip/ipip.go
index eabb630..30e80c7 100644
--- a/pkg/ipip/ipip.go
+++ b/pkg/ipip/ipip.go
@@ -54,3 +54,7 @@ func (db IPIPFree) Find(query string, params ...string) (result fmt.Stringer, er
 		return
 	}
 }
+
+func (db IPIPFree) Name() string {
+	return "ipip"
+}
diff --git a/pkg/qqwry/qqwry.go b/pkg/qqwry/qqwry.go
index d637075..2f165f0 100644
--- a/pkg/qqwry/qqwry.go
+++ b/pkg/qqwry/qqwry.go
@@ -90,6 +90,10 @@ func (db QQwry) Find(query string, params ...string) (result fmt.Stringer, err e
 	return reader.Result.DecodeGBK(), nil
 }
 
+func (db QQwry) Name() string {
+	return "qqwry"
+}
+
 func CheckFile(data []byte) bool {
 	if len(data) < 8 {
 		return false
diff --git a/pkg/zxipv6wry/zxipv6wry.go b/pkg/zxipv6wry/zxipv6wry.go
index 9816c6e..fa32115 100644
--- a/pkg/zxipv6wry/zxipv6wry.go
+++ b/pkg/zxipv6wry/zxipv6wry.go
@@ -82,6 +82,10 @@ func (db *ZXwry) Find(query string, _ ...string) (result fmt.Stringer, err error
 	return reader.Result, nil
 }
 
+func (db *ZXwry) Name() string {
+	return "xzwry"
+}
+
 func CheckFile(data []byte) bool {
 	if len(data) < 4 {
 		return false