-
Notifications
You must be signed in to change notification settings - Fork 1
/
server.go
317 lines (276 loc) · 8.77 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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
package server
import (
"context"
"errors"
"fmt"
"github.com/aldas/go-modbus-client/packet"
"io"
"log"
"net"
"os"
"sync"
"sync/atomic"
"time"
)
const (
readTimeout = 5 * time.Millisecond
writeTimeout = 50 * time.Millisecond
idleTimeout = 25 * time.Second
)
var (
// ErrServerClosed is returned when server context is ended (by shutdown)
ErrServerClosed = errors.New("modbus server closed")
)
// PacketAssembler is called when server reads data from client connection. Is responsible for assembling data read
// from the connection to whole modbus packet.
//
// return with closeConnection=true when you are done sending and want to close connection
type PacketAssembler interface {
ReceiveRead(ctx context.Context, received []byte, bytesRead int) (response []byte, closeConnection bool)
}
// RawReadTracer is PacketAssembler optional interface that is called for each Read from connection to allow tracing read results
type RawReadTracer interface {
Read(data []byte, n int, err error)
}
// ModbusHandler calls Handle method when it has received enough data to be parsed into Modbus packet.
type ModbusHandler interface {
Handle(ctx context.Context, received packet.Request) (packet.Response, error)
}
// Server simple TCP server implementation for server to serve modbus packets.
// Each connection is handled in separate goroutine, in which panics are recovered.
//
// Public fields are not designed to be goroutine safe. Do not mutate after server has been started
type Server struct {
mu sync.RWMutex
listener net.Listener // for simplicity, we only allow serving one listener
isShutdown atomic.Bool
activeConnections map[*connection]struct{}
activeConnectionCount atomic.Int64
// WriteTimeout is amount of time writing the request can take after it errors out
WriteTimeout time.Duration
// ReadTimeout is amount of time reading the response can take
ReadTimeout time.Duration
// AssemblerCreatorFunc creates Assembler for each connetion to assemble different read byte fragments into complete
// modbus packet. Could have different implementations for TCP or RTU packets
AssemblerCreatorFunc func(handler ModbusHandler) PacketAssembler
// OnServeFunc allows capturing listener address just before server starts to accepting connections. This is useful
// for testing when listener is started with random port `:0`.
OnServeFunc func(addr net.Addr)
OnErrorFunc func(err error)
// OnAcceptConnFunc is called when server accepts new connection. When method returns an error the connection will be closed.
// connectionCount indicated currently active connection count.
//
// This is where firewall rules and other limits can be implemented
OnAcceptConnFunc func(ctx context.Context, remoteAddr net.Addr, connectionCount uint64) error
// OnCloseConnFunc is called at the end of connection. isServerShutdown indicated if method is called at server shutdown.
OnCloseConnFunc func(ctx context.Context, remoteAddr net.Addr, isServerShutdown bool)
}
type connection struct {
conn net.Conn
isBeingHandled atomic.Bool
assembler PacketAssembler
writeTimeout time.Duration
readTimeout time.Duration
onErrorFunc func(error)
}
// ListenAndServe starts accepting connection on given address and handles received data with handler function.
// Method blocks until context is cancelled
func (s *Server) ListenAndServe(ctx context.Context, address string, handler ModbusHandler) error {
listener, err := net.Listen("tcp", address)
if err != nil {
return fmt.Errorf("modbus listnener creation error: %w", err)
}
return s.serve(ctx, listener, handler)
}
// Serve accepts connections from listener and handles received data with handler function.
// Method blocks until context is cancelled
func (s *Server) Serve(ctx context.Context, listener net.Listener, handler ModbusHandler) error {
return s.serve(ctx, listener, handler)
}
// ContextRemoteAddr is context.Context value containing clients remote address
type ContextRemoteAddr struct{}
func (s *Server) serve(ctx context.Context, listener net.Listener, handler ModbusHandler) error {
if s.AssemblerCreatorFunc == nil {
s.AssemblerCreatorFunc = func(handler ModbusHandler) PacketAssembler {
return &ModbusTCPAssembler{Handler: handler}
}
}
onErrorFunc := s.OnErrorFunc
if onErrorFunc == nil {
onErrorFunc = func(err error) {
log.Printf("modbus server connection error: %v", err)
}
}
if s.OnServeFunc != nil {
// when listener is started with ":0" (random port) this will be helpful knowing where to connect
// and if server is listening already
s.OnServeFunc(listener.Addr())
}
s.listener = listener
l := onceCloseListener{Listener: listener}
defer l.Close()
for {
netConn, err := l.Accept()
if err != nil {
if s.isShutdown.Load() {
return ErrServerClosed
}
return err
}
if s.OnAcceptConnFunc != nil {
if err := s.OnAcceptConnFunc(ctx, netConn.RemoteAddr(), uint64(s.activeConnectionCount.Load())); err != nil {
if err := netConn.Close(); err != nil {
onErrorFunc(fmt.Errorf("connection.close error, err: %w", err))
}
continue
}
}
select {
case <-ctx.Done():
return ErrServerClosed
default:
}
cCtx := context.WithValue(ctx, ContextRemoteAddr{}, netConn.RemoteAddr())
c := &connection{
conn: netConn,
isBeingHandled: atomic.Bool{},
assembler: s.AssemblerCreatorFunc(handler),
writeTimeout: s.WriteTimeout,
readTimeout: s.ReadTimeout,
onErrorFunc: onErrorFunc,
}
s.trackConn(c, true)
go func(ctx context.Context, conn *connection) {
defer func() {
if rec := recover(); rec != nil {
conn.onErrorFunc(fmt.Errorf("recovered panic in handler, %v", rec))
}
if err := conn.conn.Close(); err != nil {
conn.onErrorFunc(fmt.Errorf("failed to close handler connection, err: %w", err))
}
s.trackConn(c, false)
if s.OnAcceptConnFunc != nil {
s.OnCloseConnFunc(ctx, conn.conn.RemoteAddr(), s.isShutdown.Load())
}
}()
conn.handle(ctx)
}(cCtx, c)
}
}
type onceCloseListener struct {
net.Listener
once sync.Once
closeErr error
}
func (oc *onceCloseListener) Close() error {
oc.once.Do(oc.close)
return oc.closeErr
}
func (oc *onceCloseListener) close() {
oc.closeErr = oc.Listener.Close()
}
func (s *Server) trackConn(c *connection, isAdd bool) {
// this is how http.Server does it
s.mu.Lock()
defer s.mu.Unlock()
if s.activeConnections == nil {
s.activeConnections = make(map[*connection]struct{})
}
if isAdd {
s.activeConnections[c] = struct{}{}
s.activeConnectionCount.Add(1)
} else {
delete(s.activeConnections, c)
s.activeConnectionCount.Add(-1)
}
}
func (c *connection) handle(ctx context.Context) {
cCtx, cCancel := context.WithCancel(ctx)
defer cCancel()
rTimeout := readTimeout
if c.readTimeout > 0 {
rTimeout = c.readTimeout
}
wTimeout := writeTimeout
if c.writeTimeout > 0 {
wTimeout = c.writeTimeout
}
rrt, debugRawRead := c.assembler.(RawReadTracer)
conn := c.conn
var lastReceived time.Time
received := make([]byte, 300)
for {
select {
case <-cCtx.Done():
return
default:
}
_ = conn.SetReadDeadline(time.Now().Add(rTimeout))
n, err := conn.Read(received)
if debugRawRead {
rrt.Read(received[0:n], n, err)
}
if err != nil && !errors.Is(err, os.ErrDeadlineExceeded) {
if !errors.Is(err, io.EOF) {
c.onErrorFunc(err)
}
return // when read fails due some unknown error we close connection
}
if n > 0 {
lastReceived = time.Now()
} else if time.Now().Sub(lastReceived) > idleTimeout {
return // close idle connection
} else {
continue // nothing read and not idle yet
}
c.isBeingHandled.Store(true)
toSend, closeConn := c.assembler.ReceiveRead(cCtx, received[0:n], n)
if toSend != nil {
_ = conn.SetWriteDeadline(time.Now().Add(wTimeout))
if _, err := conn.Write(toSend); err != nil {
c.onErrorFunc(err)
return // when write fails to client we close connection
}
}
c.isBeingHandled.Store(false)
if closeConn {
return
}
}
}
// Addr returns currently running server address
func (s *Server) Addr() net.Addr {
s.mu.RLock()
defer s.mu.RUnlock()
return s.listener.Addr()
}
// Shutdown gracefully shuts down the server without interrupting any active connections.
// Works similarly as `http.Server.Shutdown()`
func (s *Server) Shutdown(ctx context.Context) error {
s.mu.Lock()
defer s.mu.Unlock()
s.isShutdown.Store(true)
err := s.listener.Close()
timer := time.NewTimer(50 * time.Millisecond)
defer timer.Stop()
for {
allIdle := true
for c := range s.activeConnections {
if c.isBeingHandled.Load() {
allIdle = false
continue
}
(*c).conn.Close()
delete(s.activeConnections, c)
}
if allIdle {
return err
}
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
timer.Reset(50 * time.Millisecond)
}
}
}