Further refactoring and finishing of TODO's

betheroot · Aug 30, 2017 · 62cc459 · 62cc459
1 parent b7e61fd
commit 62cc459
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 15 deletions.
diff --git a/main.go b/main.go
@@ -8,6 +8,7 @@ import (
 	"os/signal"
 	"strconv"
 	"syscall"
+	"time"
 
 	log "github.com/Sirupsen/logrus"
 )
@@ -18,6 +19,7 @@ type Configuration struct {
 	PgUsers       []string `json:"pgUsers"`
 	Debug         bool     `json:"debug"`
 	Cleartext     bool     `json:"cleartext"`
+	TcpTimeout    int      `json:"server_timeout"`
 	HpFeedsConfig `json:"hpfeedsConfig"`
 }
 
@@ -51,6 +53,7 @@ func main() {
 	debug := config.Debug
 	cleartext := config.Cleartext
 	hpFeedsConfig := config.HpFeedsConfig
+	tcpTimeout := time.Duration(config.TcpTimeout) * time.Second
 
 	if debug {
 		log.SetLevel(log.DebugLevel)
@@ -66,6 +69,7 @@ func main() {
 		addr,
 		pgUsers,
 		cleartext,
+		tcpTimeout,
 		hpfeedsChannel,
 		hpFeedsConfig.Enabled,
 	)

diff --git a/pgconn.go b/pgconn.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"net"
+	"time"
 
 	log "github.com/Sirupsen/logrus"
 )
@@ -18,7 +19,8 @@ type PostgresConnection struct {
 	postgresPacket readBuf
 }
 
-func NewPostgresConnection(conn net.Conn) *PostgresConnection {
+func NewPostgresConnection(conn net.Conn, tcpTimeout time.Duration) *PostgresConnection {
+	conn.SetDeadline(time.Now().Add(tcpTimeout))
 	return &PostgresConnection{
 		buffer:     make([]byte, maxBufSize),
 		connection: conn,

diff --git a/pghoney.conf.sample b/pghoney.conf.sample
@@ -4,6 +4,7 @@
   "pgUsers":["postgres"],
   "debug":false,
   "cleartext":false,
+  "server_timeout":10,
   "hpfeedsConfig":{
     "port":10000,
     "host":"127.0.0.1",

diff --git a/server.go b/server.go
@@ -11,11 +11,6 @@ import (
 	log "github.com/Sirupsen/logrus"
 )
 
-var (
-	// TODO: Make configurable
-	tcpTimeout = 10 * time.Second
-)
-
 type PostgresServer struct {
 	listener net.Listener
 
@@ -27,11 +22,12 @@ type PostgresServer struct {
 	addr string
 	port string
 
-	pgUsers   map[string]bool
-	cleartext bool
+	pgUsers    map[string]bool
+	cleartext  bool
+	tcpTimeout time.Duration
 }
 
-func NewPostgresServer(port string, addr string, users []string, cleartext bool, hpfeedsChan chan []byte, hpfeedsEnabled bool) *PostgresServer {
+func NewPostgresServer(port string, addr string, users []string, cleartext bool, tcpTimeout time.Duration, hpfeedsChan chan []byte, hpfeedsEnabled bool) *PostgresServer {
 	listener, err := net.Listen("tcp", addr+":"+port)
 	if err != nil {
 		log.Errorf("Error listening: %s", err)
@@ -51,6 +47,7 @@ func NewPostgresServer(port string, addr string, users []string, cleartext bool,
 		addr:           addr,
 		port:           port,
 		cleartext:      cleartext,
+		tcpTimeout:     tcpTimeout,
 		pgUsers:        pgUsers,
 	}
 }
@@ -63,17 +60,16 @@ func (p *PostgresServer) Close() {
 func (p *PostgresServer) Listen() {
 	log.Infof("Starting to listening on %s:%s...", p.addr, p.port)
 	for {
-		//FIXME: conn is an example of primitive obsession
 		conn, err := p.listener.Accept()
 		if err != nil {
 			log.Warn("Error accepting: %s", err)
 			continue
 		}
 
-		conn.SetDeadline(time.Now().Add(tcpTimeout))
+		pgConn := NewPostgresConnection(conn, p.tcpTimeout)
 
 		p.waitGroup.Add(1)
-		go p.handleRequest(conn)
+		go p.handleRequest(pgConn)
 	}
 }
 
@@ -104,10 +100,8 @@ func (p *PostgresServer) sendToHpFeeds(pgConn *PostgresConnection) error {
 	return err
 }
 
-func (p *PostgresServer) handleRequest(conn net.Conn) {
+func (p *PostgresServer) handleRequest(pgConn *PostgresConnection) {
 	defer p.waitGroup.Done()
-
-	pgConn := NewPostgresConnection(conn)
 	defer pgConn.Close()
 
 	for {