Refactor keepalive logic (ortuman#105)

Author: ortuman

CHANGELOG.md

@@ -4,7 +4,7 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html).
 
-## [0.10.0] - 2020-03-18
+## [0.10.1] - 2020-03-22
 ### Changed
 - Set resource limit
 

c2s/config.go

@@ -73,7 +73,6 @@ type TransportConfig struct {
 	Type        transport.Type
 	BindAddress string
 	Port        int
-	KeepAlive   time.Duration
 	URLPath     string
 }
 
@@ -96,9 +95,6 @@ func (t *TransportConfig) UnmarshalYAML(unmarshal func(interface{}) error) error
 	case "", "socket":
 		t.Type = transport.Socket
 
-	case "websocket":
-		t.Type = transport.WebSocket
-
 	default:
 		return fmt.Errorf("c2s.TransportConfig: unrecognized transport type: %s", p.Type)
 	}
@@ -114,10 +110,6 @@ func (t *TransportConfig) UnmarshalYAML(unmarshal func(interface{}) error) error
 	if t.Port == 0 {
 		t.Port = defaultTransportPort
 	}
-	t.KeepAlive = time.Duration(p.KeepAlive) * time.Second
-	if t.KeepAlive == 0 {
-		t.KeepAlive = defaultTransportKeepAlive
-	}
 	return nil
 }
 
