forked from ortuman/jackal
-
Notifications
You must be signed in to change notification settings - Fork 0
/
server.go
156 lines (136 loc) · 3.6 KB
/
server.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
/*
* Copyright (c) 2018 Miguel Ángel Ortuño.
* See the LICENSE file for more information.
*/
package s2s
import (
"context"
"net"
"strconv"
"sync"
"sync/atomic"
"github.com/ortuman/jackal/errors"
"github.com/ortuman/jackal/log"
"github.com/ortuman/jackal/module"
"github.com/ortuman/jackal/router"
"github.com/ortuman/jackal/stream"
"github.com/ortuman/jackal/transport"
)
var listenerProvider = net.Listen
type server struct {
cfg *Config
router *router.Router
mods *module.Modules
dialer *dialer
inConns sync.Map
outConns sync.Map
ln net.Listener
listening uint32
}
func (s *server) start() {
bindAddr := s.cfg.Transport.BindAddress
port := s.cfg.Transport.Port
address := bindAddr + ":" + strconv.Itoa(port)
log.Infof("s2s_in: listening at %s", address)
if err := s.listenConn(address); err != nil {
log.Fatalf("%v", err)
}
}
func (s *server) shutdown(ctx context.Context) error {
if atomic.CompareAndSwapUint32(&s.listening, 1, 0) {
// stop listening...
if err := s.ln.Close(); err != nil {
return err
}
// close all connections...
c, err := closeConnections(&s.outConns, ctx)
if err != nil {
return err
}
log.Infof("%s: closed %d out connection(s)", s.cfg.ID, c)
c, err = closeConnections(&s.inConns, ctx)
if err != nil {
return err
}
log.Infof("%s: closed %d in connection(s)", s.cfg.ID, c)
}
return nil
}
func (s *server) listenConn(address string) error {
ln, err := listenerProvider("tcp", address)
if err != nil {
return err
}
s.ln = ln
atomic.StoreUint32(&s.listening, 1)
for atomic.LoadUint32(&s.listening) == 1 {
conn, err := ln.Accept()
if err == nil {
go s.startInStream(transport.NewSocketTransport(conn, s.cfg.Transport.KeepAlive))
continue
}
}
return nil
}
func (s *server) getOrDial(localDomain, remoteDomain string) (stream.S2SOut, error) {
domainPair := localDomain + ":" + remoteDomain
stm, loaded := s.outConns.LoadOrStore(domainPair, newOutStream(s.router))
if !loaded {
outCfg, err := s.dialer.dial(localDomain, remoteDomain)
if err != nil {
log.Error(err)
s.outConns.Delete(domainPair)
return nil, err
}
outCfg.onOutDisconnect = s.unregisterOutStream
stm.(*outStream).start(outCfg)
log.Infof("registered s2s out stream... (domainpair: %s)", domainPair)
}
return stm.(*outStream), nil
}
func (s *server) unregisterOutStream(stm stream.S2SOut) {
domainPair := stm.ID()
s.outConns.Delete(domainPair)
log.Infof("unregistered s2s out stream... (domainpair: %s)", domainPair)
}
func (s *server) startInStream(tr transport.Transport) {
stm := newInStream(&streamConfig{
keyGen: &keyGen{s.cfg.DialbackSecret},
transport: tr,
connectTimeout: s.cfg.ConnectTimeout,
maxStanzaSize: s.cfg.MaxStanzaSize,
dialer: s.dialer,
onInDisconnect: s.unregisterInStream,
}, s.mods, s.router)
s.registerInStream(stm)
}
func (s *server) registerInStream(stm stream.S2SIn) {
s.inConns.Store(stm.ID(), stm)
log.Infof("registered s2s in stream... (id: %s)", stm.ID())
}
func (s *server) unregisterInStream(stm stream.S2SIn) {
s.inConns.Delete(stm.ID())
log.Infof("unregistered s2s in stream... (id: %s)", stm.ID())
}
func closeConnections(connections *sync.Map, ctx context.Context) (count int, err error) {
connections.Range(func(_, v interface{}) bool {
stm := v.(*inStream)
select {
case <-closeConn(stm):
count++
return true
case <-ctx.Done():
err = ctx.Err()
return false
}
})
return
}
func closeConn(stm stream.InStream) <-chan bool {
c := make(chan bool, 1)
go func() {
stm.Disconnect(streamerror.ErrSystemShutdown)
c <- true
}()
return c
}