From 9f2ad13f6fa34989a5f2e58112f7f232f7ccab8a Mon Sep 17 00:00:00 2001
From: yo <yo@plopito.tk>
Date: Sat, 6 Feb 2021 14:44:53 +0100
Subject: [PATCH] Add timeout, work on connection closing (previous commit left
 connections open)

---
 mxrouter.go | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/mxrouter.go b/mxrouter.go
index 932c4e6..b8ce17f 100644
--- a/mxrouter.go
+++ b/mxrouter.go
@@ -23,6 +23,7 @@ import (
 	"os"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/go-ldap/ldap/v3"
 	"github.com/peterbourgon/ff"
@@ -30,7 +31,7 @@ import (
 )
 
 const (
-	version = "0.9.7"
+	version = "0.9.7.1"
 )
 
 var (
@@ -46,6 +47,7 @@ var (
 	ldapAttr    *string
 	defDomain   *string
 	pidFilePath *string
+	timeout     *int
 )
 
 func ldapSearch(searchReq *ldap.SearchRequest, attempt int) (*ldap.SearchResult, error) {
@@ -85,9 +87,19 @@ func handleConnection(connClt net.Conn, conLdap *ldap.Conn) {
 	var mxdom string
 	buf := make([]byte, 1024)
 
-	// It's client responsability to close connection
-	//defer connClt.Close()
+	// Close connection when this function ends
+	defer func() {
+		logstream.Info("Closing connection")
+		connClt.Close()
+	}()
+
+	timeoutDuration := time.Duration(*timeout) * time.Second
+
 	for {
+		// Set a deadline for reading. Read operation will fail if no data
+		// is received after deadline.
+		connClt.SetReadDeadline(time.Now().Add(timeoutDuration))
+
 		readlen, err := connClt.Read(buf)
 		if err != nil {
 			// Dont notice if client closed connection
@@ -229,6 +241,7 @@ func main() {
 	ldapAttr = fs.String("ldapAttr", "relayName", "LDAP attribute containing the relay name to return")
 	defDomain = fs.String("domain", "", "Domain to add to relay name if not a FQDN (also via DOMAIN env var)")
 	pidFilePath = fs.String("pidfile", "/var/run/mxrouter/mxrouter.pid", "PID File (also via PIDFILE env var)")
+	timeout = fs.Int("timeout", 5, "timeout in seconds")
 	_ = fs.String("config", "", "config file (optional)")
 
 	// Surcharge de la fonction Usage()