@@ -132,6 +124,7 @@ type Config struct {
 	ID               string
 	ConnectTimeout   time.Duration
 	Timeout          time.Duration
+	KeepAlive        time.Duration
 	MaxStanzaSize    int
 	ResourceConflict ResourceConflictPolicy
 	Transport        TransportConfig
@@ -145,6 +138,7 @@ type configProxy struct {
 	TLS              TLSConfig       `yaml:"tls"`
 	ConnectTimeout   int             `yaml:"connect_timeout"`
 	Timeout          int             `yaml:"timeout"`
+	KeepAlive        int             `yaml:"keep_alive"`
 	MaxStanzaSize    int             `yaml:"max_stanza_size"`
 	ResourceConflict string          `yaml:"resource_conflict"`
 	Transport        TransportConfig `yaml:"transport"`
@@ -167,6 +161,10 @@ func (cfg *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	if cfg.Timeout == 0 {
 		cfg.Timeout = defaultTimeout
 	}
+	cfg.KeepAlive = time.Duration(p.KeepAlive) * time.Second
+	if cfg.KeepAlive == 0 {
+		cfg.KeepAlive = defaultTransportKeepAlive
+	}
 	cfg.MaxStanzaSize = p.MaxStanzaSize
 	if cfg.MaxStanzaSize == 0 {
 		cfg.MaxStanzaSize = defaultMaxStanzaSize
@@ -200,9 +198,9 @@ func (cfg *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
 }
 
 type streamConfig struct {
-	transport        transport.Transport
 	connectTimeout   time.Duration
 	timeout          time.Duration
+	keepAlive        time.Duration
 	maxStanzaSize    int
 	resourceConflict ResourceConflictPolicy
 	sasl             []string

c2s/config_test.go

@@ -8,7 +8,6 @@ package c2s
 import (
 	"os"
 	"testing"
-	"time"
 
 	"github.com/ortuman/jackal/transport"
 	"github.com/ortuman/jackal/transport/compress"
@@ -46,14 +45,6 @@ func TestTransportConfig(t *testing.T) {
 	require.Equal(t, transport.Socket, s.Type)
 	require.Equal(t, "0.0.0.0", s.BindAddress)
 	require.Equal(t, 5222, s.Port)
-	require.Equal(t, time.Second*time.Duration(120), s.KeepAlive)
-
-	err = yaml.Unmarshal([]byte("{type: websocket, url_path: /xmpp/ws}"), &s)
-	require.Nil(t, err)
-
-	require.Equal(t, transport.WebSocket, s.Type)
-	require.Equal(t, 5222, s.Port)
-	require.Equal(t, time.Second*time.Duration(120), s.KeepAlive)
 }
 
 func TestConfig(t *testing.T) {

c2s/in.go

@@ -46,13 +46,15 @@ type inStream struct {
 	mods           *module.Modules
 	comps          *component.Components
 	sess           *session.Session
+	tr             transport.Transport
+	mu             sync.RWMutex
 	id             string
 	connectTm      *time.Timer
+	readTimeoutTm  *time.Timer
 	state          uint32
 	authenticators []auth.Authenticator
 	activeAuth     auth.Authenticator
 	runQueue       *runqueue.RunQueue
-	mu             sync.RWMutex
 	jid            *jid.JID
 	secured        bool
 	compressed     bool
@@ -63,10 +65,11 @@ type inStream struct {
 	ctxCancelFn    context.CancelFunc
 }
 
-func newStream(id string, config *streamConfig, mods *module.Modules, comps *component.Components, router router.Router, userRep repository.User, blockListRep repository.BlockList) stream.C2S {
+func newStream(id string, config *streamConfig, tr transport.Transport, mods *module.Modules, comps *component.Components, router router.Router, userRep repository.User, blockListRep repository.BlockList) stream.C2S {
 	ctx, ctxCancelFn := context.WithCancel(context.Background())
 	s := &inStream{
 		cfg:          config,
+		tr:           tr,
 		router:       router,
 		userRep:      userRep,
 		blockListRep: blockListRep,
@@ -79,7 +82,7 @@ func newStream(id string, config *streamConfig, mods *module.Modules, comps *com
 	}
 
 	// initialize stream context
-	secured := !(config.transport.Type() == transport.Socket)
+	secured := !(tr.Type() == transport.Socket)
 	s.setSecured(secured)
 	s.setJID(&jid.JID{})
 
@@ -179,14 +182,13 @@ func (s *inStream) Disconnect(ctx context.Context, err error) {
 	waitCh := make(chan struct{})
 	s.runQueue.Run(func() {
 		s.disconnect(ctx, err)
-		s.ctxCancelFn()
 		close(waitCh)
 	})
 	<-waitCh
 }
 
 func (s *inStream) initializeAuthenticators() {
-	tr := s.cfg.transport
+	tr := s.tr
 	var authenticators []auth.Authenticator
 	for _, a := range s.cfg.sasl {
 		switch a {
@@ -266,7 +268,7 @@ func (s *inStream) handleConnecting(ctx context.Context, elem xmpp.XElement) {
 func (s *inStream) unauthenticatedFeatures() []xmpp.XElement {
 	var features []xmpp.XElement
 
-	isSocketTr := s.cfg.transport.Type() == transport.Socket
+	isSocketTr := s.tr.Type() == transport.Socket
 
 	if isSocketTr && !s.IsSecured() {
 		startTLS := xmpp.NewElementName("starttls")
@@ -302,7 +304,7 @@ func (s *inStream) unauthenticatedFeatures() []xmpp.XElement {
 func (s *inStream) authenticatedFeatures() []xmpp.XElement {
 	var features []xmpp.XElement
 
-	isSocketTr := s.cfg.transport.Type() == transport.Socket
+	isSocketTr := s.tr.Type() == transport.Socket
 
 	// attach compression feature
 	compressionAvailable := isSocketTr && s.cfg.compression.Level != compress.NoCompression
@@ -444,7 +446,7 @@ func (s *inStream) proceedStartTLS(ctx context.Context, elem xmpp.XElement) {
 	s.setSecured(true)
 	s.writeElement(ctx, xmpp.NewElementNamespace("proceed", tlsNamespace))
 
-	s.cfg.transport.StartTLS(&tls.Config{Certificates: s.router.Hosts().Certificates()}, false)
+	s.tr.StartTLS(&tls.Config{Certificates: s.router.Hosts().Certificates()}, false)
 
 	log.Infof("secured stream... id: %s", s.id)
 	s.restartSession()
@@ -470,7 +472,7 @@ func (s *inStream) compress(ctx context.Context, elem xmpp.XElement) {
 	}
 	s.writeElement(ctx, xmpp.NewElementNamespace("compressed", compressProtocolNamespace))
 
-	s.cfg.transport.EnableCompression(s.cfg.compression.Level)
+	s.tr.EnableCompression(s.cfg.compression.Level)
 	s.setCompressed(true)
 
 	log.Infof("compressed stream... id: %s", s.id)
@@ -698,7 +700,9 @@ sendMessage:
 
 // Runs on it's own goroutine
 func (s *inStream) doRead() {
+	s.scheduleReadTimeout()
 	elem, sErr := s.sess.Receive()
+	s.cancelReadTimeout()
 
 	ctx, _ := context.WithTimeout(context.Background(), s.cfg.timeout)
 	if sErr == nil {
@@ -796,12 +800,14 @@ func (s *inStream) disconnectClosingSession(ctx context.Context, closeSession, u
 	if unbind {
 		s.router.Unbind(ctx, s.JID())
 	}
+	s.ctxCancelFn()
+
 	// notify disconnection
 	if s.cfg.onDisconnect != nil {
 		s.cfg.onDisconnect(s)
 	}
 	s.setState(disconnected)
-	_ = s.cfg.transport.Close()
+	_ = s.tr.Close()
 
 	s.runQueue.Stop(nil) // stop processing messages
 }
@@ -831,9 +837,8 @@ func (s *inStream) isBlockedJID(ctx context.Context, j *jid.JID) bool {
 func (s *inStream) restartSession() {
 	s.sess = session.New(s.id, &session.Config{
 		JID:           s.JID(),
-		Transport:     s.cfg.transport,
 		MaxStanzaSize: s.cfg.maxStanzaSize,
-	}, s.router.Hosts())
+	}, s.tr, s.router.Hosts())
 	s.setState(connecting)
 }
 
@@ -885,6 +890,25 @@ func (s *inStream) setSessionStarted(sessStarted bool) {
 	s.sessStarted = sessStarted
 }
 
+func (s *inStream) scheduleReadTimeout() {
+	s.mu.Lock()
+	s.readTimeoutTm = time.AfterFunc(s.cfg.keepAlive, s.readTimeout)
+	s.mu.Unlock()
+}
+
+func (s *inStream) cancelReadTimeout() {
+	s.mu.Lock()
+	s.readTimeoutTm.Stop()
+	s.mu.Unlock()
+}
+
+func (s *inStream) readTimeout() {
+	s.runQueue.Run(func() {
+		ctx, _ := context.WithTimeout(context.Background(), s.cfg.timeout)
+		s.disconnect(ctx, streamerror.ErrConnectionTimeout)
+	})
+}
+
 func (s *inStream) setState(state uint32) {
 	atomic.StoreUint32(&s.state, state)
 }

c2s/in_test.go

@@ -423,10 +423,11 @@ func tUtilStreamStartSession(conn *fakeSocketConn, t *testing.T) {
 
 func tUtilStreamInit(r router.Router, userRep repository.User, blockListRep repository.BlockList) (*inStream, *fakeSocketConn) {
 	conn := newFakeSocketConn()
-	tr := transport.NewSocketTransport(conn, 4096)
+	tr := transport.NewSocketTransport(conn)
 	stm := newStream(
 		"abcd1234",
-		tUtilInStreamDefaultConfig(tr),
+		tUtilInStreamDefaultConfig(),
+		tr,
 		tUtilInitModules(r),
 		&component.Components{},
 		r,
@@ -435,10 +436,10 @@ func tUtilStreamInit(r router.Router, userRep repository.User, blockListRep repo
 	return stm.(*inStream), conn
 }
 
-func tUtilInStreamDefaultConfig(tr transport.Transport) *streamConfig {
+func tUtilInStreamDefaultConfig() *streamConfig {
 	return &streamConfig{
 		connectTimeout:   time.Second,
-		transport:        tr,
+		keepAlive:        time.Second,
 		maxStanzaSize:    8192,
 		resourceConflict: Reject,
 		compression:      CompressConfig{Level: compress.DefaultCompression},

c2s/server.go

@@ -7,13 +7,13 @@ package c2s
 
 import (
 	"context"
-	"crypto/tls"
 	"fmt"
 	"net"
 	"net/http"
 	"strconv"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/gorilla/websocket"
 	"github.com/ortuman/jackal/component"
@@ -67,9 +67,6 @@ func (s *server) start() {
 	switch s.cfg.Transport.Type {
 	case transport.Socket:
 		err = s.listenSocketConn(address)
-	case transport.WebSocket:
-		err = s.listenWebSocketConn(address)
-		break
 	}
 	if err != nil {
 		log.Fatalf("%v", err)
@@ -87,40 +84,13 @@ func (s *server) listenSocketConn(address string) error {
 	for atomic.LoadUint32(&s.listening) == 1 {
 		conn, err := ln.Accept()
 		if err == nil {
-			go s.startStream(transport.NewSocketTransport(conn, s.cfg.Transport.KeepAlive))
+			go s.startStream(transport.NewSocketTransport(conn), s.cfg.KeepAlive)
 			continue
 		}
 	}
 	return nil
 }
 
-func (s *server) listenWebSocketConn(address string) error {
-	http.HandleFunc(s.cfg.Transport.URLPath, s.websocketUpgrade)
-
-	s.wsSrv = &http.Server{TLSConfig: &tls.Config{Certificates: s.router.Hosts().Certificates()}}
-	s.wsUpgrader = &websocket.Upgrader{
-		Subprotocols: []string{"xmpp"},
-		CheckOrigin:  func(r *http.Request) bool { return r.Header.Get("Sec-WebSocket-Protocol") == "xmpp" },
-	}
-
-	// start listening
-	ln, err := listenerProvider("tcp", address)
-	if err != nil {
-		return err
-	}
-	atomic.StoreUint32(&s.listening, 1)
-	return s.wsSrv.ServeTLS(ln, "", "")
-}
-
-func (s *server) websocketUpgrade(w http.ResponseWriter, r *http.Request) {
-	conn, err := s.wsUpgrader.Upgrade(w, r, nil)
-	if err != nil {
-		log.Error(err)
-		return
-	}
-	s.startStream(transport.NewWebSocketTransport(conn, s.cfg.Transport.KeepAlive))
-}
-
 func (s *server) shutdown(ctx context.Context) error {
 	if atomic.CompareAndSwapUint32(&s.listening, 1, 0) {
 		// stop listening
@@ -129,10 +99,6 @@ func (s *server) shutdown(ctx context.Context) error {
 			if err := s.ln.Close(); err != nil {
 				return err
 			}
-		case transport.WebSocket:
-			if err := s.wsSrv.Shutdown(ctx); err != nil {
-				return err
-			}
 		}
 		// close all connections
 		c, err := s.closeConnections(ctx)
@@ -144,18 +110,18 @@ func (s *server) shutdown(ctx context.Context) error {
 	return nil
 }
 
-func (s *server) startStream(tr transport.Transport) {
+func (s *server) startStream(tr transport.Transport, keepAlive time.Duration) {
 	cfg := &streamConfig{
-		transport:        tr,
 		resourceConflict: s.cfg.ResourceConflict,
 		connectTimeout:   s.cfg.ConnectTimeout,
+		keepAlive:        s.cfg.KeepAlive,
 		timeout:          s.cfg.Timeout,
 		maxStanzaSize:    s.cfg.MaxStanzaSize,
 		sasl:             s.cfg.SASL,
 		compression:      s.cfg.Compression,
 		onDisconnect:     s.unregisterStream,
 	}
-	stm := newStream(s.nextID(), cfg, s.mods, s.comps, s.router, s.userRep, s.blockListRep)
+	stm := newStream(s.nextID(), cfg, tr, s.mods, s.comps, s.router, s.userRep, s.blockListRep)
 	s.registerStream(stm)
 }