Merge pull request #2 from HirbodBehnam/dev

Dev

README.md

@@ -22,6 +22,8 @@ Did you download the executable for your os? Good!
 
 Edit the `rules.json` file as you wish. Here is the cheatsheet for it:
 * `SaveDuration`: The program writes the quotas to disk every `SaveDuration` seconds. Default is 600 seconds or 10 minutes.
+* `Timeout`: The time in seconds that a connection can stay alive without transmitting any data. Default is 15 minutes. Use -1 to disable the timeout.
+* `TimeoutCheck`: The time in seconds to check if connections are still alive. Default is 60 seconds.
 * `Rules` Array: Each element represents a forwarding rule and quota for it.
     * `Listen`: The local port to accept the incoming connections for proxy.
     * `Forward`: The address that the traffic must be forwarded to. Enter it like `ip:port`
@@ -52,5 +54,10 @@ As soon as the function returns, I the quota will change.
 
 And what do I mean about the softer connections? The client can use the program after the quota is reached. When the client wants to establish a new connection it will be rejected from the server. Plus you can manage how much client has used more than its quota.
 
+### Timeout
+I implemented a custom timeout method for this. It is dead simple:
+
+Save the last time each _connection_ has transmitted something. (On each write or read function). Then every minute(default), the app reads all of the last transmit time. If it is larger than (now - timeout), close the connection.
+
 ## Other Stuff
 [Persian guild to setup this with mtproto](http://www.mediafire.com/file/4u3khp5oj7ecgxk/%25D9%2585%25D8%25AD%25D8%25AF%25D9%2588%25D8%25AF_%25DA%25A9%25D8%25B1%25D8%25AF%25D9%2586_%25DA%25A9%25D8%25A7%25D8%25B1%25D8%25A8%25D8%25B1%25D8%25A7%25D9%2586.pdf/file)

main.go

@@ -21,7 +21,7 @@ var ConfigFileName = "rules.json"
 var SimultaneousConnections CSafeConnections
 var Verbose = false
 
-const Version = "1.0.0 / Build 4"
+const Version = "1.1.0 / Build 5"
 
 type CSafeConnections struct {
 	SimultaneousConnections []int
@@ -36,14 +36,26 @@ type CSafeRule struct {
 type Rule struct {
 	Listen       uint16
 	Forward      string
-	Quota        int64
+	Quota        uint64
 	Simultaneous int
 }
 type Config struct {
 	SaveDuration int
+	Timeout      int64
+	TimeoutCheck int
 	Rules        []Rule
 }
 
+type AliveStatus struct {
+	mu         sync.Mutex
+	LastAccess int64
+	con        *net.Conn
+}
+
+var LastAlive = make(map[uint64]*AliveStatus)
+var IndexerAlive = uint64(0)
+var EnableTimeOut = true
+
 func main() {
 	{ //Parse arguments
 		configFileName := flag.String("config", "rules.json", "The config filename")
@@ -67,17 +79,22 @@ func main() {
 	}
 
 	//Read config file
-	confF, err := ioutil.ReadFile(ConfigFileName)
-	if err != nil {
-		panic("Cannot read the config file. (io Error) " + err.Error())
-	}
 	var conf Config
-	err = json.Unmarshal(confF, &conf)
-	if err != nil {
-		panic("Cannot read the config file. (Parse Error) " + err.Error())
+	{
+		confF, err := ioutil.ReadFile(ConfigFileName)
+		if err != nil {
+			panic("Cannot read the config file. (io Error) " + err.Error())
+		}
+		err = json.Unmarshal(confF, &conf)
+		if err != nil {
+			panic("Cannot read the config file. (Parse Error) " + err.Error())
+		}
+		Rules.Rules = conf.Rules
+		if conf.Timeout == -1 {
+			fmt.Println("Disabled timeout")
+			EnableTimeOut = false
+		}
 	}
-	Rules.Rules = conf.Rules
-	SimultaneousConnections.SimultaneousConnections = make([]int, len(Rules.Rules))
 
 	//Start listeners
 	for index := range Rules.Rules {
@@ -116,10 +133,44 @@ func main() {
 		}(index)
 	}
 
+	//Keep alive status
+	go func() {
+		if !EnableTimeOut {
+			return
+		}
+		//At first check for timeout values; if they are zero assign defaults
+		timeout := conf.Timeout
+		if conf.Timeout == 0 {
+			timeout = 900 //15 minute timeout
+			conf.Timeout = 900
+		}
+		checkDuration := conf.TimeoutCheck
+		if checkDuration == 0 {
+			checkDuration = 60 //Every minute check if the connections are still alive
+			conf.TimeoutCheck = 60
+		}
+		for {
+			time.Sleep(time.Duration(checkDuration) * time.Second)
+			for index, i := range LastAlive {
+				if i.LastAccess+timeout < time.Now().Unix() {
+					if Verbose {
+						log.Println("Dropping a dead connection from", (*i.con).RemoteAddr(), ";", index, "; The last accessed time is", time.Unix(i.LastAccess, 0).Format("2006-01-02 15:04:05"))
+					}
+					(*i.con).Close() //Close the connection and the copyBuffer method will throw an error; So the values will be saved
+				}
+			}
+		}
+	}()
+
 	//Save config file
 	go func() {
+		sd := conf.SaveDuration
+		if sd == 0 {
+			sd = 600
+			conf.SaveDuration = 600
+		}
 		for {
-			time.Sleep(time.Duration(conf.SaveDuration) * time.Second) //Save file every x seconds
+			time.Sleep(time.Duration(sd) * time.Second) //Save file every x seconds
 			saveConfig(conf)
 		}
 	}()
@@ -179,23 +230,84 @@ func handleRequest(conn net.Conn, index int, r Rule) {
 
 	SimultaneousConnections.mu.Lock()
 	SimultaneousConnections.SimultaneousConnections[index] += 2 //Two is added; One for client to server and another for server to client
+	if Verbose {
+		log.Println("Accepting a connection from", conn.RemoteAddr(), "; Now", SimultaneousConnections.SimultaneousConnections[index], "SimultaneousConnections")
+	}
+	LastAlive[IndexerAlive] = &AliveStatus{con: &conn, LastAccess: time.Now().Unix()}
+	t := IndexerAlive
+	IndexerAlive++
 	SimultaneousConnections.mu.Unlock()
 
-	go copyIO(conn, proxy, index)
-	go copyIO(proxy, conn, index)
+	go copyIO(conn, proxy, index, t)
+	go copyIO(proxy, conn, index, t)
 }
 
-func copyIO(src, dest net.Conn, index int) {
+func copyIO(src, dest net.Conn, index int, aLiveIndex uint64) {
 	defer src.Close()
 	defer dest.Close()
 
-	r, _ := io.Copy(src, dest) //r is the amount of bytes transferred
+	var r uint64 //r is the amount of bytes transferred
+	var err error
+
+	if EnableTimeOut {
+		r, err = copyBuffer(src, dest, aLiveIndex)
+	} else {
+		var r1 int64
+		r1, err = io.Copy(src, dest)
+		r = uint64(r1)
+	}
+
+	if Verbose && err != nil {
+		log.Println("Error on copyBuffer(Usually happens):", err.Error())
+	}
 
 	Rules.mu.Lock() //Lock to read the amount of data transferred
 	Rules.Rules[index].Quota -= r
 	Rules.mu.Unlock()
 
 	SimultaneousConnections.mu.Lock()
 	SimultaneousConnections.SimultaneousConnections[index]-- //This will actually run twice
+	delete(LastAlive, aLiveIndex)
+	if Verbose {
+		log.Println("Closing a connection from", src.RemoteAddr(), "; Connections Now:", SimultaneousConnections.SimultaneousConnections[index])
+	}
 	SimultaneousConnections.mu.Unlock()
 }
+
+func copyBuffer(dst, src net.Conn, index uint64) (written uint64, err error) {
+	buf := make([]byte, 32768)
+	for {
+		go updateDate(index)
+		nr, er := src.Read(buf)
+		if nr > 0 {
+			go updateDate(index)
+			nw, ew := dst.Write(buf[0:nr])
+			if nw > 0 {
+				written += uint64(nw)
+			}
+			if ew != nil {
+				err = ew
+				break
+			}
+			if nr != nw {
+				err = io.ErrShortWrite
+				break
+			}
+		}
+		if er != nil {
+			if er != io.EOF {
+				err = er
+			}
+			break
+		}
+	}
+	return written, err
+}
+
+func updateDate(index uint64) {
+	if LastAlive[index] != nil {
+		LastAlive[index].mu.Lock()
+		LastAlive[index].LastAccess = time.Now().Unix()
+		LastAlive[index].mu.Unlock()
+	}
+}

rules.json

@@ -1,5 +1,6 @@
 {
   "SaveDuration": 600,
+  "Timeout": -1,
   "Rules": [
     {
       "Listen": 5362,
@@ -13,4 +14,4 @@
       "Quota":9223372036854775807
     }
   ]
-}
+}