Better connection management

- Send keepalives
- Reconnect when the connection is lost
- Set timeout for write operations

1 files changed, 12 insertions, 8 deletions

diff --git a/irc/states.go b/irc/states.go
index 9737565..b6db9d0 100644
--- a/irc/states.go
+++ b/irc/states.go
@@ -6,13 +6,15 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
-	"io"
+	"net"
 	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
 )
 
+const writeDeadline = 10 * time.Second
+
 type SASLClient interface {
 	Handshake() (mech string)
 	Respond(challenge string) (res string, err error)
@@ -149,7 +151,7 @@ type SessionParams struct {
 
 // Session is an IRC session/connection/whatever.
 type Session struct {
-	conn io.ReadWriteCloser
+	conn net.Conn
 	msgs chan Message // incoming messages.
 	acts chan action  // user actions.
 	evts chan Event   // events sent to the user.
@@ -183,7 +185,7 @@ type Session struct {
 //
 // It returns an error when the paramaters are invalid, or when it cannot write
 // to the connection.
-func NewSession(conn io.ReadWriteCloser, params SessionParams) (*Session, error) {
+func NewSession(conn net.Conn, params SessionParams) (*Session, error) {
 	s := &Session{
 		conn:          conn,
 		msgs:          make(chan Message, 64),
@@ -207,11 +209,6 @@ func NewSession(conn io.ReadWriteCloser, params SessionParams) (*Session, error)
 
 	s.running.Store(true)
 
-	err := s.send("CAP LS 302\r\nNICK %s\r\nUSER %s 0 * :%s\r\n", s.nick, s.user, s.real)
-	if err != nil {
-		return nil, err
-	}
-
 	go func() {
 		r := bufio.NewScanner(conn)
 
@@ -233,6 +230,11 @@ func NewSession(conn io.ReadWriteCloser, params SessionParams) (*Session, error)
 		s.Stop()
 	}()
 
+	err := s.send("CAP LS 302\r\nNICK %s\r\nUSER %s 0 * :%s\r\n", s.nick, s.user, s.real)
+	if err != nil {
+		return nil, err
+	}
+
 	go s.run()
 
 	return s, nil
@@ -1081,6 +1083,8 @@ func (s *Session) updateFeatures(features []string) {
 
 func (s *Session) send(format string, args ...interface{}) (err error) {
 	msg := fmt.Sprintf(format, args...)
+
+	s.conn.SetWriteDeadline(time.Now().Add(writeDeadline))
 	_, err = s.conn.Write([]byte(msg))
 
 	if s.debug {