From 306babc6ace320989548fd34a542668ae3050360 Mon Sep 17 00:00:00 2001
From: zu1k <i@zu1k.com>
Date: Thu, 20 Oct 2022 15:59:03 +0800
Subject: [PATCH] fix: Check wry db before query

---
 pkg/qqwry/qqwry.go         | 4 ++++
 pkg/zxipv6wry/zxipv6wry.go | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/pkg/qqwry/qqwry.go b/pkg/qqwry/qqwry.go
index 9557366..468e626 100644
--- a/pkg/qqwry/qqwry.go
+++ b/pkg/qqwry/qqwry.go
@@ -49,6 +49,10 @@ func NewQQwry(filePath string) (*QQwry, error) {
 	start := binary.LittleEndian.Uint32(header[:4])
 	end := binary.LittleEndian.Uint32(header[4:])
 
+	if uint32(len(fileData)) < end+7 {
+		log.Fatalln("纯真 IP 库存在错误，请重新下载")
+	}
+
 	return &QQwry{
 		IPDB: wry.IPDB[uint32]{
 			Data: fileData,
diff --git a/pkg/zxipv6wry/zxipv6wry.go b/pkg/zxipv6wry/zxipv6wry.go
index 12866c4..2791457 100644
--- a/pkg/zxipv6wry/zxipv6wry.go
+++ b/pkg/zxipv6wry/zxipv6wry.go
@@ -47,6 +47,10 @@ func NewZXwry(filePath string) (*ZXwry, error) {
 	counts := binary.LittleEndian.Uint64(header[8:16])
 	end := start + counts*11
 
+	if uint64(len(fileData)) < end {
+		log.Fatalln("ZX IPv6数据库存在错误，请重新下载")
+	}
+
 	return &ZXwry{
 		IPDB: wry.IPDB[uint64]{
 			Data: fileData,