Merge pull request #7 from bubble-diff/dev

修复一些bug并支持汇报部署信息

Author: peanutzhen

config.go

@@ -1,9 +1,13 @@
 package main
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
+	"net/http"
 	"os"
+	"sync"
+	"time"
 
 	"github.com/sirupsen/logrus"
 )
@@ -22,6 +26,9 @@ type config struct {
 
 	// DeviceIPv4 网卡ipv4地址
 	DeviceIPv4 string
+
+	mu            sync.Mutex
+	isTaskRunning bool
 }
 
 var configuration = config{}
@@ -47,8 +54,95 @@ func (c *config) init() {
 		logrus.Fatalf("%s: this device has no ipv4 address.", c.Device)
 	}
 
+	// 启动后台去执行一些自动更新的动作
+	go func() {
+		logrus.Info("configuration background started.")
+		for {
+			c.setIsTaskRunning()
+			c.reportDepolyed()
+			time.Sleep(5 * time.Second)
+		}
+	}()
+
 	logrus.WithField(
 		"configuration",
 		fmt.Sprintf("%+v", configuration),
 	).Debug("configuration initialized.")
 }
+
+func (c *config) getIsTaskRunning() bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.isTaskRunning
+}
+
+func (c *config) setIsTaskRunning() {
+	var err error
+	var apiResp struct {
+		Err       string `json:"err"`
+		IsRunning bool   `json:"is_running"`
+	}
+
+	logrus.Infof("[setIsTaskRunning] updating TaskID=%d status...", c.Taskid)
+	// call bubblereplay
+	api := fmt.Sprintf("http://%s%s/%d", c.ReplaySvrAddr, ApiTaskStatus, c.Taskid)
+	resp, err := http.Get(api)
+	if err != nil {
+		logrus.Errorf("[setIsTaskRunning] call api failed, %s", err)
+		return
+	}
+
+	err = json.NewDecoder(resp.Body).Decode(&apiResp)
+	if err != nil {
+		logrus.Errorf("[setIsTaskRunning] decode json response failed, %s", err)
+		return
+	}
+	resp.Body.Close()
+
+	if apiResp.Err != "" {
+		logrus.Errorf("[setIsTaskRunning] response return error, %s", err)
+		return
+	}
+
+	// update state
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.isTaskRunning = apiResp.IsRunning
+}
+
+func (c *config) reportDepolyed() {
+	var err error
+	var apiBody = struct {
+		TaskID int64 `json:"task_id"`
+		// Addr 基准服务地址
+		Addr string `json:"addr"`
+	}{
+		TaskID: c.Taskid,
+		Addr:   c.DeviceIPv4,
+	}
+	var apiResp struct {
+		Err string `json:"err"`
+	}
+
+	logrus.Infof("[reportDepolyed] reporting TaskID=%d has been deployed the bubblecopy...", c.Taskid)
+	// call bubblereplay
+	api := fmt.Sprintf("http://%s%s", c.ReplaySvrAddr, ApiSetDeployed)
+	data, err := json.Marshal(apiBody)
+	resp, err := http.Post(api, "application/json", bytes.NewReader(data))
+	if err != nil {
+		logrus.Errorf("[reportDepolyed] call api failed, %s", err)
+		return
+	}
+
+	err = json.NewDecoder(resp.Body).Decode(&apiResp)
+	if err != nil {
+		logrus.Errorf("[reportDepolyed] decode json response failed, %s", err)
+		return
+	}
+	resp.Body.Close()
+
+	if apiResp.Err != "" {
+		logrus.Errorf("[reportDepolyed] response return error, %s", err)
+		return
+	}
+}

const.go

@@ -14,4 +14,6 @@ const (
 
 const (
 	ApiAddRecord = "/record/add"
+	ApiTaskStatus = "/task_status"
+	ApiSetDeployed = "/deploy"
 )

main.go

@@ -69,26 +69,43 @@ func main() {
 
 	ticker := time.NewTicker(time.Second * 30)
 	for {
-		// todo: 等待Diff任务启动，若未启动，请勿进行抓包消耗CPU
-		// your code here...
-
+		skip := false
 		done := false
+
+		// 等待Diff任务启动，若未启动，请勿进行抓包消耗CPU
+		if configuration.getIsTaskRunning() == false {
+			logrus.Infof("TaskID=%d is not running, self sleeping...", configuration.Taskid)
+			skip = true
+		}
+
 		select {
 		case <-signalChan:
 			logrus.Info("Caught SIGINT: aborting")
 			done = true
 		case <-ticker.C:
 			// 停止监听30秒内无数据传输的连接
 			assembler.FlushCloseOlderThan(time.Now().Add(time.Second * -30))
+		default:
+			// nop
+		}
+
+		if done {
+			break
+		}
+		if skip {
+			time.Sleep(time.Second)
+			continue
+		}
+
+		select {
 		case packet := <-source.Packets():
 			tcp := packet.Layer(layers.LayerTypeTCP)
 			if tcp != nil {
 				tcp := tcp.(*layers.TCP)
 				assembler.Assemble(packet.NetworkLayer().NetworkFlow(), tcp)
 			}
-		}
-		if done {
-			break
+		default:
+			// nop
 		}
 	}