/
bolt.go
115 lines (100 loc) · 2.14 KB
/
bolt.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
package db
import (
"fmt"
"math/rand"
"os"
"strings"
"github.com/boltdb/bolt"
"github.com/sirupsen/logrus"
"github.com/wrfly/gus-proxy/utils"
)
var (
bktName = []byte("DNS")
)
// DNS database
type DNS struct {
db *bolt.DB
}
// New database file, create one if not exist
func New(dbFileName string) (*DNS, error) {
if _, err := os.Stat(dbFileName); os.IsNotExist(err) {
// create the file
_, err := os.Create(dbFileName)
if err != nil {
return nil, err
}
}
dnsDB, err := bolt.Open(dbFileName, 0600, nil)
if err != nil {
return nil, err
}
dnsDB.Update(func(tx *bolt.Tx) error {
_, err := tx.CreateBucket(bktName)
if err != nil {
return fmt.Errorf("create bucket: %s", err)
}
return nil
})
return &DNS{
db: dnsDB,
}, nil
}
// Close the DB
func (d *DNS) Close() error {
return d.db.Close()
}
// setDNS to the DB
func (d *DNS) setDNS(domain string, answer []string) error {
answerStr := strings.Join(answer, "|")
err := d.db.Update(func(tx *bolt.Tx) error {
b := tx.Bucket(bktName)
err := b.Put([]byte(domain), []byte(answerStr))
return err
})
return err
}
// query domain from DB
func (d *DNS) query(domain string) (answer []string) {
d.db.View(func(tx *bolt.Tx) error {
b := tx.Bucket(bktName)
v := b.Get([]byte(domain))
answer = strings.Split(string(v), "|")
return nil
})
if len(answer) == 0 {
return nil
}
if answer[0] == "" {
return nil
}
return answer
}
// SelectIP domain -> IP
// kfd.me:8080 -> xx.x.xx.x:8080
func (dnsDB *DNS) SelectIP(domainPort string) (host, hostPort string) {
str := strings.Split(domainPort, ":")
domain := str[0]
port := "80"
if len(str) == 2 {
port = str[1]
}
ips := dnsDB.query(domain)
logrus.Debugf("query %s IP: %v", domain, ips)
// not found in db
if len(ips) == 0 {
digIPs, err := utils.LookupHost(domain)
if err != nil {
logrus.Errorf("Dig Error: %s", err)
return "127.0.0.1", "127.0.0.1:80"
}
// set to db
logrus.Debugf("Set DNS DB: domain: %s IP: %v", domain, digIPs)
if err := dnsDB.setDNS(domain, digIPs); err != nil {
logrus.Error(err)
}
ips = digIPs
}
host = ips[rand.Int()%len(ips)]
hostPort = host + ":" + port
return
}