From 13ff3f09e87557da59874762af7eb6a6966c002a Mon Sep 17 00:00:00 2001
From: zhiqiang feng <fengzhq@it2000.com.cn>
Date: Tue, 22 Dec 2020 11:45:25 +0800
Subject: [PATCH] fix run, add connection check

---
 cmd/start.go    |  8 ++++----
 utils/socket.go | 41 +++++++++++++++++++++++++++++++++++++++--
 utils/sqlite.go | 14 +++++++++++---
 3 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/cmd/start.go b/cmd/start.go
index 197f196..0ae9821 100644
--- a/cmd/start.go
+++ b/cmd/start.go
@@ -17,12 +17,13 @@ package cmd
 
 import (
 	"fmt"
-	"it2000.com.cn/tele-recv/utils"
 	"log"
 	"os"
 	"os/signal"
 	"syscall"
 
+	"it2000.com.cn/tele-recv/utils"
+
 	"github.com/spf13/cobra"
 )
 
@@ -41,8 +42,6 @@ to quickly create a Cobra application.`,
 	},
 }
 
-var run bool
-
 func init() {
 	rootCmd.AddCommand(startCmd)
 
@@ -77,7 +76,8 @@ func start() {
 		sig := <-sigs
 		fmt.Println()
 		fmt.Println(sig)
-		run = false
+		utils.ServerRunning = false
+		utils.StopSocketServer()
 		done <- true
 	}()
 
diff --git a/utils/socket.go b/utils/socket.go
index d5b9d57..96d1004 100644
--- a/utils/socket.go
+++ b/utils/socket.go
@@ -1,17 +1,21 @@
 package utils
 
 import (
+	"io"
 	"log"
 	"net"
+	"time"
 )
 
 const ServerType = "tcp"
 
 var ServerRunning = false
 var client net.Conn
+var server net.Listener
 
 func Listen(address string) {
-	server, err := net.Listen(ServerType, address)
+	var err error
+	server, err = net.Listen(ServerType, address)
 	if err != nil {
 		log.Fatal("error in create server ", err)
 	}
@@ -25,14 +29,22 @@ func Listen(address string) {
 		}
 		if client == nil {
 			client = conn
+			log.Println("client connected ", client.RemoteAddr())
+
 		} else {
 			log.Println("client already connected, close")
+			client.Close()
+			client = nil
 			conn.Close()
 		}
+		if IsClientConnected() {
+			connCheck()
+		}
+		time.Sleep(2 * time.Second)
 	}
 }
 
-func Write(data string) {
+func WriteToClient(data string) error {
 	defer func() {
 		client.Close()
 		client = nil
@@ -43,4 +55,29 @@ func Write(data string) {
 		client.Close()
 		client = nil
 	}
+	return err
+}
+
+func connCheck() bool {
+	_, err := client.Read(make([]byte, 0))
+	if err != nil && err != io.EOF {
+		// this connection is invalid
+		log.Println("conn closed....", err)
+		client = nil
+		return false
+	}
+	return true
+}
+
+func IsClientConnected() bool {
+	return client != nil
+}
+
+func StopSocketServer() {
+	if IsClientConnected() {
+		client.Close()
+	}
+	if server != nil {
+		server.Close()
+	}
 }
diff --git a/utils/sqlite.go b/utils/sqlite.go
index 90753d1..7647833 100644
--- a/utils/sqlite.go
+++ b/utils/sqlite.go
@@ -2,10 +2,11 @@ package utils
 
 import (
 	"database/sql"
-	_ "github.com/mattn/go-sqlite3"
 	"log"
 	"os"
 	"time"
+
+	_ "github.com/mattn/go-sqlite3"
 )
 
 const (
@@ -18,7 +19,8 @@ const (
                     [tele_text] TEXT NOT NULL    
                 )
 `
-	TableInsert = "insert into telegram (tele_text) values (?)"
+	InsertNew = "insert into telegram (tele_text) values (?)"
+	InsertOld = "insert into telegram (tele_text, tele_processed) values (?, 1)"
 )
 
 var (
@@ -69,7 +71,13 @@ func InsertTelegram(teleString string) {
 	isWriting = true
 	db := getDb()
 	defer db.Close()
-	stmt, err := db.Prepare(TableInsert)
+	var insertSQL string
+	if IsClientConnected() && WriteToClient(teleString) == nil {
+		insertSQL = InsertOld
+	} else {
+		insertSQL = InsertNew
+	}
+	stmt, err := db.Prepare(insertSQL)
 	if err != nil {
 		log.Fatal(err)
 	}