From c5863a1ce9fa27cc211f0141bea6a3d92ff98962 Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sat, 13 Jan 2024 22:25:31 +0800 Subject: [PATCH 01/10] =?UTF-8?q?=E6=B7=BB=E5=8A=A0Blacklist=20Flie?= =?UTF-8?q?=EF=BC=8C=E5=B9=B6=E6=94=AF=E6=8C=81=E7=83=AD=E6=9B=B4=E6=96=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- config.go | 73 ++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 72 insertions(+), 1 deletion(-) diff --git a/config.go b/config.go index c2028ce..fcc0783 100644 --- a/config.go +++ b/config.go @@ -1,12 +1,17 @@ package main import ( + "bufio" "encoding/json" "flag" "fmt" "github.com/sirupsen/logrus" "io/ioutil" + "os" "regexp" + "strings" + "sync" + "time" ) type configStructure struct { @@ -24,10 +29,12 @@ type ruleStructure struct { Address string `json:"address"` } `json:"targets"` FirstPacketTimeout uint64 `json:"first_packet_timeout"` - Blacklist map[string]bool `json:"blacklist"` + BlacklistFile string `json:"blacklist_file"` + blacklistMap map[string]bool `json:"-"` } var config *configStructure +var blacklistMutex = &sync.Mutex{} func init() { cfgPath := flag.String("config", "config.json", "config.json file path") @@ -85,5 +92,69 @@ func (c *ruleStructure) verify() error { v.regexp = r } } + if c.BlacklistFile != "" { + err := loadBlacklist(c.BlacklistFile, &c.blacklistMap) + if err != nil { + return fmt.Errorf("failed to load blacklist: %s", err.Error()) + } + go watchBlacklist(c.BlacklistFile, &c.blacklistMap) + } + return nil +} + +func loadBlacklist(path string, blacklist *map[string]bool) error { + file, err := os.Open(path) + if err != nil { + return err + } + defer file.Close() + + newBlacklist := make(map[string]bool) + scanner := bufio.NewScanner(file) + for scanner.Scan() { + ip := strings.TrimSpace(scanner.Text()) + newBlacklist[ip] = true + } + + if err := scanner.Err(); err != nil { + return err + } + + blacklistMutex.Lock() + oldBlacklist := *blacklist + *blacklist = newBlacklist + blacklistMutex.Unlock() + + // 打印出被移除的IP地址 + for ip := range oldBlacklist { + if !newBlacklist[ip] { + logrus.Infof("At %s, IP %s move out the Blacklist", time.Now().Format(time.RFC3339), ip) + } + } + + // 打印出新添加的IP地址 + for ip := range newBlacklist { + if !oldBlacklist[ip] { + logrus.Infof("At %s, IP %s add to the Blacklist", time.Now().Format(time.RFC3339), ip) + } + } + return nil } + + +func watchBlacklist(path string, blacklist *map[string]bool) { + ticker := time.NewTicker(10 * time.Second) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + err := loadBlacklist(path, blacklist) + if err != nil { + logrus.Errorf("failed to reload blacklist: %s", err.Error()) + } + } + } +} + From b7e37c7af42df005d7b217d34627975f391ce704 Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sat, 13 Jan 2024 22:25:45 +0800 Subject: [PATCH 02/10] =?UTF-8?q?=E6=B7=BB=E5=8A=A0Blacklist=20Flie?= =?UTF-8?q?=EF=BC=8C=E5=B9=B6=E6=94=AF=E6=8C=81=E7=83=AD=E6=9B=B4=E6=96=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core.go | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/core.go b/core.go index 5d42769..4c865e1 100644 --- a/core.go +++ b/core.go @@ -28,10 +28,13 @@ func listen(rule *ruleStructure, wg *sync.WaitGroup) { continue } //判断黑名单 - if len(rule.Blacklist) != 0 { + blacklistMutex.Lock() + blacklist := rule.blacklistMap + blacklistMutex.Unlock() + if len(blacklist) != 0 { clientIP := conn.RemoteAddr().String() clientIP = clientIP[0:strings.LastIndex(clientIP, ":")] - if rule.Blacklist[clientIP] { + if blacklist[clientIP] { logrus.Infof("[%s] disconnected ip in blacklist: %s", rule.Name, clientIP) conn.Close() continue @@ -129,4 +132,4 @@ func waitFirstPacket(conn net.Conn) ([]byte, error) { return nil, err } return buf[:n], nil -} \ No newline at end of file +} From 9e88b6361ab96e29f148179c40435f0250294b09 Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sat, 13 Jan 2024 22:27:05 +0800 Subject: [PATCH 03/10] =?UTF-8?q?=E6=B7=BB=E5=8A=A0blacklist=5Ffile?= =?UTF-8?q?=E5=8F=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 53bbfaf..ee4c68d 100644 --- a/README.md +++ b/README.md @@ -74,6 +74,7 @@ "1.2.3.4":true, "114.114.114.114":true }, + "blacklist_file": "/yourblacklistfliepath", "targets": [ { "address": "127.0.0.1:80" @@ -89,6 +90,7 @@ "1.2.3.4":true, "114.114.114.114":true }, + "blacklist_file": "/yourblacklistfliepath", "targets": [ { "regexp": "^(GET|POST|HEAD|DELETE|PUT|CONNECT|OPTIONS|TRACE)", From b6e42b13d1940e31886fba08067d824cc852fe38 Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sat, 13 Jan 2024 22:27:57 +0800 Subject: [PATCH 04/10] =?UTF-8?q?=E6=B7=BB=E5=8A=A0blacklist=5Ffile?= =?UTF-8?q?=E5=8F=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index ee4c68d..3ed5881 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ "1.2.3.4":true, "114.114.114.114":true }, + "blacklist_file": "/yourblacklistfliepath", "targets": [ 目标配置 ] From 223006aa7699817b656142e695270454791fcbbd Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sat, 13 Jan 2024 22:28:38 +0800 Subject: [PATCH 05/10] Update config.json --- config.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/config.json b/config.json index d72ad7f..0aea63d 100644 --- a/config.json +++ b/config.json @@ -7,6 +7,7 @@ "enable_regexp": false, "first_packet_timeout": 5000, "blacklist": null, + "blacklist_file": "/yourblacklistfliepath", "targets": [ { "regexp": "^(GET|POST|HEAD|DELETE|PUT|CONNECT|OPTIONS|TRACE)", @@ -19,4 +20,4 @@ ] } ] -} \ No newline at end of file +} From 89ba23e2127fe62fe4463838deb2d43de3cfa35a Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sat, 13 Jan 2024 22:56:42 +0800 Subject: [PATCH 06/10] =?UTF-8?q?=E6=B7=BB=E5=8A=A0blacklist=5Ffile?= =?UTF-8?q?=E8=AF=B4=E6=98=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 3ed5881..dbc136b 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,8 @@ 3. enable_regexp为是否开启正则表达式模式，后面有解释 4. first_packet_timeout为等待客户端第一个数据包的超时时间(**毫秒**)，仅开启正则表达式模式后有效，后面有解释 5. blacklist为黑名单IP，在黑名单里面的IP且为true的时候则直接断开链接。如不需要使用黑名单可留null -5. targets为目标配置数组，看下面 +6. blacklist_file文件为单行单IP黑名单，多行多IP同时黑名单 +7. targets为目标配置数组，看下面 #### 目标配置 目标配置有两种模式：**普通模式**和**正则模式**。 From 3a22b4be7eb65083cf96f56da94dd87fb3e0dacf Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sun, 14 Jan 2024 15:29:58 +0800 Subject: [PATCH 07/10] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=B5=81=E9=87=8F?= =?UTF-8?q?=E7=BB=9F=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core.go | 109 ++++++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 90 insertions(+), 19 deletions(-) diff --git a/core.go b/core.go index 4c865e1..c3f2336 100644 --- a/core.go +++ b/core.go @@ -10,6 +10,10 @@ import ( "time" ) +// 定义一个全局的流量统计map +var trafficMap = make(map[string]int64) +var trafficMutex = &sync.Mutex{} + func listen(rule *ruleStructure, wg *sync.WaitGroup) { defer wg.Done() //监听 @@ -53,76 +57,118 @@ func handleNormal(conn net.Conn, rule *ruleStructure) { defer conn.Close() var target net.Conn - //正常模式下挨个连接直到成功连接 + var targetAddress string for _, v := range rule.Targets { c, err := net.Dial("tcp", v.Address) if err != nil { - logrus.Errorf("[%s] try to handle connection (%s) failed because target (%s) connected failed, try next target.", + logrus.Errorf("[%s] try to handle connection %s failed because target %s connected failed, try next target.", rule.Name, conn.RemoteAddr(), v.Address) continue } target = c + targetAddress = v.Address break } if target == nil { - logrus.Errorf("[%s] unable to handle connection (%s) because all targets connected failed", + logrus.Errorf("[%s] unable to handle connection %s because all targets connected failed", rule.Name, conn.RemoteAddr()) return } - logrus.Debugf("[%s] handle connection (%s) to target (%s)", rule.Name, conn.RemoteAddr(), target.RemoteAddr()) + logrus.Debugf("[%s] handle connection %s to target %s", rule.Name, conn.RemoteAddr(), target.RemoteAddr()) defer target.Close() - //io桥 - go io.Copy(conn, target) - io.Copy(target, conn) + var wg sync.WaitGroup + wg.Add(2) + + var traffic1, traffic2 int64 + ip, _, _ := net.SplitHostPort(conn.RemoteAddr().String()) + go func() { + defer wg.Done() + traffic1 = copyWithTrafficCount(conn, target) + trafficMutex.Lock() + trafficMap[ip] += traffic1 + trafficMutex.Unlock() + }() + go func() { + defer wg.Done() + traffic2 = copyWithTrafficCount(target, conn) + trafficMutex.Lock() + trafficMap[ip] += traffic2 + trafficMutex.Unlock() + }() + + wg.Wait() + + trafficMutex.Lock() + logrus.Infof("[%s] Connection from IP %s to target %s: This connection traffic: %.2f MB, Total traffic: %.2f MB", rule.Name, conn.RemoteAddr().String(), targetAddress, float64(traffic1 + traffic2) / (1024 * 1024), float64(trafficMap[ip]) / (1024 * 1024)) + trafficMutex.Unlock() } func handleRegexp(conn net.Conn, rule *ruleStructure) { defer conn.Close() - //正则模式下需要客户端的第一个数据包判断特征，所以需要设置一个超时 conn.SetReadDeadline(time.Now().Add(time.Millisecond * time.Duration(rule.FirstPacketTimeout))) - //获取第一个数据包 firstPacket, err := waitFirstPacket(conn) if err != nil { - logrus.Errorf("[%s] unable to handle connection (%s) because failed to get first packet : %s", + logrus.Errorf("[%s] unable to handle connection %s because failed to get first packet : %s", rule.Name, conn.RemoteAddr(), err.Error()) return } var target net.Conn - //挨个匹配正则 + var targetAddress string for _, v := range rule.Targets { if !v.regexp.Match(firstPacket) { continue } c, err := net.Dial("tcp", v.Address) if err != nil { - logrus.Errorf("[%s] try to handle connection (%s) failed because target (%s) connected failed, try next match target.", + logrus.Errorf("[%s] try to handle connection %s failed because target %s connected failed, try next match target.", rule.Name, conn.RemoteAddr(), v.Address) continue } target = c + targetAddress = v.Address break } if target == nil { - logrus.Errorf("[%s] unable to handle connection (%s) because no match target", + logrus.Errorf("[%s] unable to handle connection %s because no match target", rule.Name, conn.RemoteAddr()) return } - logrus.Debugf("[%s] handle connection (%s) to target (%s)", rule.Name, conn.RemoteAddr(), target.RemoteAddr()) - //匹配到了，去除掉刚才设定的超时 + logrus.Debugf("[%s] handle connection %s to target %s", rule.Name, conn.RemoteAddr(), target.RemoteAddr()) conn.SetReadDeadline(time.Time{}) - //把第一个数据包发送给目标 io.Copy(target, bytes.NewReader(firstPacket)) defer target.Close() - //io桥 - go io.Copy(conn, target) - io.Copy(target, conn) + var wg sync.WaitGroup + wg.Add(2) + + var traffic1, traffic2 int64 + ip, _, _ := net.SplitHostPort(conn.RemoteAddr().String()) + go func() { + defer wg.Done() + traffic1 = copyWithTrafficCount(conn, target) + trafficMutex.Lock() + trafficMap[ip] += traffic1 + trafficMutex.Unlock() + }() + go func() { + defer wg.Done() + traffic2 = copyWithTrafficCount(target, conn) + trafficMutex.Lock() + trafficMap[ip] += traffic2 + trafficMutex.Unlock() + }() + + wg.Wait() + + trafficMutex.Lock() + logrus.Infof("[%s] Connection from IP %s to target %s: This connection traffic: %.2f MB, Total traffic: %.2f MB", rule.Name, conn.RemoteAddr().String(), targetAddress, float64(traffic1 + traffic2) / (1024 * 1024), float64(trafficMap[ip]) / (1024 * 1024)) + trafficMutex.Unlock() } func waitFirstPacket(conn net.Conn) ([]byte, error) { @@ -133,3 +179,28 @@ func waitFirstPacket(conn net.Conn) ([]byte, error) { } return buf[:n], nil } + +func copyWithTrafficCount(dst io.Writer, src io.Reader) int64 { + buf := make([]byte, 32*1024) + var traffic int64 = 0 + for { + nr, er := src.Read(buf) + if nr > 0 { + nw, ew := dst.Write(buf[0:nr]) + if nw > 0 { + traffic += int64(nw) + } + if ew != nil { + break + } + if nr != nw { + logrus.Errorf("partial write") + break + } + } + if er != nil { + break + } + } + return traffic +} From 8fd45104fb8a83f005d0308831217ef6cb98b359 Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sun, 14 Jan 2024 15:30:16 +0800 Subject: [PATCH 08/10] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=B5=81=E9=87=8F?= =?UTF-8?q?=E7=BB=9F=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- config.go | 1 - 1 file changed, 1 deletion(-) diff --git a/config.go b/config.go index fcc0783..11a8bcc 100644 --- a/config.go +++ b/config.go @@ -157,4 +157,3 @@ func watchBlacklist(path string, blacklist *map[string]bool) { } } } - From c5a9c55a0db0d5d76a7d72a6fee2eb926a30bbad Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sun, 14 Jan 2024 15:32:18 +0800 Subject: [PATCH 09/10] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E6=B5=81=E9=87=8F?= =?UTF-8?q?=E7=BB=9F=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/main.go b/main.go index dd52ffe..9e3cac4 100644 --- a/main.go +++ b/main.go @@ -6,7 +6,7 @@ import ( ) const ( - VERSION = "2.0" + VERSION = "2.1" ) func main() { From 95f426c83978889b17e475b32af1a990c70c8472 Mon Sep 17 00:00:00 2001 From: AKEBI-NUKUI <132581172+AKEBI-NUKUI@users.noreply.github.com> Date: Sun, 14 Jan 2024 16:16:44 +0800 Subject: [PATCH 10/10] =?UTF-8?q?=E8=A7=84=E8=8C=83=E6=B5=81=E9=87=8F?= =?UTF-8?q?=E7=BB=9F=E8=AE=A1=E6=98=BE=E7=A4=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core.go b/core.go index c3f2336..149f5b4 100644 --- a/core.go +++ b/core.go @@ -101,7 +101,7 @@ func handleNormal(conn net.Conn, rule *ruleStructure) { wg.Wait() trafficMutex.Lock() - logrus.Infof("[%s] Connection from IP %s to target %s: This connection traffic: %.2f MB, Total traffic: %.2f MB", rule.Name, conn.RemoteAddr().String(), targetAddress, float64(traffic1 + traffic2) / (1024 * 1024), float64(trafficMap[ip]) / (1024 * 1024)) + logrus.Infof("[%s] %s to target %s: This connection traffic: %.2f MB, Total traffic: %.2f MB", rule.Name, conn.RemoteAddr().String(), targetAddress, float64(traffic1 + traffic2) / (1024 * 1024), float64(trafficMap[ip]) / (1024 * 1024)) trafficMutex.Unlock() } @@ -167,7 +167,7 @@ func handleRegexp(conn net.Conn, rule *ruleStructure) { wg.Wait() trafficMutex.Lock() - logrus.Infof("[%s] Connection from IP %s to target %s: This connection traffic: %.2f MB, Total traffic: %.2f MB", rule.Name, conn.RemoteAddr().String(), targetAddress, float64(traffic1 + traffic2) / (1024 * 1024), float64(trafficMap[ip]) / (1024 * 1024)) + logrus.Infof("[%s] %s to target %s: This connection traffic: %.2f MB, Total traffic: %.2f MB", rule.Name, conn.RemoteAddr().String(), targetAddress, float64(traffic1 + traffic2) / (1024 * 1024), float64(trafficMap[ip]) / (1024 * 1024)) trafficMutex.Unlock() }